Merge pull request #237 from jeromekelleher/direct-attach-zero-mutation-samples

Add special case for inserting exact matches

Author: jeromekelleher

sc2ts/core.py

@@ -17,6 +17,7 @@
 NODE_IS_MUTATION_OVERLAP = 1 << 21
 NODE_IS_REVERSION_PUSH = 1 << 22
 NODE_IS_RECOMBINANT = 1 << 23
+NODE_IS_EXACT_MATCH = 1 << 24
 
 
 __version__ = "undefined"

sc2ts/inference.py

@@ -593,9 +593,11 @@ def extend(
     match_db.create_mask_table(base_ts)
     ts = increment_time(date, base_ts)
 
+    ts = add_exact_matches(ts=ts, match_db=match_db, date=date)
+
     logger.info(f"Update ARG with low-cost samples for {date}")
     ts = add_matching_results(
-        f"match_date=='{date}' and hmm_cost<={max_hmm_cost}",
+        f"match_date=='{date}' and hmm_cost>0 and hmm_cost<={max_hmm_cost}",
         ts=ts,
         match_db=match_db,
         date=date,
@@ -674,6 +676,31 @@ def match_path_ts(samples, ts, path, reversions):
     # print(tables)
 
 
+def add_exact_matches(match_db, ts, date):
+    where_clause = f"match_date=='{date}' AND hmm_cost==0"
+    logger.info(f"Querying match DB WHERE: {where_clause}")
+    samples = list(match_db.get(where_clause))
+    if len(samples) == 0:
+        logger.info(f"No exact matches on {date}")
+        return ts
+    logger.info(f"Update ARG with {len(samples)} exact matches for {date}")
+    tables = ts.dump_tables()
+    for sample in samples:
+        assert len(sample.path) == 1
+        assert len(sample.mutations) == 0
+        node_id = tables.nodes.add_row(
+            flags=tskit.NODE_IS_SAMPLE | core.NODE_IS_EXACT_MATCH,
+            time=0,
+            metadata=sample.metadata,
+        )
+        parent = sample.path[0].parent
+        logger.debug(f"ARG add exact match {sample.strain}:{node_id}->{parent}")
+        tables.edges.add_row(0, ts.sequence_length, parent=parent, child=node_id)
+    tables.sort()
+    tables.build_index()
+    return tables.tree_sequence()
+
+
 def add_matching_results(
     where_clause,
     match_db,

sc2ts/utils.py

@@ -494,13 +494,14 @@ def summary(self):
         mc_nodes = np.sum(self.ts.nodes_flags == sc2ts.NODE_IS_MUTATION_OVERLAP)
         pr_nodes = np.sum(self.ts.nodes_flags == sc2ts.NODE_IS_REVERSION_PUSH)
         re_nodes = np.sum(self.ts.nodes_flags == sc2ts.NODE_IS_RECOMBINANT)
+        exact_matches = np.sum((self.ts.nodes_flags & sc2ts.NODE_IS_EXACT_MATCH) > 0)
 
         samples = self.ts.samples()
         nodes_with_zero_muts = np.sum(self.nodes_num_mutations == 0)
         sites_with_zero_muts = np.sum(self.sites_num_mutations == 0)
         latest_sample = self.nodes_date[samples[-1]]
         masked_sites_per_sample = self.nodes_num_masked_sites[samples]
-        non_samples = self.ts.nodes_flags != tskit.NODE_IS_SAMPLE
+        non_samples = (self.ts.nodes_flags & tskit.NODE_IS_SAMPLE) == 0
         max_non_sample_mutations = np.max(self.nodes_num_mutations[non_samples])
         insertions = np.sum(self.mutations_inherited_state == "-")
         deletions = np.sum(self.mutations_derived_state == "-")
@@ -509,6 +510,7 @@ def summary(self):
             ("latest_sample", latest_sample),
             ("samples", self.ts.num_samples),
             ("nodes", self.ts.num_nodes),
+            ("exact_matches", exact_matches),
             ("mc_nodes", mc_nodes),
             ("pr_nodes", pr_nodes),
             ("re_nodes", re_nodes),
@@ -584,7 +586,7 @@ def _node_summary(self, u, child_mutations=True):
             qc += status
         flags = self.ts.nodes_flags[u]
         strain = ""
-        if flags == tskit.NODE_IS_SAMPLE:
+        if (flags & tskit.NODE_IS_SAMPLE) != 0:
             strain = md["strain"]
         elif flags == 1 << 21:
             if "overlap" in md: