Can be hard to constrain root node to a fixed time on a large single tree

[hyanwong]

I'm finding when dating a large single tree where the root node is fixed, tsdate makes the root time change. The effect is relatively small below, (root dates before & after: 1658.0 1679.47) but gets larger with bigger sample sizes (e.g. in the larger ts of 4M samples from which this has been simplified down, the fixed root node is assigned twice the expected age)

sts = tszip.load("simpified50000.tsz")
print(f"{sts.num_trees} trees, root is sample @ time {sts.nodes_time[sts.first().root]}: ", sts.first().root)
edge_times = sts.nodes_time[sts.edges_parent] - sts.nodes_time[sts.edges_child]
av_mu = sts.num_mutations / ((sts.edges_right - sts.edges_left) * edge_times).sum()

dts = tsdate.date(
    sts,
    mutation_rate=av_mu,
    rescaling_intervals=0,
    max_iterations=1,  # single tree, so only one round needed
    time_units=sts.time_units,
    allow_unary=True,
    progress=True,
    set_metadata=False,
)
print("Root dates before & after:", sts.nodes_time[sts.first().root], dts.nodes_time[dts.first().root])

simpified50000.tsz.zip

[hyanwong]

If I increase the mutation rate to av_mu= av_mu * 2, the times become the same, but for larger sample sizes, that doesn't work

[hyanwong]

Aha, here's a MWE. This is small enough to add to the test suite, I think, if we can fix it.

import tskit
import tsdate
import msprime
mu = 0.0001
tts = tskit.Tree.generate_balanced(5000, span=1000).tree_sequence
tts = msprime.sim_mutations(tts, rate=mu, random_seed=123)
assert tts.num_trees == 1
tree = tts.first()
tables = tts.dump_tables()
tables.nodes[tree.root] = tables.nodes[tree.root].replace(flags = tskit.NODE_IS_SAMPLE)
tts = tables.tree_sequence()
assert tts.node(tree.root).is_sample()
dts = tsdate.date(tts, rescaling_intervals=0, max_iterations=1, mutation_rate=mu, progress=True)
print(f"Expected root @ {tts.node(tree.root).time}, got {dts.node(dts.first().root).time}")

Expected root @ 13.0, got 14.120520210273911

Increasing the sample size to 100_000, it gets more extreme:

Expected root @ 17.0, got 29.93303585326362

Note that this does it too:

mu = 0.00001
tts = tskit.Tree.generate_comb(1000, span=10000).tree_sequence

Expected root @ 999.0, got 1013.965404685807

[hyanwong]

From @nspope :

the only way a fixed node time would change, afaict, is via forcing the constraint. If you try using constr_iterations=100, it should make this less likely. this is known behaviour (we talked about it before for the covid ARGs) and it's not easy to fix aside from using constr_iterations. so mainly just want to be sure it is the constraint being violated

And indeed, setting e.g. constr_iterations=1000 works for the MWE, but even when set to e.g. 100000 on my larger covid tree from the Usher project, I'm still getting the root node at 2763 days when it should be fixed at 1658, which is quite a difference. I might try to set some of the other internal nodes as fixed in case that seems to help.

[nspope]

It would probably clear up what is going on to look at the posteriors for the nodes immediately beneath the root. Presumably, one (or more) of these is excessively large, and so the constraint is very hard to satisfy by constrained least squares.

[hyanwong]

I actually tried fixing all the immediate children of the root too, and that made no difference. But the contraint could be a couple of levels down.

Here are the posteriors at and below the root (as a tuple with the number of mutations above them). Some are quite large (max 2127.18), but nothing that would make the actual node be pushed to 2559, right. The one with the largest number of mutations (15) is at time 1519.92, so that shouldn't be too hard a constraint to meet ?

print("Original time:", ts.node(tree.root).time)
print("Now:", dts.node(tree.root))
mut_nodes = dts.mutations_node
times_muts = sorted(
    [(dts.node(u).metadata['mn'], int((mut_nodes==u).sum())) for u in tree.children(tree.root)]
)
print("Child posterior node times:")
while len(times_muts):
    print(", ".join([f"t{x[0]:.2f} ({x[1]})" for x in times_muts[0:10]]))
    times_muts = times_muts[10:]

