Add shift and concatenate functions

Fixes tskit-dev#3164

Update python/tskit/tables.py

Author: hyanwong

python/CHANGELOG.rst

@@ -12,6 +12,11 @@
   associated with each individual as a numpy array.
   (:user:`benjeffery`, :pr:`3153`)
 
+- Add ``shift`` and ``concatenate`` methods to both ``TableCollection`` and
+  ``TreeSequence`` classes, allowing the coordinate system to be shifted and
+  one tree sequence to be added to the right of another.
+  (:user:`hyanwong`, :pr:`3165`, :issue:`3164`)
+
 
 **Fixes**
 

python/tests/test_table_transforms.py

@@ -39,6 +39,21 @@
 # we can remove this.
 
 
+class TestShift:
+    # Most testing is done on the TreeSequence methods. Here we just check
+    # that the TableCollection methods work even if they produce an invalid ts
+    def test_too_negative(self):
+        tables = tskit.Tree.generate_comb(2).tree_sequence.dump_tables()
+        tables.shift(-1)
+        assert np.min(tables.edges.left) == -1
+
+    def test_bad_seq_len(self):
+        tables = tskit.Tree.generate_comb(2).tree_sequence.dump_tables()
+        tables.shift(1, sequence_length=0.5)
+        assert tables.sequence_length == 0.5
+        assert np.max(tables.edges.right) == 2
+
+
 def delete_older_definition(tables, time):
     node_time = tables.nodes.time
     edges = tables.edges.copy()

python/tests/test_topology.py

@@ -7074,6 +7074,136 @@ def test_failure_with_migrations(self):
             ts.trim()
 
 
+class TestShift:
+    """
+    Test the shift functionality
+    """
+
+    @pytest.mark.parametrize("shift", [-0.5, 0, 0.5])
+    def test_shift(self, shift):
+        ts = tskit.Tree.generate_comb(2, span=2).tree_sequence
+        tables = ts.dump_tables()
+        tables.delete_intervals([[0, 1]], simplify=False)
+        tables.sites.add_row(1.5, "A")
+        ts = tables.tree_sequence()
+        ts = ts.shift(shift)
+        assert ts.sequence_length == 2 + shift
+        assert np.min(ts.tables.edges.left) == 1 + shift
+        assert np.max(ts.tables.edges.right) == 2 + shift
+        assert np.all(ts.tables.sites.position == 1.5 + shift)
+
+    def test_sequence_length(self):
+        ts = tskit.Tree.generate_comb(2).tree_sequence
+        ts = ts.shift(1, sequence_length=3)
+        assert ts.sequence_length == 3
+        ts = ts.shift(-1, sequence_length=1)
+        assert ts.sequence_length == 1
+
+    def test_empty(self):
+        empty_ts = tskit.TableCollection(1.0).tree_sequence()
+        empty_ts = empty_ts.shift(1)
+        assert empty_ts.sequence_length == 2
+        empty_ts = empty_ts.shift(-1.5)
+        assert empty_ts.sequence_length == 0.5
+        assert empty_ts.num_nodes == 0
+
+    def test_provenance(self):
+        ts = tskit.Tree.generate_comb(2).tree_sequence
+        ts = ts.shift(1, record_provenance=False)
+        params = json.loads(ts.provenance(-1).record)["parameters"]
+        assert params["command"] != "shift"
+        ts = ts.shift(1, sequence_length=9)
+        params = json.loads(ts.provenance(-1).record)["parameters"]
+        assert params["command"] == "shift"
+        assert params["value"] == 1
+        assert params["sequence_length"] == 9
+
+    def test_too_negative(self):
+        ts = tskit.Tree.generate_comb(2).tree_sequence
+        with pytest.raises(tskit.LibraryError, match="TSK_ERR_BAD_SEQUENCE_LENGTH"):
+            ts.shift(-1)
+
+    def test_bad_seq_len(self):
+        ts = tskit.Tree.generate_comb(2).tree_sequence
+        with pytest.raises(
+            tskit.LibraryError, match="TSK_ERR_RIGHT_GREATER_SEQ_LENGTH"
+        ):
+            ts.shift(1, sequence_length=1)
+
+
+class TestConcatenate:
+    def test_simple(self):
+        ts1 = tskit.Tree.generate_comb(5, span=2).tree_sequence
+        ts2 = tskit.Tree.generate_balanced(5, arity=3, span=3).tree_sequence
+        assert ts1.num_samples == ts2.num_samples
+        assert ts1.num_nodes != ts2.num_nodes
+        joint_ts = ts1.concatenate(ts2)
+        assert joint_ts.num_nodes == ts1.num_nodes + ts2.num_nodes - 5
+        assert joint_ts.sequence_length == ts1.sequence_length + ts2.sequence_length
+        assert joint_ts.num_samples == ts1.num_samples
+        ts3 = joint_ts.delete_intervals([[2, 5]]).rtrim()
+        # Have to simplify here, to remove the redundant nodes
+        assert ts3.equals(ts1.simplify(), ignore_provenance=True)
+        ts4 = joint_ts.delete_intervals([[0, 2]]).ltrim()
+        assert ts4.equals(ts2.simplify(), ignore_provenance=True)
+
+    def test_empty(self):
+        empty_ts = tskit.TableCollection(10).tree_sequence()
+        ts = empty_ts.concatenate(empty_ts)
+        assert ts.num_nodes == 0
+        assert ts.sequence_length == 20
+
+    def test_samples_at_end(self):
+        ts1 = tskit.Tree.generate_comb(5, span=2).tree_sequence
+        ts2 = tskit.Tree.generate_balanced(5, arity=3, span=3).tree_sequence
+        # reverse the node order
+        ts1 = ts1.subset(np.arange(ts1.num_nodes)[::-1])
+        assert ts1.num_samples == ts2.num_samples
+        assert np.all(ts1.samples() != ts2.samples())
+        joint_ts = ts1.concatenate(ts2)
+        assert joint_ts.num_samples == ts1.num_samples
+        assert np.all(joint_ts.samples() == ts1.samples())
+
+    def test_internal_samples(self):
+        tables = tskit.Tree.generate_comb(4, span=2).tree_sequence.dump_tables()
+        nodes_flags = tables.nodes.flags
+        nodes_flags[:] = tskit.NODE_IS_SAMPLE
+        nodes_flags[-1] = 0  # Only root is not a sample
+        tables.nodes.flags = nodes_flags
+        ts = tables.tree_sequence()
+        joint_ts = ts.concatenate(ts)
+        assert joint_ts.num_samples == ts.num_samples
+        assert joint_ts.num_nodes == ts.num_nodes + 1
+        assert joint_ts.sequence_length == ts.sequence_length * 2
+
+    def test_some_shared_samples(self):
+        ts1 = tskit.Tree.generate_comb(4, span=2).tree_sequence
+        ts2 = tskit.Tree.generate_balanced(8, arity=3, span=3).tree_sequence
+        shared = np.full(ts2.num_nodes, tskit.NULL)
+        shared[0] = 1
+        shared[1] = 0
+        joint_ts = ts1.concatenate(ts2, node_mapping=shared)
+        assert joint_ts.sequence_length == ts1.sequence_length + ts2.sequence_length
+        assert joint_ts.num_samples == ts1.num_samples + ts2.num_samples - 2
+        assert joint_ts.num_nodes == ts1.num_nodes + ts2.num_nodes - 2
+
+    def test_provenance(self):
+        ts = tskit.Tree.generate_comb(2).tree_sequence
+        ts = ts.concatenate(ts, record_provenance=False)
+        params = json.loads(ts.provenance(-1).record)["parameters"]
+        assert params["command"] != "concatenate"
+
+        ts = ts.concatenate(ts)
+        params = json.loads(ts.provenance(-1).record)["parameters"]
+        assert params["command"] == "concatenate"
+
+    def test_unequal_samples(self):
+        ts1 = tskit.Tree.generate_comb(5, span=2).tree_sequence
+        ts2 = tskit.Tree.generate_balanced(4, arity=3, span=3).tree_sequence
+        with pytest.raises(ValueError, match="must have the same number of samples"):
+            ts1.concatenate(ts2)
+
+
 class TestMissingData:
     """
     Test various aspects of missing data functionality

python/tskit/tables.py

@@ -3991,6 +3991,96 @@ def trim(self, record_provenance=True):
                 record=json.dumps(provenance.get_provenance_dict(parameters))
             )
 
+    def shift(self, value, *, sequence_length=None, record_provenance=True):
+        """
+        Shift the coordinate system (used by edges and sites) of this TableCollection by
+        a given value. This is identical to :meth:`TreeSequence.shift` but acts
+        *in place* to alter the data in this :class:`TableCollection`.
+
+        .. note::
+            No attempt is made to check that the new coordinate system or sequence length
+            is valid: if you wish to do this, use {meth}`TreeSequence.shift` instead.
+
+        :param value: The amount by which to shift the coordinate system.
+        :param sequence_length: The new sequence length of the tree sequence. If
+            ``None`` (default) add `value` to the sequence length.
+        """
+        self.edges.left += value
+        self.edges.right += value
+        self.sites.position += value
+        if sequence_length is None:
+            self.sequence_length += value
+        else:
+            self.sequence_length = sequence_length
+
+        if record_provenance:
+            parameters = {
+                "command": "shift",
+                "value": value,
+                "sequence_length": sequence_length,
+            }
+            self.provenances.add_row(
+                record=json.dumps(provenance.get_provenance_dict(parameters))
+            )
+
+    def concatenate(
+        self, other, *, node_mapping=None, record_provenance=True, **kwargs
+    ):
+        """
+        Concatenate another table collection to the right of this one. This
+        {meth}`shift`s the other table coordinate rightwards, then calls
+        {meth}`union` with ``check_shared_equality=False`` and the provided
+        ``node_mapping``. If no node mapping is given, the two table
+        collections must have the same number of samples, and those are treated
+        (in numerical order) as shared between the two table collections.
+        This is identical to :meth:`TreeSequence.concatenate` but
+        acts *in place* to alter the data in this :class:`TableCollection`.
+
+        .. note::
+            To add gaps between the concatenated tables, use :meth:`shift` before
+            concatenating; to remove gaps, use :meth:`trim`.
+
+        :param TableCollection other: The other table collection to add to the right
+            of this one.
+        :param list node_mapping: An array of integers of the same length as the number
+            of nodes in ``other``, where the _k_'th element gives the id of the node in
+            the current table collection corresponding to node _k_ in the other table
+            collection (see {meth}`union`). If None (default), only the sample nodes
+            between the two node tables, in numerical order, are mapped to each other.
+        :param bool record_provenance: If True (default), record details of this call to
+            ``concatenate`` in the returned tree sequence's provenance information
+            (Default: True).
+        :param \\**kwargs: Additional keyword arguments to pass to {meth}`union`
+            (e.g. ``add_populations``).
+        """
+        if node_mapping is None:
+            samples = np.where(self.nodes.flags & tskit.NODE_IS_SAMPLE)[0]
+            other_samples = np.where(other.nodes.flags & tskit.NODE_IS_SAMPLE)[0]
+            if len(other_samples) != len(samples):
+                raise ValueError(
+                    "each `other` must have the same number of samples as `self`"
+                )
+            node_mapping = np.full(other.nodes.num_rows, tskit.NULL, dtype=np.int32)
+            node_mapping[other_samples] = samples
+        other.shift(self.sequence_length, record_provenance=False)
+        self.sequence_length = other.sequence_length
+        # NB: should we use a different default for add_populations?
+        self.union(
+            other,
+            node_mapping=node_mapping,
+            check_shared_equality=False,  # Needed as checks fail with internal samples
+            record_provenance=False,
+            **kwargs,
+        )
+        if record_provenance:
+            parameters = {
+                "command": "concatenate",
+                "TODO": "add concatenate parameters",  # tricky as both have provenances
+            }
+            self.provenances.add_row(
+                record=json.dumps(provenance.get_provenance_dict(parameters))
+            )
+
     def delete_older(self, time):
         """
         Deletes edge, mutation and migration information at least as old as

