Add table collection index to LWT interface

Author: benjeffery

c/CHANGELOG.rst

@@ -8,6 +8,9 @@
   criteria (e.g., ignore top-level metadata and schema or provenance tables).
   (:user:`mufernando`, :issue:`896`, :pr:`897`)
 
+- Exposed ``tsk_table_collection_set_index`` to the API.
+  (:user:`benjeffery`, :issue:`870`, :pr:`921`)
+
 ---------------------
 [0.99.7] - 2020-09-29
 ---------------------

c/tskit/tables.c

@@ -7783,7 +7783,7 @@ tsk_table_collection_set_metadata_schema(tsk_table_collection_t *self,
         metadata_schema, metadata_schema_length);
 }
 
-static int
+int
 tsk_table_collection_set_index(tsk_table_collection_t *self,
     tsk_id_t *edge_insertion_order, tsk_id_t *edge_removal_order)
 {

c/tskit/tables.h

@@ -2721,6 +2721,26 @@ collection. Any existing index is first dropped using
 */
 int tsk_table_collection_build_index(tsk_table_collection_t *self, tsk_flags_t options);
 
+/**
+@brief Sets the edge insertion/removal index for this table collection
+
+@rst
+This method sets the edge insertion/removal index for this table collection
+The index arrays should have the same number of edges that are in the
+edge table. The index is not checked for validity.
+
+See the :ref:`sec_c_api_table_indexes` section for details on the index
+life-cycle.
+@endrst
+
+@param self A pointer to a tsk_table_collection_t object.
+@param edge_insertion_order Array of tsk_id_t edge ids
+@param edge_removal_order Array of tsk_id_t edge ids
+@return Return 0 on success or a negative value on failure.
+*/
+int tsk_table_collection_set_index(tsk_table_collection_t *self,
+    tsk_id_t *edge_insertion_order, tsk_id_t *edge_removal_order);
+
 /**
 @brief Runs integrity checks on this table collection.
 

python/CHANGELOG.rst

@@ -17,6 +17,9 @@
 - Added ``TableCollection.index`` for access to the edge insertion/removal order indexes.
   (:user:`benjeffery`, :issue:`4`, :pr:`916`)
 
+- The dictionary representation of a TableCollection now contains its index.
+  (:user:`benjeffery`, :issue:`870`, :pr:`921`)
+
 **Breaking changes**
 
 - The argument to ``ts.dump`` and ``tskit.load`` has been renamed `file` from `path`.

python/lwt_interface/CHANGELOG.rst

@@ -0,0 +1,6 @@
+--------------------
+[0.1.2] - 2020-10-22
+--------------------
+
+ - Added optional top-level key ``index`` which has contains ``edge_insertion_order`` and
+   ``edge_removal_order``

python/lwt_interface/dict_encoding_testlib.py

@@ -148,7 +148,7 @@ def get_example_tables():
 class TestEncodingVersion:
     def test_version(self):
         lwt = lwt_module.LightweightTableCollection()
-        assert lwt.asdict()["encoding_version"] == (1, 1)
+        assert lwt.asdict()["encoding_version"] == (1, 2)
 
 
 class TestRoundTrip:
@@ -270,6 +270,7 @@ def test_missing_tables(self):
             "metadata",
             "metadata_schema",
             "encoding_version",
+            "index",
         }
         for table_name in table_names:
             d = tables.asdict()
@@ -292,6 +293,7 @@ def verify_columns(self, value):
             "metadata",
             "metadata_schema",
             "encoding_version",
+            "index",
         }
         for table_name in table_names:
             table_dict = d[table_name]
@@ -336,6 +338,7 @@ def verify(self, num_rows):
             "metadata",
             "metadata_schema",
             "encoding_version",
+            "index",
         }
         for table_name in sorted(table_names):
             table_dict = d[table_name]
@@ -354,6 +357,30 @@ def test_two_rows(self):
     def test_zero_rows(self):
         self.verify(0)
 
+    def test_bad_index_length(self):
+        tables = get_example_tables()
+        for col in ("insertion", "removal"):
+            d = tables.asdict()
+            d["index"][f"edge_{col}_order"] = d["index"][f"edge_{col}_order"][:-1]
+            lwt = lwt_module.LightweightTableCollection()
+            with pytest.raises(
+                ValueError,
+                match="^edge_insertion_order and"
+                " edge_removal_order must be the same"
+                " length$",
+            ):
+                lwt.fromdict(d)
+        d = tables.asdict()
+        for col in ("insertion", "removal"):
+            d["index"][f"edge_{col}_order"] = d["index"][f"edge_{col}_order"][:-1]
+        lwt = lwt_module.LightweightTableCollection()
+        with pytest.raises(
+            ValueError,
+            match="^edge_insertion_order and edge_removal_order must be"
+            " the same length as the number of edges$",
+        ):
+            lwt.fromdict(d)
+
 
 class TestRequiredAndOptionalColumns:
     """
@@ -563,6 +590,38 @@ def test_provenances(self):
             ["record", "record_offset", "timestamp", "timestamp_offset"],
         )
 
+    def test_index(self):
+        tables = get_example_tables()
+        d = tables.asdict()
+        lwt = lwt_module.LightweightTableCollection()
+        lwt.fromdict(d)
+        other = lwt.asdict()
+        assert np.array_equal(
+            d["index"]["edge_insertion_order"], other["index"]["edge_insertion_order"]
+        )
+        assert np.array_equal(
+            d["index"]["edge_removal_order"], other["index"]["edge_removal_order"]
+        )
+
+        # index is optional
+        d = tables.asdict()
+        del d["index"]
+        lwt = lwt_module.LightweightTableCollection()
+        lwt.fromdict(d)
+
+        # Both columns must be provided, if one is
+        for col in ("insertion", "removal"):
+            d = tables.asdict()
+            del d["index"][f"edge_{col}_order"]
+            lwt = lwt_module.LightweightTableCollection()
+            with pytest.raises(
+                TypeError,
+                match="^edge_insertion_order and "
+                "edge_removal_order must be specified "
+                "together$",
+            ):
+                lwt.fromdict(d)
+
     def test_top_level_metadata(self):
         tables = get_example_tables()
         d = tables.asdict()

python/lwt_interface/tskit_lwt_interface.h

@@ -1172,6 +1172,69 @@ parse_provenance_table_dict(
     return ret;
 }
 
+static int
+parse_index_dict(tsk_table_collection_t *tables, PyObject *dict)
+{
+    int err;
+    int ret = -1;
+    size_t insertion_length, removal_length;
+    PyObject *insertion_input = NULL;
+    PyArrayObject *insertion_array = NULL;
+    PyObject *removal_input = NULL;
+    PyArrayObject *removal_array = NULL;
+
+    /* Get the inputs */
+    insertion_input = get_table_dict_value(dict, "edge_insertion_order", false);
+    if (insertion_input == NULL) {
+        goto out;
+    }
+    removal_input = get_table_dict_value(dict, "edge_removal_order", false);
+    if (removal_input == NULL) {
+        goto out;
+    }
+
+    if ((insertion_input == Py_None) != (removal_input == Py_None)) {
+        PyErr_SetString(PyExc_TypeError,
+            "edge_insertion_order and edge_removal_order must be specified together");
+        goto out;
+    }
+
+    if (insertion_input != Py_None) {
+        insertion_array = table_read_column_array(
+            insertion_input, NPY_INT32, &insertion_length, false);
+        if (insertion_array == NULL) {
+            goto out;
+        }
+        removal_array
+            = table_read_column_array(removal_input, NPY_INT32, &removal_length, false);
+        if (removal_array == NULL) {
+            goto out;
+        }
+        if (insertion_length != removal_length) {
+            PyErr_SetString(PyExc_ValueError,
+                "edge_insertion_order and edge_removal_order must be the same length");
+            goto out;
+        }
+        if (insertion_length != tables->edges.num_rows) {
+            PyErr_SetString(PyExc_ValueError,
+                "edge_insertion_order and edge_removal_order must be "
+                "the same length as the number of edges");
+            goto out;
+        }
+        err = tsk_table_collection_set_index(
+            tables, PyArray_DATA(insertion_array), PyArray_DATA(removal_array));
+        if (err != 0) {
+            handle_tskit_error(err);
+            goto out;
+        }
+    }
+    ret = 0;
+out:
+    Py_XDECREF(insertion_array);
+    Py_XDECREF(removal_array);
+    return ret;
+}
+
 static int
 parse_table_collection_dict(tsk_table_collection_t *tables, PyObject *tables_dict)
 {
@@ -1339,6 +1402,21 @@ parse_table_collection_dict(tsk_table_collection_t *tables, PyObject *tables_dic
         goto out;
     }
 
+    /* index */
+    value = get_table_dict_value(tables_dict, "index", false);
+    if (value == NULL) {
+        goto out;
+    }
+    if (value != Py_None) {
+        if (!PyDict_Check(value)) {
+            PyErr_SetString(PyExc_TypeError, "not a dictionary");
+            goto out;
+        }
+        if (parse_index_dict(tables, value) != 0) {
+            goto out;
+        }
+    }
+
     ret = 0;
 out:
     return ret;
@@ -1480,6 +1558,14 @@ write_table_arrays(tsk_table_collection_t *tables, PyObject *dict)
         { NULL },
     };
 
