Move inherited_state array to CPython

Author: benjeffery

c/CHANGELOG.rst

@@ -14,6 +14,15 @@
   compatible with the correct mutation parent.
   (:user:`benjeffery`, :issue:`2729`, :issue:`2732`, :pr:`3212`).
 
+- Mutations returned by ``tsk_treeseq_get_mutation`` now include pre-computed
+  ``inherited_state`` and ``inherited_state_length`` fields. The inherited state
+  is computed during tree sequence initialization and represents the state that
+  existed at the site before each mutation occurred (either the ancestral state
+  if the mutation is the root mutation or the derived state of the parent mutation).
+  Note that this breaks ABI compatibility due to the addition of these fields
+  to the ``tsk_mutation_t`` struct.
+  (:user:`benjeffery`, :pr:`3277`, :issue:`2631`).
+
 **Features**
 
 - Add ``TSK_TS_INIT_COMPUTE_MUTATION_PARENTS`` to ``tsk_treeseq_init``

c/tskit/trees.c

@@ -269,6 +269,7 @@ tsk_treeseq_init_trees(tsk_treeseq_t *self)
     bool discrete_breakpoints = true;
     tsk_id_t *node_edge_map = tsk_malloc(num_nodes * sizeof(*node_edge_map));
     tsk_mutation_t *mutation;
+    tsk_id_t parent_id;
 
     self->tree_sites_length
         = tsk_malloc(num_trees_alloc * sizeof(*self->tree_sites_length));
@@ -329,7 +330,7 @@ tsk_treeseq_init_trees(tsk_treeseq_t *self)
                           - sites_ancestral_state_offset[site_id];
                 } else {
                     /* Has parent: inherited state is parent's derived state */
-                    tsk_id_t parent_id = mutation_parent[mutation_id];
+                    parent_id = mutation_parent[mutation_id];
                     mutation->inherited_state
                         = mutations_derived_state
                           + mutations_derived_state_offset[parent_id];

docs/python-api.md

@@ -92,6 +92,7 @@ for more information.
   TreeSequence.edges_parent
   TreeSequence.edges_child
   TreeSequence.sites_position
+  TreeSequence.sites_ancestral_state
   TreeSequence.mutations_site
   TreeSequence.mutations_node
   TreeSequence.mutations_parent

python/CHANGELOG.rst

@@ -30,9 +30,6 @@
 - Add ``TreeSequence.mutations_edge`` which returns the edge ID for each mutation's
   edge. (:user:`benjeffery`, :pr:`3226`, :issue:`3189`)
 
-- Add  which returns the inherited state
-  for each mutation. (:user:`benjeffery`, :pr:`3276`, :issue:`2631`)
-
 - Add ``TreeSequence.sites_ancestral_state``, ``TreeSequence.mutations_derived_state`` and
   ``TreeSequence.mutations_inherited_state`` properties  to return the ancestral state of sites,
   derived state of mutations and inherited state of mutations as NumPy arrays of

python/_tskitmodule.c

@@ -11399,6 +11399,59 @@ TreeSequence_get_mutations_derived_state(TreeSequence *self, void *closure)
 out:
     return ret;
 }
+static PyObject *
+TreeSequence_get_mutations_inherited_state(TreeSequence *self, void *closure)
+{
+    PyObject *ret = NULL;
+    tsk_treeseq_t *ts;
+    tsk_size_t num_mutations;
+    char *inherited_state_data = NULL;
+    tsk_size_t *inherited_state_offsets = NULL;
+    tsk_size_t total_length = 0;
+    tsk_size_t j, offset;
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+
+    ts = self->tree_sequence;
+    num_mutations = ts->tables->mutations.num_rows;
+
+    /* Calculate total length needed for inherited state data */
+    for (j = 0; j < num_mutations; j++) {
+        total_length += ts->site_mutations_mem[j].inherited_state_length;
+    }
+
+    /* Allocate memory for the ragged array */
+    inherited_state_data = PyMem_Malloc(total_length * sizeof(char));
+    inherited_state_offsets = PyMem_Malloc((num_mutations + 1) * sizeof(tsk_size_t));
+    if (inherited_state_data == NULL || inherited_state_offsets == NULL) {
+        PyErr_NoMemory();
+        goto out;
+    }
+
+    /* Populate the ragged array data */
+    offset = 0;
+    for (j = 0; j < num_mutations; j++) {
+        inherited_state_offsets[j] = offset;
+        memcpy(inherited_state_data + offset, ts->site_mutations_mem[j].inherited_state,
+            ts->site_mutations_mem[j].inherited_state_length);
+        offset += ts->site_mutations_mem[j].inherited_state_length;
+    }
+    inherited_state_offsets[num_mutations] = offset;
+
+    ret = TreeSequence_decode_ragged_string_column(
+        self, num_mutations, inherited_state_data, inherited_state_offsets);
+
+out:
+    if (inherited_state_data != NULL) {
+        PyMem_Free(inherited_state_data);
+    }
+    if (inherited_state_offsets != NULL) {
+        PyMem_Free(inherited_state_offsets);
+    }
+    return ret;
+}
 #endif
 
 static PyObject *
