adding jerome's suggestions

Author: petrelharp

python/tests/tsutil.py

@@ -138,26 +138,22 @@ def insert_branch_mutations(ts, mutations_per_branch=1):
 
 def remove_mutation_times(ts):
     tables = ts.tables
-    tables.mutations.clear()
-    for mut in ts.mutations():
-        tables.mutations.add_row(
-            site=mut.site,
-            node=mut.node,
-            derived_state=mut.derived_state,
-            parent=mut.parent,
-            metadata=mut.metadata,
-            time=None,
-        )
+    tables.mutations.time = np.full_like(tables.mutations.time, tskit.UNKNOWN_TIME)
     return tables.tree_sequence()
 
 
-def insert_discrete_time_mutations(ts, num_times=4):
+def insert_discrete_time_mutations(ts, num_times=4, num_sites=10):
+    """
+    Inserts mutations in the tree sequence at regularly-spaced num_sites
+    positions, at only a discrete set of times (the same for all trees): at
+    num_times times evenly spaced between 0 and the maximum tree height.
+    """
     tables = ts.tables
     tables.sites.clear()
     tables.mutations.clear()
     height = max([t.time(t.roots[0]) for t in ts.trees()])
-    for pos in range(int(tables.sequence_length)):
-        anc = "X" * pos
+    for j, pos in enumerate(np.linspace(0, tables.sequence_length, num_sites + 1)[:-1]):
+        anc = "X" * j
         tables.sites.add_row(position=pos, ancestral_state=anc)
         t = ts.at(pos)
         for k, s in enumerate(np.linspace(0, height, num_times)):
@@ -166,11 +162,12 @@ def insert_discrete_time_mutations(ts, num_times=4):
                     (t.parent(n) == tskit.NULL) or (t.time(t.parent(n)) > s)
                 ):
                     tables.mutations.add_row(
-                        site=pos, node=n, derived_state=anc + str(k), time=s
+                        site=j, node=n, derived_state=anc + str(k), time=s
                     )
                     k += 1
     tables.sort()
-    tables.compute_mutation_parents
+    tables.build_index()
+    tables.compute_mutation_parents()
     return tables.tree_sequence()
 
 
@@ -833,13 +830,14 @@ def shuffle_tables(
     shuffle_individuals=True,
     shuffle_sites=True,
     shuffle_mutations=True,
+    shuffle_migrations=True,
     keep_mutation_parent_order=False,
 ):
     """
-    Randomizes the order the Individual, Population, Edge, Site and Mutation Tables rows.
-    Note that if mutations are completely shuffled, then TableCollection.sort() will not
-    necessarily produce valid tables (unless all mutation times are present and
-    distinct), since it only puts parent mutations before children if
+    Randomizes the order of rows in (possibly) all except the Node table.  Note
+    that if mutations are completely shuffled, then TableCollection.sort() will
+    not necessarily produce valid tables (unless all mutation times are present
+    and distinct), since it only puts parent mutations before children if
     canonical=True. However, setting keep_mutation_parent_order to True will
     maintain the order of mutations within each site.
 
@@ -885,6 +883,20 @@ def shuffle_tables(
         rng.shuffle(randomised_edges)
     for e in randomised_edges:
         tables.edges.add_row(e.left, e.right, e.parent, e.child, metadata=e.metadata)
+    # migrations
+    randomised_migrations = list(orig.migrations)
+    if shuffle_migrations:
+        rng.shuffle(randomised_migrations)
+    for m in randomised_migrations:
+        tables.migrations.add_row(
+            m.left,
+            m.right,
+            m.node,
+            pop_id_map[m.source],
+            pop_id_map[m.dest],
+            m.time,
+            m.metadata,
+        )
     # sites
     randomised_sites = list(enumerate(orig.sites))
     if shuffle_sites:
@@ -934,7 +946,9 @@ def orig_cmp_site(i, j, tables):
 
 
 def orig_cmp_mutation(i, j, tables, site_order, canonical=False, num_descendants=None):
-    ret = site_order[tables.mutations.site[i]] - site_order[tables.mutations.site[j]]
+    site_i = tables.mutations.site[i]
+    site_j = tables.mutations.site[j]
+    ret = site_order[site_i] - site_order[site_j]
     if (
         ret == 0
         and (not tskit.is_unknown_time(tables.mutations.time[i]))
@@ -1574,40 +1588,39 @@ def assert_table_collections_equal(t1, t2, ignore_provenance=False):
     Checks for table collection equality, but step-by-step,
     so it's easy to see what's different.
     """
