Add optional per-seed match dates to include_samples

CHANGELOG.md

@@ -9,6 +9,11 @@ In development
   usually far diverged from the current ARG, which previously led to spurious
   recombinations.
 
+- Entries in `include_samples` may now be `(strain, match_date)` tuples (in
+  addition to bare strain IDs). A non-None `match_date` matches the seed in on
+  that date instead of its actual date; it may be *before* the actual date, in
+  which case the seed node retains its actual date via a negative node time.
+
 - Add basic support for non-SARS-CoV-2 genomes via an optional reference FASTA.
   Supply `--reference` to `import-alignments` and a `reference_fasta` key in the
   inference config; both default to the built-in SARS-CoV-2 reference, so

sc2ts/inference.py

@@ -502,6 +502,27 @@ def preprocess(
     return samples
 
 
+def normalise_include_samples(include_samples):
+    """
+    Normalise the ``include_samples`` input to a canonical list of
+    ``(strain, match_date)`` tuples, where ``match_date`` is None if not
+    specified. Each input entry may be a bare strain string, or a
+    ``(strain, match_date)`` tuple/list of length 2. Any non-None match date
+    is validated as an ISO date string.
+    """
+    normalised = []
+    for entry in include_samples:
+        if isinstance(entry, str):
+            strain, match_date = entry, None
+        else:
+            strain, match_date = entry
+        if match_date is not None:
+            # Validate; raises ValueError on a malformed date.
+            parse_date(match_date)
+        normalised.append((strain, match_date))
+    return normalised
+
+
 def extend(
     *,
     dataset,
@@ -527,7 +548,21 @@ def extend(
     num_threads=0,
     memory_limit=0,
 ):
-
+    """
+    Extend the base tree sequence by one day, matching in the samples for the
+    given date.
+
+    ``include_samples`` is an optional list of "seed" samples that are matched
+    in unconditionally and without recombination. Each entry is either a bare
+    strain ID, or a ``(strain, match_date)`` tuple. When ``match_date`` is not
+    None the seed is matched in on ``match_date`` instead of its actual date;
+    this may be *before* the actual date, in which case the seed node retains
+    its actual date and is given a negative ("in the future") node time
+    relative to the match-day time-zero.
+    A ``match_date`` must be a date that is otherwise processed by the pipeline
+    (i.e. has real samples in the dataset); otherwise the seed is never
+    injected.
+    """
     if num_mismatches is None:
         num_mismatches = 3
     if hmm_cost_threshold is None:
@@ -552,6 +587,7 @@ def extend(
         deletions_as_missing = False
     if include_samples is None:
         include_samples = []
+    include_samples = normalise_include_samples(include_samples)
     base_ts = str(base_ts)
     dataset = str(dataset)
     match_db = str(match_db)
@@ -565,6 +601,20 @@ def extend(
     base_ts = tszip.load(base_ts)
     ds = _dataset.Dataset(dataset, date_field=date_field)
 
+    missing = sorted(
+        strain for strain, _ in include_samples if strain not in ds.metadata
+    )
+    if len(missing) > 0:
+        raise ValueError(f"Seed samples not in dataset: {missing}")
+    for strain, match_date in include_samples:
+        if match_date is not None:
+            actual_date = ds.metadata[strain]["date"]
+            if match_date > actual_date:
+                logger.warning(
+                    f"Seed sample {strain} match date {match_date} is after its "
+                    f"actual date {actual_date}; this is unusual for a seed sample"
+                )
+
     with MatchDb(match_db) as matches:
         tables = _extend(
             dataset=ds,
@@ -625,10 +675,27 @@ def _extend(
         f"mutations={base_ts.num_mutations};date={previous_date}"
     )
 
+    include_strains = {strain for strain, _ in include_samples}
+    # Seeds with an explicit match date are matched in on that date instead of
+    # their actual date, so we override which day they're processed on.
+    override_dates = {
+        strain: match_date
+        for strain, match_date in include_samples
+        if match_date is not None
+    }
+
     metadata_matches = {
         strain: dataset.metadata[strain]
         for strain in dataset.metadata.samples_for_date(date)
+        # Exclude a seed with an override match date from its actual date; it
+        # is processed only on the override date.
+        if override_dates.get(strain, date) == date
     }
+    # Inject seeds whose override match date is today but which aren't
+    # naturally sampled today. Missing strains are rejected up-front in extend().
+    for strain, match_date in override_dates.items():
+        if match_date == date and strain not in metadata_matches:
+            metadata_matches[strain] = dataset.metadata[strain]
 
     logger.info(f"Got {len(metadata_matches)} metadata matches")
 
@@ -643,7 +710,6 @@ def _extend(
     pango_lineage_key = "Viridian_pangolin"
     scorpio_key = "Viridian_scorpio"
 
-    include_strains = set(include_samples)
     unconditional_include_samples = []
     samples = []
     for s in preprocessed_samples:
@@ -1455,6 +1521,9 @@ def match_tsinfer(
     num_alleles = 4 if deletions_as_missing else 5
     mu, rho = solve_num_mismatches(num_mismatches, num_alleles)
 
+    # Detach any future (negative-time) nodes so that samples can't copy from
+    # them. Node IDs are preserved, so the returned match paths stay valid.
+    ts = tree_ops.detach_future_nodes(ts)
     tsb, coord_map = make_tsb(ts, num_alleles, mirror_coordinates)
 
     work = []
@@ -1755,7 +1824,11 @@ def attach_tree(
             node = child_ts.node(u)
             sample_date = parse_date(node.metadata["date"])
             node_time[u] = (current_date - sample_date).days
-            assert node_time[u] >= 0.0
+            # Seed samples can be matched in on a date before their actual
+            # date, giving a negative ("in the future") node time.
+            assert node_time[u] >= 0.0 or (
+                node.flags & core.NODE_IS_UNCONDITIONALLY_INCLUDED
+            )
     max_sample_time = max(node_time.values())
 
     node_id_map = {}

sc2ts/tree_ops.py

@@ -657,6 +657,41 @@ def drop_vestigial_root_edge(ts):
     return tables.tree_sequence()
 
 
+def detach_future_nodes(ts):
+    """
+    Return a copy of ``ts`` in which every node with a negative ("in the
+    future") time is fully detached: all edges incident to it are removed, any
+    mutations over it are dropped, and its ``NODE_IS_SAMPLE`` flag is cleared.
+
+    Such nodes are seed samples that were matched in on a date before their
+    actual date, and so lie in the future relative to the current time-zero.
+    Detaching them prevents other samples from copying from them during
+    matching. Node IDs and times are preserved so that any match paths
+    referring to the returned tree sequence remain valid against the original.
+    """
+    future = ts.nodes_time < 0
+    if not np.any(future):
+        return ts
+    tables = ts.dump_tables()
+    keep_edges = ~(future[tables.edges.parent] | future[tables.edges.child])
+    keep_mutations = ~future[tables.mutations.node]
+    logger.debug(
+        f"Detaching {int(np.sum(future))} future nodes "
+        f"({len(tables.edges) - int(np.sum(keep_edges))} edges, "
+        f"{len(tables.mutations) - int(np.sum(keep_mutations))} mutations removed)"
+    )
+    tables.edges.keep_rows(keep_edges)
+    tables.mutations.keep_rows(keep_mutations)
+    # A detached future node must not be treated as a sample to copy from.
+    flags = tables.nodes.flags
+    flags[future] &= ~np.uint32(tskit.NODE_IS_SAMPLE)
+    tables.nodes.flags = flags
+    tables.sort()
+    tables.build_index()
+    tables.compute_mutation_parents()
+    return tables.tree_sequence()
+
+
 def insert_vestigial_root_edge(ts):
     """
     Insert an edge between node 0 and 1 at the end of the edge table, if

tests/test_cli.py

@@ -422,7 +422,7 @@ def test_include_samples(self, tmp_path, fx_ts_map, fx_dataset):
             tmp_path,
             fx_dataset,
             exclude_sites=[56, 57, 58, 59, 60],
-            include_samples=["SRR14631544", "NO_SUCH_STRAIN"],
+            include_samples=["SRR14631544"],
         )
         runner = ct.CliRunner()
         result = runner.invoke(
@@ -440,6 +440,44 @@ def test_include_samples(self, tmp_path, fx_ts_map, fx_dataset):
         assert np.sum(ts.nodes_time[ts.samples()] == 0) == 1
         assert ts.num_samples == 1
 
+    def test_include_samples_missing_strain(self, tmp_path, fx_ts_map, fx_dataset):
+        # A seed strain not in the dataset makes the run fail.
+        config_file = self.make_config(
+            tmp_path,
+            fx_dataset,
+            exclude_sites=[56, 57, 58, 59, 60],
+            include_samples=["SRR14631544", "NO_SUCH_STRAIN"],
+        )
+        runner = ct.CliRunner()
+        with pytest.raises(ValueError, match="not in dataset"):
+            runner.invoke(
+                cli.cli,
+                f"infer {config_file} --stop 2020-01-02",
+                catch_exceptions=False,
+            )
+
+    def test_include_samples_with_dates(self, tmp_path, fx_ts_map, fx_dataset):
+        # The (strain, date) tuple form is expressed in TOML as a 2-element
+        # array, mixed with bare strings.
+        config_file = self.make_config(
+            tmp_path,
+            fx_dataset,
+            exclude_sites=[56, 57, 58, 59, 60],
+            include_samples=[["SRR14631544", "2020-01-01"]],
+        )
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.cli,
+            f"infer {config_file} --stop 2020-01-02",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        date = "2020-01-01"
+        ts_path = tmp_path / "results" / "test" / f"test_{date}.ts"
+        ts = tskit.load(ts_path)
+        assert "SRR14631544" in ts.metadata["sc2ts"]["samples_strain"]
+        assert ts.num_samples == 1
+
     def test_override(self, tmp_path, fx_ts_map, fx_dataset):
         hmm_cost_threshold = 47
         config_file = self.make_config(