Bugfixes and refinements to retro groups

sc2ts/cli.py

@@ -400,6 +400,16 @@ def summarise_base(ts, date, progress):
         "group tree"
     ),
 )
+@click.option(
+    "--max-recurrent-mutations",
+    default=10,
+    show_default=True,
+    type=int,
+    help=(
+        "Maximum number of recurrent mutations in an inferred retrospective "
+        "group tree"
+    ),
+)
 @click.option(
     "--retrospective-window",
     default=30,
@@ -465,6 +475,7 @@ def extend(
     min_group_size,
     min_root_mutations,
     max_mutations_per_sample,
+    max_recurrent_mutations,
     retrospective_window,
     deletions_as_missing,
     max_daily_samples,
@@ -509,6 +520,7 @@ def extend(
             min_group_size=min_group_size,
             min_root_mutations=min_root_mutations,
             max_mutations_per_sample=max_mutations_per_sample,
+            max_recurrent_mutations=max_recurrent_mutations,
             retrospective_window=retrospective_window,
             deletions_as_missing=deletions_as_missing,
             max_daily_samples=max_daily_samples,

sc2ts/inference.py

@@ -550,6 +550,7 @@ def extend(
     min_root_mutations=None,
     min_different_dates=None,
     max_mutations_per_sample=None,
+    max_recurrent_mutations=None,
     deletions_as_missing=None,
     max_daily_samples=None,
     show_progress=False,
@@ -568,6 +569,8 @@ def extend(
         min_root_mutations = 2
     if max_mutations_per_sample is None:
         max_mutations_per_sample = 100
+    if max_recurrent_mutations is None:
+        max_recurrent_mutations = 100
     if min_different_dates is None:
         min_different_dates = 3
     if retrospective_window is None:
@@ -673,12 +676,15 @@ def extend(
         min_different_dates=min_different_dates,
         min_root_mutations=min_root_mutations,
         max_mutations_per_sample=max_mutations_per_sample,
+        max_recurrent_mutations=max_recurrent_mutations,
         additional_node_flags=core.NODE_IN_RETROSPECTIVE_SAMPLE_GROUP,
         show_progress=show_progress,
         phase="retro",
     )
     for group in groups:
-        logger.warning(f"Add retro group {dict(group.pango_count)}: {group.tree_quality_metrics.summary()}")
+        logger.warning(
+            f"Add retro group {dict(group.pango_count)}: {group.tree_quality_metrics.summary()}"
+        )
     return update_top_level_metadata(ts, date, groups)
 
 
@@ -693,6 +699,9 @@ def update_top_level_metadata(ts, date, retro_groups):
         samples_strain.append(node.metadata["strain"])
     md["sc2ts"]["samples_strain"] = samples_strain
     existing_retro_groups = md["sc2ts"].get("retro_groups", [])
+    if isinstance(existing_retro_groups, dict):
+        # Hack to implement metadata format change
+        existing_retro_groups = []
     for group in retro_groups:
         d = group.tree_quality_metrics.asdict()
         d["group_id"] = group.sample_hash
@@ -797,6 +806,7 @@ class GroupTreeQualityMetrics:
     """
     Set of metrics used to assess the quality of an in inferred sample group tree.
     """
+
     strains: List[str]
     pango_lineages: List[str]
     dates: List[str]
@@ -805,6 +815,7 @@ class GroupTreeQualityMetrics:
     num_mutations: int
     num_recurrent_mutations: int
     depth: int
+    date_added: str
 
     def asdict(self):
         return dataclasses.asdict(self)
@@ -876,7 +887,7 @@ def summary(self):
             f"strains={self.strains}"
         )
 
-    def add_tree_quality_metrics(self, ts):
+    def add_tree_quality_metrics(self, ts, date):
         tree = ts.first()
         assert ts.num_trees == 1
         self.tree_quality_metrics = GroupTreeQualityMetrics(
@@ -888,6 +899,7 @@ def add_tree_quality_metrics(self, ts):
             num_root_mutations=int(np.sum(ts.mutations_node == tree.root)),
             num_recurrent_mutations=int(np.sum(ts.mutations_parent != -1)),
             depth=max(tree.depth(u) for u in ts.samples()),
+            date_added=date,
         )
         return self.tree_quality_metrics
 
@@ -901,6 +913,7 @@ def add_matching_results(
     min_different_dates=1,
     min_root_mutations=0,
     max_mutations_per_sample=np.inf,
+    max_recurrent_mutations=np.inf,
     additional_node_flags=None,
     show_progress=False,
     additional_group_metadata_keys=list(),
@@ -964,17 +977,23 @@ def add_matching_results(
                 binary_ts = tree_ops.infer_binary(flat_ts)
                 poly_ts = tree_ops.trim_branches(binary_ts)
             assert poly_ts.num_samples == flat_ts.num_samples
-            tqm = group.add_tree_quality_metrics(poly_ts)
+            tqm = group.add_tree_quality_metrics(poly_ts, date)
             if tqm.num_root_mutations < min_root_mutations:
                 logger.debug(
-                    f"Skipping root_mutations={tqm.num_root_mutations}: "
+                    f"Skipping root_mutations={tqm.num_root_mutations} < threshold "
                     f"{group.summary()}"
                 )
                 continue
             if tqm.mean_mutations_per_sample > max_mutations_per_sample:
                 logger.debug(
-                    f"Skipping mutation_per_sample={tqm.mutations_per_sample}: exceeds threshold "
-                    f"{group.summary()}"
+                    f"Skipping mean_mutations_per_sample={tqm.mean_mutations_per_sample} "
+                    f"exceeds threshold {group.summary()}"
+                )
+                continue
+            if tqm.num_recurrent_mutations > max_recurrent_mutations:
+                logger.debug(
+                    f"Skipping num_recurrent_mutations={tqm.num_recurrent_mutations} "
+                    f"exceeds threshold: {group.summary()}"
                 )
                 continue
             nodes = attach_tree(ts, tables, group, poly_ts, date, additional_node_flags)

tests/test_inference.py

@@ -1,5 +1,6 @@
 import collections
 import hashlib
+import logging
 
 import numpy as np
 import numpy.testing as nt
@@ -555,8 +556,123 @@ def test_2020_02_14_all_matches(
             "num_root_mutations": 0,
             "pango_lineages": ["A.5"],
             "strains": ["SRR15736313"],
+            "date_added": "2020-02-15",
         }
 
+    def test_2020_02_14_skip_recurrent(
+        self,
+        tmp_path,
+        fx_ts_map,
+        fx_alignment_store,
+        fx_metadata_db,
+        fx_match_db,
+        caplog,
+    ):
+        date = "2020-02-14"
+        assert len(list(fx_metadata_db.get(date))) == 0
+        with caplog.at_level("DEBUG", logger="sc2ts.inference"):
+            ts = sc2ts.extend(
+                alignment_store=fx_alignment_store,
+                metadata_db=fx_metadata_db,
+                base_ts=fx_ts_map["2020-02-13"],
+                date="2020-02-15",
+                match_db=fx_match_db,
+                # This should allow everything in but exclude on max_recurrent
+                min_root_mutations=0,
+                min_group_size=1,
+                min_different_dates=1,
+                max_recurrent_mutations=-1,
+            )
+            retro_groups = ts.metadata["sc2ts"]["retro_groups"]
+            assert len(retro_groups) == 0
+            assert "Skipping num_recurrent_mutations=0 exceeds threshold" in caplog.text
+
+    def test_2020_02_14_skip_max_mutations(
+        self,
+        tmp_path,
+        fx_ts_map,
+        fx_alignment_store,
+        fx_metadata_db,
+        fx_match_db,
+        caplog,
+    ):
+        date = "2020-02-14"
+        assert len(list(fx_metadata_db.get(date))) == 0
+        with caplog.at_level("DEBUG", logger="sc2ts.inference"):
+            ts = sc2ts.extend(
+                alignment_store=fx_alignment_store,
+                metadata_db=fx_metadata_db,
+                base_ts=fx_ts_map["2020-02-13"],
+                date="2020-02-15",
+                match_db=fx_match_db,
+                min_root_mutations=0,
+                min_group_size=1,
+                min_different_dates=1,
+                max_recurrent_mutations=100,
+                # This should allow everything in but exclude on max_mtuations
+                max_mutations_per_sample=-1,
+            )
+            retro_groups = ts.metadata["sc2ts"]["retro_groups"]
+            assert len(retro_groups) == 0
+            assert (
+                "Skipping mean_mutations_per_sample=1.0 exceeds threshold"
+                in caplog.text
+            )
+
+    def test_2020_02_14_skip_root_mutations(
+        self,
+        tmp_path,
+        fx_ts_map,
+        fx_alignment_store,
+        fx_metadata_db,
+        fx_match_db,
+        caplog,
+    ):
+        date = "2020-02-14"
+        assert len(list(fx_metadata_db.get(date))) == 0
+        with caplog.at_level("DEBUG", logger="sc2ts.inference"):
+            ts = sc2ts.extend(
+                alignment_store=fx_alignment_store,
+                metadata_db=fx_metadata_db,
+                base_ts=fx_ts_map["2020-02-13"],
+                date="2020-02-15",
+                match_db=fx_match_db,
+                # This should allow everything in but exclude on min_root_mutations
+                min_root_mutations=100,
+                min_group_size=1,
+                min_different_dates=1,
+            )
+            retro_groups = ts.metadata["sc2ts"]["retro_groups"]
+            assert len(retro_groups) == 0
+            assert "Skipping root_mutations=0 < threshold" in caplog.text
+
+    def test_2020_02_14_skip_group_size(
+        self,
+        tmp_path,
+        fx_ts_map,
+        fx_alignment_store,
+        fx_metadata_db,
+        fx_match_db,
+        caplog,
+    ):
+        date = "2020-02-14"
+        assert len(list(fx_metadata_db.get(date))) == 0
+        with caplog.at_level("DEBUG", logger="sc2ts.inference"):
+            ts = sc2ts.extend(
+                alignment_store=fx_alignment_store,
+                metadata_db=fx_metadata_db,
+                base_ts=fx_ts_map["2020-02-13"],
+                date="2020-02-15",
+                match_db=fx_match_db,
+                min_root_mutations=0,
+                # This should allow everything in but exclude on group size
+                min_group_size=100,
+                min_different_dates=1,
+            )
+            retro_groups = ts.metadata["sc2ts"]["retro_groups"]
+            assert len(retro_groups) == 0
+            assert "Skipping size=" in caplog.text
+
     @pytest.mark.parametrize("date", dates)
     def test_date_metadata(self, fx_ts_map, date):
         ts = fx_ts_map[date]