Add genotype_values() method

See tskit-dev/tsinfer#739

python/tests/test_genotypes.py

@@ -655,6 +655,33 @@ def test_snipped_tree_sequence_mutations_over_isolated(self):
         assert non_missing_found
         assert missing_found
 
+    def test_genotype_values(self):
+        tables = tskit.TableCollection(1.0)
+        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
+        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
+        s = tables.sites.add_row(0, "C")
+        tables.mutations.add_row(site=s, derived_state="G", node=0)
+        tables.mutations.add_row(site=s, derived_state="T", node=1)
+        s = tables.sites.add_row(0.5, "")
+        tables.mutations.add_row(site=s, derived_state="A long string", node=0)
+        ts = tables.tree_sequence()
+
+        v = ts.variants(isolated_as_missing=False)
+        vals = next(v).genotype_values()
+        assert vals.dtype.type == np.str_
+        assert np.array_equal(vals, np.array(["G", "T"]))
+        vals = next(v).genotype_values()
+        assert vals.dtype.type == np.str_
+        assert np.array_equal(vals, np.array(["A long string", ""]))
+
+        v = ts.variants(isolated_as_missing=True)
+        vals = next(v).genotype_values()
+        assert vals.dtype.type == np.str_
+        assert np.array_equal(vals, np.array(["G", "T"]))
+        vals = next(v).genotype_values()
+        assert vals.dtype.type == np.object_
+        assert np.array_equal(vals, np.array(["A long string", None]))
+
 
 class TestLimitInterval:
     def test_simple_case(self, ts_fixture):

python/tskit/genotypes.py

@@ -245,6 +245,19 @@ def copy(self) -> Variant:
         variant_copy._ll_variant = self._ll_variant.restricted_copy()
         return variant_copy
 
+    def genotype_values(self) -> np.ndarray:
+        """
+        Returns the genotypes at this site as an numpy array of strings (if
+        there is no missing data) or objects (if the genotypes contain missing data,
+        in which case some elements will be equal to ``None``),
+        rather than an array of integer indexes. Note that this is inefficient
+        compared to working with the underlying integer representation as
+        returned by the :attr:`.genotypes`` property.
+
+        :return: An array of length ``num_sites`` containing strings or objects.
+        """
+        return np.array(self.alleles)[self.genotypes]
+
     def counts(self) -> typing.Counter[str | None]:
         """
         Returns a :class:`python:collections.Counter` object providing counts for each