python/tskit/trees.py

@@ -7060,6 +7060,68 @@ def trim(self, record_provenance=True):
         tables.trim(record_provenance)
         return tables.tree_sequence()
 
+    def shift(self, value, sequence_length=None, record_provenance=True):
+        """
+        Shift the coordinate system (used by edges and sites) of this TableCollection by
+        a given value. Positive values shift the coordinate system to the right, negative
+        values to the left. The sequence length of the tree sequence will be changed by
+        ``value``, unless ``sequence_length`` is given, in which case this will be used
+        for the new sequence length.
+
+        :param value: The amount by which to shift the coordinate system.
+        :param sequence_length: The new sequence length of the tree sequence. If
+            ``None`` (default) add ``value`` to the sequence length.
+        :raises ValueError: If the new coordinate system is invalid (e.g., if
+            shifting the coordinate system results in negative coordinates).
+        """
+        tables = self.dump_tables()
+        tables.shift(
+            value=value,
+            sequence_length=sequence_length,
+            record_provenance=record_provenance,
+        )
+        try:
+            ts = tables.tree_sequence()
+        except ValueError as e:
+            raise ValueError("Cannot shift due to bad coordinate values") from e
+        return ts
+
+    def concatenate(
+        self, other, *, node_mapping=None, record_provenance=True, **kwargs
+    ):
+        """
+        Concatenate another tree sequence to the right of this one. This shifts the
+        coordinate system of the other tree sequence rightwards, then calls
+        {meth}`union` with the provided ``node_mapping``. If no node mapping
+        is given, matches sample nodes only, in numerical order.
+
+        .. note::
+            To add gaps between the concatenated tables, use :meth:`shift` or
+            to remove gaps, use :meth:`trim` before concatenating.
+
+        :param TableCollection other: The other table collection to add to the right
+            of this one.
+        :param list node_mapping: An array of integers of the same length as the number
+            of nodes in ``other``, where the _k_'th element gives the id of the node in
+            the current table collection corresponding to node _k_ in the other table
+            collection (see :meth:`union`). If None (default), only the sample nodes
+            between the two node tables, in numerical order, are mapped to each other.
+        :param bool record_provenance: If True (default), record details of this call to
+            ``concatenate`` in the returned tree sequence's provenance information
+            (Default: True).
+        :param \\**kwargs: Additional keyword arguments to pass to :meth:`union`,
+            e.g. ``add_populations``.
+        """
+
+        tables = self.dump_tables()
+        tables.concatenate(
+            other.tables,
+            node_mapping=node_mapping,
+            record_provenance=record_provenance,
+            **kwargs,
+        )
+        return tables.tree_sequence()
+
     def split_edges(self, time, *, flags=None, population=None, metadata=None):
         """
         Returns a copy of this tree sequence in which we replace any