Check tree sequences topologically equal (i.e. ignoring node times)

[hyanwong]

Is there a way of checking whether one tree sequence is topologically identical to another (i.e. if node times are ignored and, therefore, assuming nodes are in time order, ignoring internal node ids)? I can’t think of an easy way off the top of my head. We can compare tree.rank() for trees, but there’s no obvious way to do this for an entire tree sequence, is there?

[hyanwong]

It's not enough to simply check that the breakpoints in the two tree sequences are the same and that each tree is topologically identical between the two tree sequences, because a tree could be topologically identical but the node identities could have changed between trees (which would imply a different topology). Perhaps a brute force approach would be to coiterate over the edge diffs (and trees) and check that each newly introduced parent node has the same set of samples underneath it?

[hyanwong]

Here's my first attempt, with some rough tests. I'm not sure if the logic is correct, however. E.g. do I only need to test the edges_in, as I'm doing here? It's also pretty inefficient, as I checking all the descendant samples for every single node in the TS (possible multiple times). But at least I don't do every node in every tree.

def topologies_equal(ts_a, ts_b):
    if any(ts_a.samples() != ts_b.samples()):
        return False
    if ts_a.num_trees != ts_b.num_trees:
        return False
    for tree_a, diffs_a, tree_b, diffs_b in zip(
        ts_a.trees(sample_lists=True),
        ts_a.edge_diffs(),
        ts_b.trees(sample_lists=True),
        ts_b.edge_diffs()
    ):
        if tree_a.interval != tree_b.interval:
            return False
        a_in = diffs_a[2]
        b_in = diffs_b[2]
        a_tips = set(frozenset(tree_a.samples(e.parent)) for e in a_in)
        b_tips = set(frozenset(tree_b.samples(e.parent)) for e in b_in)
        if a_tips != b_tips:
            return False
    return True
    
### TEST the topologies_equal() function
ts1 =  msprime.simulate(10, recombination_rate=1, random_seed = 2)

tables = ts1.dump_tables()
tables.nodes.time += 1
ts2 = tables.tree_sequence()

# Jiggle the nonsample node ids
tables = ts1.dump_tables()
nodes = tables.nodes
edges = tables.edges
internal_nodes = (nodes.flags & tskit.NODE_IS_SAMPLE == 0)
shuffled_internal_nodes = np.where(internal_nodes)[0]
np.random.shuffle(shuffled_internal_nodes)
reorder = np.arange(nodes.num_rows)
reorder[internal_nodes] = shuffled_internal_nodes
nodes.set_columns(
    flags=nodes.flags[reorder],
    time=nodes.time[reorder],
    population=nodes.population[reorder],
    individual=nodes.individual[reorder],
)
_, mapping = np.unique(reorder, return_index=True)  # make a reverse mapping of new->old ids
mapping = mapping.astype(edges.parent.dtype)
tables.edges.child = mapping[edges.child]
tables.edges.parent = mapping[edges.parent]
tables.sort()
ts3 = tables.tree_sequence()

assert not all(ts1.tables.nodes.time == ts2.tables.nodes.time)
assert not all(ts1.tables.nodes.time == ts3.tables.nodes.time)
assert topologies_equal(ts1, ts2)
assert topologies_equal(ts1, ts3)
assert not topologies_equal(ts1, msprime.simulate(10, recombination_rate=1, random_seed = 1))

[petrelharp]

You could put a mutation on every branch of every tree and check that the genotype matrix is the same!

[hyanwong]

Would that be enough? What if a breakpoint caused only a shift in height of a node? I'm not sure if I'd count this as "topologically identical" or not, as there would be different node identities in the 2 tree sequences. The ARG implied by the TS is different, certainly.

[petrelharp]

Oh, I see. Well, could you compare the distances between all the nodes for all the trees? Obtained by:

ts.divergence([[s] for s in ts.samples()], indexes=[(i, j) for i in range(num_samples) for j in range(i, num_samples)], mode="branch", windows="trees")

[hyanwong]

