Change to a single concatenate function

Author: hyanwong

python/tests/test_table_transforms.py

@@ -39,21 +39,6 @@
 # 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

@@ -7091,6 +7091,7 @@ def test_shift(self, 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)
+        assert len(list(ts.trees())) == ts.num_trees
 
     def test_sequence_length(self):
         ts = tskit.Tree.generate_comb(2).tree_sequence
@@ -7147,11 +7148,31 @@ def test_simple(self):
         ts4 = joint_ts.delete_intervals([[0, 2]]).ltrim()
         assert ts4.equals(ts2.simplify(), ignore_provenance=True)
 
+    def test_multiple(self):
+        np.random.seed(42)
+        ts3 = [
+            tskit.Tree.generate_comb(5, span=2).tree_sequence,
+            tskit.Tree.generate_balanced(5, arity=3, span=3).tree_sequence,
+            tskit.Tree.generate_star(5, span=5).tree_sequence,
+        ]
+        for i in range(1, len(ts3)):
+            # shuffle the sample nodes so they don't have the same IDs
+            ts3[i] = ts3[i].subset(np.random.permutation(ts3[i].num_nodes))
+        assert not np.all(ts3[0].samples() == ts3[1].samples())
+        assert not np.all(ts3[0].samples() == ts3[2].samples())
+        assert not np.all(ts3[1].samples() == ts3[2].samples())
+        ts = ts3[0].concatenate(*ts3[1:])
+        assert ts.sequence_length == sum([t.sequence_length for t in ts3])
+        assert ts.num_nodes - ts.num_samples == sum(
+            [t.num_nodes - t.num_samples for t in ts3]
+        )
+        assert np.all(ts.samples() == ts3[0].samples())
+
     def test_empty(self):
         empty_ts = tskit.TableCollection(10).tree_sequence()
-        ts = empty_ts.concatenate(empty_ts)
+        ts = empty_ts.concatenate(empty_ts, empty_ts, empty_ts)
         assert ts.num_nodes == 0
-        assert ts.sequence_length == 20
+        assert ts.sequence_length == 40
 
     def test_samples_at_end(self):
         ts1 = tskit.Tree.generate_comb(5, span=2).tree_sequence
@@ -7182,7 +7203,7 @@ def test_some_shared_samples(self):
         shared = np.full(ts2.num_nodes, tskit.NULL)
         shared[0] = 1
         shared[1] = 0
-        joint_ts = ts1.concatenate(ts2, node_mapping=shared)
+        joint_ts = ts1.concatenate(ts2, node_mappings=[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

python/tskit/tables.py

@@ -4005,6 +4005,7 @@ def shift(self, value, *, sequence_length=None, record_provenance=True):
         :param sequence_length: The new sequence length of the tree sequence. If
             ``None`` (default) add `value` to the sequence length.
         """
+        self.drop_index()
         self.edges.left += value
         self.edges.right += value
         self.sites.position += value
@@ -4023,64 +4024,6 @@ def shift(self, value, *, sequence_length=None, record_provenance=True):
                 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

@@ -32,6 +32,7 @@
 import functools
 import io
 import itertools
+import json
 import math
 import numbers
 import warnings
@@ -46,6 +47,7 @@
 import tskit.combinatorics as combinatorics
 import tskit.drawing as drawing
 import tskit.metadata as metadata_module
+import tskit.provenance as provenance
 import tskit.tables as tables
 import tskit.text_formats as text_formats
 import tskit.util as util
@@ -7091,39 +7093,75 @@ def shift(self, value, sequence_length=None, record_provenance=True):
         return ts
 
     def concatenate(
-        self, other, *, node_mapping=None, record_provenance=True, **kwargs
+        self, *args, node_mappings=None, record_provenance=True, add_populations=None
     ):
-        """
-        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.
+        r"""
+        Concatenate a set of tree sequences to the right of this one, by repeatedly
+        calling {meth}`union` with an (optional)
+        node mapping for each of the ``others``. If any node mapping is ``None``
+        only map the sample nodes between the input tree sequence and this one,
+        based on the numerical order of sample node IDs.
 
         .. 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``.
-        """
+        :param TreeSequence \*args: A list of other tree sequences to append to
+            the right of this one.
+        :param Union[list, None] node_mappings: An list of node mappings for each
+            input tree sequence in ``args``. Each should either be an array of
+            integers of the same length as the number of nodes in the equivalent
+            input tree sequence (see :meth:`union` for details), or ``None``.
+            If ``None``, only sample nodes are mapped to each other.
+            Default: ``None``, treated as ``[None] * len(args)``.
+        :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 bool add_populations: If True (default), nodes new to ``self`` will
+            be assigned new population IDs (see :meth:`union`)
+        """
+        if node_mappings is None:
+            node_mappings = [None] * len(args)
+        if add_populations is None:
+            add_populations = True
+        if len(node_mappings) != len(args):
+            raise ValueError(
+                "You must provide the same number of node_mappings as args"
+            )
 
+        samples = self.samples()
         tables = self.dump_tables()
-        tables.concatenate(
-            other.tables,
-            node_mapping=node_mapping,
-            record_provenance=record_provenance,
-            **kwargs,
-        )
+        tables.drop_index()
+
+        for node_mapping, other in zip(node_mappings, args):
+            if node_mapping is None:
+                other_samples = other.samples()
+                if len(other_samples) != len(samples):
+                    raise ValueError(
+                        "each `other` must have the same number of samples as `self`"
+                    )
+                node_mapping = np.full(other.num_nodes, tskit.NULL, dtype=np.int32)
+                node_mapping[other_samples] = samples
+            other_tables = other.dump_tables()
+            other_tables.shift(tables.sequence_length, record_provenance=False)
+            tables.sequence_length = other_tables.sequence_length
+            # NB: should we use a different default for add_populations?
+            tables.union(
+                other_tables,
+                node_mapping=node_mapping,
+                check_shared_equality=False,  # Else checks fail with internal samples
+                record_provenance=False,
+                add_populations=add_populations,
+            )
+        if record_provenance:
+            parameters = {
+                "command": "concatenate",
+                "TODO": "add concatenate parameters",  # tricky as both have provenances
+            }
+            tables.provenances.add_row(
+                record=json.dumps(provenance.get_provenance_dict(parameters))
+            )
+
         return tables.tree_sequence()
 
     def split_edges(self, time, *, flags=None, population=None, metadata=None):