tweak to bincount

jeromekelleher · Jul 20, 2022 · c12c160 · c12c160
1 parent 33df3ca
commit c12c160
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/tskit/genotypes.py b/python/tskit/genotypes.py
@@ -256,9 +256,9 @@ def counts(self) -> typing.Counter[str | None]:
             for i, allele in enumerate(self.alleles[:-1]):
                 counts[allele] = np.sum(self.genotypes == i)
         else:
-            bincounts = np.bincount(self.genotypes)
+            bincounts = np.bincount(self.genotypes, minlength=self.num_alleles)
             for i, allele in enumerate(self.alleles):
-                counts[allele] = bincounts[i] if i < len(bincounts) else 0
+                counts[allele] = bincounts[i]
         return counts
 
     def frequencies(self, remove_missing=None) -> dict[str, float]: