Add highlevel API

Author: benjeffery

python/tests/test_highlevel.py

@@ -5321,6 +5321,53 @@ def test_mixed_sample_status(self):
         assert_array_equal(result, expected)
 
 
+class TestRaggedArrays:
+    @pytest.mark.parametrize("num_rows", [0, 1, 100])
+    @pytest.mark.parametrize("column", ["ancestral_state", "derived_state"])
+    def test_site_ancestral_state(self, num_rows, column):
+        tables = tskit.TableCollection(sequence_length=100)
+        rng = random.Random(42)
+        for i in range(num_rows):
+            state_length = rng.randint(0, 10)
+            state = "".join(
+                chr(rng.randint(0x1F300, 0x1F6FF)) for _ in range(state_length)
+            )
+            if column == "ancestral_state":
+                tables.sites.add_row(position=i, ancestral_state=state)
+            elif column == "derived_state":
+                tables.nodes.add_row()
+                tables.sites.add_row(position=i, ancestral_state="A")
+                tables.mutations.add_row(site=i, node=0, derived_state=state)
+        ts = tables.tree_sequence()
+        a = getattr(
+            ts,
+            (
+                "sites_ancestral_state"
+                if column == "ancestral_state"
+                else "mutations_derived_state"
+            ),
+        )
+        assert isinstance(a, np.ndarray)
+        assert a.shape == (num_rows,)
+        assert a.dtype == np.dtype("T")
+        assert a.size == num_rows
+
+        # Check that the value is cached
+        assert a is getattr(
+            ts,
+            (
+                "sites_ancestral_state"
+                if column == "ancestral_state"
+                else "mutations_derived_state"
+            ),
+        )
+
+        for state, row in itertools.zip_longest(
+            a, ts.sites() if column == "ancestral_state" else ts.mutations()
+        ):
+            assert state == getattr(row, column)
+
+
 class TestSampleNodesByPloidy:
     @pytest.mark.parametrize(
         "n_samples,ploidy,expected",

python/tskit/trees.py

@@ -4138,6 +4138,8 @@ def __init__(self, ll_tree_sequence):
         self._individuals_location = None
         self._individuals_nodes = None
         self._mutations_edge = None
+        self._sites_ancestral_state = None
+        self._mutations_derived_state = None
         # NOTE: when we've implemented read-only access via the underlying
         # tables we can replace these arrays with reference to the read-only
         # tables here (and remove the low-level boilerplate).
@@ -5952,6 +5954,26 @@ def sites_position(self):
         """
         return self._sites_position
 
+    @property
+    def sites_ancestral_state(self):
+        """
+        The ``ancestral_state`` column in the
+        :ref:`sec_site_table_definition` as a numpy array (dtype=StringDtype).
+        """
+        if self._sites_ancestral_state is None:
+            try:
+                self._sites_ancestral_state = (
+                    self._ll_tree_sequence.sites_ancestral_state
+                )
+            except AttributeError:
+                if not _tskit.HAS_NUMPY_2:
+                    raise RuntimeError(
+                        "The sites_ancestral_state property requires numpy 2.0 or later."
+                    )
+                else:
+                    raise
+        return self._sites_ancestral_state
+
     @property
     def sites_metadata(self):
         """
@@ -6009,6 +6031,27 @@ def mutations_time(self):
         """
         return self._mutations_time
 
+    @property
+    def mutations_derived_state(self):
+        """
+        Access to the ``derived_state`` column in the
+        :ref:`sec_mutation_table_definition` as a numpy array (dtype=StringDtype).
+        """
+        if self._mutations_derived_state is None:
+            try:
+                self._mutations_derived_state = (
+                    self._ll_tree_sequence.mutations_derived_state
+                )
+            except AttributeError:
+                if not _tskit.HAS_NUMPY_2:
+                    raise RuntimeError(
+                        "The mutations_derived_state property requires "
+                        "numpy 2.0 or later."
+                    )
+                else:
+                    raise
+        return self._mutations_derived_state
+
     @property
     def mutations_metadata(self):
         """