Tskit: Save sample-relevant intermediate (unary) nodes during simplify()

[hyanwong]

@pkalbers is interested in saving information about those mutations that make it into the samples in a (forwards-simulated) tree sequence. These mutations are often likely to belong to a node that is "intermediate" (i.e. it is not a coalescence node - it has a single child and a single parent at the focal position). Normally these nodes are cut out during simplify(), even if they their descendants are present in the samples. It would be useful to be able to mark those nodes during simplify(), such that they aren't cut out, unless they have no descendant samples any more. Also see a short discussion at https://github.com/MesserLab/SLiM/issues/25#issuecomment-437515243 and @petrelharp 's reply to the point just below that point in the thread.

p.s. I have put this as a (first!) issue in ts-kit, as it really isn't relevant to msprime, which has no concept of a unary node with an associated mutation (mutations are always placed above binary nodes)

[petrelharp]

There's a couple of possibilities for what exactly you'd want to do here. Thinking about it a bit more, I think all of them are addressable by first going through the tree sequence to decide which nodes you want to keep around, then passing those to simplify(): seems like you should start with that, at the python level; if it turns out to be a common use case we can make the case for adding this functionality to simplify().

[hyanwong]

Yes, I think you are right: I can call simplify() with the samples plus the internal nodes to keep, and work out whether a node has descendants before deciding to keep it. Although I don’t know how I would call simplify() with bespoke parameters like this within a SLiM simulation: I guess I would have to hack the SLiM codebase?

[hyanwong]

I'm assuming that @jeromekelleher might have some nifty fast algorithms for working out whether a node has any sample descendants anywhere in the TS. In fact, I suspect this must be coded in C already, as part of simplify(), so it may just be a question of exposing that functionality to python.

[petrelharp]

If you want to do this in SLiM, as you go along, then yes, you need to go into the C++.

But, before we think about addressing this, you need to make a quantitative case that this is necessary: it is not clear to me that there exists a use case for which this would give you substantial savings over just Remembering everyone that you might want to know about later. To weigh in on this any further, I'm going to need to see something concrete.

[hyanwong]

Thanks @petrelharp - this seems like a reasonable request, especially as I haven't actually coded anything up yet. My general thinking is that if we want to Remember all neutral mutations in a spatial model, there are likely to be a large number of individuals in which mutations occur, but which never spread (~ 1/Ne will go to fixation). So I assume remembering all individuals in which mutations occur will mean remembering roughly Ne times the number of individuals that would be required in a properly simplified TS.

I've probably not thought it through properly, though. I guess it's easy enough to test in SLiM, with a post-processing simplification stage to compare the two approaches.

[hyanwong]

@petrelharp: attached is a python3 script to test this (in .txt form, to allow uploading to github). It is based on the continuous space model in section 14.1 of the manual (the weeding out of unconnected individuals is very hacky, but it does work).

For a population of size 500, run for 2000 generations, I get 20727 Remembered individuals, including the 500 ALIVE, which is weeded down to 749 individuals if we only keep those that are intermediates in the simplified tree sequence. That reduces the number of edges from 72409 to 2121, essentially a 34X reduction in tree sequence size. I see this as a reasonable saving (although nothing like my naive reduction by Ne times).

test_remember_mutants.txt (see below for updated script)

[hyanwong]

Another related thought - presumably the simplify() function in SLiM has to do some jiggery-pokery with moving mutations to a node that won't be simplified away. But in fact these are the very nodes that I want to keep. So there might be an argument to provide a switch in the simplify call for forwards simulations (or other equivalent models) that doesn't carry out this particular part of the process. But I suspect the details are a bit more complex than I have been assuming. I will ask @jeromekelleher when he's back.

[hyanwong]

And an updated script with a much more efficient trigger for treeSeqRememberIndividuals(), within a fitness callback, thanks to @bhaller.

test_remember_mutants.txt

[bhaller]

Ah, I see you got the setMutationType() call in there too, excellent.

[jeromekelleher]

Catching up on this, it seems like this has evolved into something that would be useful feature for forward simulators rather than something that tskit itself needs to worry about. If this is the case, perhaps we could move the discuss elsewhere?