+    struct table_col index_cols[] = {
+        { "edge_insertion_order", (void *) tables->indexes.edge_insertion_order,
+            tables->indexes.num_edges, NPY_INT32 },
+        { "edge_removal_order", (void *) tables->indexes.edge_removal_order,
+            tables->indexes.num_edges, NPY_INT32 },
+        { NULL },
+    };
+
     struct table_desc table_descs[] = {
         { "individuals", individual_cols, tables->individuals.metadata_schema,
             tables->individuals.metadata_schema_length },
@@ -1496,6 +1582,7 @@ write_table_arrays(tsk_table_collection_t *tables, PyObject *dict)
         { "populations", population_cols, tables->populations.metadata_schema,
             tables->populations.metadata_schema_length },
         { "provenances", provenance_cols, NULL, 0 },
+        { "index", index_cols, NULL, 0 },
     };
 
     for (j = 0; j < sizeof(table_descs) / sizeof(*table_descs); j++) {
@@ -1557,7 +1644,7 @@ dump_tables_dict(tsk_table_collection_t *tables)
     }
 
     /* Dict representation version */
-    val = Py_BuildValue("ll", 1, 1);
+    val = Py_BuildValue("ll", 1, 2);
     if (val == NULL) {
         goto out;
     }

python/tests/test_tables.py

