NaN's in `ld_matrix` output with "ghost" alleles

[nspope]

Here's a MWE:

# /// script                                                                                                                                                   
# dependencies = ["tskit==1.0.3"]
# requires-python = ">=3.13"
# ///
import tskit

tables = tskit.TableCollection(sequence_length=10)
for _ in range(4):
    tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
tables.nodes.add_row(flags=0, time=1)  # node 4: parent of samples 0,1
tables.nodes.add_row(flags=0, time=2)  # node 5: root
for parent, child in [(4, 0), (4, 1), (5, 4), (5, 2), (5, 3)]:
    tables.edges.add_row(left=0, right=10, parent=parent, child=child)

# Site 0: A -> G then G -> T stacked on node 4 
# (so "G" is in the allele list but carried by zero samples)
s0 = tables.sites.add_row(position=2, ancestral_state="A")
mG = tables.mutations.add_row(site=s0, node=4, derived_state="G", time=1.5)
tables.mutations.add_row(site=s0, node=4, derived_state="T", time=1.2, parent=mG)

# Site 1: plain biallelic A -> C on node 4
s1 = tables.sites.add_row(position=6, ancestral_state="A")
tables.mutations.add_row(site=s1, node=4, derived_state="C", time=1.5)

# Site 2: plain biallelic A -> G on node 4
s2 = tables.sites.add_row(position=8, ancestral_state="A")
tables.mutations.add_row(site=s2, node=4, derived_state="G", time=1.5)

tables.sort()
ts = tables.tree_sequence()

print(f"\n---\ntree sequence:")
print(ts.draw_text())

print("\n---\nsites:")
for v in ts.variants():
    print(f"site {v.site.id}: alleles={v.alleles} genotypes={v.genotypes.tolist()}")

print("\n---\nstats (sites 0, 1):")
for stat in ["D", "r", "r2", "D_prime", "D2"]:
    print(f"{stat:8s} = {ts.ld_matrix(sites=[[0], [1]], stat=stat)[0, 0]}")

print("\n---\nstats (sites 1, 2):")
for stat in ["D", "r", "r2", "D_prime", "D2"]:
    print(f"{stat:8s} = {ts.ld_matrix(sites=[[1], [2]], stat=stat)[0, 0]}")   

Output:

#---                                                                                                                                                            
#tree sequence:                                                                                                                                                 
#2.00┊     5   ┊                                                                                                                                                
#    ┊  ┏━━╋━┓ ┊                                                                                                                                                
#1.00┊  4  ┃ ┃ ┊                                                                                                                                                
#    ┊ ┏┻┓ ┃ ┃ ┊                                                                                                                                                
#0.00┊ 0 1 2 3 ┊                                                                                                                                                
#    0        10                                                                                                                                                                                                                                                                                                               
#                                                                                                                                                               
---
sites:
site 0: alleles=('A', 'G', 'T') genotypes=[2, 2, 0, 0]
site 1: alleles=('A', 'C') genotypes=[1, 1, 0, 0]
site 2: alleles=('A', 'G') genotypes=[1, 1, 0, 0]

---
stats (sites 0, 1):
D        = 0.125
r        = nan
r2       = nan
D_prime  = nan
D2       = 0.041666666666666664

---
stats (sites 1, 2):
D        = 0.25
r        = 1.0
r2       = 1.0
D_prime  = 1.0
D2       = 0.0625

Given the stats that get NaN'd are the normalised ones, I'm assuming this is a division-by-zero somewhere in the normalisation.

[apragsdale]

Thank @nspope - yes, there is a division by zero. For the multi-allelic site, it computes the statistic for each allele (even those that have frequency zero or one), and those statistics that are giving NaN values divide by p*(1-p), which causes the division by zero.

This is related to #3441. Note that your example doesn't even require a multi-allelic sites, but a back mutation (e.g., A->T, followed by T->A on the same branch) would also trigger this, since we would compute those statistics for a site that is monomorphic.

So this isn't a bug in the sense that the code is doing something unexpected based on what we wrote up and initially intended, but it probably is a bug in the sense that it's an unexpected behavior for a user and probably not the behavior we want for these statistics. The fix requires some careful thought (there are quite a few edge cases with different statistics, especially once you consider multi-population stats), and I haven't had the bandwidth so far unfortunately.

[nspope]

Got it --- thanks @apragsdale. I'll close as there's another issue tracking this.

[apragsdale]

I think this issue is somewhat separate, though the same root cause, since one has to do with alleles only found in one subset of the full set of samples, and this issue is due to multiallelic sites where there are alleles with zero copies in the sample. (So it could be worth keeping this open.)