On a potentially related note, I do think we need another mode for simplify which keeps unary nodes though. We need this for tsinfer, because the current behaviour of removing all unary nodes actually results in larger tree sequences than we would have, and also obscures some ancestral relationships. If we kept all unary nodes reachable from samples, would this resolve your issue above @hyanwong?

[hyanwong]

I assumed that ts-kit might want to support generic features that might be used by forwards simulators too (i.e. ts-kit was about tree sequences in general, whereas e.g. msprime was the repo for backwards simulator features in particular). Also, I can imagine this being a potentially useful feature in a backwards simulator that (inefficiently) allocates mutations to a specific intermediate node, rather than picking a Poisson number above each coalescent node.

Re: keeping unary nodes - the main problem I have is working out which nodes (both unary and binary) do not lead to any samples in the final TS. What would be useful to me is to have a flag on simplify() that says to drop these sorts of nodes, even if they are passed in as the first parameter. Or alternatively (and probably better), a fast way of identifying such nodes from a list of nodes, some of which might be binary and some unary. Does that make sense, @jeromekelleher ?

[jeromekelleher]

I'm not fully following the thread here @hyanwong; probably best if we talk it through in person. But yes, your intial assertion is right: tskit provides the basic tools that others use to make interesting features.

[hyanwong]

@petrelharp - having chatted to @jeromekelleher he thinks that the first (and possibly only necessary) step is to have a less aggressive simplify() method that leaves all unary nodes as long as they are on the path to a sample. He thinks this should only increase the size of the TS by ~ 1 order of magnitude (and certainly only by a constant amount). I have updated the title of this issue accordingly - I suspect it's something that you are also interested in, and might have some ideas about how to implement it in the guts of simplify?

[petrelharp]

I am interested in this! I suspect it will increase the size of the tree sequence by a factor of N rather than an order of magnitude, but I could easily be wrong and am not coming up with a concise reason at the moment. Modifying simplify in this way should be reasonably straightforward; we could do some experiments with the python version.

[jeromekelleher]

Thinking about this again from the tsinfer perspective, I realised that simplify won't remove any nodes in this case, because every ancestor that we generate is copied from at least twice in one position (i.e., at the focal site). So, in in principle the simplified tree sequence containing all the unary nodes would be exactly the same size as the raw output of the copying process. Hmmm.

Still, it's worth thinking about. The trees we output are very cumbersome as there's lots of stuff reachable from the root that is non ancestry to samples. There must be a way of getting rid of this without chopping up edges too much.

[petrelharp]

The trees we output are very cumbersome as there's lots of stuff reachable from the root that is non ancestry to samples. There must be a way of getting rid of this without chopping up edges too much.

My suspicion is that if we keep information along all of the edges, instead of at the nodes, then we'll be in N^2 land, not N log N land. But this would be great to verify.

[gtsambos]

Hello @bhaller! It's now possible to retain all unary nodes using simplify() with an additional argument keep_unary=True. I think it would be very useful to give SLiM users the option of simplifying in this way during simulation.

Although this would increase memory needs somewhat (see some of the comments above), it would also allow users to retain extra useful information that is currently 'thrown out' by the default implementation of simplify(). For example:

1. You could retain information about the exact meioses at which de novo mutations appear, and thus the ages of the mutations (see @hyanwong's comment that started this thread.)

2. You could retain information about all the populations that are ancestral to a sample at a given location. This would make SLiM a more effective tool for simulating local ancestry. (I'm actively working on this now, and am happy to explain further elsewhere on a more suitable 'issue page').

[bhaller]

Hi @gtsambos. Yes, I think this will be coming over to SLiM. We just need to merge the latest tskit into SLiM, and then adding the necessary API support should be trivial.

[bhaller]

@hyanwong, @petrelharp : I see this issue has been closed, and I'm not sure whether the through-line of actually fixing this issue got dropped as a result. When this issue was closed, should an issue have been opened on the SLiM side to represent the need to now add a flag for this behavior in SLiM? Does the issue at https://github.com/MesserLab/SLiM/issues/25 represent this, or is that a separate issue? (I thought it was separate, but I get a bit lost in the details of the discussions...)