Add mutation.inherited_state

Author: benjeffery

c/tests/test_trees.c

@@ -3998,24 +3998,31 @@ test_single_tree_good_mutations(void)
     CU_ASSERT_EQUAL(other_mutations[0].id, 0);
     CU_ASSERT_EQUAL(other_mutations[0].node, 2);
     CU_ASSERT_NSTRING_EQUAL(other_mutations[0].derived_state, "1", 1);
+    CU_ASSERT_NSTRING_EQUAL(other_mutations[0].inherited_state, "0", 1);
     CU_ASSERT_EQUAL(other_mutations[1].id, 1);
     CU_ASSERT_EQUAL(other_mutations[1].node, 4);
     CU_ASSERT_NSTRING_EQUAL(other_mutations[1].derived_state, "1", 1);
+    CU_ASSERT_NSTRING_EQUAL(other_mutations[1].inherited_state, "0", 1);
     CU_ASSERT_EQUAL(other_mutations[2].id, 2);
     CU_ASSERT_EQUAL(other_mutations[2].node, 0);
     CU_ASSERT_NSTRING_EQUAL(other_mutations[2].derived_state, "0", 1);
+    CU_ASSERT_NSTRING_EQUAL(other_mutations[2].inherited_state, "1", 1);
     CU_ASSERT_EQUAL(other_mutations[3].id, 3);
     CU_ASSERT_EQUAL(other_mutations[3].node, 0);
     CU_ASSERT_NSTRING_EQUAL(other_mutations[3].derived_state, "1", 1);
+    CU_ASSERT_NSTRING_EQUAL(other_mutations[3].inherited_state, "0", 1);
     CU_ASSERT_EQUAL(other_mutations[4].id, 4);
     CU_ASSERT_EQUAL(other_mutations[4].node, 1);
     CU_ASSERT_NSTRING_EQUAL(other_mutations[4].derived_state, "1", 1);
+    CU_ASSERT_NSTRING_EQUAL(other_mutations[4].inherited_state, "0", 1);
     CU_ASSERT_EQUAL(other_mutations[5].id, 5);
     CU_ASSERT_EQUAL(other_mutations[5].node, 2);
     CU_ASSERT_NSTRING_EQUAL(other_mutations[5].derived_state, "1", 1);
+    CU_ASSERT_NSTRING_EQUAL(other_mutations[5].inherited_state, "0", 1);
     CU_ASSERT_EQUAL(other_mutations[6].id, 6);
     CU_ASSERT_EQUAL(other_mutations[6].node, 3);
     CU_ASSERT_NSTRING_EQUAL(other_mutations[6].derived_state, "1", 1);
+    CU_ASSERT_NSTRING_EQUAL(other_mutations[6].inherited_state, "0", 1);
 
     tsk_treeseq_free(&ts);
 }

c/tskit/tables.h

@@ -158,6 +158,10 @@ typedef struct {
     /** @brief The ID of the edge that this mutation lies on, or TSK_NULL
       if there is no corresponding edge.*/
     tsk_id_t edge;
+    /** @brief Inherited state. */
+    const char *inherited_state;
+    /** @brief Size of the inherited state in bytes. */
+    tsk_size_t inherited_state_length;
 } tsk_mutation_t;
 
 /**

c/tskit/trees.c

@@ -253,6 +253,13 @@ tsk_treeseq_init_trees(tsk_treeseq_t *self)
     const tsk_size_t num_nodes = self->tables->nodes.num_rows;
     const double *restrict site_position = self->tables->sites.position;
     const tsk_id_t *restrict mutation_site = self->tables->mutations.site;
+    const tsk_id_t *restrict mutation_parent = self->tables->mutations.parent;
+    const char *restrict sites_ancestral_state = self->tables->sites.ancestral_state;
+    const tsk_size_t *restrict sites_ancestral_state_offset
+        = self->tables->sites.ancestral_state_offset;
+    const char *restrict mutations_derived_state = self->tables->mutations.derived_state;
+    const tsk_size_t *restrict mutations_derived_state_offset
+        = self->tables->mutations.derived_state_offset;
     const tsk_id_t *restrict I = self->tables->indexes.edge_insertion_order;
     const tsk_id_t *restrict O = self->tables->indexes.edge_removal_order;
     const double *restrict edge_right = self->tables->edges.right;
@@ -311,6 +318,26 @@ tsk_treeseq_init_trees(tsk_treeseq_t *self)
                 mutation_id < num_mutations && mutation_site[mutation_id] == site_id) {
                 mutation = self->site_mutations_mem + mutation_id;
                 mutation->edge = node_edge_map[mutation->node];
+
+                /* Compute inherited state */
+                if (mutation_parent[mutation_id] == TSK_NULL) {
+                    /* No parent: inherited state is the site's ancestral state */
+                    mutation->inherited_state
+                        = sites_ancestral_state + sites_ancestral_state_offset[site_id];
+                    mutation->inherited_state_length
+                        = sites_ancestral_state_offset[site_id + 1]
+                          - 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];
+                    mutation->inherited_state
+                        = mutations_derived_state
+                          + mutations_derived_state_offset[parent_id];
+                    mutation->inherited_state_length
+                        = mutations_derived_state_offset[parent_id + 1]
+                          - mutations_derived_state_offset[parent_id];
+                }
+
                 mutation_id++;
             }
             site_id++;
@@ -5228,6 +5255,9 @@ tsk_treeseq_get_mutation(
         goto out;
     }
     mutation->edge = self->site_mutations_mem[index].edge;
+    mutation->inherited_state = self->site_mutations_mem[index].inherited_state;
+    mutation->inherited_state_length
+        = self->site_mutations_mem[index].inherited_state_length;
 out:
     return ret;
 }

docs/c-api.rst

@@ -65,7 +65,7 @@ arbitrarily between releases.
     stability, since the primary use-case is for tskit to be embedded
     within another application rather than used as a shared library. If you
     do intend to use tskit as a shared library and ABI stability is
-    therefore imporant to you, please let us know and we can plan
+    therefore important to you, please let us know and we can plan
     accordingly.
 
 .. _sec_c_api_overview_structure:

python/CHANGELOG.rst

@@ -30,13 +30,14 @@
 - Add ``TreeSequence.mutations_edge`` which returns the edge ID for each mutation's
   edge. (:user:`benjeffery`, :pr:`3226`, :issue:`3189`)
 
-- Add ``TreeSequence.mutations_inherited_state`` which returns the inherited state
+- Add  which returns the inherited state
   for each mutation. (:user:`benjeffery`, :pr:`3276`, :issue:`2631`)
 
-- Add ``TreeSequence.sites_ancestral_state`` and ``TreeSequence.mutations_derived_state`` properties
-  to return the ancestral state of sites and derived state of mutations as NumPy arrays of
+- 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
   the new numpy 2.0 StringDType.
-  (:user:`benjeffery`, :pr:`3228`, :issue:`2632`)
+  (:user:`benjeffery`, :pr:`3228`, :issue:`2632`, :pr:`3276`, :issue:`2631`)
 
 - Tskit now distributes with a requirement of numpy version 2 or greater. However, you can still use
   tskit with numpy 1.X by building tskit from source with numpy 1.X using ``pip install tskit --no-binary tskit``.
@@ -45,6 +46,9 @@
   the error "A module that was compiled using NumPy 1.x cannot be run in NumPy 2.0.0 as it may crash.".
   If no newer version of the module is available you will have to use the Numpy 1.X build as above.
 
+- Add ``Mutation.inherited_state`` property which returns the inherited state
+  for a single mutation. (:user:`benjeffery`, :pr:`3277`, :issue:`2631`)
+
 **Bugfixes**
 
 - In some tables with mutations out-of-order `TableCollection.sort` did not re-order

python/_tskitmodule.c

@@ -363,9 +363,10 @@ make_mutation_object(const tsk_mutation_t *mutation)
     if (metadata == NULL) {
         goto out;
     }
-    ret = Py_BuildValue("iis#iOdi", mutation->site, mutation->node,
+    ret = Py_BuildValue("iis#iOdis#", mutation->site, mutation->node,
         mutation->derived_state, (Py_ssize_t) mutation->derived_state_length,
-        mutation->parent, metadata, mutation->time, mutation->edge);
+        mutation->parent, metadata, mutation->time, mutation->edge,
+        mutation->inherited_state, (Py_ssize_t) mutation->inherited_state_length);
 out:
     Py_XDECREF(metadata);
     return ret;