Original time: 1658.0
Now: Node(id=5312042, flags=0, time=2559.453162237328, population=-1, individual=-1, metadata={'Country': '', 'Date_tree': '', 'Genbank_pangolin': '', 'Genbank_scorpio': '', 'Pangolin_usher': '', 'Platform': '', 'Viridian_amplicon_scheme': '', 'Viridian_pangolin_1.29': '', 'Viridian_scorpio': '', 'mn': 1585.10833331689, 'strain': 'node_1', 'vr': 2512.568428350649})
Child posterior node times:
t674.67 (2), t948.00 (1), t967.00 (0), t1040.00 (0), t1040.00 (0), t1040.00 (0), t1040.00 (0), t1040.00 (0), t1058.00 (2), t1115.00 (0)
t1253.00 (0), t1253.00 (0), t1253.00 (0), t1253.00 (1), t1325.00 (0), t1337.00 (0), t1340.00 (0), t1360.00 (0), t1367.00 (0), t1451.00 (3)
t1485.00 (1), t1496.00 (6), t1504.00 (1), t1511.00 (0), t1511.00 (0), t1512.00 (4), t1513.00 (0), t1514.00 (1), t1515.00 (2), t1516.00 (0)
t1518.00 (0), t1519.00 (0), t1519.00 (0), t1519.00 (1), t1519.93 (15), t1520.00 (0), t1520.00 (0), t1520.00 (1), t1522.00 (0), t1522.00 (0)
t1522.00 (0), t1522.00 (1), t1523.00 (1), t1524.00 (3), t1525.00 (0), t1525.00 (0), t1525.00 (0), t1525.00 (0), t1525.00 (4), t1526.00 (3)
t1527.00 (0), t1527.00 (1), t1527.00 (2), t1527.00 (3), t1527.00 (3), t1527.00 (4), t1528.00 (1), t1528.00 (2), t1529.00 (0), t1529.00 (0)
t1529.00 (0), t1529.00 (0), t1529.00 (0), t1529.00 (0), t1529.00 (3), t1529.00 (3), t1530.00 (0), t1530.00 (0), t1530.00 (1), t1530.00 (1)
t1530.00 (2), t1530.00 (3), t1531.00 (0), t1531.00 (1), t1531.00 (1), t1531.00 (1), t1531.00 (3), t1532.00 (0), t1532.00 (2), t1532.00 (2)
t1533.00 (0), t1533.00 (0), t1533.00 (0), t1533.00 (1), t1533.00 (1), t1533.00 (2), t1533.00 (2), t1534.00 (0), t1534.00 (0), t1534.00 (0)
t1534.00 (1), t1534.00 (3), t1535.00 (0), t1535.00 (0), t1535.00 (0), t1535.00 (1), t1535.00 (1), t1535.00 (1), t1536.00 (0), t1536.00 (0)
t1536.00 (0), t1536.00 (0), t1536.00 (0), t1536.00 (1), t1536.00 (1), t1536.00 (1), t1536.00 (2), t1536.00 (3), t1536.00 (3), t1537.00 (0)
t1537.00 (0), t1537.00 (0), t1537.00 (0), t1537.00 (1), t1537.00 (1), t1537.00 (2), t1538.00 (0), t1538.00 (0), t1538.00 (0), t1538.00 (0)
t1538.00 (0), t1538.00 (1), t1538.00 (1), t1538.00 (3), t1539.00 (0), t1539.00 (0), t1539.00 (0), t1540.00 (0), t1540.00 (1), t1540.00 (1)
t1541.00 (0), t1541.00 (0), t1541.00 (0), t1541.00 (1), t1541.00 (1), t1541.00 (2), t1541.00 (2), t1541.00 (2), t1542.00 (0), t1542.00 (0)
t1542.00 (0), t1542.00 (1), t1543.00 (0), t1543.00 (0), t1543.00 (0), t1543.00 (0), t1543.00 (0), t1543.00 (1), t1543.00 (1), t1543.00 (1)
t1543.00 (3), t1544.00 (0), t1544.00 (0), t1545.00 (1), t1546.00 (1), t1547.00 (0), t1547.00 (0), t1547.00 (0), t1547.00 (0), t1547.00 (0)
t1547.00 (0), t1547.00 (1), t1548.00 (0), t1548.00 (0), t1548.00 (0), t1548.00 (0), t1548.00 (0), t1548.00 (1), t1548.00 (1), t1548.00 (2)
t1549.00 (0), t1549.00 (0), t1549.00 (0), t1549.00 (0), t1549.00 (1), t1550.00 (0), t1550.00 (1), t1551.00 (1), t1551.00 (1), t1552.00 (0)
t1552.00 (0), t1552.00 (0), t1553.00 (0), t1553.00 (0), t1553.00 (0), t1553.00 (0), t1554.00 (1), t1554.66 (4), t1555.00 (0), t1555.00 (0)
t1555.00 (0), t1555.00 (0), t1555.00 (0), t1555.00 (0), t1556.00 (0), t1556.00 (0), t1558.00 (0), t1558.00 (2), t1559.00 (0), t1559.00 (0)
t1561.00 (1), t1562.00 (0), t1563.00 (0), t1563.00 (0), t1564.00 (0), t1564.00 (0), t1564.00 (0), t1564.00 (1), t1566.00 (0), t1567.00 (0)
t1567.00 (2), t1567.00 (2), t1570.31 (3), t1571.35 (4), t1578.00 (0), t1578.85 (2), t1579.00 (0), t1580.00 (2), t1580.46 (4), t1582.00 (3)
t1582.00 (4), t1583.00 (0), t1584.00 (1), t1584.00 (2), t1586.00 (1), t1586.00 (1), t1587.00 (1), t1587.00 (4), t1590.00 (0), t1591.00 (1)
t1592.16 (2), t1592.45 (1), t1593.54 (0), t1594.22 (3), t1595.29 (6), t1597.67 (1), t1600.12 (1), t1600.50 (1), t1601.60 (1), t1601.73 (1)
t1601.88 (2), t1602.42 (1), t1602.48 (1), t1603.07 (0), t1604.10 (1), t1604.15 (1), t1606.92 (1), t1608.36 (1), t1608.84 (4), t1609.46 (1)
t1611.24 (2), t1614.71 (1), t1614.83 (2), t1617.68 (1), t1619.55 (1), t1632.45 (1), t1634.48 (1), t1634.68 (1), t1636.37 (1), t1640.98 (4)
t1641.88 (1), t1650.38 (1), t1653.68 (1), t1653.89 (1), t1654.47 (1), t1656.50 (2), t1661.18 (1), t1663.86 (1), t1664.41 (1), t1665.08 (4)
t1665.41 (1), t1671.63 (1), t1671.88 (4), t1671.89 (1), t1672.53 (1), t1673.72 (2), t1674.37 (2), t1675.95 (1), t1684.64 (1), t1685.07 (1)
t1687.42 (1), t1695.87 (1), t1699.79 (1), t1702.54 (1), t1706.37 (1), t1707.97 (2), t1710.11 (2), t1722.84 (1), t1724.73 (2), t1734.70 (1)
t1736.17 (1), t1742.42 (1), t1755.16 (1), t1761.81 (2), t1762.21 (1), t1763.35 (1), t1767.57 (1), t1775.98 (2), t1778.63 (1), t1780.88 (1)
t1792.54 (2), t1796.08 (1), t1831.89 (1), t1862.01 (1), t1864.46 (1), t1897.68 (1), t1925.76 (1), t1926.22 (1), t1939.64 (1), t1984.35 (1)
t1990.02 (1), t2009.30 (1), t2117.88 (1), t2127.18 (1)

