Add tests for two trees

jeromekelleher · Oct 7, 2024 · aa4f58c · aa4f58c
1 parent 9a27181
commit aa4f58c
Showing 1 changed file with 68 additions and 0 deletions.
diff --git a/tests/test_info.py b/tests/test_info.py
@@ -2,6 +2,7 @@
 import numpy as np
 import pandas as pd
 
+import msprime
 import tskit
 
 from sc2ts import info
@@ -78,3 +79,70 @@ def test_1tree_2mut_reversion(self):
         expected[4] = 1
         actual = info.get_num_muts(ts)
         np.testing.assert_equal(expected, actual)
+
+    def test_2trees_0mut(self):
+        ts = msprime.sim_ancestry(
+            2,
+            recombination_rate=1e6, # Nearly guarantee recomb.
+            sequence_length=2,
+        )
+        assert ts.num_trees == 2
+        expected = np.zeros(ts.num_nodes, dtype=np.int32)
+        actual = info.get_num_muts(ts)
+        np.testing.assert_equal(expected, actual)
+
+    def test_2trees_1mut(self):
+        ts = msprime.sim_ancestry(
+            4,
+            ploidy=1,
+            recombination_rate=1e6, # Nearly guarantee recomb.
+            sequence_length=2,
+        )
+        tables = ts.dump_tables()
+        tables.sites.add_row(0, "A")
+        tables.mutations.add_row(site=0, node=0, derived_state="T")
+        ts = tables.tree_sequence()
+        assert ts.num_trees == 2
+        expected = np.zeros(ts.num_nodes, dtype=np.int32)
+        expected[0] = 1
+        actual = info.get_num_muts(ts)
+        np.testing.assert_equal(expected, actual)
+
+    def test_2trees_2mut_diff_trees(self):
+        ts = msprime.sim_ancestry(
+            4,
+            ploidy=1,
+            recombination_rate=1e6, # Nearly guarantee recomb.
+            sequence_length=2,
+        )
+        tables = ts.dump_tables()
+        tables.sites.add_row(0, "A")
+        tables.sites.add_row(1, "A")
+        tables.mutations.add_row(site=0, node=0, derived_state="T")
+        tables.mutations.add_row(site=1, node=0, derived_state="T")
+        ts = tables.tree_sequence()
+        assert ts.num_trees == 2
+        expected = np.zeros(ts.num_nodes, dtype=np.int32)
+        expected[0] = 2
+        actual = info.get_num_muts(ts)
+        np.testing.assert_equal(expected, actual)
+
+    def test_2trees_2mut_same_tree(self):
+        ts = msprime.sim_ancestry(
+            4,
+            ploidy=1,
+            recombination_rate=1e6, # Nearly guarantee recomb.
+            sequence_length=2,
+        )
+        tables = ts.dump_tables()
+        tables.sites.add_row(0, "A")
+        tables.sites.add_row(1, "A")
+        tables.mutations.add_row(site=0, node=0, derived_state="T")
+        tables.mutations.add_row(site=1, node=3, derived_state="T")
+        ts = tables.tree_sequence()
+        assert ts.num_trees == 2
+        expected = np.zeros(ts.num_nodes, dtype=np.int32)
+        expected[0] = 1
+        expected[3] = 1
+        actual = info.get_num_muts(ts)
+        np.testing.assert_equal(expected, actual)