Storing multiple similar tree sequences, e.g. for Singer / ARG MCMC output

[hyanwong]

Prompted by @YunDeng98, I thought I would have a go at prototyping a Python class for storing multiple similar tree sequences. I call this a MultiTableCollection. The main idea is to keep everything similar to a tskit TableCollection, but change the edge table so that it contains 2 extra columns, a start_index and an end_index. There's no need to make different node tables: we can easily have nodes present in some of the tree sequences and not in others, as tskit doesn't require nodes to be referenced by any edge in a tree sequence.

The general idea is that you add to the MultiTableCollection by calling .add_row() with the required start_index, and remove an edge from all subsequent tree sequences by calling .end_row() or .end_rows(). At any time you can call MultiTableCollection.tree_sequence(i) to extract the ith tree sequence from the MultiTableCollection (checking consistency is only done when you grab out a tree sequence like this). In my code below, a separate MultiEdgeTable is stored, and subset down to a selection of edges when the .tree_sequence(i) method is called. The implementation below is a bit scrappy, and missing lots of functionality (e.g. you can't save one of these things to disk at the moment, which would require implementing dump and 'load') but it seems to work!

I guess it could form the base for testing e.g. a full C implementation, if required. Implementing it first in python makes it a lot quicker to test the basic principle.

import tskit
import numpy as np

class MultiEdgeTable(tskit.EdgeTable):
    """
    Like a normal edge table, except we have 2 extra columns: start_index and end_index.
    For a prototype, we hack these columns so they are simply stored as python lists of
    integers but we should probably store them as more memory-efficient structures
    (e.g. numpy arrays with space reserved for expanding them, if we stay in Python)
    """
    def __init__(self, max_rows_increment=0, ll_table=None):
        nrows = 0 if ll_table is None else ll_table.num_rows
        # Extra storage 
        self.start_index = [0] * nrows
        # Use tskit.NULL to indicate the edge persists for the rest of the indexes
        self.end_index = [tskit.NULL] * nrows
        super().__init__(max_rows_increment, ll_table)

    def clear(self):
        self.start_index = []
        self.end_index = []
        super().clear()
    
    def asdict(self):
        # Used to create a new MultiEdgeTable
        ret = super().asdict()
        ret["start_index"] = np.array(self.start_index, dtype=int)
        ret["end_index"] = np.array(self.end_index, dtype=int)
        return ret
        
    def _text_header_and_rows(self, *args, **kwargs):
        # useful to be able to print a MultiEdgeTable for debug purposes
        headers, rows = super()._text_header_and_rows(*args, **kwargs)
        headers = headers[:-1] + ("start_index", "end_index") + headers[-1:]
        for r in rows:
            if len(r) > 1:  # Don't alter the "__skipped__" row
                id_ = int(r[0])
                md = r.pop()  # Hack: stick metadata on the end of the printout
                r.append(str(self.start_index[id_]))
                r.append(str(self.end_index[id_]))
                r.append(md)
        return headers, rows
        
    def set_columns(
        self,
        left=None,
        right=None,
        parent=None,
        child=None,
        start_index=None,
        end_index=None,
        metadata=None,
        metadata_offset=None,
        metadata_schema=None,
    ):
        # required for replace_with to work
        super().set_columns(
            left, right, parent, child,
            metadata=metadata, metadata_offset=metadata_offset, metadata_schema=metadata_schema
        )
        for idx, name, val in [(start_index, "start_index", 0), (end_index, "end_index", tskit.NULL)]:
            if idx is None:
                idx = [val] * self.num_rows
            else:
                if len(idx) != self.num_rows:
                    raise ValueError(f"{name} must be same length as other fields")
                try:
                    idx = idx.tolist()  # convert from numpy array
                except AttributeError:
                    pass
            setattr(self, name, idx)

    
    def replace_with(self, other):
        # Overwrite the contents of this table with a copy of the other table
        # This will also work if passed a standard tskit.EdgeTable
        self.set_columns(**other.asdict())

    def __getattr__(self, name):
        if name == "start_index":
            return np.array(self.start_index, dtype=int)
        if name == "end_index":
            return np.array(self.start_index, dtype=int)
        return super().__getattr__(name)
    
    def add_row(self, left, right, parent, child, start_index=0, metadata=None):
        self.start_index.append(start_index)
        self.end_index.append(tskit.NULL)
        return super().add_row(left, right, parent, child, metadata)

    def end_row(self, id_, end_index):
        """
        A MultiEdgeTable-only function, to set the last index over which an edge is valid
        """
        self.end_index[id_] = end_index

    def end_rows(self, ids_, end_index):
        """
        A MultiEdgeTable-only function, to set the last index over which a set of edges is valid
        """
        # TODO, if we store self.end_index as a numpy array, we shouldn't need to go via tmp
        tmp = np.array(self.end_index)
        tmp[ids_] = end_index
        self.end_index = tmp.tolist()


class MultiTableCollection(tskit.TableCollection):
    def __init__(self, sequence_length=0, *, ll_tables=None):
        super().__init__(sequence_length, ll_tables=ll_tables)
        # Make a new memory store for the edges, not associated with the underlying _ll_tables
        # (we clear the existing edges, and fill them when returning a TS)
        self.edges.clear()
        self._edges = MultiEdgeTable()

    @staticmethod
    def from_TableCollection(other):
        # Make a completely new copy of all the data
        tc = MultiTableCollection(other.sequence_length)
        tc._individuals.replace_with(other._individuals)
        tc._nodes.replace_with(other._nodes)
        tc._edges.replace_with(other._edges)
        tc._migrations.replace_with(other._migrations)
        tc._sites.replace_with(other._sites)
        tc._mutations.replace_with(other._mutations)
        tc._populations.replace_with(other._populations)
        tc._provenances.replace_with(other._provenances)
        tc.time_units = other.time_units
        return tc

    def _to_TableCollection(self, index):
        # This can simply take the lower level table data, writing into the Edge Table
        tables = tskit.TableCollection(sequence_length=self.sequence_length, ll_tables=self._ll_tables)
        self._set_ll_edge_table(index)
        return tables

    def _set_ll_edge_table(self, index):
        # overwrite the lower level edge table (which is not being used to store the
        # MultiEdgeTable data), so that we can return a standard tskit.TableCollection
        # backed by the underlying lower level memory storage.
        # TODO: use a (numpy-like) `view` into the MultiEdge table, to avoid having to copy memory around

        # Finding the array of edges to use (below) should probably be sped up using some sort of index
        np_end = np.array(self.edges.end_index)
        use = np.logical_and(
            index >= np.array(self.edges.start_index),
            np.logical_or(index < np_end, np_end == tskit.NULL)
        )
        # Partly copied from BaseTable.__getitem__
        use = tskit.util.safe_np_int_cast(np.flatnonzero(use), np.int32)
        self._ll_tables.edges.clear()
        self._ll_tables.edges.extend(self.edges.ll_table, row_indexes=use)

    def tree_sequence(self, index):
        tables = self._to_TableCollection(index)
        tables.sort()
        return tables.tree_sequence()

    #### not implemented ####
    def freeze(self):
        raise NotImplementedError("Should have something to make a MultiTableCollection immutable")

    def sort(self):
        raise NotImplementedError("Should sort edges as in tskit, but break ties by start_index")

Here's a quick test:

import msprime
tables = msprime.simulate(5, recombination_rate=1, random_seed=1).dump_tables()
multi_tables = MultiTableCollection.from_TableCollection(tables)
ts0 = multi_tables.tree_sequence(0)
print("First tree sequence")
print(ts0.draw_text())

# Now make a new tree sequence by changing an edge from the first
edge_above_4 = np.where(ts0.edges_child == 4)[0]
multi_tables.edges.end_rows(edge_above_4, 1)  # End this edge after the first TS
multi_tables.edges.add_row(left=0, right=1, child=4, parent=8, start_index=1)
#  For speed it would be nice to do the following:
# multi_tables.sort()
# multi_tables.freeze()

print("Second tree sequence")
print(multi_tables.tree_sequence(1).draw_text())
print(multi_tables.edges)

First tree sequence
2.50┊     9     ┊      9    ┊  
    ┊   ┏━┻━━┓  ┊   ┏━━┻━━┓ ┊  
0.55┊   8    ┃  ┊   8     ┃ ┊  
    ┊ ┏━┻┓   ┃  ┊ ┏━┻━┓   ┃ ┊  
0.21┊ ┃  ┃   7  ┊ ┃   ┃   ┃ ┊  
    ┊ ┃  ┃  ┏┻┓ ┊ ┃   ┃   ┃ ┊  
0.14┊ ┃  6  ┃ ┃ ┊ ┃   6   ┃ ┊  
    ┊ ┃ ┏┻┓ ┃ ┃ ┊ ┃  ┏┻━┓ ┃ ┊  
0.11┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃  5  ┃ ┃ ┊  
    ┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┏┻┓ ┃ ┃ ┊  
0.00┊ 0 3 4 1 2 ┊ 0 1 3 4 2 ┊  
  0.00        0.60        1.00 

Second tree sequence
2.50┊     9     ┊       9   ┊  
    ┊   ┏━┻━━┓  ┊    ┏━━┻━┓ ┊  
0.55┊   8    ┃  ┊    8    ┃ ┊  
    ┊ ┏━╋━┓  ┃  ┊ ┏━━╋━━┓ ┃ ┊  
0.21┊ ┃ ┃ ┃  7  ┊ ┃  ┃  ┃ ┃ ┊  
    ┊ ┃ ┃ ┃ ┏┻┓ ┊ ┃  ┃  ┃ ┃ ┊  
0.14┊ ┃ 6 ┃ ┃ ┃ ┊ ┃  6  ┃ ┃ ┊  
    ┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃  ┃  ┃ ┃ ┊  
0.11┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃  5  ┃ ┃ ┊  
    ┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┏┻┓ ┃ ┃ ┊  
0.00┊ 0 3 4 1 2 ┊ 0 1 3 4 2 ┊  
  0.00        0.60        1.00 

╔══╤══════════╤══════════╤══════╤═════╤═══════════╤═════════╤════════╗
║id│left      │right     │parent│child│start_index│end_index│metadata║
╠══╪══════════╪══════════╪══════╪═════╪═══════════╪═════════╪════════╣
║0 │0.60162245│1.00000000│     5│    1│          0│       -1│        ║
║1 │0.60162245│1.00000000│     5│    3│          0│       -1│        ║
║2 │0.00000000│0.60162245│     6│    3│          0│       -1│        ║
║3 │0.00000000│1.00000000│     6│    4│          0│        1│        ║
║4 │0.60162245│1.00000000│     6│    5│          0│       -1│        ║
║5 │0.00000000│0.60162245│     7│    1│          0│       -1│        ║
║6 │0.00000000│0.60162245│     7│    2│          0│       -1│        ║
║7 │0.00000000│1.00000000│     8│    0│          0│       -1│        ║
║8 │0.00000000│1.00000000│     8│    6│          0│       -1│        ║
║9 │0.60162245│1.00000000│     9│    2│          0│       -1│        ║
║10│0.00000000│0.60162245│     9│    7│          0│       -1│        ║
║11│0.00000000│1.00000000│     9│    8│          0│       -1│        ║
║12│0.00000000│1.00000000│     8│    4│          1│       -1│        ║
╚══╧══════════╧══════════╧══════╧═════╧═══════════╧═════════╧════════╝

[hyanwong]

FWIW, I implemented a data store of numpy values for giglib, which could serve as a template for converting the self.start_index and self.end_index values to a single 2xN numpy array of ints, here: https://github.com/hyanwong/giglib/blob/2139326ea1da6c5e1265757484ad4b76db6028f2/giglib/tables.py#L152

I think 32 bit ints would be enough because we are unlikely to want to store more than 2147483648 separate ARGs, I would have thought?

[hyanwong]

From a comment on slack, re efficiency of adding new ARGs:

Other than re-implementing this in a different language, the two things that could speed up my prototype are (a) storing the start_index and end_index fields as dynamically-growing numpy arrays, and (b) efficient indexing systems to extract the required subset of edges from the lower-level tskit.EdgeTable stored within a MultiEdgeTable.

Adding new edges is probably fast. Extracting the subset to make a TreeSequence is probably slow.

A bigger change to substantially improve efficiency of traversing through the ARGs would be a way to avoid having to do all the .tree_sequence() checks after extracting the relevant edges. I presume this would involve sorting, indexing, and freezing the entire MultiTableCollection into an immutable form (a "MultiTreeSequence"), in which you are guaranteed to be able to extract the edges in a pre-defined and valid order for each extracted TreeSequence. Implementation is left to the reader...

[hyanwong]

Also some more thoughts: we could do a MultiTableCollection.sort() and simply keep the MultiEdges in the normal order, with an additional (last) sort requirement by start_index. Then we would be guaranteed that in the extraction phase, the edges would be in the required order for a tree sequence, and we wouldn't need to check sorting validity when extracting the tree sequence. Presumably this would save time during extraction, especially if we could request the .tree_sequence() method not to check validity. We could used a freeze() method (unimplemented above) to the MultiTableCollection to make it immutable (as I have done for GIGs), and then we might be able drop to the C level and ask https://tskit.dev/tskit/docs/stable/c-api.html#c.tsk_treeseq_init to take ownership when making a tree sequence, so that we don't have to make a copy of all the data.

I'm not sure how this would work at to extract a "view" of the MultiEdgeTable: I suspect this is only a numpy/python thing, so would need some reworking for a C method to take subsets to make a valid edge table.

All these are elaborations to make the general idea faster, though.