Assertion error in split_disjoint_nodes

[hyanwong]

There's something odd going on with the logic in split_disjoint_nodes. In the following example I get an assertion error:

in split_disjoint_nodes(ts, record_provenance)
    583 tables.build_index()
    584 tables.compute_mutation_parents()
--> 586 assert np.array_equal(
    587     tables.nodes.time[tables.mutations.node], ts.nodes_time[ts.mutations_node]
    588 )

Here's the code to reproduce

import numpy as np
import stdpopsim
import tsinfer
import tsdate

species = stdpopsim.get_species('HomSap')
model = species.get_demographic_model('PapuansOutOfAfrica_10J19')
contig = species.get_contig('1', right=1000000, mutation_rate=model.mutation_rate)

engine = stdpopsim.get_engine('msprime')
samples = {
    "YRI": 20,
    "DenA": 2,
    "NeaA": 2
}

sim_ts = engine.simulate(model, contig, samples=samples, seed=123)
sd = tsinfer.SampleData.from_tree_sequence(ts)
sample_times = ts.nodes_time[ts.samples()]

ancestors = tsinfer.generate_ancestors(sd, progress_monitor=True)
a_ts = tsinfer.match_ancestors(sd, ancestors, progress_monitor=True)
inferred_ts = tsinfer.match_samples(
    sd,
    a_ts,
    indexes = np.where(sample_times == 0)[0],
    progress_monitor=True,
)
tsdate.preprocess_ts(inferred_ts)  # Assertion error

Have you any idea what's going on here @nspope ? I dug a bit but got confused by the flow of logic. It appears as if there are swaps in the node assignments, e.g.

Node 29 has time 0.0, but mutation 509 is assigned to node 885 with time 0.08333333333333333
Node 885 has time 0.08333333333333333, but mutation 510 is assigned to node 29 with time 0.0

It can be avoided by fully simplifying first:

tsdate.preprocess_ts(inferred_ts.simplify())

[nspope]

Hm, I'll take a look!

[nspope]

This seems to have been introduced with tsdate v0.2.5 and tskit v1.0.0 --- the assertion does not get hit after downgrading to tskit 0.6.4 and tsdate v0.2.4 / tsinfer v0.5.0 (the last versions to support tskit <1). We need a test for the entire pipeline in the test suite.

I'll get this sorted ASAP, as it pops up in all the simulated examples I've looked at. I'm surprised we haven't run into it before now; guess I'm in the habit of always immediately simplifying output from tsinfer.

[nspope]

I think what is going on here is that, as of tskit v1, the mutation order may be changed upon tables.sort(). This makes the triggered assertion invalid, as it assumes that mutations in the input/output tree sequences are in the same order. I'm struggling to understand why this only happens with unsimplified input, though (might simply be much rarer in simplified output w/o unary nodes)

[nspope]

Poking at this further: I can only trigger the assertion through preprocess_ts, not with split_disjoint_nodes run separately. In general, if the tables are sorted, then split_disjoint_nodes will maintain the order of the mutation table. However, the telomere trimming in preprocess_ts occasionally puts the mutation table in not-canonical order (i.e. mutations at the same site may not be in time order). I'm not sure yet if this is a bug or expected.

To reproduce, uv run --script split_disjoint_bug.py

# split_disjoint_bug.py
# /// script
# requires-python = ">=3.11"
# dependencies = ["stdpopsim==0.3.0", "tsinfer==0.5.1", "tsdate==0.2.6", "tskit==1.0.2"]
# ///

import tsinfer
import tsdate
import numpy as np
import stdpopsim

engine = stdpopsim.get_engine("msprime")
homsap = stdpopsim.get_species("HomSap")
demogr = homsap.get_demographic_model("OutOfAfrica_3G09")
contig = homsap.get_contig(length=1e6)
                                                                                                                                                               
samples = { 
    population.name: 500 
    for population in demogr.populations
}
ts = engine.simulate(
    samples=samples,
    contig=contig,
    demographic_model=demogr,
    seed=1024,
)
sample_data = tsinfer.SampleData.from_tree_sequence(ts)
print("-> tsinfer.infer")
ts = tsinfer.infer(sample_data)
print("-> tsdate.util.split_disjoint_nodes")
tsdate.util.split_disjoint_nodes(ts)
print("-> tsdate.util.preprocess_ts")
tsdate.util.preprocess_ts(ts)

Adding a tables.sort() between the telomere removal and the disjoint node splitting fixes the issue with preprocess_ts.

[nspope]

After looking into it---the issue is that preprocess_ts uses tables.simplify(), which removes unary nodes by default. This can result in the mutation table getting out of canonical order; we just need to sort afterwards to put it back into order. (For example, we can force the code above to pass by explicitly passing keep_unary=True to preprocess_ts.) We don't expect the assertion in split_disjoint_nodes to hold for all inputs, and it is not required (more a leftover sanity check), so should remove it. So I think we can close this issue by:

1. Adding a sort after tables.simplify in preprocess_ts
2. Removing assertion from split_disjoint_nodes
3. Adding a test to check that preprocess_ts returns a tree sequence with the mutation table in the right order

I'm not clear on whether or not tree sequence conversion should fail when the mutation table is out of sorted order, but that's a separate issue. (Probably not, if the mutation parents are valid?)