Expose low-level individual and provenance arrays

Author: benjeffery

python/_tskitmodule.c

@@ -11074,6 +11074,104 @@ TreeSequence_get_individuals_metadata_offset(TreeSequence *self, void *closure)
     return ret;
 }
 
+static PyObject *
+TreeSequence_get_individuals_location(TreeSequence *self, void *closure)
+{
+    PyObject *ret = NULL;
+    tsk_individual_table_t individuals;
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+    individuals = self->tree_sequence->tables->individuals;
+    ret = TreeSequence_make_array(
+        self, individuals.location_length, NPY_FLOAT64, individuals.location);
+out:
+    return ret;
+}
+
+static PyObject *
+TreeSequence_get_individuals_location_offset(TreeSequence *self, void *closure)
+{
+    PyObject *ret = NULL;
+    tsk_individual_table_t individuals;
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+    individuals = self->tree_sequence->tables->individuals;
+    ret = TreeSequence_make_array(
+        self, individuals.num_rows + 1, NPY_UINT64, individuals.location_offset);
+out:
+    return ret;
+}
+
+static PyObject *
+TreeSequence_get_individuals_parents(TreeSequence *self, void *closure)
+{
+    PyObject *ret = NULL;
+    tsk_individual_table_t individuals;
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+    individuals = self->tree_sequence->tables->individuals;
+    ret = TreeSequence_make_array(
+        self, individuals.parents_length, NPY_INT32, individuals.parents);
+out:
+    return ret;
+}
+
+static PyObject *
+TreeSequence_get_individuals_parents_offset(TreeSequence *self, void *closure)
+{
+    PyObject *ret = NULL;
+    tsk_individual_table_t individuals;
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+    individuals = self->tree_sequence->tables->individuals;
+    ret = TreeSequence_make_array(
+        self, individuals.num_rows + 1, NPY_UINT64, individuals.parents_offset);
+out:
+    return ret;
+}
+
+#if HAVE_NUMPY_2
+static PyObject *
+TreeSequence_get_provenances_timestamp(TreeSequence *self, void *closure)
+{
+    PyObject *ret = NULL;
+    tsk_provenance_table_t provenances;
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+    provenances = self->tree_sequence->tables->provenances;
+    ret = TreeSequence_decode_ragged_string_column(
+        self, provenances.num_rows, provenances.timestamp, provenances.timestamp_offset);
+out:
+    return ret;
+}
+
+static PyObject *
+TreeSequence_get_provenances_record(TreeSequence *self, void *closure)
+{
+    PyObject *ret = NULL;
+    tsk_provenance_table_t provenances;
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+    provenances = self->tree_sequence->tables->provenances;
+    ret = TreeSequence_decode_ragged_string_column(
+        self, provenances.num_rows, provenances.record, provenances.record_offset);
+out:
+    return ret;
+}
+#endif
+
 static PyObject *
 TreeSequence_get_nodes_time(TreeSequence *self, void *closure)
 {
@@ -12005,6 +12103,18 @@ static PyGetSetDef TreeSequence_getsetters[] = {
     { .name = "individuals_metadata_offset",
         .get = (getter) TreeSequence_get_individuals_metadata_offset,
         .doc = "The individual metadata offset array" },
+    { .name = "individuals_location",
+        .get = (getter) TreeSequence_get_individuals_location,
+        .doc = "The individual location array" },
+    { .name = "individuals_location_offset",
+        .get = (getter) TreeSequence_get_individuals_location_offset,
+        .doc = "The individual location offset array" },
+    { .name = "individuals_parents",
+        .get = (getter) TreeSequence_get_individuals_parents,
+        .doc = "The individual parents array" },
+    { .name = "individuals_parents_offset",
+        .get = (getter) TreeSequence_get_individuals_parents_offset,
+        .doc = "The individual parents offset array" },
     { .name = "nodes_time",
         .get = (getter) TreeSequence_get_nodes_time,
         .doc = "The node time array" },
@@ -12117,6 +12227,14 @@ static PyGetSetDef TreeSequence_getsetters[] = {
     { .name = "indexes_edge_removal_order",
         .get = (getter) TreeSequence_get_indexes_edge_removal_order,
         .doc = "The edge removal order array" },
+#if HAVE_NUMPY_2
+    { .name = "provenances_timestamp",
+        .get = (getter) TreeSequence_get_provenances_timestamp,
+        .doc = "The provenance timestamp array" },
+    { .name = "provenances_record",
+        .get = (getter) TreeSequence_get_provenances_record,
+        .doc = "The provenance record array" },
+#endif
     { NULL } /* Sentinel */
 };
 

python/tests/test_lowlevel.py

@@ -1178,6 +1178,10 @@ class TestTreeSequence(LowLevelTestCase, MetadataTestMixin):
 
     ARRAY_NAMES = [
         "individuals_flags",
+        "individuals_location",
+        "individuals_location_offset",
+        "individuals_parents",
+        "individuals_parents_offset",
         "individuals_metadata",
         "individuals_metadata_offset",
         "nodes_time",
@@ -2197,6 +2201,8 @@ def test_generated_columns(self, ts_fixture, name):
             "sites_ancestral_state",
             "mutations_derived_state",
             "mutations_inherited_state",
+            "provenances_timestamp",
+            "provenances_record",
         ],
     )
     @pytest.mark.parametrize(
@@ -2218,6 +2224,12 @@ def test_string_arrays(self, ts_fixture, str_lengths, string_array):
                 elif string_array == "mutations_inherited_state":
                     assert ts.num_mutations > 0
                     assert {len(mut.inherited_state) for mut in ts.mutations()} == {1}
+                elif string_array == "provenances_timestamp":
+                    assert ts.num_provenances > 0
+                    assert len(ts.provenance(3).timestamp) == 1
+                elif string_array == "provenances_record":
+                    assert ts.num_provenances > 0
+                    assert len(ts.provenance(3).record) == 1
             else:
                 tables = ts_fixture.dump_tables()
 
@@ -2266,6 +2278,22 @@ def test_string_arrays(self, ts_fixture, str_lengths, string_array):
                                 derived_state=get_derived_state(i, mutation)
                             )
                         )
+                elif string_array == "provenances_timestamp":
+                    provenances = tables.provenances.copy()
+                    tables.provenances.clear()
+                    get_timestamp = str_map[str_lengths]
+                    for i, provenance in enumerate(provenances):
+                        tables.provenances.append(
+                            provenance.replace(timestamp=get_timestamp(i, provenance))
+                        )
+                elif string_array == "provenances_record":
+                    provenances = tables.provenances.copy()
+                    tables.provenances.clear()
+                    get_record = str_map[str_lengths]
+                    for i, provenance in enumerate(provenances):
+                        tables.provenances.append(
+                            provenance.replace(record=get_record(i, provenance))
+                        )
 
                 ts = tables.tree_sequence()
         ll_ts = ts.ll_tree_sequence
@@ -2285,6 +2313,12 @@ def test_string_arrays(self, ts_fixture, str_lengths, string_array):
             elif string_array == "mutations_inherited_state":
                 for mutation in ts.mutations():
                     assert a[mutation.id] == mutation.inherited_state
+            elif string_array == "provenances_timestamp":
+                for provenance in ts.provenances():
+                    assert a[provenance.id] == provenance.timestamp
+            elif string_array == "provenances_record":
+                for provenance in ts.provenances():
+                    assert a[provenance.id] == provenance.record
 
         # Read only
         with pytest.raises(AttributeError, match="not writable"):