Merge pull request #499 from jeromekelleher/deletion-sites

Switch map-deletions to require explicit list

Author: jeromekelleher

sc2ts/cli.py

@@ -589,28 +589,16 @@ def minimise_metadata(
 @click.command()
 @click.argument("dataset", type=click.Path(exists=True, dir_okay=False))
 @click.argument("ts_in", type=click.Path(exists=True, dir_okay=False))
+@click.argument("sites", type=click.Path(exists=True, dir_okay=False))
 @click.argument("ts_out", type=click.Path(exists=False, dir_okay=False))
-@click.option(
-    "--frequency-threshold",
-    type=float,
-    default=0.01,
-    help="Frequency threshold for deletions to get mapped back",
-)
-@click.option(
-    "--mutations-threshold",
-    type=int,
-    default=None,
-    help="Maximum number of mutations at a site after parsimony",
-)
 @click.option("--progress/--no-progress", default=True)
 @click.option("-v", "--verbose", count=True)
 @click.option("-l", "--log-file", default=None, type=click.Path(dir_okay=False))
 def map_deletions(
-    ts_in,
     dataset,
+    ts_in,
+    sites,
     ts_out,
-    frequency_threshold,
-    mutations_threshold,
     progress,
     verbose,
     log_file,
@@ -621,13 +609,8 @@ def map_deletions(
     setup_logging(verbose, log_file)
     ds = sc2ts.Dataset(dataset)
     ts = tszip.load(ts_in)
-    ts = sc2ts.map_deletions(
-        ts,
-        ds,
-        frequency_threshold=frequency_threshold,
-        mutations_threshold=mutations_threshold,
-        show_progress=progress,
-    )
+    sites = np.loadtxt(sites, dtype=int)
+    ts = sc2ts.map_deletions(ts, ds, sites, show_progress=progress)
     ts.dump(ts_out)
 
 

sc2ts/inference.py

@@ -1954,38 +1954,13 @@ def get_recombinant_strains(ts):
     return ret
 
 
-def map_deletions(
-    ts, ds, *, frequency_threshold, mutations_threshold=None, show_progress=False
-):
+def map_deletions(ts, ds, sites, *, show_progress=False):
     """
-    Map deletions at sites that exceed the specified frequency threshold onto the
-    ARG using parsimony (excluding flanks), and insert the mutations for any sites
-    that have less than the specified threshold number of mutations.
+    Map deletions at the specified set of site positions ARG using parsimony.
     """
-    mutations_threshold = 2**64 if mutations_threshold is None else mutations_threshold
     start_time = time.time()  # wall time
-    genes = core.get_gene_coordinates()
-    start = genes["ORF1ab"][0]
-    end = genes["ORF10"][1]
-
+    sites = np.array(sites, dtype=int)
     md = ts.metadata
-    num_samples = ts.num_samples
-    ts_contains_exact_matches = md["sc2ts"].get("includes_exact_matches")
-    if ts_contains_exact_matches:
-        total_exact_matches = sum(
-            md["sc2ts"]["cumulative_stats"]["exact_matches"]["pango"].values()
-        )
-        logger.info(
-            f"Exact matches included; adjusting num_samples to remove {total_exact_matches}"
-        )
-        num_samples -= total_exact_matches
-
-    del_sites = []
-    for site in ts.sites():
-        if start <= site.position < end:
-            deletion_samples = site.metadata["sc2ts"]["deletion_samples"]
-            if deletion_samples / ts.num_samples >= frequency_threshold:
-                del_sites.append(site.id)
 
     sample_id = md["sc2ts"]["samples_strain"]
     assert not sample_id[0].startswith("Wuhan")
@@ -1994,26 +1969,27 @@ def map_deletions(
     tree = ts.first()
 
     variants = get_progress(
-        ds.variants(sample_id, ts.sites_position[del_sites]),
+        ds.variants(sample_id, sites),
         title="Map deletions",
         phase="",
         show_progress=show_progress,
-        total=len(del_sites),
+        total=len(sites),
     )
 
-    logger.info(f"Remapping {len(del_sites)} sites")
+    logger.info(f"Remapping {len(sites)} sites")
 
     mut_metadata = {"sc2ts": {"type": "post_parsimony"}}
     site_metadata = {}
     keep_mutations = np.ones(ts.num_mutations, dtype=bool)
     for var in variants:
         tree.seek(var.position)
-        site = ts.site(position=var.position)
+        try:
+            site = ts.site(position=var.position)
+        except ValueError:
+            logger.warning(f"No site at position {var.position}; skipping")
+            continue
 
         g = mask_ambiguous(var.genotypes)
-        deletion_samples = site.metadata["sc2ts"]["deletion_samples"]
-        if not ts_contains_exact_matches:
-            assert deletion_samples == np.sum(g == DELETION)
         _, mutations = tree.map_mutations(
             g, list(var.alleles), ancestral_state=site.ancestral_state
         )
@@ -2022,31 +1998,28 @@ def map_deletions(
             f"Site {int(site.position)} "
             f"mapped mutations = {len(mutations)}; current = {len(site.mutations)}"
         )
-        if len(mutations) < mutations_threshold:
-
-            old_mutations = []
-            for mut in site.mutations:
-                keep_mutations[mut.id] = False
-                old_mutations.append(
-                    {
-                        "node": mut.node,
-                        "derived_state": mut.derived_state,
-                        "metadata": mut.metadata,
-                    }
-                )
-            md = dict(site.metadata)
-            md["sc2ts"]["original_mutations"] = old_mutations
-            site_metadata[site.id] = md
-            for m in mutations:
-                tables.mutations.add_row(
-                    site=site.id,
-                    node=m.node,
-                    derived_state=m.derived_state,
-                    time=ts.nodes_time[m.node],
-                    metadata=mut_metadata,
-                )
-        else:
-            logger.debug(f"Skipping site {int(site.position)} ")
+
+        old_mutations = []
+        for mut in site.mutations:
+            keep_mutations[mut.id] = False
+            old_mutations.append(
+                {
+                    "node": mut.node,
+                    "derived_state": mut.derived_state,
+                    "metadata": mut.metadata,
+                }
+            )
+        md = dict(site.metadata)
+        md["sc2ts"]["original_mutations"] = old_mutations
+        site_metadata[site.id] = md
+        for m in mutations:
+            tables.mutations.add_row(
+                site=site.id,
+                node=m.node,
+                derived_state=m.derived_state,
+                time=ts.nodes_time[m.node],
+                metadata=mut_metadata,
+            )
 
     added_mutations = len(tables.mutations) - ts.num_mutations
     logger.info(
@@ -2064,11 +2037,7 @@ def map_deletions(
     tables.sort()
     tables.build_index()
     tables.compute_mutation_parents()
-    params = {
-        "dataset": str(ds.path),
-        "frequency_threshold": float(frequency_threshold),
-        "mutations_threshold": int(mutations_threshold),
-    }
+    params = {"dataset": str(ds.path), "sites": sites.tolist()}
     prov = get_provenance_dict("map_deletions", params, start_time)
     tables.provenances.add_row(json.dumps(prov))
     return tables.tree_sequence()

tests/test_cli.py

@@ -417,11 +417,15 @@ class TestMapDeletions:
     def test_example(self, tmp_path, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
         out_ts_path = tmp_path / "ts.ts"
+        del_sites_path = tmp_path / "deletion_sites.txt"
+        with open(del_sites_path, "w") as f:
+            print("1547 3951 3952 3953", file=f)
+
         runner = ct.CliRunner()
         result = runner.invoke(
             cli.cli,
-            f"map-deletions {fx_dataset.path} {ts.path} {out_ts_path} "
-            "--frequency-threshold=0.0001",
+            f"map-deletions {fx_dataset.path} {ts.path} {del_sites_path} "
+            f"{out_ts_path} ",
             catch_exceptions=False,
         )
         assert result.exit_code == 0

tests/test_inference.py

@@ -1398,13 +1398,14 @@ def fx_ts_exact_matches(fx_ts_map, fx_match_db):
 class TestMapDeletions:
     def test_example(self, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
-        new_ts = sc2ts.map_deletions(ts, fx_dataset, frequency_threshold=0.001)
+        sites = [1547, 3951, 3952, 3953]
+        new_ts = sc2ts.map_deletions(ts, fx_dataset, sites)
         remapped_sites = [
             j
             for j in range(ts.num_sites)
             if "original_mutations" in new_ts.site(j).metadata["sc2ts"]
         ]
-        assert remapped_sites == [1541, 3945, 3946, 3947]
+        assert remapped_sites == list(np.searchsorted(ts.sites_position, sites))
 
         for site_id in remapped_sites:
             site = new_ts.site(site_id)
@@ -1425,11 +1426,22 @@ def test_example(self, fx_ts_map, fx_dataset):
             for mut in site.mutations:
                 assert mut.metadata["sc2ts"]["type"] == "post_parsimony"
 
-    def test_filter_all(self, fx_ts_map, fx_dataset):
+    def test_empty(self, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
-        new_ts = sc2ts.map_deletions(
-            ts, fx_dataset, frequency_threshold=0.001, mutations_threshold=0
-        )
+        new_ts = sc2ts.map_deletions(ts, fx_dataset, [])
+        remapped_sites = [
+            j
+            for j in range(ts.num_sites)
+            if "original_mutations" in new_ts.site(j).metadata["sc2ts"]
+        ]
+        assert remapped_sites == []
+
+    def test_missing_site(self, fx_ts_map, fx_dataset):
+        ts = fx_ts_map["2020-02-13"]
+        missing_positions = [56, 57, 58, 59, 60]
+        assert len(set(missing_positions) & set(ts.sites_position.astype(int))) == 0
+
+        new_ts = sc2ts.map_deletions(ts, fx_dataset, missing_positions)
         remapped_sites = [
             j
             for j in range(ts.num_sites)
@@ -1439,29 +1451,31 @@ def test_filter_all(self, fx_ts_map, fx_dataset):
 
     def test_example_exact_matches(self, fx_ts_exact_matches, fx_dataset):
         ts = fx_ts_exact_matches
-        new_ts = sc2ts.map_deletions(ts, fx_dataset, frequency_threshold=0.001)
+        sites = [1547, 3951, 3952, 3953]
+        new_ts = sc2ts.map_deletions(ts, fx_dataset, sites)
         remapped_sites = [
             j
             for j in range(ts.num_sites)
             if "original_mutations" in new_ts.site(j).metadata["sc2ts"]
         ]
-        assert remapped_sites == [1541, 3945, 3946, 3947]
+        assert remapped_sites == list(np.searchsorted(ts.sites_position, sites))
 
     def test_validate(self, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
-        new_ts = sc2ts.map_deletions(ts, fx_dataset, frequency_threshold=0.001)
+        sites = [1547, 3951, 3952, 3953]
+        new_ts = sc2ts.map_deletions(ts, fx_dataset, sites)
         sc2ts.validate(new_ts, fx_dataset, deletions_as_missing=False)
 
     def test_provenance(self, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
-        tsp = sc2ts.map_deletions(ts, fx_dataset, frequency_threshold=0.125)
+        sites = [1547, 3951, 3952, 3953]
+        tsp = sc2ts.map_deletions(ts, fx_dataset, sites)
         assert tsp.num_provenances == ts.num_provenances + 1
         prov = tsp.provenance(-1)
         assert json.loads(prov.record)["parameters"] == {
             "command": "map_deletions",
             "dataset": str(fx_dataset.path),
-            "frequency_threshold": 0.125,
-            "mutations_threshold": 2**64,
+            "sites": sites,
         }