python/tests/__init__.py

@@ -150,6 +150,7 @@ def make_mutation(id_):
                 metadata,
                 time,
                 edge,
+                inherited_state,
             ) = ll_ts.get_mutation(id_)
             return tskit.Mutation(
                 id=id_,
@@ -160,6 +161,7 @@ def make_mutation(id_):
                 parent=parent,
                 metadata=metadata,
                 edge=edge,
+                inherited_state=inherited_state,
                 metadata_decoder=tskit.metadata.parse_metadata_schema(
                     ll_ts.get_table_metadata_schemas().mutation
                 ).decode_row,

python/tests/test_highlevel.py

@@ -1218,6 +1218,20 @@ def test_trees(self, ts):
     def test_mutations(self, ts):
         self.verify_mutations(ts)
 
+    @pytest.mark.parametrize("ts", tsutil.get_example_tree_sequences())
+    def test_mutation_inherited_state_property(self, ts):
+        inherited_states = ts.mutations_inherited_state
+        for mut in ts.mutations():
+            expected = inherited_states[mut.id]
+            actual = mut.inherited_state
+            assert actual == expected
+
+            if mut.parent == tskit.NULL:
+                expected_direct = ts.site(mut.site).ancestral_state
+            else:
+                expected_direct = ts.mutation(mut.parent).derived_state
+            assert actual == expected_direct
+
     def verify_pairwise_diversity(self, ts):
         haplotypes = ts.genotype_matrix(isolated_as_missing=False).T
         if ts.num_samples == 0:

python/tests/test_lowlevel.py

@@ -1336,14 +1336,24 @@ def test_get_mutation_interface(self):
             assert len(mutations) == ts.get_num_mutations()
             # Check the form of the mutations
             for packed in mutations:
-                site, node, derived_state, parent, metadata, time, edge = packed
+                (
+                    site,
+                    node,
+                    derived_state,
+                    parent,
+                    metadata,
+                    time,
+                    edge,
+                    inherited_state,
+                ) = packed
                 assert isinstance(site, int)
                 assert isinstance(node, int)
                 assert isinstance(derived_state, str)
                 assert isinstance(parent, int)
                 assert isinstance(metadata, bytes)
                 assert isinstance(time, float)
                 assert isinstance(edge, int)
+                assert isinstance(inherited_state, str)
 
     def test_get_edge_interface(self):
         for ts in self.get_example_tree_sequences():
@@ -3867,6 +3877,7 @@ def test_sites(self):
                             metadata,
                             time,
                             edge,
+                            inherited_state,
                         ) = ts.get_mutation(mut_id)
                         assert site == index
                         assert mutation_id == mut_id

python/tskit/trees.py

@@ -425,6 +425,7 @@ class Mutation(util.Dataclass):
         "metadata",
         "time",
         "edge",
+        "inherited_state",
     ]
     id: int  # noqa A003
     """
@@ -473,6 +474,13 @@ class Mutation(util.Dataclass):
     """
     The ID of the edge that this mutation is on.
     """
+    inherited_state: str
+    """
+    The inherited state for this mutation. This is the state that existed at the site
+    before this mutation occurred. This is either the ancestral state of the site
+    (if the mutation has no parent) or the derived state of the mutation's
+    parent mutation (if it has a parent).
+    """
 
     # To get default values on slots we define a custom init
     def __init__(
@@ -485,6 +493,7 @@ def __init__(
         parent=NULL,
         metadata=b"",
         edge=NULL,
+        inherited_state=None,
     ):
         self.id = id
         self.site = site
@@ -494,6 +503,7 @@ def __init__(
         self.parent = parent
         self.metadata = metadata
         self.edge = edge
+        self.inherited_state = inherited_state
 
     # We need a custom eq to compare unknown times.
     def __eq__(self, other):
@@ -6333,6 +6343,7 @@ def mutation(self, id_):
             metadata,
             time,
             edge,
+            inherited_state,
         ) = self._ll_tree_sequence.get_mutation(id_)
         return Mutation(
             id=id_,
@@ -6343,6 +6354,7 @@ def mutation(self, id_):
             metadata=metadata,
             time=time,
             edge=edge,
+            inherited_state=inherited_state,
             metadata_decoder=self.table_metadata_schemas.mutation.decode_row,
         )