-    if ignore_provenance:
-        t1.provenances.clear()
-        t2.provenances.clear()
     assert_tables_equal(t1.populations, t2.populations, "populations")
     assert_tables_equal(t1.individuals, t2.individuals, "individuals")
     assert_tables_equal(t1.nodes, t2.nodes, "nodes")
     assert_tables_equal(t1.edges, t2.edges, "edges")
     assert_tables_equal(t1.sites, t2.sites, "sites")
     assert_tables_equal(t1.mutations, t2.mutations, "mutations")
     assert_tables_equal(t1.migrations, t2.migrations, "migrations")
-    assert_tables_equal(t1.provenances, t2.provenances, "provenances")
+    if not ignore_provenance:
+        assert_tables_equal(t1.provenances, t2.provenances, "provenances")
     assert t1.metadata_schema == t2.metadata_schema
     assert t1.metadata == t2.metadata
     assert t1.metadata_bytes == t2.metadata_bytes
     assert t1.sequence_length == t2.sequence_length
-    assert t1 == t2
+    assert t1.equals(t2, ignore_provenance=ignore_provenance)
 
 
 def assert_tables_equal(t1, t2, label=""):
     if hasattr(t1, "metadata_schema"):
         if t1.metadata_schema != t2.metadata_schema:
-            print(f"{label} :::::::::: t1 ::::::::::::")
-            print(t1.metadata_schema)
-            print(f"{label} :::::::::: t2 ::::::::::::")
-            print(t2.metadata_schema)
-        assert t1.metadata_schema == t2.metadata_schema
-    if t1.num_rows != t2.num_rows:
-        print(f"{label}: t1.num_rows {t1.num_rows} != {t2.num_rows} t2.num_rows")
+            msg = (
+                f"{label} :::::::::: t1 ::::::::::::\n{t1.metadata_schema}"
+                f"{label} :::::::::: t2 ::::::::::::\n{t1.metadata_schema}"
+            )
+            raise AssertionError(msg)
     for k, (e1, e2) in enumerate(zip(t1, t2)):
         if e1 != e2:
-            print(f"{label} :::::::::: t1 (row {k}) ::::::::::::")
-            print(e1)
-            print(f"{label} :::::::::: t2 (row {k}) ::::::::::::")
-            print(e2)
-        assert e1 == e2
-    assert t1.num_rows == t2.num_rows
+            msg = (
+                f"{label} :::::::::: t1 (row {k}) ::::::::::::\n{e1}"
+                f"{label} :::::::::: t2 (row {k}) ::::::::::::\n{e2}"
+            )
+            raise AssertionError(msg)
+    if t1.num_rows != t2.num_rows:
+        raise AssertionError(
+            f"{label}: t1.num_rows {t1.num_rows} != {t2.num_rows} t2.num_rows"
+        )
     assert t1 == t2

python/tskit/tables.py

@@ -2603,7 +2603,7 @@ def sort(self, edge_start=0):
 
     def canonicalise(self):
         """
-        This puts the tables in *canonical* order - to do this, the individual
+        This puts the tables in *canonical* form - to do this, the individual
         and population tables are sorted by the first node that refers to each
         (see :meth:`TreeSequence.subset`, and note that individuals and
         populations not referred to by any nodes will be put at the end of the
@@ -2667,8 +2667,8 @@ def deduplicate_sites(self):
         site), and renumbering the ``site`` column of the mutation table
         appropriately.  This requires the site table to be sorted by position.
 
-        This method does not sort the tables afterwards, so mutations may no longer
-        be sorted by time.
+        ..warning:: This method does not sort the tables afterwards, so
+            mutations may no longer be sorted by time.
         """
         self._ll_tables.deduplicate_sites()
         # TODO add provenance
@@ -2956,20 +2956,22 @@ def subset(
     ):
         """
         Modifies the tables in place to contain only the entries referring to
-        the provided list of nodes, with nodes reordered according to the order
-        they appear in the list. See :meth:`TreeSequence.subset` for a more
-        detailed description.
+        the provided list of node IDs, with nodes reordered according to the
+        order they appear in the list. See :meth:`TreeSequence.subset` for a
+        more detailed description.
 
-        Note: the tables can be completely unsorted.
+        Note: there are no sortedness requirements on the tables.
 
         :param list nodes: The list of nodes for which to retain information. This
             may be a numpy array (or array-like) object (dtype=np.int32).
         :param bool record_provenance: Whether to record a provenance entry
             in the provenance table for this operation.
         :param bool filter_populations: Whether to remove populations not referenced by
-            retained nodes. If False, the population table will remain unchanged.
+            retained nodes. If False, the population table will not be altered
+            in any way.
         :param bool filter_individuals: Whether to remove individuals not referenced by
-            retained nodes. If False, the individuals table will remain unchanged.
+            retained nodes. If False, the individuals table will not be altered
+            in any way.
         :param bool filter_sites: Whether to remove sites not referenced by
             retained mutations. If False, the site table will remain unchanged.
         :param bool canonicalise: If True, retains all unused entries, putting