Tree array access.

Closes #1299

Author: jeromekelleher

python/_tskitmodule.c

@@ -8786,6 +8786,121 @@ Tree_set_root_threshold(Tree *self, PyObject *args)
     return ret;
 }
 
+/* The x_array properties are the high-performance zero-copy interface to the
+ * corresponding arrays in the tsk_tree object. We use properties and
+ * return a new array each time rather than trying to create a single array
+ * at Tree initialisation time to avoid a circular reference counting loop,
+ * which (it seems) the even cyclic garbage collection support can't resolve.
+ */
+static PyObject *
+Tree_make_array(Tree *self, int dtype, void *data)
+{
+    PyObject *ret = NULL;
+    PyArrayObject *array = NULL;
+    npy_intp dims = self->tree->num_nodes;
+
+    array = (PyArrayObject *) PyArray_SimpleNewFromData(1, &dims, dtype, data);
+    if (array == NULL) {
+        goto out;
+    }
+    PyArray_CLEARFLAGS(array, NPY_ARRAY_WRITEABLE);
+    if (PyArray_SetBaseObject(array, (PyObject *) self) != 0) {
+        goto out;
+    }
+    /* PyArray_SetBaseObject steals a reference, so we have to incref the tree
+     * object. This makes sure that the Tree instance will stay alive if there
+     * are any arrays that refer to its memory. */
+    Py_INCREF(self);
+    ret = (PyObject *) array;
+    array = NULL;
+out:
+    Py_XDECREF(array);
+    return ret;
+}
+
+static PyObject *
+Tree_get_parent_array(Tree *self, void *closure)
+{
+    PyObject *ret = NULL;
+
+    if (Tree_check_state(self) != 0) {
+        goto out;
+    }
+    ret = Tree_make_array(self, NPY_INT32, self->tree->parent);
+out:
+    return ret;
+}
+
+static PyObject *
+Tree_get_left_child_array(Tree *self, void *closure)
+{
+    PyObject *ret = NULL;
+
+    if (Tree_check_state(self) != 0) {
+        goto out;
+    }
+    ret = Tree_make_array(self, NPY_INT32, self->tree->left_child);
+out:
+    return ret;
+}
+
+static PyObject *
+Tree_get_right_child_array(Tree *self, void *closure)
+{
+    PyObject *ret = NULL;
+
+    if (Tree_check_state(self) != 0) {
+        goto out;
+    }
+    ret = Tree_make_array(self, NPY_INT32, self->tree->right_child);
+out:
+    return ret;
+}
+
+static PyObject *
+Tree_get_left_sib_array(Tree *self, void *closure)
+{
+    PyObject *ret = NULL;
+
+    if (Tree_check_state(self) != 0) {
+        goto out;
+    }
+    ret = Tree_make_array(self, NPY_INT32, self->tree->left_sib);
+out:
+    return ret;
+}
+
+static PyObject *
+Tree_get_right_sib_array(Tree *self, void *closure)
+{
+    PyObject *ret = NULL;
+
+    if (Tree_check_state(self) != 0) {
+        goto out;
+    }
+    ret = Tree_make_array(self, NPY_INT32, self->tree->right_sib);
+out:
+    return ret;
+}
+
+static PyGetSetDef Tree_getsetters[]
+    = { { .name = "parent_array",
+            .get = (getter) Tree_get_parent_array,
+            .doc = "The parent array in the quintuply linked tree." },
+          { .name = "left_child_array",
+              .get = (getter) Tree_get_left_child_array,
+              .doc = "The left_child array in the quintuply linked tree." },
+          { .name = "right_child_array",
+              .get = (getter) Tree_get_right_child_array,
+              .doc = "The right_child array in the quintuply linked tree." },
+          { .name = "left_sib_array",
+              .get = (getter) Tree_get_left_sib_array,
+              .doc = "The left_sib array in the quintuply linked tree." },
+          { .name = "right_sib_array",
+              .get = (getter) Tree_get_right_sib_array,
+              .doc = "The right_sib array in the quintuply linked tree." },
+          { NULL } };
+
 static PyMethodDef Tree_methods[] = {
     { .ml_name = "first",
         .ml_meth = (PyCFunction) Tree_first,
@@ -8961,6 +9076,7 @@ static PyTypeObject TreeType = {
     .tp_flags = Py_TPFLAGS_DEFAULT,
     .tp_doc = "Tree objects",
     .tp_methods = Tree_methods,
+    .tp_getset = Tree_getsetters,
     .tp_init = (initproc) Tree_init,
     .tp_new = PyType_GenericNew,
     // clang-format on

python/tests/test_highlevel.py

@@ -2692,6 +2692,39 @@ def test_interval(self):
                 breakpoints[i + 1] - breakpoints[i]
             )
 
+    def verify_tree_arrays(self, tree):
+        ts = tree.tree_sequence
+        assert tree.parent_array.shape == (ts.num_nodes,)
+        assert tree.left_child_array.shape == (ts.num_nodes,)
+        assert tree.right_child_array.shape == (ts.num_nodes,)
+        assert tree.left_sib_array.shape == (ts.num_nodes,)
+        assert tree.right_sib_array.shape == (ts.num_nodes,)
+        for u in range(ts.num_nodes):
+            assert tree.parent(u) == tree.parent_array[u]
+            assert tree.left_child(u) == tree.left_child_array[u]
+            assert tree.right_child(u) == tree.right_child_array[u]
+            assert tree.left_sib(u) == tree.left_sib_array[u]
+            assert tree.right_sib(u) == tree.right_sib_array[u]
+
+    def test_tree_arrays(self):
+        ts = msprime.simulate(10, recombination_rate=1, random_seed=1)
+        assert ts.num_trees > 1
+        for tree in ts.trees():
+            self.verify_tree_arrays(tree)
+
+    @pytest.mark.parametrize(
+        "array", ["parent", "left_child", "right_child", "left_sib", "right_sib"]
+    )
+    def test_tree_array_properties(self, array):
+        name = array + "_array"
+        ts = msprime.simulate(10, random_seed=1)
+        tree = ts.first()
+        a = getattr(tree, name)
+        assert getattr(tree, name) is a
+        assert a.base is tree._ll_tree
+        with pytest.raises(AttributeError):
+            setattr(tree, name, None)
+
     def verify_empty_tree(self, tree):
         ts = tree.tree_sequence
         assert tree.index == -1
@@ -2711,6 +2744,7 @@ def verify_empty_tree(self, tree):
                 assert tree.left_sib(samples[j]) == samples[j - 1]
             if j < ts.num_samples - 1:
                 assert tree.right_sib(samples[j]) == samples[j + 1]
+        self.verify_tree_arrays(tree)
 
     def test_empty_tree(self):
         ts = msprime.simulate(10, recombination_rate=3, length=3, random_seed=42)
@@ -2741,21 +2775,11 @@ def test_clear(self):
     def verify_trees_identical(self, t1, t2):
         assert t1.tree_sequence is t2.tree_sequence
         assert t1.num_nodes is t2.num_nodes
-        assert [t1.parent(u) for u in range(t1.num_nodes)] == [
-            t2.parent(u) for u in range(t2.num_nodes)
-        ]
-        assert [t1.left_child(u) for u in range(t1.num_nodes)] == [
-            t2.left_child(u) for u in range(t2.num_nodes)
-        ]
-        assert [t1.right_child(u) for u in range(t1.num_nodes)] == [
-            t2.right_child(u) for u in range(t2.num_nodes)
-        ]
-        assert [t1.left_sib(u) for u in range(t1.num_nodes)] == [
-            t2.left_sib(u) for u in range(t2.num_nodes)
-        ]
-        assert [t1.right_sib(u) for u in range(t1.num_nodes)] == [
-            t2.right_sib(u) for u in range(t2.num_nodes)
-        ]
+        assert np.all(t1.parent_array == t2.parent_array)
+        assert np.all(t1.left_child_array == t2.left_child_array)
+        assert np.all(t1.right_child_array == t2.right_child_array)
+        assert np.all(t1.left_sib_array == t2.left_sib_array)
+        assert np.all(t1.right_sib_array == t2.right_sib_array)
         assert list(t1.sites()) == list(t2.sites())
 
     def test_copy_seek(self):

python/tests/test_lowlevel.py

@@ -2550,6 +2550,76 @@ def test_map_mutations_errors(self):
             with pytest.raises(_tskit.LibraryError):
                 tree.map_mutations(genotypes)
 
+    @pytest.mark.parametrize(
+        "array", ["parent", "left_child", "right_child", "left_sib", "right_sib"]
+    )
+    def test_array_read_only(self, array):
+        name = array + "_array"
+        ts1 = self.get_example_tree_sequence(10)
+        t1 = _tskit.Tree(ts1)
+        t1.first()
+        with pytest.raises(AttributeError, match="not writable"):
+            setattr(t1, name, None)
+        a = getattr(t1, name)
+        with pytest.raises(ValueError, match="assignment destination"):
+            a[:] = 0
+        with pytest.raises(ValueError, match="assignment destination"):
+            a[0] = 0
+        with pytest.raises(ValueError, match="cannot set WRITEABLE"):
+            a.setflags(write=True)
+
+    @pytest.mark.parametrize(
+        "array", ["parent", "left_child", "right_child", "left_sib", "right_sib"]
+    )
+    def test_array_properties(self, array):
+        ts1 = self.get_example_tree_sequence(10)
+        t1 = _tskit.Tree(ts1)
+        a = getattr(t1, array + "_array")
+        t1.first()
+        a = getattr(t1, array + "_array")
+        assert a.dtype == np.int32
+        assert a.shape == (ts1.get_num_nodes(),)
+        assert a.base == t1
+        assert not a.flags.writeable
+        assert a.flags.aligned
+        assert a.flags.c_contiguous
+        b = getattr(t1, array + "_array")
+        assert a is not b
+        assert np.all(a == b)
+        a_copy = a.copy()
+        # This checks that the underlying pointer to memory is the same in
+        # both arrays.
+        assert a.__array_interface__ == b.__array_interface__
+        t1.next()
+        # NB! Because we are pointing to the underlying memory, the arrays
+        # will change as we iterate along the trees! This is a gotcha, but
+        # it's just something we have to document as it's a consequence of the
+        # zero copy semantics.
+        b = getattr(t1, array + "_array")
+        assert np.all(a == b)
+        assert np.any(a_copy != b)
+
+    @pytest.mark.parametrize(
+        "array", ["parent", "left_child", "right_child", "left_sib", "right_sib"]
+    )
+    def test_array_lifetime(self, array):
+        ts1 = self.get_example_tree_sequence(10)
+        t1 = _tskit.Tree(ts1)
+        t1.first()
+        a1 = getattr(t1, array + "_array")
+        a2 = a1.copy()
+        assert a1 is not a2
+        del t1
+        # Do some memory operations
+        a3 = np.ones(10 ** 6)
+        assert np.all(a1 == a2)
+        del ts1
+        assert np.all(a1 == a2)
+        del a1
+        # Just do something to touch memory
+        a2[:] = 0
+        assert a3 is not a2
+
 
 class TestTableMetadataSchema(MetadataTestMixin):
     def test_metadata_schema_attribute(self):

python/tskit/trees.py

@@ -669,6 +669,14 @@ def __init__(
         self._tree_sequence = tree_sequence
         self._ll_tree = _tskit.Tree(tree_sequence.ll_tree_sequence, **kwargs)
         self._ll_tree.set_root_threshold(root_threshold)
+        # Store the low-level arrays for efficiency. There's no real overhead
+        # in this, because the refer to the same underlying memory as the
+        # tree object.
+        self._parent_array = self._ll_tree.parent_array
+        self._left_child_array = self._ll_tree.left_child_array
+        self._right_child_array = self._ll_tree.right_child_array
+        self._left_sib_array = self._ll_tree.left_sib_array
+        self._right_sib_array = self._ll_tree.right_sib_array
 
     def copy(self):
         """
@@ -1014,6 +1022,10 @@ def parent(self, u):
         """
         return self._ll_tree.get_parent(u)
 
+    @property
+    def parent_array(self):
+        return self._parent_array
+
     # Quintuply linked tree structure.
 
     def left_child(self, u):
@@ -1035,6 +1047,10 @@ def left_child(self, u):
         """
         return self._ll_tree.get_left_child(u)
 
+    @property
+    def left_child_array(self):
+        return self._left_child_array
+
     def right_child(self, u):
         """
         Returns the rightmost child of the specified node. Returns
@@ -1054,6 +1070,10 @@ def right_child(self, u):
         """
         return self._ll_tree.get_right_child(u)
 
+    @property
+    def right_child_array(self):
+        return self._right_child_array
+
     def left_sib(self, u):
         """
         Returns the sibling node to the left of u, or :data:`tskit.NULL`
@@ -1069,6 +1089,10 @@ def left_sib(self, u):
         """
         return self._ll_tree.get_left_sib(u)
 
+    @property
+    def left_sib_array(self):
+        return self._left_sib_array
+
     def right_sib(self, u):
         """
         Returns the sibling node to the right of u, or :data:`tskit.NULL`
@@ -1084,6 +1108,10 @@ def right_sib(self, u):
         """
         return self._ll_tree.get_right_sib(u)
 
+    @property
+    def right_sib_array(self):
+        return self._right_sib_array
+
     # Sample list.
 
     def left_sample(self, u):