Merge pull request #505 from jeromekelleher/insert-missing-sites

Insert missing sites

Author: jeromekelleher

sc2ts/cli.py

@@ -610,7 +610,7 @@ def minimise_metadata(
 @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(
+def map_parsimony(
     dataset,
     ts_in,
     ts_out,
@@ -621,14 +621,14 @@ def map_deletions(
     log_file,
 ):
     """
-    Map deletions onto the specified ARG.
+    Map variation at the specified set of sites to the ARG using parsimony.
     """
     setup_logging(verbose, log_file)
     ds = sc2ts.Dataset(dataset)
     ts = tszip.load(ts_in)
     if sites is not None:
         sites = np.loadtxt(sites, dtype=int)
-    result = sc2ts.map_deletions(ts, ds, sites, show_progress=progress)
+    result = sc2ts.map_parsimony(ts, ds, sites, show_progress=progress)
     if report is not None:
         result.report.to_csv(report)
     result.tree_sequence.dump(ts_out)
@@ -670,5 +670,5 @@ def cli():
 cli.add_command(validate)
 cli.add_command(postprocess)
 cli.add_command(minimise_metadata)
-cli.add_command(map_deletions)
+cli.add_command(map_parsimony)
 cli.add_command(run_hmm)

sc2ts/inference.py

@@ -1956,14 +1956,16 @@ def get_recombinant_strains(ts):
 
 
 @dataclasses.dataclass
-class MapDeletionsResult:
+class MapParsimonyResult:
     tree_sequence: tskit.TreeSequence
     report: pd.DataFrame
 
 
-def map_deletions(ts, ds, sites=None, *, show_progress=False):
+def map_parsimony(ts, ds, sites=None, *, show_progress=False):
     """
-    Map deletions at the specified set of site positions ARG using parsimony.
+    Map variation at the specified set of site positions to the ARG using parsimony
+    and return the resulting tree sequence. States ACGT- are treated as alleles
+    and all other states are treated as missing data.
     """
     start_time = time.time()  # wall time
     if sites is not None:
@@ -1977,9 +1979,11 @@ def map_deletions(ts, ds, sites=None, *, show_progress=False):
     # will have been dropped.
     if "sc2ts" in md:
         mut_metadata = {"sc2ts": {"type": "post_parsimony"}}
+        site_metadata = {"sc2ts": {"type": "post_parsimony"}}
         sample_id = md["sc2ts"]["samples_strain"]
     else:
         mut_metadata = None
+        site_metadata = None
         dfn = stats.node_data(ts, inheritance_stats=False)
         sample_id = dfn[dfn.is_sample].sample_id.values
 
@@ -2000,13 +2004,25 @@ def map_deletions(ts, ds, sites=None, *, show_progress=False):
 
     keep_mutations = np.ones(ts.num_mutations, dtype=bool)
     report_data = []
+    current_sites = set(ts.sites_position)
     for var in variants:
         tree.seek(var.position)
-        try:
+        if var.position in current_sites:
             site = ts.site(position=var.position)
-        except ValueError:
-            logger.warning(f"No site at position {var.position}; skipping")
-            continue
+            site_id = site.id
+            old_mutations = set((mut.node, mut.derived_state) for mut in site.mutations)
+            for mut in site.mutations:
+                keep_mutations[mut.id] = False
+        else:
+            assert ts.reference_sequence.data[0] == "X"
+            site_id = tables.sites.add_row(
+                position=var.position,
+                ancestral_state=ts.reference_sequence.data[var.position],
+                metadata=site_metadata,
+            )
+            site = tables.sites[site_id]
+            old_mutations = set()
+
         g = mask_ambiguous(var.genotypes)
 
         if np.all(g == -1):
@@ -2016,7 +2032,6 @@ def map_deletions(ts, ds, sites=None, *, show_progress=False):
                 g, list(var.alleles), ancestral_state=site.ancestral_state
             )
 
-        old_mutations = set((mut.node, mut.derived_state) for mut in site.mutations)
         new_mutations = set((mut.node, mut.derived_state) for mut in mutations)
         inter = old_mutations & new_mutations
         report_data.append(
@@ -2032,11 +2047,9 @@ def map_deletions(ts, ds, sites=None, *, show_progress=False):
             f"Site {int(site.position)} "
             f"old={len(old_mutations)} new={len(new_mutations)} inter={len(inter)}"
         )
-        for mut in site.mutations:
-            keep_mutations[mut.id] = False
         for m in mutations:
             tables.mutations.add_row(
-                site=site.id,
+                site=site_id,
                 node=m.node,
                 derived_state=m.derived_state,
                 time=ts.nodes_time[m.node],
@@ -2055,10 +2068,10 @@ def map_deletions(ts, ds, sites=None, *, show_progress=False):
     tables.build_index()
     tables.compute_mutation_parents()
     params = {"dataset": str(ds.path), "sites": sites.tolist()}
-    prov = get_provenance_dict("map_deletions", params, start_time)
+    prov = get_provenance_dict("map_parsimony", params, start_time)
     tables.provenances.add_row(json.dumps(prov))
 
-    return MapDeletionsResult(tables.tree_sequence(), pd.DataFrame(report_data))
+    return MapParsimonyResult(tables.tree_sequence(), pd.DataFrame(report_data))
 
 
 def append_exact_matches(ts, match_db, show_progress=False):

tests/test_cli.py

@@ -433,7 +433,7 @@ def test_example_args(self, tmp_path, fx_ts_map, fx_match_db):
         assert "pango" in dfn
 
 
-class TestMapDeletions:
+class TestMapParsimony:
 
     def test_example(self, tmp_path, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
@@ -446,7 +446,7 @@ def test_example(self, tmp_path, fx_ts_map, fx_dataset):
         runner = ct.CliRunner()
         result = runner.invoke(
             cli.cli,
-            f"map-deletions {fx_dataset.path} {ts.path} --sites={del_sites_path} "
+            f"map-parsimony {fx_dataset.path} {ts.path} --sites={del_sites_path} "
             f"{out_ts_path} --report={report_path}",
             catch_exceptions=False,
         )
@@ -466,7 +466,7 @@ def test_example_no_report(self, tmp_path, fx_ts_map, fx_dataset):
         runner = ct.CliRunner()
         result = runner.invoke(
             cli.cli,
-            f"map-deletions {fx_dataset.path} {ts.path} --sites={del_sites_path} "
+            f"map-parsimony {fx_dataset.path} {ts.path} --sites={del_sites_path} "
             f"{out_ts_path}",
             catch_exceptions=False,
         )
@@ -483,7 +483,7 @@ def test_example_no_sites(self, tmp_path, fx_ts_map, fx_dataset):
         runner = ct.CliRunner()
         result = runner.invoke(
             cli.cli,
-            f"map-deletions {fx_dataset.path} {ts.path} {out_ts_path} "
+            f"map-parsimony {fx_dataset.path} {ts.path} {out_ts_path} "
             f"--report={report_path}",
             catch_exceptions=False,
         )

tests/test_inference.py

@@ -1402,11 +1402,11 @@ def fx_ts_exact_matches(fx_ts_map, fx_match_db):
     return tsp
 
 
-class TestMapDeletions:
+class TestMapParsimony:
     def test_example(self, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
         sites = [1547, 3951, 3952, 3953]
-        result = sc2ts.map_deletions(ts, fx_dataset, sites)
+        result = sc2ts.map_parsimony(ts, fx_dataset, sites)
         report = result.report
         assert list(report.site.values) == sites
         assert np.all(report.old == 1)
@@ -1422,7 +1422,7 @@ def test_example(self, fx_ts_map, fx_dataset):
     def test_example_minimised_md(self, fx_ts_min_2020_02_15, fx_dataset):
         ts = fx_ts_min_2020_02_15
         sites = [1547, 3951, 3952, 3953]
-        result = sc2ts.map_deletions(ts, fx_dataset, sites)
+        result = sc2ts.map_parsimony(ts, fx_dataset, sites)
         report = result.report
         assert list(report.site.values) == sites
         assert np.all(report.old == 1)
@@ -1438,7 +1438,7 @@ def test_example_minimised_md(self, fx_ts_min_2020_02_15, fx_dataset):
     def test_example_minimised_zero_sites(self, fx_ts_min_2020_02_15, fx_dataset):
         ts = fx_ts_min_2020_02_15
         sites = []
-        result = sc2ts.map_deletions(ts, fx_dataset, sites)
+        result = sc2ts.map_parsimony(ts, fx_dataset, sites)
         report = result.report
         assert report.shape[0] == 0
         ts.tables.assert_equals(
@@ -1447,33 +1447,60 @@ def test_example_minimised_zero_sites(self, fx_ts_min_2020_02_15, fx_dataset):
 
     def test_empty(self, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
-        result = sc2ts.map_deletions(ts, fx_dataset, [])
+        result = sc2ts.map_parsimony(ts, fx_dataset, [])
         ts.tables.assert_equals(
             result.tree_sequence.tables, ignore_metadata=True, ignore_provenance=True
         )
         assert result.report.shape[0] == 0
 
     def test_empty_all_sites(self, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
-        result = sc2ts.map_deletions(ts, fx_dataset)
+        result = sc2ts.map_parsimony(ts, fx_dataset)
         assert result.report.shape[0] == ts.num_sites
 
-    def test_missing_site(self, fx_ts_map, fx_dataset):
+    def test_missing_sites(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
 
-        result = sc2ts.map_deletions(ts, fx_dataset, missing_positions)
+        result = sc2ts.map_parsimony(ts, fx_dataset, missing_positions)
+        assert (
+            len(
+                set(missing_positions)
+                & set(result.tree_sequence.sites_position.astype(int))
+            )
+            == 5
+        )
+        assert list(result.report.site) == missing_positions
+        assert np.all(result.report.old == 0)
+        assert np.all(result.report.new == 0)
+        assert np.all(result.report.intersection == 0)
+
+    def test_readd_existing_site(self, fx_ts_map, fx_dataset):
+        ts = fx_ts_map["2020-02-13"]
+        site = ts.site(position=203)
+        assert len(site.mutations) == 1
+        ts_del = ts.delete_sites([site.id])
+
+        result = sc2ts.map_parsimony(ts_del, fx_dataset, [site.position])
+        new_site = result.tree_sequence.site(position=site.position)
+        assert new_site.ancestral_state == site.ancestral_state
+        assert new_site.position == site.position
+        assert new_site.id == site.id
+
         ts.tables.assert_equals(
             result.tree_sequence.tables, ignore_metadata=True, ignore_provenance=True
         )
-        assert result.report.shape[0] == 0
+        assert result.report.shape[0] == 1
+        assert np.all(result.report.old == 0)
+        assert np.all(result.report.new == 1)
+        assert np.all(result.report.intersection == 0)
 
     def test_example_exact_matches(self, fx_ts_exact_matches, fx_dataset):
         ts = fx_ts_exact_matches
         sites = [1547, 3951, 3952, 3953]
-        result = sc2ts.map_deletions(ts, fx_dataset, sites)
-        result = sc2ts.map_deletions(ts, fx_dataset, sites)
+        result = sc2ts.map_parsimony(ts, fx_dataset, sites)
+        result = sc2ts.map_parsimony(ts, fx_dataset, sites)
         report = result.report
         assert list(report.site.values) == sites
         assert np.all(report.old == 1)
@@ -1483,18 +1510,18 @@ def test_example_exact_matches(self, fx_ts_exact_matches, fx_dataset):
     def test_validate(self, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
         sites = [1547, 3951, 3952, 3953]
-        result = sc2ts.map_deletions(ts, fx_dataset, sites)
+        result = sc2ts.map_parsimony(ts, fx_dataset, sites)
         sc2ts.validate(result.tree_sequence, fx_dataset, deletions_as_missing=False)
 
     def test_provenance(self, fx_ts_map, fx_dataset):
         ts = fx_ts_map["2020-02-13"]
         sites = [1547, 3951, 3952, 3953]
-        result = sc2ts.map_deletions(ts, fx_dataset, sites)
+        result = sc2ts.map_parsimony(ts, fx_dataset, sites)
         tsp = result.tree_sequence
         assert tsp.num_provenances == ts.num_provenances + 1
         prov = tsp.provenance(-1)
         assert json.loads(prov.record)["parameters"] == {
-            "command": "map_deletions",
+            "command": "map_parsimony",
             "dataset": str(fx_dataset.path),
             "sites": sites,
         }