Thanks for the replies. But doesn't that take node heights into account, though? I'm looking for a topology-only measure - i.e. do the two TSes define identical DAGs from the sample nodes to the root(s). Otherwise I guess I could use e.g. kc_distance too. (KC on topology-only suffers from the same problem I mentioned above).

[hyanwong]

I think any measure that based on calculations on successive trees independently is likely to fail here, because the topological structure of the tree sequence includes nodes that are shared between trees, and you need to know about the node sharing to check for topological identity.

[petrelharp]

Hm. Well, how about adding mutations to all branches and then checking for equality of genotypes of all nodes (not just samples)?

[jeromekelleher]

This is a Hard Problem I think. We would need to have a combinatorial concept of what a tree sequence topology is to be able to do this properly. I.e., how many topologically distinct tree sequences are there (under some reasonable assumptions), how do we define an ordering over them, etc. This is a PhD project (at least!) as far as I'm concerned...

What do you want this for? (Perhaps there's a less difficult problem that actually needs to be solved?)

[hyanwong]

In this case, I want to run msprime (say) 10 million times, pick out all the tree sequences with a given ARG topology, then look at the distribution of node times, so that I can test that tsdate will accurately converge on the empirical distribution of node times for a fixed topology. It's easy for a single tree, but not so easy for an entire TS.

I have the additional problem that I would like to ignore the sample node labelling too, which reduces the number of topologies quite drastically, but will make the coding harder.

[hyanwong]

I'm not sure it is hard to define topological equivalence, if we think of the tree sequence as a DAG. It simply turns into a question of whether 2 graphs have an identical topology (and since it's a tree sequence, with an additional constraint that any edges that only span part of the genome must span exactly the same region to be considered identical). Comparing 2 DAGs for equality must surely be a classic graph theory question? So it's just a question of what an efficient implementation is.

I think my code above does this, but I'm not 100% sure. I'm finding it quite tricky to think about.

Edit - for graphs this is known as isomorphism, and e.g. networkx has some implemented algorithms at https://networkx.org/documentation/stable/reference/algorithms/isomorphism.html

[jeromekelleher]

I think I'd do something different (as in, evaluate whatever you're trying to evaluate some other way) - this is a near bottomless rabbithole. There's years of research involved in doing this properly.

[hyanwong]

Well, I'm only doing this for very small tree sequences (like, 5 tips, 2 trees), and I think my algorithm above is correct, so I can actually check examples by hand easily enough.

[hyanwong]

I've just realised that there is a trivial definition of tree sequence topological equality: is there a way of simultaneously remapping the edge.parent and edge.child IDs to create an exactly identical edge table (after sorting).

At least this defines the problem.

[hyanwong]

I've just PR'ed a method to port whole tree sequences (i.e. the ARG, if we have recombination nodes) into NetworkX. So now to check for tree sequence topological identity we can just read them into NetworkX and use nx.is_isomorphic while ensuring that the edges have identical left and right attributes:

import msprime
import networkx as nx

def edges_identical(d1, d2):
    return set((v["left"], v["right"]) for v in d1.values()) == set((v["left"], v["right"]) for v in d2.values())

def as_dict_of_dicts_reorder(ts, mapping=None):
    # like ts.as_dict_of_dicts PRed in #1296 but also allows the node IDs
    # to be randomly shuffled, so that the mappings are unbiased
    # (only needed if you are sampling from the mappings)
    dod = {}
    for edge in ts.edges():
        if mapping is not None:
            parent = mapping[edge.parent]
            child = mapping[edge.child]
        else:
            parent = edge.parent
            child = edge.child
        if parent not in dod:
            dod[parent] = {}
        if child not in dod[parent]:
            dod[parent][child] = {}
        dod[parent][child][edge.id] = {
            "left": edge.left,
            "right": edge.right,
            "branch_length": ts.node(parent).time - ts.node(child).time,
        }
    return dod

target_ts = msprime.sim_ancestry(3, recombination_rate=0.1, sequence_length=2, random_seed=1)
target_graph = nx.from_dict_of_dicts(
    as_dict_of_dicts_reorder(target_ts),
    create_using=nx.MultiDiGraph(),
    multigraph_input=True)