@@ -12018,6 +12071,9 @@ static PyGetSetDef TreeSequence_getsetters[] = {
     { .name = "mutations_derived_state",
         .get = (getter) TreeSequence_get_mutations_derived_state,
         .doc = "The mutation derived state array" },
+    { .name = "mutations_inherited_state",
+        .get = (getter) TreeSequence_get_mutations_inherited_state,
+        .doc = "The mutation inherited state array" },
 #endif
     { .name = "mutations_metadata",
         .get = (getter) TreeSequence_get_mutations_metadata,

python/tests/test_highlevel.py

@@ -5560,6 +5560,14 @@ def test_equality_mutations_derived_state(self, ts):
             [mutation.derived_state for mutation in ts.mutations()],
         )
 
+    @pytest.mark.skipif(not _tskit.HAS_NUMPY_2, reason="Requires NumPy 2.0 or higher")
+    @pytest.mark.parametrize("ts", tsutil.get_example_tree_sequences())
+    def test_equality_mutations_inherited_state(self, ts):
+        assert_array_equal(
+            ts.mutations_inherited_state,
+            [mutation.inherited_state for mutation in ts.mutations()],
+        )
+
     @pytest.mark.skipif(not _tskit.HAS_NUMPY_2, reason="Requires NumPy 2.0 or higher")
     @pytest.mark.parametrize("ts", tsutil.get_example_tree_sequences())
     def test_mutations_inherited_state(self, ts):

python/tests/test_jit.py

@@ -420,6 +420,10 @@ def test_jitwrap_properties(ts):
     assert numba_ts.mutations_time.dtype == np.float64
     nt.assert_array_equal(numba_ts.mutations_derived_state, ts.mutations_derived_state)
     assert numba_ts.mutations_derived_state.dtype.kind == "U"  # Unicode string
+    nt.assert_array_equal(
+        numba_ts.mutations_inherited_state, ts.mutations_inherited_state
+    )
+    assert numba_ts.mutations_inherited_state.dtype.kind == "U"  # Unicode string
     nt.assert_array_equal(
         numba_ts.indexes_edge_insertion_order, ts.indexes_edge_insertion_order
     )

python/tests/test_lowlevel.py

@@ -2192,7 +2192,12 @@ def test_generated_columns(self, ts_fixture, name):
 
     @pytest.mark.skipif(not _tskit.HAS_NUMPY_2, reason="Requires NumPy 2.0+")
     @pytest.mark.parametrize(
-        "string_array", ["sites_ancestral_state", "mutations_derived_state"]
+        "string_array",
+        [
+            "sites_ancestral_state",
+            "mutations_derived_state",
+            "mutations_inherited_state",
+        ],
     )
     @pytest.mark.parametrize(
         "str_lengths",
@@ -2210,6 +2215,9 @@ def test_string_arrays(self, ts_fixture, str_lengths, string_array):
                 elif string_array == "mutations_derived_state":
                     assert ts.num_mutations > 0
                     assert {len(mut.derived_state) for mut in ts.mutations()} == {1}
+                elif string_array == "mutations_inherited_state":
+                    assert ts.num_mutations > 0
+                    assert {len(mut.inherited_state) for mut in ts.mutations()} == {1}
             else:
                 tables = ts_fixture.dump_tables()
 
@@ -2239,6 +2247,25 @@ def test_string_arrays(self, ts_fixture, str_lengths, string_array):
                                 derived_state=get_derived_state(i, mutation)
                             )
                         )
+                elif string_array == "mutations_inherited_state":
+                    # For inherited state, we modify sites and mutations to create
+                    # varied lengths
+                    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))
+                        )
+                    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
@@ -2255,6 +2282,9 @@ def test_string_arrays(self, ts_fixture, str_lengths, string_array):
             elif string_array == "mutations_derived_state":
                 for mutation in ts.mutations():
                     assert a[mutation.id] == mutation.derived_state
