Make sure that variant.count() returns alleles in an order

And document/test it.

Author: hyanwong

python/CHANGELOG.rst

@@ -7,6 +7,9 @@
 - CLI commands that load a tree sequence now read from stdin when the input
   path argument is omitted. (:user:`chris-a-talbot`, :user:`jeromekelleher`,
   :issue:`3468`, :pr:`3469`)
+- The returned object from ``variant.counts()`` and ``variant.frequencies()``
+  now stores alleles in the order defined in ``variant.alleles``.
+  (:user:`hyanwong`, :pr:`3471`)
 
 --------------------
 [1.0.3] - 2026-05-14

python/tests/test_genotypes.py

@@ -2609,7 +2609,7 @@ def test_variant_counts(self, ts_fixture):
         assert len(variant.alleles) > 2
         assert None in variant.alleles
         counts = variant.counts()
-        assert len(counts) == len(variant.alleles)
+        assert list(counts.keys()) == list(variant.alleles)
         assert np.sum(list(counts.values())) == ts_fixture.num_samples
         assert counts[None] == variant.num_missing
         assert ts_fixture.num_samples > variant.num_missing

python/tskit/genotypes.py

@@ -283,16 +283,19 @@ def counts(self) -> typing.Counter[str | None]:
         possible :attr:`allele <Variant.alleles>` at this site: i.e. the number of
         samples possessing that allele among the set of samples specified when creating
         this Variant (by default, this is all the sample nodes in the tree sequence).
-        Missing data is represented by an allelic state of ``None``.
+        Missing data is represented by an allelic state of ``None``. The order of
+        alleles in the Counter is the same as the order of alleles in the
+        :attr:`alleles <Variant.alleles>` tuple.
 
         :return: A counter of the number of samples associated with each allele.
         """
         counts = collections.Counter()
         if self.alleles[-1] is None:
             # we have to treat the last element of the genotypes array as special
-            counts[None] = np.sum(self.genotypes == tskit.MISSING_DATA)
             for i, allele in enumerate(self.alleles[:-1]):
                 counts[allele] = np.sum(self.genotypes == i)
+            # Add the count of missing data as the last item
+            counts[None] = np.sum(self.genotypes == tskit.MISSING_DATA)
         else:
             bincounts = np.bincount(self.genotypes, minlength=self.num_alleles)
             for i, allele in enumerate(self.alleles):
@@ -308,6 +311,8 @@ def frequencies(self, remove_missing=None) -> dict[str, float]:
         sample nodes in the tree sequence). Note, therefore, that if a restricted set
         of samples was specified on creation, the allele frequencies returned here
         will *not* be the global allele frequencies in the whole tree sequence.
+        The order of alleles in the returned dictionary is the same as the order
+        of alleles in the :attr:`alleles <Variant.alleles>` tuple.
 
         :param bool remove_missing: If True, only samples with non-missing data will
             be counted in the total number of samples used to calculate the frequency,