Merge pull request #278 from jeromekelleher/imr-2

Refine conditions for immediate reversion flag

run.sh

@@ -8,7 +8,8 @@ max_daily_samples=1000
 num_threads=8
 
 # Paths
-datadir=testrun
+# datadir=testrun
+datadir=data/viridian
 run_id=tmp-dev-hp
 # run_id=upgma-mds-$max_daily_samples-md-$max_submission_delay-mm-$mismatches
 resultsdir=results/$run_id
@@ -19,22 +20,25 @@ alignments=$datadir/alignments.db
 metadata=$datadir/metadata.db
 matches=$resultsdir/matches.db
 
-dates=`python3 -m sc2ts list-dates $metadata | grep -v 2021-12-31 | head -n 14`
-
-options="--num-threads $num_threads -vv -l $logfile "
+options="--num-threads $num_threads -vv " #-l $logfile "
 # options+="--max-submission-delay $max_submission_delay "
 # options+="--max-daily-samples $max_daily_samples "
 options+="--num-mismatches $mismatches"
 
 mkdir -p $resultsdir
 
-last_ts=$resultsdir/initial.ts
+# date=2000-01-01
+# last_ts=$resultsdir/initial.ts
+# python3 -m sc2ts initialise $last_ts $matches
 
-python3 -m sc2ts initialise $last_ts $matches
+date=2020-03-01
+last_ts="$results_prefix$date".ts
 
+dates=`python3 -m sc2ts list-dates --after $date $metadata | grep -v 2021-12-31`
 for date in $dates; do
     out_ts="$results_prefix$date".ts
     python3 -m sc2ts extend $last_ts $date $alignments $metadata \
         $matches $out_ts $options
     last_ts=$out_ts
+    # break
 done

sc2ts/inference.py

@@ -1265,6 +1265,7 @@ def delete_immediate_reversion_nodes(ts, attach_nodes):
     tables = ts.dump_tables()
     edges_to_delete = []
     for u in nodes_to_delete:
+        logger.debug(f"Deleting immediate reversion node {u}")
         edges_to_delete.append(tree.edge(u))
         parent = tree.parent(u)
         assert tree.num_children(u) > 0
@@ -1697,8 +1698,9 @@ def attach_tree(
         raise ValueError("Incompatible sequence length")
 
     tree = child_ts.first()
+    has_root_mutations = np.any(child_ts.mutations_node == tree.root)
     condition = (
-        np.any(child_ts.mutations_node == tree.root)
+        has_root_mutations
         or len(attach_path) > 1
         or len(group.immediate_reversions) > 0
     )
@@ -1783,9 +1785,12 @@ def attach_tree(
         assert tree.num_children(tree.root) == 1
         u = node_id_map[node]
         row = parent_tables.nodes[u]
-        parent_tables.nodes[u] = row.replace(
-            flags=core.NODE_IS_IMMEDIATE_REVERSION_MARKER
-        )
+        if not has_root_mutations:
+            # We only really want to remove this if there are no nodes
+            logger.debug(f"Flagging reversion at node {u} for {group.summary()}")
+            parent_tables.nodes[u] = row.replace(
+                flags=core.NODE_IS_IMMEDIATE_REVERSION_MARKER
+            )
         # print("attaching reversions at ", node, node_id_map[node])
         # print(child_ts.draw_text())
         for site_id, derived_state in group.immediate_reversions:

sc2ts/info.py

@@ -280,7 +280,7 @@ def summary(self):
         re_nodes = np.sum(self.ts.nodes_flags == core.NODE_IS_RECOMBINANT)
         exact_matches = np.sum((self.ts.nodes_flags & core.NODE_IS_EXACT_MATCH) > 0)
         imr_nodes = np.sum(
-            (self.ts.nodes_flags & core.NODE_IS_IMMEDIATE_REVERSION_MARKER) > 0
+            (self.ts.nodes_flags == core.NODE_IS_IMMEDIATE_REVERSION_MARKER)
         )
 
         samples = self.ts.samples()