[jeromekelleher]

Just to note that I'm not sure this is critical - just dating the internal nodes crudely and forcing the root to be at a sensible time is probably good enough for anything we'll do with this tree.

[nspope]

That's good to know @jeromekelleher -- though I'm happy to spend a little bit of time trying to figure out if this is a bug or just suboptimal behaviour

[nspope]

@hyanwong hmm this is a bit tricky, because the constraint works cumulatively. Say there's one subtree of the root, where, for whatever reason, one node gets bumped up substantially. Then many of its parents will get bumped up, and grandparents, etc and this'll cascade up to the root. But I think we should have one posterior somewhere that is ~2559, for this to be happening?

Also the constraint doesn't care about mutations (it's just using the posteriors) so I think that's a red herring until we trace this back to an EP update.

[hyanwong]

@hyanwong hmm this is a bit tricky, because the constraint works cumulatively. Say there's one subtree of the root, where, for whatever reason, one node gets bumped up substantially. Then many of its parents will get bumped up, and grandparents, etc and this'll cascade up to the root. But I think we should have one posterior somewhere that is ~2559, for this to be happening?

Yes, good point. Here are the 20 nodes with the oldest posterior times, plus their times in the dated and original TS

print("Posterior:", post_mn[np.argsort(post_mn)[-20:]])
print("Final dated:", dts.nodes_time[np.argsort(post_mn)[-20:]])
print("Original dates:", ts.nodes_time[np.argsort(post_mn)[-20:]])
print(f"Hops to root (@ time : {ts.nodes_time[tree.root]})", [len(list(tree.ancestors(u))) for u in np.argsort(post_mn)[-20:]])
print("Muts above", [(ts.mut_nodes==u).sum() for u in np.argsort(post_mn)[-20:]])```

Posterior: [2699.74661004 2723.05063179 2731.20393664 2735.58902176 2737.70149581
 2742.09482063 2744.81373805 2747.50833239 2762.95865797 2770.08300993
 2778.93670399 2821.20101346 2842.65449721 2845.8847962  2911.24774113
 2919.89663357 3039.04476697 3052.20277163 3130.53783579 3193.6573234 ]
Final dated: [2443.80201359 2443.80201359 2559.45316224 2403.77061284 2443.80201359
 2559.45316224 2559.45316224 2559.45316224 2559.45316224 2559.45316224
 2559.45316224 2559.45316224 2559.45316224 2559.45316224 2559.45316224
 2559.45316224 2559.45316224 2559.45316224 2559.45316224 2559.45316224]
Original dates: [1357. 1479. 1650. 1263. 1358. 1305. 1258. 1306. 1264. 1615.  871. 1501.
 1273.  730. 1260. 1272. 1649.  872.  874.  873.]
Hops to root (@ time : 1658.0) [11, 9, 8, 7, 10, 10, 11, 9, 7, 8, 13, 6, 7, 13, 9, 8, 9, 12, 10, 11]
Muts above [2, 1, 7, 2, 1, 1, 3, 2, 1, 4, 8, 1, 2, 5, 2, 1, 11, 3, 6, 3]

By my reading, the problem is the node fourth from the end, which is at time 1649 (close to the root at 1658), but has 11 mutations above it (contributing to the EP update) and therefore a posterior of 3039. Will treating this as a sample help to fix it?

Also the constraint doesn't care about mutations (it's just using the posteriors) so I think that's a red herring until we trace this back to an EP update.

Doh, yes, of course.

[hyanwong]

Will treating this as a sample help to fix it?

I'm not sure that the non-zero fixed nodes are actually fixed until the constraint step, right?

Also, if I increase the mutation rate (I presume making mutation-full branches less unlikely), then the result converges more quickly to what I expect. I presume that increases biases elsewhere, but it might be OK for the sort other thing we are proposing, on a single tree.

[nspope]

fixed nodes are actually fixed until the constraint step

On the contrary, everything that is fixed should be fixed through all steps, except the last forcing of the constraint.

Also, if I increase the mutation rate (I presume making mutation-full branches less unlikely)

Presumably (as this is from real viral data) the effective mutation rate could vary through time and across branches? Like, due to selection?

[hyanwong]

On the contrary, everything that is fixed should be fixed through all steps, except the last forcing of the constraint.

Oh yes, doh. Sorry. I was having a senior moment.

In our specific use-case, we can work around this by using a tree that has been subsampled to only those samples which we have identified (via sc2ts) as not being "time-travellers".

So solving this issue has slightly less priority for me now.