for s in range(2, 1000):
    ts = msprime.sim_ancestry(3, recombination_rate=0.1, sequence_length=2, random_seed=s)
    graph = nx.from_dict_of_dicts(
        as_dict_of_dicts_reorder(ts),
        create_using=nx.MultiDiGraph(),
        multigraph_input=True)
    DiGM = nx.isomorphism.DiGraphMatcher(graph, target_graph, edge_match=edges_identical)
    if DiGM.is_isomorphic():
        break

print(target_ts.draw_text())
print(ts.draw_text())
print("corresponding nodes:", DiGM.mapping)

Which indeed works: the following tree sequences are what I would define as "topologically identical"

1.83┊       11    ┊      11     ┊  
    ┊     ┏━━┻━━┓ ┊    ┏━━┻━━┓  ┊  
1.38┊    10     ┃ ┊   10     ┃  ┊  
    ┊   ┏━┻━━┓  ┃ ┊  ┏━┻━┓   ┃  ┊  
1.01┊   ┃    ┃  ┃ ┊  ┃   ┃   9  ┊  
    ┊   ┃    ┃  ┃ ┊  ┃   ┃  ┏┻┓ ┊  
0.85┊   ┃    8  ┃ ┊  ┃   8  ┃ ┃ ┊  
    ┊   ┃   ┏┻┓ ┃ ┊  ┃  ┏┻┓ ┃ ┃ ┊  
0.80┊   7   ┃ ┃ ┃ ┊  ┃  ┃ ┃ ┃ ┃ ┊  
    ┊  ┏┻━┓ ┃ ┃ ┃ ┊  ┃  ┃ ┃ ┃ ┃ ┊  
0.78┊  6  ┃ ┃ ┃ ┃ ┊  6  ┃ ┃ ┃ ┃ ┊  
    ┊ ┏┻┓ ┃ ┃ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┃ ┃ ┊  
0.00┊ 0 1 5 2 3 4 ┊ 0 1 2 3 4 5 ┊  
  0.00          1.00          2.00 

6.51┊   11        ┊    11       ┊  
    ┊ ┏━━┻━━┓     ┊  ┏━━┻━━┓    ┊  
5.25┊ ┃    10     ┊  ┃    10    ┊  
    ┊ ┃   ┏━┻━━┓  ┊  ┃   ┏━┻━┓  ┊  
1.43┊ ┃   ┃    ┃  ┊  9   ┃   ┃  ┊  
    ┊ ┃   ┃    ┃  ┊ ┏┻┓  ┃   ┃  ┊  
0.87┊ ┃   8    ┃  ┊ ┃ ┃  ┃   ┃  ┊  
    ┊ ┃  ┏┻━┓  ┃  ┊ ┃ ┃  ┃   ┃  ┊  
0.49┊ ┃  ┃  ┃  7  ┊ ┃ ┃  ┃   7  ┊  
    ┊ ┃  ┃  ┃ ┏┻┓ ┊ ┃ ┃  ┃  ┏┻┓ ┊  
0.03┊ ┃  6  ┃ ┃ ┃ ┊ ┃ ┃  6  ┃ ┃ ┊  
    ┊ ┃ ┏┻┓ ┃ ┃ ┃ ┊ ┃ ┃ ┏┻┓ ┃ ┃ ┊  
0.00┊ 0 1 3 5 2 4 ┊ 0 5 1 3 2 4 ┊  
  0.00          1.00          2.00 

corresponding nodes: {6: 6, 1: 0, 3: 1, 8: 7, 5: 5, 9: 9, 0: 4, 10: 10, 7: 8, 2: 2, 4: 3, 11: 11}

[hyanwong]

