Rough first-pass at site-divmat

Author: jeromekelleher

python/tests/test_divmat.py

@@ -22,6 +22,8 @@
 """
 Test cases for divergence matrix based pairwise stats
 """
+import collections
+
 import msprime
 import numpy as np
 import pytest
@@ -258,6 +260,60 @@ def stats_api_divergence_matrix(ts, windows=None, samples=None):
     return out
 
 
+def rootward_path(tree, u, v):
+    while u != v:
+        yield u
+        u = tree.parent(u)
+
+
+def site_divergence_matrix_naive(ts, windows=None, samples=None):
+    windows_specified = windows is not None
+    windows = [0, ts.sequence_length] if windows is None else windows
+    num_windows = len(windows) - 1
+    samples = ts.samples() if samples is None else samples
+
+    n = len(samples)
+    D = np.zeros((num_windows, n, n))
+    tree = tskit.Tree(ts)
+    for i in range(num_windows):
+        left = windows[i]
+        right = windows[i + 1]
+        tree.seek(left)
+        # Iterate over the trees in this window
+        while tree.interval.left < right and tree.index != -1:
+            span_left = max(tree.interval.left, left)
+            span_right = min(tree.interval.right, right)
+            mutations_per_node = collections.Counter()
+            for site in tree.sites():
+                if span_left <= site.position < span_right:
+                    for mutation in site.mutations:
+                        mutations_per_node[mutation.node] += 1
+            for j in range(n):
+                u = samples[j]
+                for k in range(j + 1, n):
+                    v = samples[k]
+                    w = tree.mrca(u, v)
+                    if w != tskit.NULL:
+                        wu = w
+                        wv = w
+                    else:
+                        wu = local_root(tree, u)
+                        wv = local_root(tree, v)
+                    du = sum(mutations_per_node[x] for x in rootward_path(tree, u, wu))
+                    dv = sum(mutations_per_node[x] for x in rootward_path(tree, v, wv))
+                    # NOTE: we're just accumulating the raw mutation counts, not
+                    # multiplying by span
+                    D[i, j, k] += du + dv
+            tree.next()
+        # Fill out symmetric triangle in the matrix
+        for j in range(n):
+            for k in range(j + 1, n):
+                D[i, k, j] = D[i, j, k]
+    if not windows_specified:
+        D = D[0]
+    return D
+
+
 # NOTE: the internal_checks argument is left over from an older version that
 # used a more complex algorithm - not removing for now so we can reuse the
 # tests that were designed to exercise the code paths of that algorithm.
@@ -696,6 +752,15 @@ def test_threads_windows(self, ts):
         self.check(ts, num_threads=5, windows=windows)
 
 
+class TestSiteDivergence:
+    def test_simulation_example(self):
+        ts = msprime.sim_ancestry(4, sequence_length=100, random_seed=2)
+        ts = msprime.sim_mutations(ts, rate=1)
+        print(ts.num_mutations)
+        D1 = site_divergence_matrix_naive(ts)
+        print(D1)
+
+
 class TestThreadsNoWindows:
     def check(self, ts, num_threads, samples=None):
         D1 = ts.divergence_matrix(num_threads=0, samples=samples)