+            elif string_array == "mutations_inherited_state":
+                for mutation in ts.mutations():
+                    assert a[mutation.id] == mutation.inherited_state
 
         # Read only
         with pytest.raises(AttributeError, match="not writable"):

python/tskit/jit/numba.py

@@ -402,9 +402,11 @@ def __init__(
         mutations_parent,
         mutations_time,
         mutations_derived_state,
+        mutations_inherited_state,
         breakpoints,
         max_ancestral_length,
         max_derived_length,
+        max_inherited_length,
     ):
         self.num_trees = num_trees
         self.num_nodes = num_nodes
@@ -431,9 +433,11 @@ def __init__(
         self.mutations_parent = mutations_parent
         self.mutations_time = mutations_time
         self.mutations_derived_state = mutations_derived_state
+        self.mutations_inherited_state = mutations_inherited_state
         self.breakpoints = breakpoints
         self.max_ancestral_length = max_ancestral_length
         self.max_derived_length = max_derived_length
+        self.max_inherited_length = max_inherited_length
 
     def tree_index(self):
         """
@@ -526,7 +530,7 @@ def parent_index(self):
 
 # We cache these classes to avoid repeated JIT compilation
 @functools.lru_cache(None)
-def _jitwrap(max_ancestral_length, max_derived_length):
+def _jitwrap(max_ancestral_length, max_derived_length, max_inherited_length):
     # We have a circular dependency in JIT compilation between NumbaTreeSequence
     # and NumbaTreeIndex so we used a deferred type to break it
     tree_sequence_type = numba.deferred_type()
@@ -576,9 +580,14 @@ def _jitwrap(max_ancestral_length, max_derived_length):
         ("mutations_parent", numba.int32[:]),
         ("mutations_time", numba.float64[:]),
         ("mutations_derived_state", numba.types.UnicodeCharSeq(max_derived_length)[:]),
+        (
+            "mutations_inherited_state",
+            numba.types.UnicodeCharSeq(max_inherited_length)[:],
+        ),
         ("breakpoints", numba.float64[:]),
         ("max_ancestral_length", numba.int32),
         ("max_derived_length", numba.int32),
+        ("max_inherited_length", numba.int32),
     ]
 
     # The `tree_index` method on NumbaTreeSequence uses NumbaTreeIndex
@@ -614,8 +623,13 @@ def jitwrap(ts):
     """
     max_ancestral_length = max(1, max(map(len, ts.sites_ancestral_state), default=1))
     max_derived_length = max(1, max(map(len, ts.mutations_derived_state), default=1))
+    max_inherited_length = max(
+        1, max(map(len, ts.mutations_inherited_state), default=1)
+    )
 
-    JittedTreeSequence = _jitwrap(max_ancestral_length, max_derived_length)
+    JittedTreeSequence = _jitwrap(
+        max_ancestral_length, max_derived_length, max_inherited_length
+    )
 
     # Create the tree sequence instance
     numba_ts = JittedTreeSequence(
@@ -648,9 +662,13 @@ def jitwrap(ts):
         mutations_derived_state=ts.mutations_derived_state.astype(
             f"U{max_derived_length}"
         ),
+        mutations_inherited_state=ts.mutations_inherited_state.astype(
+            f"U{max_inherited_length}"
+        ),
         breakpoints=ts.breakpoints(as_array=True),
         max_ancestral_length=max_ancestral_length,
         max_derived_length=max_derived_length,
+        max_inherited_length=max_inherited_length,
     )
 
     return numba_ts

python/tskit/trees.py

@@ -6092,14 +6092,14 @@ def mutations_inherited_state(self):
         :return: Array of shape (num_mutations,) containing inherited states.
         :rtype: numpy.ndarray
         """
+        if not _tskit.HAS_NUMPY_2:
+            raise RuntimeError(
+                "The mutations_inherited_state property requires numpy 2.0 or later."
+            )
         if self._mutations_inherited_state is None:
-            inherited_state = self.sites_ancestral_state[self.mutations_site]
-            mutations_with_parent = self.mutations_parent != -1
-            parent = self.mutations_parent[mutations_with_parent]
-            inherited_state[mutations_with_parent] = self.mutations_derived_state[
-                parent
-            ]
-            self._mutations_inherited_state = inherited_state
+            self._mutations_inherited_state = (
+                self._ll_tree_sequence.mutations_inherited_state
+            )
         return self._mutations_inherited_state
 
     @property