Merge pull request #476 from jeromekelleher/postprocess

Postprocess

Author: jeromekelleher

sc2ts/cli.py

@@ -543,6 +543,41 @@ def run_hmm(
         print(run.asjson())
 
 
+@click.command()
+@click.argument("ts_in", type=click.Path(exists=True, dir_okay=False))
+@click.argument("ts_out", type=click.Path(exists=False, dir_okay=False))
+@click.option("--match-db", type=click.Path(exists=True, dir_okay=False))
+@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 postprocess(
+    ts_in,
+    ts_out,
+    match_db,
+    progress,
+    verbose,
+    log_file,
+):
+    """
+    Perform final postprocessing steps to the specified ARG.
+    """
+    setup_logging(verbose, log_file)
+    ts = tszip.load(ts_in)
+    if match_db is not None:
+        with sc2ts.MatchDb(match_db) as db:
+            ts = sc2ts.append_exact_matches(ts, db, show_progress=progress)
+
+    ts = sc2ts.trim_metadata(ts, show_progress=progress)
+
+    ts = sc2ts.push_up_unary_recombinant_mutations(ts)
+
+    # See if we can remove some of the reversions in a straightforward way.
+    mutations_is_reversion = sc2ts.find_reversions(ts)
+    mutations_before = ts.num_mutations
+    ts = sc2ts.push_up_reversions(ts, ts.mutations_node[mutations_is_reversion])
+    ts.dump(ts_out)
+
+
 def find_previous_date_path(date, path_pattern):
     """
     Find the path with the most-recent date to the specified one
@@ -577,6 +612,7 @@ def cli():
 
 cli.add_command(infer)
 cli.add_command(validate)
+cli.add_command(postprocess)
 cli.add_command(run_hmm)
 
 cli.add_command(tally_lineages)

sc2ts/inference.py

@@ -814,7 +814,7 @@ def update_top_level_metadata(ts, date, retro_groups, samples):
     return tables
 
 
-def add_sample_to_tables(sample, tables, group_id=None):
+def add_sample_to_tables(sample, tables, group_id=None, time=0):
     sc2ts_md = {
         "hmm_match": sample.hmm_match.asdict(),
         "alignment_composition": dict(sample.alignment_composition),
@@ -825,7 +825,7 @@ def add_sample_to_tables(sample, tables, group_id=None):
     if group_id is not None:
         sc2ts_md["group_id"] = group_id
     metadata = {**sample.metadata, "sc2ts": sc2ts_md}
-    return tables.nodes.add_row(flags=sample.flags, metadata=metadata)
+    return tables.nodes.add_row(flags=sample.flags, metadata=metadata, time=time)
 
 
 def match_path_ts(group):
@@ -2031,3 +2031,133 @@ def map_deletions(ts, ds, *, frequency_threshold, show_progress=False):
     tables.build_index()
     tables.compute_mutation_parents()
     return tables.tree_sequence()
+
+
+def append_exact_matches(ts, match_db, show_progress=False):
+    """
+    Update the specified tree sequence to include all exact matches
+    from the specified match DB.
+    """
+    md = ts.metadata
+    date = md["sc2ts"]["date"]
+    total_exact_matches = sum(
+        md["sc2ts"]["cumulative_stats"]["exact_matches"]["pango"].values()
+    )
+    samples_strain = md["sc2ts"]["samples_strain"]
+    tables = ts.dump_tables()
+    L = tables.sequence_length
+    time_zero = parse_date(date)
+    with match_db.conn:
+        sql = f"SELECT * FROM samples WHERE hmm_cost == 0 AND match_date <= '{date}'"
+        rows = tqdm.tqdm(
+            match_db.conn.execute(sql),
+            total=total_exact_matches,
+            desc="Exact matches",
+            disable=not show_progress,
+        )
+        for row in rows:
+            pkl = row.pop("pickle")
+            sample = pickle.loads(bz2.decompress(pkl))
+            sample.flags |= core.NODE_IS_EXACT_MATCH
+            delta = time_zero - parse_date(sample.date)
+            assert delta.days >= 0
+            u = add_sample_to_tables(sample, tables, time=delta.days)
+            parent = sample.hmm_match.path[0].parent
+            tables.edges.add_row(0, L, parent=parent, child=u)
+            samples_strain.append(sample.strain)
+
+    assert total_exact_matches == len(tables.nodes) - ts.num_nodes
+    md["sc2ts"]["samples_strain"] = samples_strain
+    tables.metadata = md
+    tables.sort()
+    return tables.tree_sequence()
+
+
+def trim_metadata(ts, show_progress=False):
+    tables = ts.dump_tables()
+
+    tables.nodes.clear()
+
+    nodes = tqdm.tqdm(
+        ts.nodes(),
+        total=ts.num_nodes,
+        desc="Trim node metadata",
+        disable=not show_progress,
+    )
+    for node in nodes:
+        md = node.metadata
+        if node.is_sample():
+            # Note it would be nice to trim down the name of the pango field here
+            # but it's too tedious to test.
+            md = {k: md[k] for k in ["strain", "date", "Viridian_pangolin"]}
+        tables.nodes.append(node.replace(metadata=md))
+    return tables.tree_sequence()
+
+
+def find_reversions(ts):
+    """
+    Return a boolean array with True for all mutations in which the
+    inherited_state of the parent is equal to the derived_state of the
+    child.
+    """
+    tables = ts.tables
+    assert np.all(
+        tables.mutations.derived_state_offset == np.arange(ts.num_mutations + 1)
+    )
+    derived_state = tables.mutations.derived_state.view("S1").astype(str)
+    assert np.all(tables.sites.ancestral_state_offset == np.arange(ts.num_sites + 1))
+    ancestral_state = tables.sites.ancestral_state.view("S1").astype(str)
+    del tables
+    inherited_state = ancestral_state[ts.mutations_site]
+    mutations_with_parent = ts.mutations_parent != -1
+    parent = ts.mutations_parent[mutations_with_parent]
+    assert np.all(parent >= 0)
+    inherited_state[mutations_with_parent] = derived_state[parent]
+
+    assert np.all(inherited_state != derived_state)
+
+    is_reversion = np.zeros(ts.num_mutations, dtype=bool)
+    is_reversion[mutations_with_parent] = (
+        derived_state[mutations_with_parent] == inherited_state[parent]
+    )
+    return is_reversion
+
+
+def push_up_unary_recombinant_mutations(ts):
+    """
+    Find any mutations that occur on unary children of a recombinant node,
+    and push those mutations onto the recombinant node itself. The
+    rationale for this is that, due to technical details of tree building,
+    we sometimes get a single child of a recombinant node, which can have
+    a large number of mutations. It is more parsimonious to assume that the
+    mutations occured on the branch(es) *leading to* the recombinant than
+    to have succeeded it.
+    """
+    recomb_parent_edges = np.where(
+        ts.nodes_flags[ts.edges_parent] & core.NODE_IS_RECOMBINANT > 0
+    )[0]
+    by_parent = collections.defaultdict(list)
+    logger.info(f"Found {len(recomb_parent_edges)} edges with recombinant parent")
+    for e in recomb_parent_edges:
+        edge = ts.edge(e)
+        if edge.left == 0 and edge.right == ts.sequence_length:
+            by_parent[edge.parent].append(edge)
+
+    # We're only interested in full-span edges with a single child.
+    child_to_parent = {
+        e[0].child: e[0].parent for e in by_parent.values() if len(e) == 1
+    }
+    logger.info(f"Of which {len(child_to_parent)} are unary")
+    mutations_to_move = np.isin(
+        ts.mutations_node, np.array(list(child_to_parent.keys()), dtype=np.int32)
+    )
+    tables = ts.dump_tables()
+    for m in np.where(mutations_to_move)[0]:
+        row = tables.mutations[m]
+        node = child_to_parent[row.node]
+        # We're only changing the node and time, which are fixed size so we
+        # don't rewrite the table for each of these.
+        tables.mutations[m] = row.replace(node=node, time=ts.nodes_time[node])
+    logger.info(f"Moved up {np.sum(mutations_to_move)} mutations")
+    tables.sort()
+    return tables.tree_sequence()

tests/conftest.py

@@ -364,6 +364,19 @@ def recombinant_example_3(tmp_path, fx_ts_map, fx_dataset, ds_path):
     return rts
 
 
+def recombinant_example_4(tmp_path, fx_recombinant_example_2):
+    """
+    Same as recombinant_ex2 but with two mutations below the recombinant.
+    """
+    tables = fx_recombinant_example_2.dump_tables()
+    u = 55
+    tables.mutations.add_row(
+        site=2500, derived_state="G", node=u, time=tables.nodes.time[u], metadata={}
+    )
+    tables.sort()
+    return tables.tree_sequence()
+
+
 @pytest.fixture
 def fx_recombinant_example_1(tmp_path, fx_data_cache, fx_ts_map, fx_dataset):
     cache_path = fx_data_cache / "recombinant_ex1.ts"
@@ -395,3 +408,13 @@ def fx_recombinant_example_3(tmp_path, fx_data_cache, fx_ts_map, fx_dataset):
         ts = recombinant_example_3(tmp_path, fx_ts_map, fx_dataset, ds_cache_path)
         ts.dump(cache_path)
     return tskit.load(cache_path)
+
+
+@pytest.fixture
+def fx_recombinant_example_4(tmp_path, fx_data_cache, fx_recombinant_example_2):
+    cache_path = fx_data_cache / "recombinant_ex4.ts"
+    if not cache_path.exists():
+        print(f"Generating {cache_path}")
+        ts = recombinant_example_4(tmp_path, fx_recombinant_example_2)
+        ts.dump(cache_path)
+    return tskit.load(cache_path)

tests/test_cli.py

@@ -376,6 +376,22 @@ def test_multiple_override(self, tmp_path, fx_ts_map, fx_dataset):
         assert ts.num_samples == 0
 
 
+class TestPostprocess:
+
+    def test_example(self, tmp_path, fx_ts_map, fx_match_db):
+        ts = fx_ts_map["2020-02-13"]
+        out_ts_path = tmp_path / "ts.ts"
+        runner = ct.CliRunner(mix_stderr=False)
+        result = runner.invoke(
+            cli.cli,
+            f"postprocess {ts.path} {out_ts_path} --match-db={fx_match_db.path}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        out = tskit.load(out_ts_path)
+        assert out.num_samples == ts.num_samples + 8
+
+
 class TestValidate:
 
     @pytest.mark.parametrize("date", ["2020-01-01", "2020-02-11"])

tests/test_inference.py

@@ -71,6 +71,10 @@ def test_get_recombinant_strains_ex2(self, fx_recombinant_example_2):
         d = sc2ts.get_recombinant_strains(fx_recombinant_example_2)
         assert d == {56: ["recombinant_114:29825"]}
 
+    def test_get_recombinant_strains_ex4(self, fx_recombinant_example_4):
+        d = sc2ts.get_recombinant_strains(fx_recombinant_example_4)
+        assert d == {56: ["recombinant_114:29825"]}
+
     def test_recombinant_example_1(self, fx_recombinant_example_1):
         ts = fx_recombinant_example_1
         samples_strain = ts.metadata["sc2ts"]["samples_strain"]
@@ -1412,3 +1416,101 @@ 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)
+
+
+@pytest.fixture
+def fx_ts_exact_matches(fx_ts_map, fx_match_db):
+    ts = fx_ts_map["2020-02-13"]
+    tsp = sc2ts.append_exact_matches(ts, fx_match_db)
+    return tsp
+
+
+class TestAppendExactMatches:
+    def test_validate(self, fx_ts_exact_matches, fx_dataset):
+        sc2ts.validate(fx_ts_exact_matches, fx_dataset)
+
+    def test_example_properties(self, fx_ts_exact_matches):
+        ts = fx_ts_exact_matches
+        samples_strain = ts.metadata["sc2ts"]["samples_strain"]
+        assert [ts.node(u).metadata["strain"] for u in ts.samples()] == samples_strain
+        assert ts.num_nodes == 61
+        tree = ts.first()
+        assert tree.num_roots == 1
+
+    def test_times_agree(self, fx_ts_exact_matches):
+        ts = fx_ts_exact_matches
+        date_to_time = {}
+        time_to_date = {}
+        for u in ts.samples():
+            node = ts.node(u)
+            time = node.time
+            date = node.metadata["date"]
+            if date not in date_to_time:
+                date_to_time[date] = time
+            assert date_to_time[date] == time
+            if time not in time_to_date:
+                time_to_date[time] = date
+            assert time_to_date[time] == date
+
+    def test_flags(self, fx_ts_exact_matches):
+        ts = fx_ts_exact_matches
+        assert np.all(
+            ts.nodes_flags[-8:] == sc2ts.NODE_IS_EXACT_MATCH | tskit.NODE_IS_SAMPLE
+        )
+
+    def test_exact_match_counts(self, fx_ts_exact_matches):
+        ts = fx_ts_exact_matches
+        tree = ts.first()
+        node_count = ts.metadata["sc2ts"]["cumulative_stats"]["exact_matches"]["node"]
+        for u in tree.nodes():
+            num_exact_matches = 0
+            for v in tree.children(u):
+                if (ts.nodes_flags[v] & sc2ts.NODE_IS_EXACT_MATCH) > 0:
+                    num_exact_matches += 1
+            assert node_count.get(str(u), 0) == num_exact_matches
+
+
+class TestTrimMetadata:
+    def test_validate(self, fx_ts_map, fx_dataset):
+        ts = fx_ts_map["2020-02-13"]
+        tsp = sc2ts.trim_metadata(ts)
+        sc2ts.validate(tsp, fx_dataset)
+
+    def test_fields(self, fx_ts_map):
+        ts = fx_ts_map["2020-02-13"]
+        tsp = sc2ts.trim_metadata(ts)
+        for u in tsp.samples():
+            node = tsp.node(u)
+            assert set(node.metadata.keys()) == {"strain", "date", "Viridian_pangolin"}
+
+
+class TestPushUpRecombinantMutations:
+
+    def test_no_recombinants(self, fx_ts_map):
+        ts = fx_ts_map["2020-02-13"]
+        tsp = sc2ts.push_up_unary_recombinant_mutations(ts)
+        ts.tables.assert_equals(tsp.tables)
+
+    def test_recombinant_example_1(self, fx_recombinant_example_1):
+        ts = fx_recombinant_example_1
+        tsp = sc2ts.push_up_unary_recombinant_mutations(ts)
+        ts.tables.assert_equals(tsp.tables)
+
+    def test_recombinant_example_2(self, fx_recombinant_example_2):
+        ts = fx_recombinant_example_2
+        tsp = sc2ts.push_up_unary_recombinant_mutations(ts)
+        ts.tables.assert_equals(tsp.tables)
+
+    def test_recombinant_example_3(self, fx_recombinant_example_3):
+        ts = fx_recombinant_example_3
+        tsp = sc2ts.push_up_unary_recombinant_mutations(ts)
+        ts.tables.assert_equals(tsp.tables)
+
+    def test_recombinant_example_4(self, fx_recombinant_example_4):
+        ts = fx_recombinant_example_4
+        site = 2500
+        mut = ts.site(site).mutations[0]
+        assert mut.node == 55
+        tsp = sc2ts.push_up_unary_recombinant_mutations(ts)
+        mut = tsp.site(site).mutations[0]
+        assert mut.node == 56