Add max_pango_lineages retro group filter.

docs/example_config.toml

@@ -55,8 +55,12 @@ deletions_as_missing=true
 
 # The minimum number of samples in a retro group
 min_group_size=10
+# The minimum number of distinct dates the samples were collected on
+min_different_dates=3
 # The minimum number of mutations shared by all samples
 min_root_mutations=2
+# The maximum number of distinct pango lineages in the group
+max_pango_lineages=10
 # The maxmimum number of recurrent mutations in the group tree
 max_recurrent_mutations=2
 # The maxmimum number of mutations per sample, overall

sc2ts/inference.py

@@ -513,6 +513,7 @@ def extend(
     min_different_dates=None,
     max_mutations_per_sample=None,
     max_recurrent_mutations=None,
+    max_pango_lineages=None,
     deletions_as_missing=None,
     max_daily_samples=None,
     show_progress=False,
@@ -535,6 +536,8 @@ def extend(
         max_mutations_per_sample = 100
     if max_recurrent_mutations is None:
         max_recurrent_mutations = 100
+    if max_pango_lineages is None:
+        max_pango_lineages = 10
     if min_different_dates is None:
         min_different_dates = 3
     if retrospective_window is None:
@@ -572,6 +575,7 @@ def extend(
             min_different_dates=min_different_dates,
             max_mutations_per_sample=max_mutations_per_sample,
             max_recurrent_mutations=max_recurrent_mutations,
+            max_pango_lineages=max_pango_lineages,
             deletions_as_missing=deletions_as_missing,
             max_daily_samples=max_daily_samples,
             retrospective_window=retrospective_window,
@@ -601,6 +605,7 @@ def _extend(
     min_different_dates,
     max_mutations_per_sample,
     max_recurrent_mutations,
+    max_pango_lineages,
     deletions_as_missing,
     max_daily_samples,
     show_progress,
@@ -724,6 +729,7 @@ def _extend(
         min_root_mutations=min_root_mutations,
         max_mutations_per_sample=max_mutations_per_sample,
         max_recurrent_mutations=max_recurrent_mutations,
+        max_pango_lineages=max_pango_lineages,
         show_progress=show_progress,
         phase="retro",
     )
@@ -985,6 +991,7 @@ def add_matching_results(
     min_root_mutations=0,
     max_mutations_per_sample=np.inf,
     max_recurrent_mutations=np.inf,
+    max_pango_lineages=np.inf,
     show_progress=False,
     phase=None,
 ):
@@ -1039,6 +1046,13 @@ def add_matching_results(
                     f"{group.summary()}"
                 )
                 continue
+            num_pango_lineages = len(group.pango_count)
+            if num_pango_lineages > max_pango_lineages:
+                logger.debug(
+                    f"Skipping num_pango_lineages={num_pango_lineages} exceeds "
+                    f"threshold: {group.summary()}"
+                )
+                continue
             flat_ts = match_path_ts(group, ts.sequence_length)
             if flat_ts.num_mutations == 0 or flat_ts.num_samples == 1:
                 poly_ts = flat_ts

tests/test_inference.py

@@ -805,6 +805,52 @@ def test_2020_02_14_all_matches(self, tmp_path, fx_ts_map, fx_dataset, fx_match_
             "date_added": "2020-02-15",
         }
 
+    def test_2020_02_14_allow_pango_lineages(
+        self, tmp_path, fx_ts_map, fx_dataset, fx_match_db
+    ):
+        date = "2020-02-14"
+        assert len(list(fx_dataset.metadata.samples_for_date(date))) == 0
+        ts = run_extend(
+            dataset=fx_dataset,
+            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,
+            # The largest group here has 2 lineages, so this lets everything in
+            max_pango_lineages=2,
+        )
+        retro_groups = ts.metadata["sc2ts"]["retro_groups"]
+        assert len(retro_groups) == 6
+
+    def test_2020_02_14_skip_pango_lineages(
+        self,
+        tmp_path,
+        fx_ts_map,
+        fx_dataset,
+        fx_match_db,
+        caplog,
+    ):
+        date = "2020-02-14"
+        assert len(list(fx_dataset.metadata.samples_for_date(date))) == 0
+        with caplog.at_level("DEBUG", logger="sc2ts.inference"):
+            ts = run_extend(
+                dataset=fx_dataset,
+                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,
+                # Excludes the one group spanning the B and B.4 lineages
+                max_pango_lineages=1,
+            )
+            retro_groups = ts.metadata["sc2ts"]["retro_groups"]
+            assert len(retro_groups) == 5
+            assert all(len(set(g["pango_lineages"])) == 1 for g in retro_groups)
+            assert "Skipping num_pango_lineages=2 exceeds threshold" in caplog.text
+
     def test_2020_02_14_skip_recurrent(
         self,
         tmp_path,