Add mutations derived state

Author: benjeffery

python/_tskitmodule.c

@@ -11195,6 +11195,24 @@ TreeSequence_get_mutations_time(TreeSequence *self, void *closure)
     return ret;
 }
 
+#if HAVE_NUMPY_2
+static PyObject *
+TreeSequence_get_mutations_derived_state(TreeSequence *self, void *closure)
+{
+    PyObject *ret = NULL;
+    tsk_mutation_table_t mutations;
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+    mutations = self->tree_sequence->tables->mutations;
+    ret = TreeSequence_decode_ragged_string_column(
+        self, mutations.num_rows, mutations.derived_state, mutations.derived_state_offset);
+out:
+    return ret;
+}
+#endif
+
 static PyObject *
 TreeSequence_get_mutations_metadata(TreeSequence *self, void *closure)
 {
@@ -11792,6 +11810,11 @@ static PyGetSetDef TreeSequence_getsetters[] = {
     { .name = "mutations_time",
         .get = (getter) TreeSequence_get_mutations_time,
         .doc = "The mutation time array" },
+#if HAVE_NUMPY_2
+    { .name = "mutations_derived_state",
+        .get = (getter) TreeSequence_get_mutations_derived_state,
+        .doc = "The mutation derived state array" },
+#endif
     { .name = "mutations_metadata",
         .get = (getter) TreeSequence_get_mutations_metadata,
         .doc = "The mutation metadata array" },

python/tests/test_lowlevel.py

@@ -1976,69 +1976,91 @@ def test_individuals_nodes(self, ts_fixture):
 
     @pytest.mark.skipif(not _tskit.HAS_NUMPY_2, reason="Requires NumPy 2.0+")
     @pytest.mark.parametrize(
-        "site_lengths",
+        "string_array", ["sites_ancestral_state", "mutations_derived_state"]
+    )
+    @pytest.mark.parametrize(
+        "str_lengths",
         ["none", "all-0", "all-1", "all-2", "mixed", "very_long", "unicode"],
     )
-    def test_sites_ancestral_state(self, ts_fixture, site_lengths):
-        if site_lengths == "none":
+    def test_string_arrays(self, ts_fixture, str_lengths, string_array):
+        if str_lengths == "none":
             ts = tskit.TableCollection(1.0).tree_sequence()
         else:
-            if site_lengths == "all-1":
+            if str_lengths == "all-1":
                 ts = ts_fixture
-                assert {len(site.ancestral_state) for site in ts.sites()} == {1}
+                if string_array == "sites_ancestral_state":
+                    assert ts.num_sites > 0
+                    assert {len(site.ancestral_state) for site in ts.sites()} == {1}
+                elif string_array == "mutations_derived_state":
+                    assert ts.num_mutations > 0
+                    assert {len(mut.derived_state) for mut in ts.mutations()} == {1}
             else:
                 tables = ts_fixture.dump_tables()
-                sites = tables.sites.copy()
-                tables.sites.clear()
-
-                ancestral_state_map = {
-                    "all-0": lambda i, site: "",
-                    "all-2": lambda i, site: chr(ord("A") + (i % 26)) * 2,
-                    "mixed": lambda i, site: chr(ord("A") + (i % 26)) * (i % 20),
-                    "very_long": lambda i, site: "A" * 100_000_000 if i == 1 else "A",
-                    "unicode": lambda i, site: "💩" * (i + 1),
+
+                str_map = {
+                    "all-0": lambda i, item: "",
+                    "all-2": lambda i, item: chr(ord("A") + (i % 26)) * 2,
+                    "mixed": lambda i, item: chr(ord("A") + (i % 26)) * (i % 20),
+                    "very_long": lambda i, item: "A" * 100_000_000 if i == 1 else "T",
+                    "unicode": lambda i, item: "🧬" * (i + 1),
                 }
 
-                get_ancestral_state = ancestral_state_map[site_lengths]
-                for i, site in enumerate(sites):
-                    tables.sites.append(
-                        site.replace(ancestral_state=get_ancestral_state(i, site))
-                    )
+                if string_array == "sites_ancestral_state":
+                    sites = tables.sites.copy()
+                    tables.sites.clear()
+                    get_ancestral_state = str_map[str_lengths]
+                    for i, site in enumerate(sites):
+                        tables.sites.append(
+                            site.replace(ancestral_state=get_ancestral_state(i, site))
+                        )
+                elif string_array == "mutations_derived_state":
+                    mutations = tables.mutations.copy()
+                    tables.mutations.clear()
+                    get_derived_state = str_map[str_lengths]
+                    for i, mutation in enumerate(mutations):
+                        tables.mutations.append(
+                            mutation.replace(
+                                derived_state=get_derived_state(i, mutation)
+                            )
+                        )
+
                 ts = tables.tree_sequence()
         ll_ts = ts.ll_tree_sequence
 
-        a = ll_ts.sites_ancestral_state
+        a = getattr(ll_ts, string_array)
+
         # Contents
-        if site_lengths == "none":
+        if str_lengths == "none":
             assert a.size == 0
         else:
-            for site in ts.sites():
-                assert a[site.id] == site.ancestral_state
+            if string_array == "sites_ancestral_state":
+                for site in ts.sites():
+                    assert a[site.id] == site.ancestral_state
+            elif string_array == "mutations_derived_state":
+                for mutation in ts.mutations():
+                    assert a[mutation.id] == mutation.derived_state
 
         # Read only
         with pytest.raises(AttributeError, match="not writable"):
-            ll_ts.sites_ancestral_state = None
+            setattr(ll_ts, string_array, None)
         with pytest.raises(AttributeError, match="not writable"):
-            del ll_ts.sites_ancestral_state
+            delattr(ll_ts, string_array)
 
         with pytest.raises(ValueError, match="assignment destination"):
             a[:] = 0
         with pytest.raises(ValueError, match="assignment destination"):
             a[0] = 0
-        if site_lengths in [("all-1",)]:
-            with pytest.raises(ValueError, match="cannot set WRITEABLE"):
-                a.setflags(write=True)
 
         # Properties
         assert a.dtype == np.dtypes.StringDType()
         assert a.flags.aligned
         assert a.flags.c_contiguous
-        b = ll_ts.sites_ancestral_state
+        b = getattr(ll_ts, string_array)
         assert a is not b
         assert np.all(a == b)
 
         # Lifetime
-        a1 = ll_ts.sites_ancestral_state
+        a1 = getattr(ll_ts, string_array)
         a2 = a1.copy()
         assert a1 is not a2
         del ll_ts