nans in `ld_matrix`

[petrelharp]

I'm moving discussion of nans in ts.ld_matrix here from #3416.

Consider the almost-simplest application of branch-mode LD:

>>> ts = msprime.sim_ancestry(6)
>>> ts.ld_matrix([0,1, 2, 3], mode='branch')
array([[nan]])
>>> ts.ld_matrix([[0,1,2,3], [4,5]], mode='branch', indexes=(0,1))
array([[nan]])

This is because there are in both cases branches that aren't ancestral to some of the sample sets (of course).
In other words - unless I'm missing something - r2 will always be nan unless you're doing it on all the samples in a simplified ARG. This seems unhelpful.

The same thing happens in site mode if there's mutations on those branches, even under infinite sites:

>>> ts = msprime.sim_mutations(ts, rate=0.3, discrete_genome=False)
>>> ts.ld_matrix([0,1,2,3])
array([[       nan,        nan,        nan,        nan],
       [       nan,        nan,        nan,        nan],
       [       nan,        nan, 1.        , 0.11111111],
       [       nan,        nan, 0.11111111, 1.        ]])
>>> ts.ld_matrix([[0,1,2,3], [4,5]], indexes=(0,1))
array([[nan, nan, nan, nan],
       [nan, nan, nan, nan],
       [nan, nan, nan, nan],
       [nan, nan, nan, nan]])

I think this probably ought to be explained more prominantly, or maybe have the behavior changed, because this is probably the first thing that people will run into when trying some of this stuff out. So, question for @apragsdale:

Is this the desired behavior?

If it is, then I think the explanation of this (and more importantly, suggestion of what to do instead) should be put semi-prominantly in the docs.

I think this question boils down to "what is the natural thing to do with genomic data". In other words, if you compute LD from a subset of samples, what do you do with the sites that are monomorphic in your sample? I think you ignore them, which suggests that branch mode should just ignore those branches that are not polymorphic. This is what the other statistics do, mostly: see this not note (that could be better). However, I could also see a good answer here being "r2 is a ratio and so it's not a good thing to average across a bunch of sites; nan is the right answer here and you should just use a different, non-ratio statistic instead". But then I might suggest either not making r2 the default or throwing a warning for r2 in branch mode.

[apragsdale]

Thanks @petrelharp - @lkirk and I had also discussed this quite a bit a while ago last year, and I was also concerned about the nans popping up in the ratio statistics.

r2 will always be nan unless you're doing it on all the samples in a simplified ARG. This seems unhelpful.

Is this the desired behavior?

In branch mode specifically, I agree that this is unhelpful, and no it is not the desired behavior. The "quick fix" is to return zero instead of nan for those branches that would result in a non-polymorphism within a sample set. Then when you add up the stat over pairs of branches (weighted by branch lengths), those pairs that are monomorphic at one or both loci will just add zero to the total.

But we don't want to just change the summary function to return zero instead of nan, because I think we actually do want nan values when computing r2 in site mode with monomorphic sites. These are sites for which the statistics is undefined.

So, if this is the desired behavior, then one possible solution is to define a separate summary function for r2 and r that is called in branch mode, which return zeros instead of nans. Then in site mode, we still call the existing summary function that returns nan for those stats with monomorphisms. Perhaps there are some issues with this approach, though.

Whatever we decide here should be thoroughly documented, but I think it should be its own PR rather than fitting it into #3416.

[petrelharp]

But we don't want to just change the summary function to return zero instead of nan, because I think we actually do want nan values when computing r2 in site mode with monomorphic sites. These are sites for which the statistics is undefined.

I mostly agree. However, consider the following: suppose a site has more than two alleles. Computing r2-LD with sample sets that don't inherit from some of those alleles returns nan, whereas if you computed it using the simplified tree sequence it'd not return nan. In other words, it's being affected by the presence of alleles outside the samples that you're passing in. Here's an example of this:

t = tskit.Tree.generate_star(3).tree_sequence.dump_tables()
t.sites.add_row(position=0, ancestral_state='A')
t.mutations.add_row(site=0, node=0, derived_state='C')
t.mutations.add_row(site=0, node=2, derived_state='T')
sts = t.tree_sequence()

sts.ld_matrix()  # 1.0
sts.ld_matrix([0,1])  # nan
sts.simplify([0,1]).ld_matrix()  # 1.0

The "quick fix" is to return zero instead of nan for those branches that would result in a non-polymorphism within a sample set.

The fix I was thinking of was along these lines, but to not mess with the summary functions, but to define ld_matrix so that:

1. alleles or branches not inherited by anyone in the sample sets are just plain not included (the summary function is never evaluated there)
2. same for those ancestral to all of them for polarised stats
3. if then the summary function gets evaluated zero times (as for instance if it's a monomorphic site) the answer is nan.

I haven't checked carefully, so maybe this doesn't do the right thing for some of the stats. (In particular, I'm not sure about (2).) But the prinicple is that if we can't "see" something from the samples we're using, then it shouldn't affect the calculation.

[petrelharp]

Notes with @apragsdale:

1. All the non-ratio stats return 0 for monomorphic alleles (freq 0 or freq 1).
2. Ratio statistics are nan for monomorphic alleles. (this is 0/0)
3. In practice, when summarizing ratio stats (like r2) we average the values for sites at which the value isn't nan.
4. So, we should do the same thing: for branch stats, and when dealing with multiallelic sites; that is, average over only the branches/alleles that are polymorphic.
5. What does "polymorphic" mean? In a one-way stat, it means "in the sample set we're looking at". In a two-way stat it could mean either "in both the sample sets" or "in their union"; we think for the same reason it should mean "in both the sample sets".
6. The maybe confusing edge case then is "what if there are no branches or alleles that are polymorphic in both"; then we want to return either nan (for ratio stats) or zero for non-ratio stats. (This may or may not require special-casing; maybe this just ends up at 0/0 in the code, naturally; if we "initialize both numerator and denominator sums to be zero".)

[petrelharp]

Takeaway: we think "all" we have to do is to only let branches/alleles contribute that are polymorphic in the sample set(s) under consideration.