Merge pull request #471 from jeromekelleher/map-deletions

Mapping deletions with metadata

Author: jeromekelleher

Diff for: sc2ts/dataset.py

@@ -317,44 +317,6 @@ def reorder(self, path, additional_fields=list(), show_progress=False):
         index = np.lexsort([self.metadata.fields[f] for f in sort_key[::-1]])
         self.copy(path, sample_id=sample_id[index], show_progress=show_progress)
 
-    def map_deletions_to_ts(self, ts, start, end):
-        """
-        Map a region of the dataset onto the specified tree sequence,
-        adding deletions to the tree sequence where necessary.
-        The tree sequence samples must all be present in the dataset.
-        A new tree sequence is returned whose original mutations in the
-        region have been removed and replaced using parsimony
-        """
-        sample_id = ts.metadata["sc2ts"]["samples_strain"]
-        if sample_id[0].startswith("Wuhan"):
-            # Note: a lot of fiddling around here is due to the potential for sample 0
-            # to be the reference, which is not in the viridian dataset.
-            sample_id = sample_id[1:]
-        tables = ts.dump_tables()
-        n_tm = ts.nodes_time
-        tree = ts.first()
-        del_sites = [ts.site(position=p).id for p in range(start, end)]
-        tables.mutations.keep_rows(np.logical_not(np.isin(ts.mutations_site, del_sites)))
-        for var in self.variants(sample_id, np.arange(start, end)):
-            tree.seek(var.position)
-            site = ts.site(position=var.position)
-            # treat non-nucleotide alleles (other IUPAC codes) as missing , but keep "-"
-            keep = np.isin(var.alleles, np.array(['A', 'C', 'G', 'T', '-']))
-            g = var.genotypes.copy()
-            g[keep[g] == False] = tskit.MISSING_DATA
-            anc = site.ancestral_state
-            # Pad the start with the anc state (non-viridian) Wuhan sample, so add that
-            if len(g) < ts.num_samples:
-                g = np.append([list(var.alleles).index(anc)], g)
-            _, mutations = tree.map_mutations(g, list(var.alleles), ancestral_state=anc)
-            m_id_map = {tskit.NULL: tskit.NULL}
-            for list_id, m in enumerate(mutations):
-                m_id_map[list_id] = tables.mutations.append(
-                    m.replace(site=site.id, parent=m_id_map[m.parent], time=n_tm[m.node])
-                )
-        tables.sort()
-        return tables.tree_sequence()
-
     @staticmethod
     def new(path, samples_chunk_size=None, variants_chunk_size=None):
 

Diff for: sc2ts/inference.py

@@ -1945,3 +1945,83 @@ def get_recombinant_strains(ts):
         group_id = node.metadata["sc2ts"]["group_id"]
         ret[u] = groups[group_id]
     return ret
+
+
+def map_deletions(ts, ds, *, frequency_threshold, show_progress=False):
+    """
+    Map deletions at sites that exceed the specified frequency threshold onto the
+    ARG using parsimony (excluding flanks).
+    """
+    genes = core.get_gene_coordinates()
+    start = genes["ORF1ab"][0]
+    end = genes["ORF10"][1]
+
+    del_sites = []
+    keep_mutations = np.ones(ts.num_mutations, dtype=bool)
+    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)
+                site_start = np.searchsorted(ts.mutations_site, site.id, side="left")
+                site_stop = np.searchsorted(ts.mutations_site, site.id, side="right")
+                keep_mutations[site_start:site_stop] = False
+
+    sample_id = ts.metadata["sc2ts"]["samples_strain"]
+    assert not sample_id[0].startswith("Wuhan")
+
+    tables = ts.dump_tables()
+    tree = ts.first()
+
+    tables.mutations.keep_rows(keep_mutations)
+
+    variants = get_progress(
+        ds.variants(sample_id, ts.sites_position[del_sites]),
+        title="Map deletions",
+        phase="",
+        show_progress=show_progress,
+        total=len(del_sites),
+    )
+
+    mut_metadata = {"sc2ts": {"type": "post_parsimony"}}
+    site_metadata = {}
+    for var in variants:
+        tree.seek(var.position)
+        site = ts.site(position=var.position)
+        old_mutations = []
+        for mut in site.mutations:
+            assert not keep_mutations[mut.id]
+            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
+        g = mask_ambiguous(var.genotypes)
+        deletion_samples = site.metadata["sc2ts"]["deletion_samples"]
+        assert deletion_samples == np.sum(g == DELETION)
+        _, mutations = tree.map_mutations(
+            g, list(var.alleles), ancestral_state=site.ancestral_state
+        )
+        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,
+            )
+
+    tables.sites.clear()
+    for site in ts.sites():
+        md = site_metadata.get(site.id, site.metadata)
+        tables.sites.append(site.replace(metadata=md))
+
+    tables.sort()
+    tables.build_index()
+    tables.compute_mutation_parents()
+    return tables.tree_sequence()

Diff for: tests/test_inference.py

@@ -1376,3 +1376,39 @@ def test_one_leaf_mutation(self, samples, result):
         tables.mutations.add_row(site=0, node=3, derived_state="T")
         ts = tables.tree_sequence()
         nt.assert_array_equal(sc2ts.extract_haplotypes(ts, samples), result)
+
+
+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)
+        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]
+
+        for site_id in remapped_sites:
+            site = new_ts.site(site_id)
+            old_site = ts.site(site_id)
+            original_mutations = site.metadata["sc2ts"]["original_mutations"]
+            assert original_mutations == [
+                {
+                    "node": mutation.node,
+                    "derived_state": mutation.derived_state,
+                    "metadata": mutation.metadata,
+                }
+                for mutation in old_site.mutations
+            ]
+            d = site.metadata["sc2ts"]
+            del d["original_mutations"]
+            assert old_site.metadata["sc2ts"] == d
+
+            for mut in site.mutations:
+                assert mut.metadata["sc2ts"]["type"] == "post_parsimony"
+
+    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)
+        sc2ts.validate(new_ts, fx_dataset, deletions_as_missing=False)