Dump pickled objects of samples filtered by HMM cost

Author: szhan

sc2ts/cli.py

@@ -8,6 +8,7 @@
 import contextlib
 import dataclasses
 import datetime
+import pickle
 
 import tqdm
 import tskit
@@ -147,6 +148,11 @@ def add_provenance(ts, output_file):
     tables.dump(output_file)
 
 
+def dump_samples(samples, output_file):
+    with open(output_file, "wb") as f:
+        pickle.dump(samples, file=f)
+
+
 @click.command()
 @click.argument("alignments", type=click.Path(exists=True, dir_okay=False))
 @click.argument("metadata", type=click.Path(exists=True, dir_okay=False))
@@ -229,9 +235,11 @@ def daily_extend(
             num_threads=num_threads,
             show_progress=not no_progress,
         )
-        for ts, date in ts_iter:
-            output = output_prefix + date + ".ts"
-            add_provenance(ts, output)
+        for ts, excluded_samples, date in ts_iter:
+            output_ts = output_prefix + date + ".ts"
+            add_provenance(ts, output_ts)
+            output_excluded_samples = output_prefix + date + ".excluded_samples.pickle"
+            dump_samples(excluded_samples, output_excluded_samples)
 
 
 @click.command()

sc2ts/inference.py

@@ -263,7 +263,7 @@ def daily_extend(
     start_day = last_date(base_ts)
     last_ts = base_ts
     for date in metadata_db.get_days(start_day):
-        ts = extend(
+        ts, excluded_samples = extend(
             alignment_store=alignment_store,
             metadata_db=metadata_db,
             date=date,
@@ -277,7 +277,7 @@ def daily_extend(
             precision=precision,
             rng=rng,
         )
-        yield ts, date
+        yield ts, excluded_samples, date
         last_ts = ts
 
 
@@ -380,7 +380,7 @@ def extend(
     )
     ts = increment_time(date, base_ts)
 
-    return add_matching_results(
+    ts, excluded_samples = add_matching_results(
         samples=samples,
         ts=ts,
         date=date,
@@ -389,6 +389,8 @@ def extend(
         show_progress=show_progress,
     )
 
+    return ts, excluded_samples
+
 
 def match_path_ts(samples, ts, path, reversions):
     """
@@ -449,9 +451,11 @@ def add_matching_results(
 
     # Group matches by path and set of immediate reversions.
     grouped_matches = collections.defaultdict(list)
+    excluded_samples = []
     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:
+            excluded_samples.append(sample)
             continue
         site_masked_samples[sample.masked_sites] += 1
         path = tuple(sample.path)
@@ -532,7 +536,8 @@ def add_matching_results(
     # print("AFTER")
     # print(ts.draw_text())
     ts = coalesce_mutations(ts, attach_nodes)
-    return ts
+
+    return ts, excluded_samples
 
 
 def solve_num_mismatches(ts, k):

tests/test_inference.py

@@ -12,12 +12,12 @@ class TestAddMatchingResults:
     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,
+        ts2, _ = sc2ts.add_matching_results(
+            samples=samples,
+            ts=ts,
+            date=date,
+            num_mismatches=num_mismatches,
+            max_hmm_cost=max_hmm_cost,
         )
         assert ts2.num_samples == len(samples) + ts.num_samples
         for u, sample in zip(ts2.samples()[-len(samples) :], samples):
@@ -76,7 +76,7 @@ def test_one_sample_recombinant_filtered(self):
         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(
+        ts2, _ = sc2ts.add_matching_results(
             samples, ts, "2021", num_mismatches=1e3, max_hmm_cost=1e3 - 1
         )
         assert ts2.num_trees == 1
@@ -97,7 +97,7 @@ def test_two_samples_recombinant_one_filtered(self):
             ],  # Filtered
         ]
         samples = util.get_samples(ts, new_paths)
-        ts2 = sc2ts.add_matching_results(
+        ts2, _ = sc2ts.add_matching_results(
             samples, ts, "2021", num_mismatches=3, max_hmm_cost=4
         )
         assert ts2.num_trees == 2
@@ -124,7 +124,7 @@ def test_one_sample_one_mutation_filtered(self):
         samples = util.get_samples(
             ts, [[(0, ts.sequence_length, 1)]], mutations=[[(0, "X")]]
         )
-        ts2 = sc2ts.add_matching_results(
+        ts2, _ = sc2ts.add_matching_results(
             samples, ts, "2021", num_mismatches=0.0, max_hmm_cost=0.0
         )
         assert ts2.num_trees == ts.num_trees
@@ -148,7 +148,7 @@ def test_two_samples_one_mutation_one_filtered(self):
             paths=new_paths,
             mutations=new_mutations,
         )
-        ts2 = sc2ts.add_matching_results(
+        ts2, _ = sc2ts.add_matching_results(
             samples, ts, "2021", num_mismatches=3, max_hmm_cost=1
         )
         assert ts2.num_trees == ts.num_trees

tests/test_utils.py

@@ -34,7 +34,7 @@ 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(
+        ts_rec, _ = sc2ts.add_matching_results(
             samples, ts, "2021", num_mismatches=None, max_hmm_cost=None
         )
         assert ts_rec.num_trees == 2