@@ -1309,6 +1309,16 @@ def test_add_row_bad_data(self):
             t.add_row(metadata=[0])
 
 
+class TestTableCollectionIndex:
+    def test_index(self):
+        i = np.arange(20)
+        r = np.arange(20)[::-1]
+        index = tskit.TableCollectionIndex(edge_insertion_order=i, edge_removal_order=r)
+        assert index.edge_insertion_order is i
+        assert index.edge_removal_order is r
+        assert index.asdict() == {"edge_insertion_order": i, "edge_removal_order": r}
+
+
 class TestSortTables:
     """
     Tests for the TableCollection.sort() method.
@@ -2254,6 +2264,7 @@ def test_asdict(self):
             "mutations": t.mutations.asdict(),
             "migrations": t.migrations.asdict(),
             "provenances": t.provenances.asdict(),
+            "index": t.index.asdict(),
         }
         d2 = t.asdict()
         assert set(d1.keys()) == set(d2.keys())
@@ -2278,6 +2289,7 @@ def test_from_dict(self):
             "mutations": t1.mutations.asdict(),
             "migrations": t1.migrations.asdict(),
             "provenances": t1.provenances.asdict(),
+            "index": t1.index.asdict(),
         }
         t2 = tskit.TableCollection.fromdict(d)
         assert t1 == t2

python/tskit/tables.py

@@ -145,6 +145,9 @@ class TableCollectionIndex:
     edge_insertion_order: np.ndarray
     edge_removal_order: np.ndarray
 
+    def asdict(self):
+        return attr.asdict(self)
+
 
 def keep_with_offset(keep, data, offset):
     """
@@ -2110,6 +2113,7 @@ def asdict(self):
             "mutations": self.mutations.asdict(),
             "populations": self.populations.asdict(),
             "provenances": self.provenances.asdict(),
+            "index": self.index.asdict(),
         }
 
     @property

python/tskit/util.py

@@ -248,7 +248,7 @@ def intervals_to_np_array(intervals, start, end):
 def negate_intervals(intervals, start, end):
     """
     Returns the set of intervals *not* covered by the specified set of
-    disjoint intervals in the specfied range.
+    disjoint intervals in the specified range.
     """
     intervals = intervals_to_np_array(intervals, start, end)
     other_intervals = []