Merge pull request #189 from szhan/filter_by_hmm_cost_without_group

Filter samples by HMM cost without match grouping

Author: jeromekelleher

sc2ts/inference.py

@@ -447,10 +447,12 @@ def add_matching_results(
         # By default, arbitraily high.
         max_hmm_cost = 1e6
 
-    # Group matches by path and set of reversion mutations
+    # Group matches by path and set of immediate reversions.
     grouped_matches = collections.defaultdict(list)
     site_masked_samples = np.zeros(int(ts.sequence_length), dtype=int)
     for sample in samples:
+        if sample.get_hmm_cost(num_mismatches) > max_hmm_cost:
+            continue
         site_masked_samples[sample.masked_sites] += 1
         path = tuple(sample.path)
         reversions = tuple(
@@ -460,17 +462,6 @@ def add_matching_results(
         )
         grouped_matches[(path, reversions)].append(sample)
 
-    # Exclude single samples with "high-HMM cost" attachment paths.
-    tmp = {}
-    for k, v in grouped_matches.items():
-        if len(v) == 1:
-            # Exclude sample if it's HMM cost exceeds a maximum.
-            sample = v[0]
-            if sample.get_hmm_cost(num_mismatches) > max_hmm_cost:
-                continue
-        tmp[k] = v
-    grouped_matches = tmp
-
     tables = ts.dump_tables()
     logger.info(f"Got {len(grouped_matches)} distinct paths")
 

tests/test_inference.py

@@ -83,24 +83,25 @@ def test_one_sample_recombinant_filtered(self):
         assert ts2.num_nodes == ts.num_nodes
         assert ts2.num_samples == ts.num_samples
 
-    def test_two_samples_recombinant_not_filtered(self):
-        """
-        Test case where two identical recombinant samples get added
-        but not excluded despite HMM costs above the max threshold.
-        """
+    def test_two_samples_recombinant_one_filtered(self):
         ts = util.example_binary(2)
         L = ts.sequence_length
         x = L / 2
-        new_sample_paths = [
-            [(0, x, 2), (x, L, 3)],
-            [(0, x, 2), (x, L, 3)],
+        new_paths = [
+            [(0, x, 2), (x, L, 3)],  # Added
+            [
+                (0, L / 4, 2),
+                (L / 4, L / 2, 3),
+                (L / 2, 3 / 4 * L, 4),
+                (3 / 4 * L, L, 2),
+            ],  # Filtered
         ]
-        samples = util.get_samples(ts, new_sample_paths)
+        samples = util.get_samples(ts, new_paths)
         ts2 = sc2ts.add_matching_results(
-            samples, ts, "2021", num_mismatches=1e3, max_hmm_cost=1e3 - 1
+            samples, ts, "2021", num_mismatches=3, max_hmm_cost=4
         )
         assert ts2.num_trees == 2
-        assert ts2.num_samples == ts.num_samples + len(new_sample_paths)
+        assert ts2.num_samples == ts.num_samples + 1
 
     def test_one_sample_one_mutation(self):
         ts = sc2ts.initial_ts()
@@ -130,24 +131,25 @@ def test_one_sample_one_mutation_filtered(self):
         assert ts2.site(0).ancestral_state == ts.site(0).ancestral_state
         assert ts2.num_mutations == 0
 
-    def test_two_samples_one_mutation_not_filtered(self):
+    def test_two_samples_one_mutation_one_filtered(self):
         ts = sc2ts.initial_ts()
         ts = sc2ts.increment_time("2020-01-01", ts)
-        new_sample_paths = [
+        x = int(ts.sequence_length / 2)
+        new_paths = [
             [(0, ts.sequence_length, 1)],
             [(0, ts.sequence_length, 1)],
         ]
-        new_sample_mutations = [
-            [(0, "X")],
-            [(0, "X")],
+        new_mutations = [
+            [(0, "X")],  # Added
+            [(0, "X"), (x, "X")],  # Filtered
         ]
         samples = util.get_samples(
             ts,
-            paths=new_sample_paths,
-            mutations=new_sample_mutations,
+            paths=new_paths,
+            mutations=new_mutations,
         )
         ts2 = sc2ts.add_matching_results(
-            samples, ts, "2021", num_mismatches=0.0, max_hmm_cost=0.0
+            samples, ts, "2021", num_mismatches=3, max_hmm_cost=1
         )
         assert ts2.num_trees == ts.num_trees
         assert ts2.site(0).ancestral_state == ts.site(0).ancestral_state

tests/test_utils.py

@@ -35,7 +35,7 @@ def make_recombinant_tree(self, num_samples=1):
         x = L / 2
         samples = util.get_samples(ts, [[(0, x, 2), (x, L, 3)]] * num_samples)
         ts_rec = sc2ts.add_matching_results(
-            samples, ts, "2021", num_mismatches=None, num_samples=None
+            samples, ts, "2021", num_mismatches=None, max_hmm_cost=None
         )
         assert ts_rec.num_trees == 2
         return ts_rec