Note that the above answer uses a bespoke "as_dict_of_dicts_reorder" function which allows the node IDs to be randomly reshuffled. I needed this for the main use of this function, which was to sample topologies, because there are often multiple possible mappings compatible with an isomorphism (for example, the case about has {2:2, 4: 3 ...} but could equally have {2:3, 4: 2 ...}. This can be a particular issue when e.g. a pair of internal nodes in the ts map to a another pair in the target_ts. In this case, the isomorphism detector can end up always mapping the youngest nodes to each other, so that the isomorphism match is dependent on the nod ID ordering. Here's a more fully-worked-out example of sampling node times for each possible tree sequence topology produced by msprime (it also has a simple optimisation to not bother running the full isomorphism algorithm if the first tree in both tree sequences is not the same shape)

import collections
import time

import msprime
import networkx as nx
import numpy as np

Topology = collections.namedtuple("Topology", "example_ts, graph, first_tree_shape, num_trees, node_times")

results = []

def edges_identical(d1, d2):
    return set((v["left"], v["right"]) for v in d1.values()) == set((v["left"], v["right"]) for v in d2.values())

def as_dict_of_dicts_reorder(ts, mapping=None):
    dod = {}
    for edge in ts.edges():
        if mapping is not None:
            parent = mapping[edge.parent]
            child = mapping[edge.child]
        else:
            parent = edge.parent
            child = edge.child
        if parent not in dod:
            dod[parent] = {}
        if child not in dod[parent]:
            dod[parent][child] = {}
        dod[parent][child][edge.id] = {
            "left": edge.left,
            "right": edge.right,
            "branch_length": ts.node(parent).time - ts.node(child).time,
        }
    return dod

start_time = time.time()
num_reps = 100000
np.random.seed(1)
for rep, ts in enumerate(
    msprime.sim_ancestry(
        6,
        ploidy=1,
        sequence_length=3,
        recombination_rate=0.3,
        record_full_arg=True,  # So we can detect the number of breakpoints
        num_replicates=num_reps,
        random_seed=123)
):
    if rep % 10000 == 0:
        print("Rep", rep, ":", len(results), "topologies")

    num_recomb_events = np.sum(ts.tables.nodes.flags & msprime.NODE_IS_RE_EVENT > 0)//2
    if num_recomb_events > 1:
        continue
    ts = ts.simplify()
    first_tree_shape = ts.first().rank()[0]
    # need to provide the nodes in a random order, otherwise the isomorphism func will always
    # pick a specific mapping from among the several possible, which biases the node labels
    random_map = np.random.permutation(ts.num_nodes)
    graph = nx.from_dict_of_dicts(
        as_dict_of_dicts_reorder(ts, random_map),
        create_using=nx.MultiDiGraph,
        multigraph_input=True)

    for i, target in enumerate(results):
        if target.num_trees==ts.num_trees and first_tree_shape==target.first_tree_shape:
            # It's quicker to weed out some of the non-matching
            # tree seqs by checking on the number of trees and
            # the shape rank of the first tree before actually converting into a graph
            DiGM = nx.isomorphism.DiGraphMatcher(graph, target.graph, edge_match=edges_identical)
            if DiGM.is_isomorphic():
                # Must map the nodes from the random_mapping used via the graph mapping
                reverse_random_map = {v:k for k, v in enumerate(random_map)}
                g_map = DiGM.mapping
                mapping = np.array([reverse_random_map[k] for k in sorted(g_map, key=g_map.get)])
                results[i].node_times.append(
                    ts.tables.nodes.time[mapping]
                )
                break
    else:
        # Haven't seen this before: make sure graph is labelled in the same way as the ts
        # then save it, so it is treated as providing the "canonical" labelling
        graph = nx.from_dict_of_dicts(
            as_dict_of_dicts_reorder(ts),
            create_using=nx.MultiDiGraph,
            multigraph_input=True)
        results.append(
            Topology(
                ts, graph, first_tree_shape, ts.num_trees, [ts.tables.nodes.time]))

print("Final:", [len(r.node_times) for r in results])
print("Time taken", time.time()-start_time, "seconds for", num_reps, "reps")
print((time.time()-start_time) / num_reps, "per replicate")

np.savez_compressed("node_times_per_topology.npz", *[np.vstack(r.node_times) for r in results])
# Stream them all into a single ts file
with open("topologies.ts", "w") as filehandle:
    for r in results:
        r.example_ts.dump(filehandle)