Merge pull request #187 from szhan/hmm_cost_filter

Implement HMM cost to filter samples

Author: jeromekelleher

sc2ts/cli.py

@@ -162,6 +162,7 @@ def add_provenance(ts, output_file):
     ),
 )
 @click.option("--num-mismatches", default=None, type=float, help="num-mismatches")
+@click.option("--max-hmm-cost", default=None, type=float, help="max-hmm-cost")
 @click.option(
     "--max-submission-delay",
     default=None,
@@ -192,6 +193,7 @@ def daily_extend(
     output_prefix,
     base,
     num_mismatches,
+    max_hmm_cost,
     max_submission_delay,
     max_daily_samples,
     num_threads,
@@ -219,6 +221,7 @@ def daily_extend(
             metadata_db=metadata_db,
             base_ts=base_ts,
             num_mismatches=num_mismatches,
+            max_hmm_cost=max_hmm_cost,
             max_submission_delay=max_submission_delay,
             max_daily_samples=max_daily_samples,
             rng=rng,

sc2ts/inference.py

@@ -222,6 +222,11 @@ def submission_date(self):
     def submission_delay(self):
         return (self.submission_date - self.date).days
 
+    def get_hmm_cost(self, num_mismatches):
+        # Note that Recombinant objects have total_cost.
+        # This bit of code is sort of repeated.
+        return num_mismatches * (len(self.path) - 1) + len(self.mutations)
+
     def asdict(self):
         return {
             "strain": self.strain,
@@ -247,6 +252,7 @@ def daily_extend(
     metadata_db,
     base_ts,
     num_mismatches=None,
+    max_hmm_cost=None,
     show_progress=False,
     max_submission_delay=None,
     max_daily_samples=None,
@@ -263,6 +269,7 @@ def daily_extend(
             date=date,
             base_ts=last_ts,
             num_mismatches=num_mismatches,
+            max_hmm_cost=max_hmm_cost,
             show_progress=show_progress,
             max_submission_delay=max_submission_delay,
             max_daily_samples=max_daily_samples,
@@ -340,14 +347,14 @@ def extend(
     date,
     base_ts,
     num_mismatches=None,
+    max_hmm_cost=None,
     show_progress=False,
     max_submission_delay=None,
     max_daily_samples=None,
     num_threads=None,
     precision=None,
     rng=None,
 ):
-
     date_samples = [Sample(md) for md in metadata_db.get(date)]
     samples = filter_samples(date_samples, alignment_store, max_submission_delay)
 
@@ -361,6 +368,7 @@ def extend(
 
     logger.info(f"Got {len(samples)} samples")
 
+    # Note num_mismatches is assigned a default value in match_tsinfer.
     samples = match(
         samples,
         alignment_store=alignment_store,
@@ -371,7 +379,15 @@ def extend(
         precision=precision,
     )
     ts = increment_time(date, base_ts)
-    return add_matching_results(samples, ts, date, show_progress)
+
+    return add_matching_results(
+        samples=samples,
+        ts=ts,
+        date=date,
+        num_mismatches=num_mismatches,
+        max_hmm_cost=max_hmm_cost,
+        show_progress=show_progress,
+    )
 
 
 def match_path_ts(samples, ts, path, reversions):
@@ -394,7 +410,7 @@ def match_path_ts(samples, ts, path, reversions):
                 "qc": sample.alignment_qc,
                 "path": [x.asdict() for x in sample.path],
                 "mutations": [x.asdict() for x in sample.mutations],
-            }
+            },
         }
         node_id = tables.nodes.add_row(
             flags=tskit.NODE_IS_SAMPLE, time=0, metadata=metadata
@@ -407,7 +423,7 @@ def match_path_ts(samples, ts, path, reversions):
 
     # Now add the mutations
     for node_id, sample in enumerate(samples, first_sample):
-        #metadata = {**sample.metadata, "sc2ts_qc": sample.alignment_qc}
+        # metadata = {**sample.metadata, "sc2ts_qc": sample.alignment_qc}
         for mut in sample.mutations:
             tables.mutations.add_row(
                 site=site_id_map[mut.site_id],
@@ -420,7 +436,16 @@ def match_path_ts(samples, ts, path, reversions):
     # print(tables)
 
 
-def add_matching_results(samples, ts, date, show_progress=False):
+def add_matching_results(
+    samples, ts, date, num_mismatches, max_hmm_cost, show_progress=False
+):
+    if num_mismatches is None:
+        # Note that this is the default assigned in match_tsinfer.
+        num_mismatches = 1e3
+
+    if max_hmm_cost is None:
+        # By default, arbitraily high.
+        max_hmm_cost = 1e6
 
     # Group matches by path and set of reversion mutations
     grouped_matches = collections.defaultdict(list)
@@ -435,6 +460,17 @@ def add_matching_results(samples, ts, date, show_progress=False):
         )
         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")
 
@@ -981,6 +1017,7 @@ def match_tsinfer(
     show_progress=False,
     mirror_coordinates=False,
 ):
+    # TODO: Should this default be assigned elsewhere?
     if num_mismatches is None:
         # Default to no recombination
         num_mismatches = 1000

sc2ts/utils.py

@@ -179,14 +179,14 @@ def data_summary(self):
         return d
 
     @property
-    def total_cost(self):
+    def total_cost(self, num_mismatches):
         """
         How different is the causal sequence from the rest, roughly?
         """
         fwd = self.hmm_runs["forward"]
         bck = self.hmm_runs["backward"]
-        cost_fwd = 3 * (len(fwd.parents) - 1) + len(fwd.mutations)
-        cost_bck = 3 * (len(bck.parents) - 1) + len(bck.mutations)
+        cost_fwd = num_mismatches * (len(fwd.parents) - 1) + len(fwd.mutations)
+        cost_bck = num_mismatches * (len(bck.parents) - 1) + len(bck.mutations)
         assert cost_fwd == cost_bck
         return cost_fwd
 

tests/test_inference.py

@@ -9,8 +9,16 @@
 
 
 class TestAddMatchingResults:
-    def add_matching_results(self, samples, ts, date="2020-01-01"):
-        ts2 = sc2ts.add_matching_results(samples, ts, date)
+    def add_matching_results(
+        self, samples, ts, date="2020-01-01", num_mismatches=None, max_hmm_cost=None
+    ):
+        ts2 = sc2ts.add_matching_results(
+            samples,
+            ts,
+            date,
+            num_mismatches,
+            max_hmm_cost,
+        )
         assert ts2.num_samples == len(samples) + ts.num_samples
         for u, sample in zip(ts2.samples()[-len(samples) :], samples):
             node = ts2.node(u)
@@ -54,6 +62,46 @@ def test_one_sample_recombinant(self):
         assert ts2.node(6).flags == sc2ts.NODE_IS_RECOMBINANT
         assert ts2.node(6).metadata == {"date_added": "2021"}
 
+    def test_one_sample_recombinant_filtered(self):
+        # 4.00┊  0  ┊
+        #     ┊  ┃  ┊
+        # 3.00┊  1  ┊
+        #     ┊  ┃  ┊
+        # 2.00┊  4  ┊
+        #     ┊ ┏┻┓ ┊
+        # 1.00┊ 2 3 ┊
+        #     0   29904
+        ts = util.example_binary(2)
+        L = ts.sequence_length
+        x = L / 2
+        samples = util.get_samples(ts, [[(0, x, 2), (x, L, 3)]])
+        # Note that it is calling the function in the main module.
+        ts2 = sc2ts.add_matching_results(
+            samples, ts, "2021", num_mismatches=1e3, max_hmm_cost=1e3 - 1
+        )
+        assert ts2.num_trees == 1
+        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.
+        """
+        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)],
+        ]
+        samples = util.get_samples(ts, new_sample_paths)
+        ts2 = sc2ts.add_matching_results(
+            samples, ts, "2021", num_mismatches=1e3, max_hmm_cost=1e3 - 1
+        )
+        assert ts2.num_trees == 2
+        assert ts2.num_samples == ts.num_samples + len(new_sample_paths)
+
     def test_one_sample_one_mutation(self):
         ts = sc2ts.initial_ts()
         ts = sc2ts.increment_time("2020-01-01", ts)
@@ -69,6 +117,44 @@ def test_one_sample_one_mutation(self):
         var = next(ts2.variants())
         assert var.alleles[var.genotypes[0]] == "X"
 
+    def test_one_sample_one_mutation_filtered(self):
+        ts = sc2ts.initial_ts()
+        ts = sc2ts.increment_time("2020-01-01", ts)
+        samples = util.get_samples(
+            ts, [[(0, ts.sequence_length, 1)]], mutations=[[(0, "X")]]
+        )
+        ts2 = sc2ts.add_matching_results(
+            samples, ts, "2021", num_mismatches=0.0, max_hmm_cost=0.0
+        )
+        assert ts2.num_trees == ts.num_trees
+        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):
+        ts = sc2ts.initial_ts()
+        ts = sc2ts.increment_time("2020-01-01", ts)
+        new_sample_paths = [
+            [(0, ts.sequence_length, 1)],
+            [(0, ts.sequence_length, 1)],
+        ]
+        new_sample_mutations = [
+            [(0, "X")],
+            [(0, "X")],
+        ]
+        samples = util.get_samples(
+            ts,
+            paths=new_sample_paths,
+            mutations=new_sample_mutations,
+        )
+        ts2 = sc2ts.add_matching_results(
+            samples, ts, "2021", num_mismatches=0.0, max_hmm_cost=0.0
+        )
+        assert ts2.num_trees == ts.num_trees
+        assert ts2.site(0).ancestral_state == ts.site(0).ancestral_state
+        assert ts2.num_mutations == 1
+        var = next(ts2.variants())
+        assert var.alleles[var.genotypes[0]] == "X"
+
 
 class TestMatchTsinfer:
     def match_tsinfer(self, samples, ts, haplotypes, **kwargs):

tests/test_utils.py

@@ -34,15 +34,17 @@ def make_recombinant_tree(self, num_samples=1):
         L = ts.sequence_length
         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")
+        ts_rec = sc2ts.add_matching_results(
+            samples, ts, "2021", num_mismatches=None, num_samples=None
+        )
         assert ts_rec.num_trees == 2
         return ts_rec
 
     @pytest.mark.parametrize(
         "ts",
         # Should probably add sc2ts.initial_ts() here, but see
         # https://github.com/jeromekelleher/sc2ts/issues/152
-        [util.example_binary(1), util.example_binary(2), util.example_binary(3)]
+        [util.example_binary(1), util.example_binary(2), util.example_binary(3)],
     )
     def test_no_recombinants(self, ts):
         ts2 = utils.detach_singleton_recombinants(ts)
@@ -52,9 +54,7 @@ def test_one_sample_recombinant(self):
         ts = self.make_recombinant_tree()
         assert ts.num_samples == 3
         re_nodes = [
-            node.id
-            for node in ts.nodes()
-            if node.flags & sc2ts.NODE_IS_RECOMBINANT
+            node.id for node in ts.nodes() if node.flags & sc2ts.NODE_IS_RECOMBINANT
         ]
         assert len(re_nodes) == 1
         re_node = re_nodes[0]