From 56352465eef8753012fd21834cdf948ec0640890 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Tue, 28 Jul 2026 17:05:12 +0100 Subject: [PATCH 01/10] Add optional per-seed match dates to include_samples Each include_samples entry may now be a (strain, match_date) tuple in addition to a bare strain ID. 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 ("in the future") node time relative to the match-day time-zero. This lets widely-diverged seeds (e.g. Omicron BA.1/BA.2/BA.3) be matched in on the same early day so a rudimentary ancestry tree forms before other tree building, instead of matching to derived samples with many reversions. Also detach future (negative-time) nodes when setting up match_tsinfer, so that other samples cannot copy from a node that lies in the future relative to the current time-zero. --- CHANGELOG.md | 7 ++ sc2ts/inference.py | 83 ++++++++++++++++++- sc2ts/tree_ops.py | 27 ++++++ tests/test_cli.py | 22 +++++ tests/test_inference.py | 177 +++++++++++++++++++++++++++++++++++++++- tests/test_tree_ops.py | 61 ++++++++++++++ 6 files changed, 372 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ea3f9b..1537713 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,13 @@ 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. + This lets widely-diverged seeds (e.g. Omicron BA.1/BA.2/BA.3) be matched in + on the same early day. + - 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 diff --git a/sc2ts/inference.py b/sc2ts/inference.py index a543388..855ceb8 100644 --- a/sc2ts/inference.py +++ b/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,22 @@ 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. This allows widely-diverged seeds + (e.g. Omicron BA.1/BA.2/BA.3) to be matched in on the same early day. + 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 +588,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 +602,15 @@ def extend( base_ts = tszip.load(base_ts) ds = _dataset.Dataset(dataset, date_field=date_field) + for strain, match_date in include_samples: + if match_date is not None and strain in ds.metadata: + 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 +671,30 @@ 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. + for strain, match_date in override_dates.items(): + if match_date == date and strain not in metadata_matches: + if strain in dataset.metadata: + metadata_matches[strain] = dataset.metadata[strain] + else: + logger.warning(f"Seed sample {strain} not in dataset; cannot include") logger.info(f"Got {len(metadata_matches)} metadata matches") @@ -643,12 +709,14 @@ 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: if s.haplotype is None: - logger.debug(f"No alignment stored for {s.strain}") + if s.strain in include_strains: + logger.warning(f"No alignment stored for seed sample {s.strain}") + else: + logger.debug(f"No alignment stored for {s.strain}") continue md = metadata_matches[s.strain] s.metadata = md @@ -1455,6 +1523,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 +1826,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 = {} diff --git a/sc2ts/tree_ops.py b/sc2ts/tree_ops.py index 6e22fba..1c8f5c4 100644 --- a/sc2ts/tree_ops.py +++ b/sc2ts/tree_ops.py @@ -657,6 +657,33 @@ 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 detached by removing all edges incident to it. + + 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 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 = ~(future[tables.edges.parent] | future[tables.edges.child]) + num_detached = int(np.sum(future)) + logger.debug( + f"Detaching {num_detached} future nodes " + f"({len(tables.edges) - int(np.sum(keep))} edges removed)" + ) + tables.edges.keep_rows(keep) + tables.sort() + tables.build_index() + 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 diff --git a/tests/test_cli.py b/tests/test_cli.py index dd066c5..b0ea126 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -440,6 +440,28 @@ 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_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"], "NO_SUCH_STRAIN"], + ) + 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( diff --git a/tests/test_inference.py b/tests/test_inference.py index 5d838bb..3f045a9 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -350,6 +350,56 @@ def test_match_reference_all_same(self, mirror, allele): assert mut.site_id == site_id assert mut.derived_state == sc2ts.IUPAC_ALLELES[allele] + def _extra_node_setup(self, node_time, site_id=5): + # A truncated reference tree sequence with one extra node hanging off + # the reference (node 1), carrying a single distinguishing mutation. + ts = util.initial_ts() + tables = ts.dump_tables() + tables.sites.truncate(20) + # Lift the base times so the reference sits well above time zero, + # mimicking a base_ts after several days of time increments. This + # leaves room to attach both present (positive-time) and future + # (negative-time) nodes below the reference. + tables.nodes.time += 10 + ancestral = tables.sites[site_id].ancestral_state + derived = "A" if ancestral != "A" else "C" + u = tables.nodes.add_row( + flags=tskit.NODE_IS_SAMPLE, + time=node_time, + metadata={"strain": "extra", "date": "2030-01-01"}, + ) + tables.edges.add_row(0, ts.sequence_length, parent=1, child=u) + tables.mutations.add_row(site=site_id, node=u, derived_state=derived) + tables.sort() + tables.build_index() + ts = tables.tree_sequence() + # A sample whose haplotype exactly matches the extra node. + alignment = util.reference_array() + alignment[0] = "A" + h = jit.encode_alleles(alignment)[ts.sites_position.astype(int)] + h[site_id] = sc2ts.IUPAC_ALLELES.index(derived) + sample = si.Sample("test", "2020-01-01", haplotype=h) + return ts, u, sample, site_id + + def test_matches_to_present_node(self): + # Sanity check: an ordinary (non-future) node that is an exact match IS + # copied from. This is the behaviour we suppress for future nodes. + ts, u, sample, site_id = self._extra_node_setup(node_time=0.5) + matches = self.match_tsinfer([sample], ts) + assert matches[0].parents == [u] + assert len(matches[0].mutations) == 0 + + def test_no_match_to_future_node(self): + # A future (negative-time) node that is an exact match must NOT be + # copied from; the sample falls back to the reference and carries the + # differing site as a mutation. + ts, u, sample, site_id = self._extra_node_setup(node_time=-5) + matches = self.match_tsinfer([sample], ts) + assert u not in matches[0].parents + assert matches[0].parents == [1] + assert len(matches[0].mutations) == 1 + assert matches[0].mutations[0].site_id == site_id + class TestMirrorTsCoords: def test_dense_sites_example(self): @@ -432,6 +482,22 @@ def test_high_recomb_mutation(self): self.check_double_mirror(ts) +class TestNormaliseIncludeSamples: + def test_empty(self): + assert si.normalise_include_samples([]) == [] + + def test_bare_strings(self): + assert si.normalise_include_samples(["a", "b"]) == [("a", None), ("b", None)] + + def test_tuples_and_lists(self): + result = si.normalise_include_samples([("a", "2021-01-01"), ["b", None], "c"]) + assert result == [("a", "2021-01-01"), ("b", None), ("c", None)] + + def test_malformed_date_raises(self): + with pytest.raises(ValueError, match="isoformat"): + si.normalise_include_samples([("a", "not-a-date")]) + + class TestRealData: dates = [ "2020-01-01", @@ -551,7 +617,15 @@ def test_2020_02_02(self, tmp_path, fx_ts_map, fx_dataset, num_threads): ts.tables.assert_equals(fx_ts_map["2020-02-02"].tables, ignore_provenance=True) @pytest.mark.parametrize( - "include_samples", (["SRR11597115"], ["SRR11597115", "NOSUCHSTRAIN"]) + "include_samples", + ( + ["SRR11597115"], + ["SRR11597115", "NOSUCHSTRAIN"], + # The tuple form with a None match date is equivalent to the + # bare-string form. + [("SRR11597115", None)], + [("SRR11597115", "2020-02-02")], + ), ) def test_2020_02_02_include_samples( self, @@ -585,6 +659,107 @@ def test_2020_02_02_include_samples( assert edges[0].left == 0 assert edges[0].right == ts.sequence_length + def test_seed_early_match_date_negative_time(self, tmp_path, fx_ts_map, fx_dataset): + # SRR11597115's actual date is 2020-02-02; match it in two days early. + strain = "SRR11597115" + ts = run_extend( + dataset=fx_dataset, + base_ts=fx_ts_map["2020-01-30"], + date="2020-01-31", + match_db=si.MatchDb.initialise(tmp_path / "match.db"), + include_samples=[(strain, "2020-01-31")], + ) + assert strain in ts.metadata["sc2ts"]["samples_strain"] + u = ts.samples()[ts.metadata["sc2ts"]["samples_strain"].index(strain)] + assert ts.nodes_flags[u] & sc2ts.NODE_IS_UNCONDITIONALLY_INCLUDED > 0 + # The node retains its actual date, two days in the future relative to + # the match-day time-zero, so its time is -2. + assert ts.nodes_time[u] == -2 + # Matched without recombination: a single full-span parent edge, and + # the parent must be older (larger time) than the future-dated seed. + assert ts.nodes_flags[u] & sc2ts.NODE_IS_RECOMBINANT == 0 + edges = [e for e in ts.edges() if e.child == u] + assert len(edges) == 1 + assert edges[0].left == 0 + assert edges[0].right == ts.sequence_length + assert ts.nodes_time[edges[0].parent] > ts.nodes_time[u] + + def test_seed_early_match_date_evolves_over_days( + self, tmp_path, fx_ts_map, fx_dataset + ): + # Inject the seed two days before its actual date, then keep extending + # past the actual date. Its node time must rise by +1/day and hit 0 on + # the actual date, exercising the low-level matcher against a base_ts + # that contains negative ("future") node times. + strain = "SRR11597115" + include_samples = [(strain, "2020-01-31")] + base_path = tmp_path / "base.ts" + fx_ts_map["2020-01-30"].dump(base_path) + match_db = si.MatchDb.initialise(tmp_path / "match.db") + dates = ["2020-01-31", "2020-02-01", "2020-02-02", "2020-02-03"] + expected_time = { + "2020-01-31": -2, + "2020-02-01": -1, + "2020-02-02": 0, + "2020-02-03": 1, + } + for date in dates: + ts = si.extend( + dataset=fx_dataset.path, + base_ts=base_path, + date=date, + match_db=match_db.path, + include_samples=include_samples, + ) + ts.dump(base_path) + strains = ts.metadata["sc2ts"]["samples_strain"] + # The seed is added exactly once and never double-processed on its + # natural date. + assert strains.count(strain) == 1 + u = ts.samples()[strains.index(strain)] + assert ts.nodes_time[u] == expected_time[date] + + def test_seed_excluded_on_natural_date(self, tmp_path, fx_ts_map, fx_dataset): + # A seed with an override match date must NOT be processed on its + # actual date. SRR11597115 is naturally sampled on 2020-02-02; with an + # override of 2020-01-31 it should be absent when we extend 2020-02-02. + strain = "SRR11597115" + ts = run_extend( + dataset=fx_dataset, + base_ts=fx_ts_map["2020-02-01"], + date="2020-02-02", + match_db=si.MatchDb.initialise(tmp_path / "match.db"), + include_samples=[(strain, "2020-01-31")], + ) + assert strain not in ts.metadata["sc2ts"]["samples_strain"] + + def test_seed_override_date_no_samples(self, tmp_path, fx_ts_map, fx_dataset): + # An override match date with no naturally-sampled strains that day is + # never reached by the driver, so the seed is silently not injected. + # Here we simply confirm extend on such a date doesn't crash and the + # seed isn't added. + strain = "SRR11597115" + ts = run_extend( + dataset=fx_dataset, + base_ts=fx_ts_map["2020-02-01"], + date="2020-02-02", + match_db=si.MatchDb.initialise(tmp_path / "match.db"), + include_samples=[(strain, "2020-02-12")], + ) + assert strain not in ts.metadata["sc2ts"]["samples_strain"] + + def test_seed_unknown_strain_with_override(self, tmp_path, fx_ts_map, fx_dataset): + # An unknown seed strain (even with an override date) is tolerated + # silently, as bare unknown strains already are. + ts = run_extend( + dataset=fx_dataset, + base_ts=fx_ts_map["2020-02-01"], + date="2020-02-02", + match_db=si.MatchDb.initialise(tmp_path / "match.db"), + include_samples=[("NOSUCHSTRAIN", "2020-02-02")], + ) + assert "NOSUCHSTRAIN" not in ts.metadata["sc2ts"]["samples_strain"] + def test_2020_02_02_mutation_overlap( self, tmp_path, diff --git a/tests/test_tree_ops.py b/tests/test_tree_ops.py index c8bb5bd..a64c917 100644 --- a/tests/test_tree_ops.py +++ b/tests/test_tree_ops.py @@ -957,3 +957,64 @@ def test_msprime_input_fails(self): ts = msprime.sim_ancestry(2) with pytest.raises(ValueError, match="Oldest edge"): tree_ops.insert_vestigial_root_edge(ts) + + +def _future_node_ts(future_times=(-1,)): + """ + A tiny tree sequence: a root (time 2) over a present-day node (time 1), + with a chain of ``future_times`` nodes hanging below it. + """ + tables = tskit.TableCollection(sequence_length=10) + tables.nodes.add_row(time=2) # 0: root + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=1) # 1: present + tables.edges.add_row(0, 10, parent=0, child=1) + parent = 1 + for t in future_times: + child = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=t) + tables.edges.add_row(0, 10, parent=parent, child=child) + parent = child + tables.sort() + tables.build_index() + return tables.tree_sequence() + + +class TestDetachFutureNodes: + def _incident_edges(self, ts, node): + return [e for e in ts.edges() if e.parent == node or e.child == node] + + def test_no_future_nodes_is_noop(self): + ts = _future_node_ts(future_times=()) + result = tree_ops.detach_future_nodes(ts) + # Returned unchanged (same object) when there's nothing to do. + assert result is ts + + def test_detaches_single_future_node(self): + ts = _future_node_ts(future_times=(-1,)) + assert len(self._incident_edges(ts, 2)) == 1 + result = tree_ops.detach_future_nodes(ts) + # Node preserved, but now isolated. + assert result.num_nodes == ts.num_nodes + assert result.nodes_time[2] == -1 + assert self._incident_edges(result, 2) == [] + # The non-future edge is retained. + assert result.num_edges == ts.num_edges - 1 + retained = list(result.edges()) + assert len(retained) == 1 + assert retained[0].parent == 0 + assert retained[0].child == 1 + + def test_detaches_chain_of_future_nodes(self): + ts = _future_node_ts(future_times=(-1, -2, -3)) + result = tree_ops.detach_future_nodes(ts) + for node in (2, 3, 4): + assert self._incident_edges(result, node) == [] + # Only the root -> present edge survives. + assert result.num_edges == 1 + assert result.nodes_time[2] == -1 + assert result.nodes_time[4] == -3 + + def test_node_ids_preserved(self): + ts = _future_node_ts(future_times=(-5,)) + result = tree_ops.detach_future_nodes(ts) + nt.assert_array_equal(result.nodes_time, ts.nodes_time) + nt.assert_array_equal(result.nodes_flags, ts.nodes_flags) From 2656af62b5b9e9e047911e91c8ab2385fd5e0f2e Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Tue, 28 Jul 2026 21:18:22 +0100 Subject: [PATCH 02/10] Harden seed handling: error on missing seeds, fully detach future nodes A seed strain in include_samples that is not present in the dataset now raises a ValueError up-front in extend() instead of being silently warned-and-skipped, so typos and bad configs fail fast. The no-alignment corner case is handled the same as for any other sample (no seed-specific path). detach_future_nodes now fully detaches negative-time nodes: in addition to removing incident edges, it drops mutations over those nodes and clears their NODE_IS_SAMPLE flag so the matcher cannot treat them as ancestors, recomputing mutation parents afterwards. Node IDs and times are preserved. Rewrite the detach_future_nodes tests with hand-drawn topologies (including recombination) following the conventions in test_tree_ops.py. --- CHANGELOG.md | 2 - sc2ts/inference.py | 22 ++--- sc2ts/tree_ops.py | 24 +++-- tests/test_cli.py | 20 +++- tests/test_inference.py | 57 ++++++++--- tests/test_tree_ops.py | 212 +++++++++++++++++++++++++++++++--------- 6 files changed, 257 insertions(+), 80 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1537713..ae15f7a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,8 +13,6 @@ In development 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. - This lets widely-diverged seeds (e.g. Omicron BA.1/BA.2/BA.3) be matched in - on the same early day. - 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 diff --git a/sc2ts/inference.py b/sc2ts/inference.py index 855ceb8..fc86863 100644 --- a/sc2ts/inference.py +++ b/sc2ts/inference.py @@ -558,8 +558,7 @@ def extend( 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. This allows widely-diverged seeds - (e.g. Omicron BA.1/BA.2/BA.3) to be matched in on the same early day. + 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. @@ -602,8 +601,13 @@ 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 and strain in ds.metadata: + if match_date is not None: actual_date = ds.metadata[strain]["date"] if match_date > actual_date: logger.warning( @@ -688,13 +692,10 @@ def _extend( if override_dates.get(strain, date) == date } # Inject seeds whose override match date is today but which aren't - # naturally sampled today. + # 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: - if strain in dataset.metadata: - metadata_matches[strain] = dataset.metadata[strain] - else: - logger.warning(f"Seed sample {strain} not in dataset; cannot include") + metadata_matches[strain] = dataset.metadata[strain] logger.info(f"Got {len(metadata_matches)} metadata matches") @@ -713,10 +714,7 @@ def _extend( samples = [] for s in preprocessed_samples: if s.haplotype is None: - if s.strain in include_strains: - logger.warning(f"No alignment stored for seed sample {s.strain}") - else: - logger.debug(f"No alignment stored for {s.strain}") + logger.debug(f"No alignment stored for {s.strain}") continue md = metadata_matches[s.strain] s.metadata = md diff --git a/sc2ts/tree_ops.py b/sc2ts/tree_ops.py index 1c8f5c4..63d701c 100644 --- a/sc2ts/tree_ops.py +++ b/sc2ts/tree_ops.py @@ -660,27 +660,35 @@ def drop_vestigial_root_edge(ts): def detach_future_nodes(ts): """ Return a copy of ``ts`` in which every node with a negative ("in the - future") time is detached by removing all edges incident to it. + 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 are preserved so that any match paths referring to the - returned tree sequence remain valid against the original. + 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 = ~(future[tables.edges.parent] | future[tables.edges.child]) - num_detached = int(np.sum(future)) + keep_edges = ~(future[tables.edges.parent] | future[tables.edges.child]) + keep_mutations = ~future[tables.mutations.node] logger.debug( - f"Detaching {num_detached} future nodes " - f"({len(tables.edges) - int(np.sum(keep))} edges removed)" + 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) + 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() diff --git a/tests/test_cli.py b/tests/test_cli.py index b0ea126..b104a3c 100644 --- a/tests/test_cli.py +++ b/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,22 @@ 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. @@ -447,7 +463,7 @@ def test_include_samples_with_dates(self, tmp_path, fx_ts_map, fx_dataset): tmp_path, fx_dataset, exclude_sites=[56, 57, 58, 59, 60], - include_samples=[["SRR14631544", "2020-01-01"], "NO_SUCH_STRAIN"], + include_samples=[["SRR14631544", "2020-01-01"]], ) runner = ct.CliRunner() result = runner.invoke( diff --git a/tests/test_inference.py b/tests/test_inference.py index 3f045a9..b378b71 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -620,7 +620,6 @@ def test_2020_02_02(self, tmp_path, fx_ts_map, fx_dataset, num_threads): "include_samples", ( ["SRR11597115"], - ["SRR11597115", "NOSUCHSTRAIN"], # The tuple form with a None match date is equivalent to the # bare-string form. [("SRR11597115", None)], @@ -684,6 +683,30 @@ def test_seed_early_match_date_negative_time(self, tmp_path, fx_ts_map, fx_datas assert edges[0].right == ts.sequence_length assert ts.nodes_time[edges[0].parent] > ts.nodes_time[u] + def test_seed_match_date_equals_actual_date(self, tmp_path, fx_ts_map, fx_dataset): + # An override match date equal to the sample's actual date behaves like + # an ordinary seed: the node is present-dated (time 0), not in the + # future. SRR11597115's actual date is 2020-02-02. + strain = "SRR11597115" + ts = run_extend( + dataset=fx_dataset, + base_ts=fx_ts_map["2020-02-01"], + date="2020-02-02", + match_db=si.MatchDb.initialise(tmp_path / "match.db"), + include_samples=[(strain, "2020-02-02")], + ) + assert strain in ts.metadata["sc2ts"]["samples_strain"] + u = ts.samples()[ts.metadata["sc2ts"]["samples_strain"].index(strain)] + assert ts.nodes_flags[u] & sc2ts.NODE_IS_UNCONDITIONALLY_INCLUDED > 0 + # match date == actual date, so the node sits at time zero. + assert ts.nodes_time[u] == 0 + # Still matched without recombination: a single full-span parent edge. + assert ts.nodes_flags[u] & sc2ts.NODE_IS_RECOMBINANT == 0 + edges = [e for e in ts.edges() if e.child == u] + assert len(edges) == 1 + assert edges[0].left == 0 + assert edges[0].right == ts.sequence_length + def test_seed_early_match_date_evolves_over_days( self, tmp_path, fx_ts_map, fx_dataset ): @@ -748,17 +771,27 @@ def test_seed_override_date_no_samples(self, tmp_path, fx_ts_map, fx_dataset): ) assert strain not in ts.metadata["sc2ts"]["samples_strain"] - def test_seed_unknown_strain_with_override(self, tmp_path, fx_ts_map, fx_dataset): - # An unknown seed strain (even with an override date) is tolerated - # silently, as bare unknown strains already are. - ts = run_extend( - dataset=fx_dataset, - base_ts=fx_ts_map["2020-02-01"], - date="2020-02-02", - match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=[("NOSUCHSTRAIN", "2020-02-02")], - ) - assert "NOSUCHSTRAIN" not in ts.metadata["sc2ts"]["samples_strain"] + @pytest.mark.parametrize( + "include_samples", + ( + ["SRR11597115", "NOSUCHSTRAIN"], + [("NOSUCHSTRAIN", "2020-02-02")], + [("NOSUCHSTRAIN", None)], + ), + ) + def test_seed_missing_strain_raises( + self, tmp_path, fx_ts_map, fx_dataset, include_samples + ): + # A seed strain that isn't in the dataset is an error, whether it's a + # bare strain or carries an override match date. + with pytest.raises(ValueError, match="not in dataset"): + run_extend( + dataset=fx_dataset, + base_ts=fx_ts_map["2020-02-01"], + date="2020-02-02", + match_db=si.MatchDb.initialise(tmp_path / "match.db"), + include_samples=include_samples, + ) def test_2020_02_02_mutation_overlap( self, diff --git a/tests/test_tree_ops.py b/tests/test_tree_ops.py index a64c917..3c934cb 100644 --- a/tests/test_tree_ops.py +++ b/tests/test_tree_ops.py @@ -959,62 +959,186 @@ def test_msprime_input_fails(self): tree_ops.insert_vestigial_root_edge(ts) -def _future_node_ts(future_times=(-1,)): - """ - A tiny tree sequence: a root (time 2) over a present-day node (time 1), - with a chain of ``future_times`` nodes hanging below it. - """ - tables = tskit.TableCollection(sequence_length=10) - tables.nodes.add_row(time=2) # 0: root - tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=1) # 1: present - tables.edges.add_row(0, 10, parent=0, child=1) - parent = 1 - for t in future_times: - child = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=t) - tables.edges.add_row(0, 10, parent=parent, child=child) - parent = child - tables.sort() - tables.build_index() - return tables.tree_sequence() +def _incident_edges(ts, node): + return [(e.left, e.right, e.parent) for e in ts.edges() if e.child == node] -class TestDetachFutureNodes: - def _incident_edges(self, ts, node): - return [e for e in ts.edges() if e.parent == node or e.child == node] +def _is_sample(ts, node): + return bool(ts.nodes_flags[node] & tskit.NODE_IS_SAMPLE) + +class TestDetachFutureNodes: def test_no_future_nodes_is_noop(self): - ts = _future_node_ts(future_times=()) + # 2.00┊ 2 ┊ + # ┊ ┏━┻━┓ ┊ + # 0.00┊ 0 1 ┊ both present samples + ts = tskit.Tree.generate_balanced(2, span=10).tree_sequence result = tree_ops.detach_future_nodes(ts) # Returned unchanged (same object) when there's nothing to do. assert result is ts - def test_detaches_single_future_node(self): - ts = _future_node_ts(future_times=(-1,)) - assert len(self._incident_edges(ts, 2)) == 1 + def test_single_future_leaf(self): + # 2.00┊ 0 ┊ root (non-sample) + # ┊ ┏━┻━┓ ┊ + # 0.00┊ 1 ┃ ┊ present sample, mutation at site 0 + # -2.00┊ 2 ┊ future sample, mutation at site 1 + tables = tskit.TableCollection(sequence_length=10) + tables.nodes.add_row(time=2) # 0 root + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) # 1 present + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=-2) # 2 future + tables.edges.add_row(0, 10, parent=0, child=1) + tables.edges.add_row(0, 10, parent=0, child=2) + tables.sites.add_row(1, "A") + tables.sites.add_row(2, "A") + tables.mutations.add_row(site=0, node=1, derived_state="T", time=0) + tables.mutations.add_row(site=1, node=2, derived_state="G", time=-2) + ts = prepare(tables) + result = tree_ops.detach_future_nodes(ts) - # Node preserved, but now isolated. + # Nodes preserved (ids and times), future node now isolated. assert result.num_nodes == ts.num_nodes - assert result.nodes_time[2] == -1 - assert self._incident_edges(result, 2) == [] - # The non-future edge is retained. - assert result.num_edges == ts.num_edges - 1 - retained = list(result.edges()) - assert len(retained) == 1 - assert retained[0].parent == 0 - assert retained[0].child == 1 + nt.assert_array_equal(result.nodes_time, ts.nodes_time) + assert _incident_edges(result, 2) == [] + assert result.num_edges == 1 + assert _incident_edges(result, 1) == [(0, 10, 0)] + # The mutation over the future node is dropped; the present one stays. + assert result.num_mutations == 1 + assert result.mutation(0).node == 1 + # The future node is no longer a sample; the present one is untouched. + assert not _is_sample(result, 2) + assert _is_sample(result, 1) + + def test_chain_of_future_nodes(self): + # 3.00┊ 0 ┊ root (non-sample) + # ┊ ┏━┻━┓ ┊ + # 0.00┊ 1 2 ┊ present sample (2 has mutation), future chain below 2 + # -1.00┊ 3 ┊ future sample, mutation + # -3.00┊ 4 ┊ future sample, mutation + tables = tskit.TableCollection(sequence_length=10) + tables.nodes.add_row(time=3) # 0 root + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) # 1 present + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) # 2 present + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=-1) # 3 future + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=-3) # 4 future + tables.edges.add_row(0, 10, parent=0, child=1) + tables.edges.add_row(0, 10, parent=0, child=2) + tables.edges.add_row(0, 10, parent=2, child=3) + tables.edges.add_row(0, 10, parent=3, child=4) + for pos, node in [(1, 2), (2, 3), (3, 4)]: + tables.sites.add_row(pos, "A") + tables.mutations.add_row(site=0, node=2, derived_state="T", time=0) + tables.mutations.add_row(site=1, node=3, derived_state="C", time=-1) + tables.mutations.add_row(site=2, node=4, derived_state="G", time=-3) + ts = prepare(tables) + + result = tree_ops.detach_future_nodes(ts) + # Both future edges gone; only the two present edges remain. + assert result.num_edges == 2 + assert _incident_edges(result, 3) == [] + assert _incident_edges(result, 4) == [] + assert _incident_edges(result, 2) == [(0, 10, 0)] + # Only the mutation over the present node survives. + assert result.num_mutations == 1 + assert result.mutation(0).node == 2 + assert not _is_sample(result, 3) + assert not _is_sample(result, 4) + + def test_future_internal_node(self): + # 4.00┊ 0 ┊ root (non-sample) + # ┊ ┏━┻━┓ ┊ + # 0.00┊ 1 ┃ ┊ present sample + # -1.00┊ 2 ┊ FUTURE internal node (non-sample), mutation + # ┊ ┃ ┊ + # -3.00┊ 3 ┊ future sample, mutation + tables = tskit.TableCollection(sequence_length=10) + tables.nodes.add_row(time=4) # 0 root + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) # 1 present + tables.nodes.add_row(time=-1) # 2 future internal (non-sample) + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=-3) # 3 future + tables.edges.add_row(0, 10, parent=0, child=1) + tables.edges.add_row(0, 10, parent=0, child=2) + tables.edges.add_row(0, 10, parent=2, child=3) + tables.sites.add_row(1, "A") + tables.sites.add_row(2, "A") + tables.mutations.add_row(site=0, node=2, derived_state="T", time=-1) + tables.mutations.add_row(site=1, node=3, derived_state="G", time=-3) + ts = prepare(tables) - def test_detaches_chain_of_future_nodes(self): - ts = _future_node_ts(future_times=(-1, -2, -3)) result = tree_ops.detach_future_nodes(ts) - for node in (2, 3, 4): - assert self._incident_edges(result, node) == [] - # Only the root -> present edge survives. assert result.num_edges == 1 - assert result.nodes_time[2] == -1 - assert result.nodes_time[4] == -3 + assert _incident_edges(result, 2) == [] + assert _incident_edges(result, 3) == [] + # Both future mutations dropped. + assert result.num_mutations == 0 + # The already-non-sample internal node's flags are unchanged. + assert result.nodes_flags[2] == ts.nodes_flags[2] + assert not _is_sample(result, 3) + + def test_future_recombinant_node(self): + tables = tskit.TableCollection(sequence_length=10) + tables.nodes.add_row(time=3) # 0 root + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) # 1 present + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) # 2 present + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=-1) # 3 future recomb + tables.edges.add_row(0, 10, parent=0, child=1) + tables.edges.add_row(0, 10, parent=0, child=2) + tables.edges.add_row(0, 5, parent=1, child=3) + tables.edges.add_row(5, 10, parent=2, child=3) + tables.sites.add_row(2, "A") + tables.mutations.add_row(site=0, node=3, derived_state="T", time=-1) + ts = prepare(tables) - def test_node_ids_preserved(self): - ts = _future_node_ts(future_times=(-5,)) result = tree_ops.detach_future_nodes(ts) - nt.assert_array_equal(result.nodes_time, ts.nodes_time) - nt.assert_array_equal(result.nodes_flags, ts.nodes_flags) + # Both partial-span parent edges of the future recombinant are removed. + assert _incident_edges(result, 3) == [] + assert result.num_edges == 2 + assert result.num_mutations == 0 + assert not _is_sample(result, 3) + + def test_present_recombinant_preserved(self): + tables = tskit.TableCollection(sequence_length=10) + tables.nodes.add_row(time=3) # 0 root + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=1) # 1 present + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=1) # 2 present + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) # 3 present recomb + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=-2) # 4 future leaf + tables.edges.add_row(0, 10, parent=0, child=1) + tables.edges.add_row(0, 10, parent=0, child=2) + tables.edges.add_row(0, 5, parent=1, child=3) + tables.edges.add_row(5, 10, parent=2, child=3) + tables.edges.add_row(0, 10, parent=0, child=4) + ts = prepare(tables) + + result = tree_ops.detach_future_nodes(ts) + # Only the future leaf's edge is removed. + assert _incident_edges(result, 4) == [] + assert not _is_sample(result, 4) + # The present recombinant keeps both of its partial-span parent edges. + assert _incident_edges(result, 3) == [(0, 5, 1), (5, 10, 2)] + assert result.num_edges == ts.num_edges - 1 + + def test_single_tree(self): + # 2.00┊ 6 ┊ + # ┊ ┏━┻━┓ ┊ + # 1.00┊ ┃ 5 ┊ + # ┊ ┃ ┏━┻┓ ┊ + # 0.00┊ ┃ ┃ 4 ┊ + # ┊ ┃ ┃ ┏┻┓ ┊ + # -1.00┊ 0 1 2 3 ┊ + # 0 1 + # -> + # 2.00┊ 6 ┊ + # ┊ ┻━┓ ┊ + # 1.00┊ 5 ┊ + # ┊ ┻┓ ┊ + # 0.00┊ 4 ┊ + # ┊ ┊ + # -1.00┊ 0 1 2 3 ┊ + tables = tskit.Tree.generate_comb(4, span=10).tree_sequence.dump_tables() + t = tables.nodes.time + t -= 1 + tables.nodes.time = t + ts = tables.tree_sequence() + result = tree_ops.detach_future_nodes(ts) + parent_dict = result.first().parent_dict + assert parent_dict == {4: 5, 5: 6} From 71fd049d286b556fccafbdba6ffd86605ee0ed8c Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 30 Jul 2026 08:44:24 +0000 Subject: [PATCH 03/10] Reimplement seeding as tree-based sample groups include_samples changes from a list of strains (optionally with a match_date) to a list of groups, each a list of strain IDs inserted into the ARG together as a single local tree. Rather than matching each seed individually, infer a tree over the group's haplotypes, match the haplotype of the group's inferred ancestor against the ARG with recombination disallowed, and attach the whole group at that single placement. Individual seeds from a major saltation are so far diverged from the contemporaneous ARG that matching them one at a time gives either spurious recombination or a match to some derived sample with a pile of reversions; the inferred ancestor is much closer. A group is inserted on the minimum date over its members; later-dated members keep their real dates and so get negative ("future") node times, using the same mechanism the old match_date override relied on. Seeds no longer go through the MatchDb. Their matches are constructed directly by forcing them down the ancestor's path, which removes the hmm_match.cost = 0.5 hack and the risk of a group being split apart by add_matching_results' (path, immediate_reversions) grouping key. Main changes: - tree_ops.infer_binary_topology(ts) now returns the (pi, tau) oriented forest instead of mutating a table collection, and returns None for fewer than two samples. tree_ops.infer_binary(ts, topology=None) accepts a pre-inferred topology, so a group's topology can be inferred once against the reference and reused when its mutations are re-mapped against the matched path. - New inference helpers: node_haplotypes, reference_haplotype, path_haplotype, flat_group_ts and force_match. - normalise_include_samples is replaced by check_include_samples and seed_group_dates. - add_matching_results is split into the MatchDb query/grouping and add_sample_groups, which add_seed_groups also uses. - create_mask_table is now called unconditionally in _extend, since a day can consist solely of seeds and the retrospective query must not run against a stale (or absent) used_samples table. Backwards compatibility with the old include_samples format is deliberately not preserved. --- CHANGELOG.md | 21 +- docs/example_config.toml | 14 +- sc2ts/inference.py | 500 ++++++++++++++++++++++++++++++--------- sc2ts/tree_ops.py | 36 ++- tests/test_cli.py | 29 +-- tests/test_inference.py | 262 +++++++++++--------- tests/test_tree_ops.py | 64 +++++ 7 files changed, 671 insertions(+), 255 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ae15f7a..78504e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,15 +4,18 @@ In development -- Seed samples listed in `include_samples` are now matched separately from the - daily samples, with recombination effectively disallowed. These samples are - 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. +- `include_samples` is now a list of seed *groups*, each a list of strain IDs + that are inserted into the ARG together as a single local tree. A tree is + inferred over the group's haplotypes, the haplotype of the group's inferred + ancestor is matched against the ARG with recombination disallowed, and the + whole group is then attached at that single placement. Seed samples are + usually far diverged from the current ARG, so matching each one individually + gave spurious recombinations or matches to derived samples with large + numbers of reversions; the inferred ancestor is much closer to the + contemporaneous ARG. A group is inserted on the minimum date over its + members, and members with later dates keep their real dates via negative + ("in the future") node times. The previous format (bare strain IDs, or + `(strain, match_date)` tuples) is no longer supported. - 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 diff --git a/docs/example_config.toml b/docs/example_config.toml index 0ee7863..e9aa637 100644 --- a/docs/example_config.toml +++ b/docs/example_config.toml @@ -82,8 +82,18 @@ num_threads=-1 # limit. memory_limit=32 -# A list of sample IDs (strings) for unconditional inclusion (e.g., to -# help seed major saltation events). +# A list of "seed" groups for unconditional inclusion (e.g., to help seed +# major saltation events). Each entry is a list of sample IDs that are +# inserted into the ARG together as a single local tree: a tree is inferred +# over the group's haplotypes, the haplotype of the group's inferred ancestor +# is matched against the ARG without recombination, and the whole group is +# attached at that single placement. +# +# A group is inserted on the *minimum* date over its members. Members with +# later dates keep their real dates, and so are given negative ("in the +# future") node times until their real date is reached. +# +# include_samples=[["BA.1_strain_a", "BA.1_strain_b"], ["BA.2_strain_c"]] include_samples=[] # Override specific parameter values over a time period. diff --git a/sc2ts/inference.py b/sc2ts/inference.py index fc86863..959f14c 100644 --- a/sc2ts/inference.py +++ b/sc2ts/inference.py @@ -502,25 +502,54 @@ def preprocess( return samples -def normalise_include_samples(include_samples): +def check_include_samples(include_samples, metadata): """ - 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. + Check the ``include_samples`` seed specification and return it as a list of + tuples of strain IDs, one tuple per seed group. + + Each entry must be a non-empty list of strain IDs which are inserted into + the ARG together as a single group. Every strain must be present in + ``metadata``, and no strain may appear in more than one group. """ - normalised = [] + groups = [] + seen = {} 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 + if isinstance(entry, str) or not isinstance(entry, (list, tuple)): + raise ValueError( + f"Each include_samples entry must be a list of strain IDs, not {entry!r}" + ) + if len(entry) == 0: + raise ValueError("include_samples groups must not be empty") + for strain in entry: + if not isinstance(strain, str): + raise ValueError(f"Seed strain IDs must be strings, not {strain!r}") + group = tuple(entry) + for strain in group: + if strain in seen: + raise ValueError( + f"Seed sample {strain} appears in more than one " + "include_samples group" + ) + seen[strain] = group + groups.append(group) + + missing = sorted(strain for strain in seen if strain not in metadata) + if len(missing) > 0: + raise ValueError(f"Seed samples not in dataset: {missing}") + return groups + + +def seed_group_dates(include_samples, metadata): + """ + Return a list of ``(date, strains)`` tuples, one per seed group, where + ``date`` is the minimum date over the group's members. The whole group is + inserted into the ARG on that date; members with later dates keep their + real dates and so get negative ("in the future") node times. + """ + return [ + (min(metadata[strain]["date"] for strain in group), group) + for group in include_samples + ] def extend( @@ -552,16 +581,21 @@ def extend( 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. + ``include_samples`` is an optional list of "seed" groups that are inserted + unconditionally and without recombination. Each entry is a list of strain + IDs which are inserted together as a single local tree: a tree is inferred + over the group's haplotypes, the haplotype of the group's inferred ancestor + is matched against the ARG, and the whole group is then attached at that + single placement. This places major saltations far better than matching + each seed individually, since the inferred ancestor is much closer to the + contemporaneous ARG than any of the group's leaves. + + A group is inserted on the *minimum* date over its members. Members with + later dates keep their real dates and so are given negative ("in the + future") node times relative to that day's time-zero, becoming visible to + the matcher once their real date is reached. Because the insertion date is + always one of the group's own dataset dates, it is a date the pipeline + processes, unless it falls outside the run's date window. """ if num_mismatches is None: num_mismatches = 3 @@ -587,7 +621,6 @@ 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) @@ -601,19 +634,7 @@ 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" - ) + include_samples = check_include_samples(include_samples, ds.metadata) with MatchDb(match_db) as matches: tables = _extend( @@ -675,26 +696,26 @@ 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 + # A seed group is inserted on the minimum date over its members, so a seed + # is processed on its group's date rather than on its own date. + group_dates = seed_group_dates(include_samples, dataset.metadata) + seed_date = { + strain: group_date for group_date, group in group_dates for strain in group } + todays_groups = [group for group_date, group in group_dates if group_date == date] + todays_seeds = {strain for group in todays_groups for strain in group} 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 + # Exclude a seed from any date other than its group's date; it is + # processed only once, on that date. + if seed_date.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: + # Inject today's seeds that aren't naturally sampled today. Missing strains + # are rejected up-front in extend(). + for strain in todays_seeds: + if strain not in metadata_matches: metadata_matches[strain] = dataset.metadata[strain] logger.info(f"Got {len(metadata_matches)} metadata matches") @@ -710,7 +731,7 @@ def _extend( pango_lineage_key = "Viridian_pangolin" scorpio_key = "Viridian_scorpio" - unconditional_include_samples = [] + seeds = {} samples = [] for s in preprocessed_samples: if s.haplotype is None: @@ -727,9 +748,10 @@ def _extend( f"Encoded {s.strain} {s.scorpio} {s.pango} missing={num_missing_sites} " f"deletions={num_deletion_sites}" ) - if s.strain in include_strains: + if s.strain in todays_seeds: + # Seeds bypass the max_missing_sites and max_daily_samples filters. s.flags |= core.NODE_IS_UNCONDITIONALLY_INCLUDED - unconditional_include_samples.append(s) + seeds[s.strain] = s elif num_missing_sites <= max_missing_sites: samples.append(s) else: @@ -737,6 +759,16 @@ def _extend( f"Filter {s.strain}: missing={num_missing_sites} > {max_missing_sites}" ) + # Seeds without a stored alignment are dropped from their group, and a + # group left with no members is skipped entirely. + todays_seed_groups = [] + for group in todays_groups: + group_samples = [seeds[strain] for strain in group if strain in seeds] + if len(group_samples) == 0: + logger.warning(f"Skipping seed group with no usable samples: {list(group)}") + else: + todays_seed_groups.append(group_samples) + if max_daily_samples is not None: if max_daily_samples < len(samples): seed_prefix = bytes(np.array([random_seed], dtype=int).data) @@ -746,48 +778,24 @@ def _extend( samples = rng.sample(samples, max_daily_samples) ts = increment_time(date, base_ts) - if len(samples) + len(unconditional_include_samples) > 0: - if len(samples) > 0: - match_samples( - date, - samples, - base_ts=base_ts, - num_mismatches=num_mismatches, - deletions_as_missing=deletions_as_missing, - show_progress=show_progress, - num_threads=num_threads, - memory_limit=memory_limit, - ) - if len(unconditional_include_samples) > 0: - # Seed samples are usually far diverged from the current ARG, and - # matching them with the standard num_mismatches gives spurious - # recombinations. Match them without recombination instead. - match_samples( - date, - unconditional_include_samples, - base_ts=base_ts, - num_mismatches=NO_RECOMBINATION_NUM_MISMATCHES, - deletions_as_missing=deletions_as_missing, - show_progress=show_progress, - num_threads=num_threads, - memory_limit=memory_limit, - ) - - samples = samples + unconditional_include_samples - samples.sort(key=lambda s: s.strain) - + # A day can consist solely of seeds, but the retrospective query below must + # not run against a stale mask table, so this is unconditional. + match_db.create_mask_table(base_ts) + if len(samples) > 0: + match_samples( + date, + samples, + base_ts=base_ts, + num_mismatches=num_mismatches, + deletions_as_missing=deletions_as_missing, + show_progress=show_progress, + num_threads=num_threads, + memory_limit=memory_limit, + ) characterise_match_mutations(base_ts, samples) characterise_recombinants(base_ts, samples) - for sample in unconditional_include_samples: - # We want this sample to included unconditionally, so we set the - # hmm cost to 0 < hmm_cost < hmm_cost_threshold. We use 0.5 - # arbitrarily here to distinguish it from real one-mutation - sample.hmm_match.cost = 0.5 - logger.warning(f"Unconditionally including {sample.summary()}") - match_db.add(samples, date, show_progress) - match_db.create_mask_table(base_ts) ts = add_exact_matches(ts=ts, match_db=match_db, date=date) @@ -802,6 +810,24 @@ def _extend( phase="close", ) + if len(todays_seed_groups) > 0: + # Seeds are attached directly rather than via the MatchDb, so that a + # group can't be split up and can't be added a second time. + ts, seed_samples = add_seed_groups( + ts, + base_ts, + todays_seed_groups, + date, + num_mismatches=num_mismatches, + deletions_as_missing=deletions_as_missing, + show_progress=show_progress, + num_threads=num_threads, + memory_limit=memory_limit, + ) + samples = samples + seed_samples + + samples.sort(key=lambda s: s.strain) + logger.info("Looking for retrospective matches") assert min_group_size is not None earliest_date = parse_date(date) - datetime.timedelta(days=retrospective_window) @@ -935,6 +961,46 @@ def match_path_ts(group, sequence_length): return tables.tree_sequence() +def flat_group_ts(ts, haplotypes, ancestral_haplotype, deletions_as_missing=False): + """ + Return a "flat" (star) tree sequence in which each of the specified + haplotypes is a sample node hanging off a root carrying + ``ancestral_haplotype``. + + Site positions and ancestral states are taken from ``ts``, and a mutation + is added wherever a haplotype's non-missing allele differs from + ``ancestral_haplotype``. Missing data therefore reads as the ancestral + state, exactly as it does in the trees built by ``match_path_ts``. + """ + tables = tskit.TableCollection(ts.sequence_length) + tables.nodes.metadata_schema = tskit.MetadataSchema.permissive_json() + sites_position = ts.sites_position + sites_ancestral_state = ts.sites_ancestral_state + root = len(haplotypes) + site_id_map = {} + for h in haplotypes: + node_id = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) + tables.edges.add_row(0, tables.sequence_length, parent=root, child=node_id) + if deletions_as_missing: + h = np.where(h == DELETION, MISSING, h) + (sites,) = np.where((h != ancestral_haplotype) & (h != MISSING)) + for site_id in sites: + if site_id not in site_id_map: + site_id_map[site_id] = tables.sites.add_row( + sites_position[site_id], sites_ancestral_state[site_id] + ) + tables.mutations.add_row( + site=site_id_map[site_id], + node=node_id, + time=0, + derived_state=core.IUPAC_ALLELES[h[site_id]], + ) + # add the root + tables.nodes.add_row(time=1) + tables.sort() + return tables.tree_sequence() + + 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}") @@ -1015,6 +1081,10 @@ class SampleGroup: immediate_reversions: List = None sample_hash: str = None tree_quality_metrics: GroupTreeQualityMetrics = None + # Optionally, a pre-inferred (pi, tau) topology to use when building the + # group's local tree instead of inferring one. Must be trailing, because + # add_matching_results constructs SampleGroups positionally. + topology: tuple = None def __post_init__(self): m = hashlib.md5() @@ -1067,26 +1137,15 @@ def add_tree_quality_metrics(self, ts, date): return self.tree_quality_metrics -def add_matching_results( - where_clause, - match_db, - ts, - date, - min_group_size=1, - min_different_dates=1, - 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, -): +def add_matching_results(where_clause, match_db, ts, date, **kwargs): + """ + Query the MatchDb, group the resulting matches by their path and set of + immediate reversions, and add each group into the ARG as a local tree. + """ logger.info(f"Querying match DB WHERE: {where_clause}") # Group matches by path and set of immediate reversions. grouped_matches = collections.defaultdict(list) - site_missing_samples = np.zeros(ts.num_sites, dtype=int) - site_deletion_samples = np.zeros(ts.num_sites, dtype=int) num_samples = 0 for sample in match_db.get(where_clause): assert all(mut.is_reversion is not None for mut in sample.hmm_match.mutations) @@ -1116,7 +1175,29 @@ def add_matching_results( for key, samples in grouped_matches.items() ] logger.info(f"Got {len(groups)} groups for {num_samples} samples") + return add_sample_groups(ts, groups, date, **kwargs) + +def add_sample_groups( + ts, + groups, + date, + *, + min_group_size=1, + min_different_dates=1, + 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, +): + """ + Add each of the specified SampleGroups into the ARG as a local tree, + subject to the given quality thresholds. + """ + site_missing_samples = np.zeros(ts.num_sites, dtype=int) + site_deletion_samples = np.zeros(ts.num_sites, dtype=int) tables = ts.dump_tables() attach_nodes = [] @@ -1143,7 +1224,7 @@ def add_matching_results( if flat_ts.num_mutations == 0 or flat_ts.num_samples == 1: poly_ts = flat_ts else: - binary_ts = tree_ops.infer_binary(flat_ts) + binary_ts = tree_ops.infer_binary(flat_ts, group.topology) 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, date) @@ -1207,6 +1288,123 @@ def add_matching_results( return ts, added_groups +def add_seed_groups( + ts, + base_ts, + seed_groups, + date, + *, + num_mismatches, + deletions_as_missing=False, + show_progress=False, + num_threads=0, + memory_limit=-1, +): + """ + Add the specified groups of seed samples into the ARG, one local tree per + group. + + Rather than matching each seed against the ARG individually, we infer a + tree over each group's haplotypes, match the haplotype of the group's + inferred ancestor (with recombination disallowed), and then rebuild the + group's tree against that path so that the whole group is attached in one + piece. Seeds never go through the MatchDb: their matches are constructed + directly by forcing them down the ancestor's path. + + Returns the updated tree sequence and the list of seed samples added. + """ + reference = reference_haplotype(base_ts) + root_samples = [] + pending = [] + for samples in seed_groups: + flat_ts = flat_group_ts( + base_ts, + [s.haplotype for s in samples], + reference, + deletions_as_missing=deletions_as_missing, + ) + topology = None + if flat_ts.num_samples > 1 and flat_ts.num_mutations > 0: + topology = tree_ops.infer_binary_topology(flat_ts) + + if topology is None: + # Either a single sample, or a group that is identical to the + # reference at every non-missing site. + if len(samples) == 1: + root_haplotype = samples[0].haplotype.copy() + else: + root_haplotype = reference.copy() + else: + binary_ts = tree_ops.infer_binary(flat_ts, topology) + tree = binary_ts.first() + # The single child of the outgroup root is the group's MRCA. + mrca = tree.children(tree.root)[0] + root_haplotype = reference.copy() + site_index = np.searchsorted( + base_ts.sites_position, binary_ts.sites_position + ) + root_haplotype[site_index] = node_haplotypes(binary_ts, [mrca])[0] + + # Sites at which every member is missing are unknown at the root, so + # let the HMM ignore them rather than asserting the reference allele. + H = np.array([s.haplotype for s in samples]) + root_haplotype[np.all(H == MISSING, axis=0)] = MISSING + + group = SampleGroup(samples, topology=topology) + root_sample = Sample( + strain=f"seed_root_{group.sample_hash}", + date=date, + haplotype=root_haplotype, + ) + root_samples.append(root_sample) + pending.append((group, root_sample)) + + # Seed samples are usually far diverged from the current ARG, and matching + # them with the standard num_mismatches gives spurious recombinations. + # Match the inferred ancestors without recombination instead. + match_samples( + date, + root_samples, + base_ts=base_ts, + num_mismatches=NO_RECOMBINATION_NUM_MISMATCHES, + deletions_as_missing=deletions_as_missing, + show_progress=show_progress, + num_threads=num_threads, + memory_limit=memory_limit, + ) + + groups = [] + seed_samples = [] + for group, root_sample in pending: + path = root_sample.hmm_match.path + parent_haplotype = path_haplotype(base_ts, path) + for sample in group: + force_match( + sample, + path, + parent_haplotype, + base_ts.sites_position, + deletions_as_missing, + num_mismatches, + ) + logger.warning(f"Unconditionally including {sample.summary()}") + group.path = tuple(path) + group.immediate_reversions = () + groups.append(group) + seed_samples.extend(group.samples) + + characterise_recombinants(base_ts, seed_samples) + ts, _ = add_sample_groups( + ts, + groups, + date, + min_group_size=1, + show_progress=show_progress, + phase="seed", + ) + return ts, seed_samples + + def solve_num_mismatches(k, num_alleles=5): """ Return the low-level LS parameters corresponding to accepting @@ -1761,6 +1959,84 @@ def extract_haplotypes(ts, samples): return ret +def node_haplotypes(ts, nodes): + """ + Return the haplotype of each of the specified nodes, using the integer + allele encoding defined by ``core.IUPAC_ALLELES``. + + This is the encoding used by ``Sample.haplotype``. Note that + ``extract_haplotypes`` is *not* equivalent: it returns per-site allele + indexes local to each site rather than IUPAC codes. + """ + # Annoyingly tskit doesn't allow us to specify duplicate samples, which can + # happen perfectly well here, so we must work around. + unique_nodes = list(set(nodes)) + H = ts.genotype_matrix( + samples=unique_nodes, + alleles=tuple(core.IUPAC_ALLELES), + isolated_as_missing=False, + ).T + return [H[unique_nodes.index(node_id)] for node_id in nodes] + + +def reference_haplotype(ts): + """ + Return the haplotype of the ARG's ancestral (reference) sequence, i.e. the + ancestral state of each site, in the ``core.IUPAC_ALLELES`` encoding. + """ + return jit.encode_alleles(np.asarray(ts.sites_ancestral_state, dtype="U1")) + + +def path_haplotype(ts, path): + """ + Return the haplotype copied along the specified HMM path, in the + ``core.IUPAC_ALLELES`` encoding. + """ + H = node_haplotypes(ts, [seg.parent for seg in path]) + h = np.zeros(ts.num_sites, dtype=np.int8) + for seg, parent_haplotype in zip(path, H): + left, right = np.searchsorted(ts.sites_position, [seg.left, seg.right]) + h[left:right] = parent_haplotype[left:right] + return h + + +def force_match( + sample, + path, + parent_haplotype, + sites_position, + deletions_as_missing, + num_mismatches, +): + """ + Set the sample's ``hmm_match`` to the match obtained by forcing it down the + specified path, with one mutation at each site where the sample's + non-missing haplotype differs from the haplotype copied along the path. + """ + h = sample.haplotype + if deletions_as_missing: + h = np.where(h == DELETION, MISSING, h) + (sites,) = np.where( + (h != parent_haplotype) & (h != MISSING) & (parent_haplotype != MISSING) + ) + mutations = [ + MatchMutation( + derived_state=core.IUPAC_ALLELES[h[site_id]], + inherited_state=core.IUPAC_ALLELES[parent_haplotype[site_id]], + site_id=int(site_id), + site_position=int(sites_position[site_id]), + # A forced path has no HMM-match artefacts, so there are no + # reversions to mark up for tree building. Parsimony over the full + # haplotypes handles reversions on the group's root branch instead. + is_reversion=False, + is_immediate_reversion=False, + ) + for site_id in sites + ] + sample.hmm_match = HmmMatch(list(path), mutations) + sample.hmm_match.compute_cost(num_mismatches) + + def characterise_recombinants(ts, samples): """ Update the metadata for any recombinants to add interval information to the metadata. diff --git a/sc2ts/tree_ops.py b/sc2ts/tree_ops.py index 63d701c..749a8ca 100644 --- a/sc2ts/tree_ops.py +++ b/sc2ts/tree_ops.py @@ -148,12 +148,23 @@ def add_tree_to_tables(tables, pi, tau): tables.edges.add_row(0, L, parent, u) -def infer_binary_topology(ts, tables): +def infer_binary_topology(ts): + """ + Infer a strictly binary topology from the variation data in the specified + tree sequence, returning it as an oriented forest ``(pi, tau)``. + + The root of ``ts`` is included as an outgroup in the tree building so that + the unrooted NJ tree can be rooted; node ``ts.num_samples`` in the returned + arrays is that outgroup, and is the root of the returned topology. + + Returns ``None`` if there are fewer than two samples, in which case there + is no topology to infer. + """ assert ts.num_trees == 1 assert ts.num_mutations > 0 if ts.num_samples < 2: - return tables.tree_sequence() + return None samples = ts.samples() tree = ts.first() @@ -176,20 +187,23 @@ def infer_binary_topology(ts, tables): # Node n - 1 is the pre-specified root, so force rerooting around that. reroot(pi, n - 1) - assert n == len(tables.nodes) + 1 + assert n == ts.num_samples + 1 tau = max_leaf_distance(pi, n) tau /= max(1, np.max(tau)) - add_tree_to_tables(tables, pi, tau) - tables.sort() - - return tables.tree_sequence() + return pi, tau # TODO rename this to infer_sample_group_tree -def infer_binary(ts): +def infer_binary(ts, topology=None): """ Infer a strictly binary tree from the variation data in the specified tree sequence. + + If ``topology`` is not None it must be an oriented forest ``(pi, tau)`` as + returned by :func:`infer_binary_topology`, and is used directly instead of + inferring the topology from ``ts``. This lets a group's topology be + inferred once and then reused when the mutations are re-mapped against a + different ancestral haplotype. """ assert ts.num_trees == 1 assert list(ts.samples()) == list(range(ts.num_samples)) @@ -202,7 +216,11 @@ def infer_binary(ts): tables.nodes.truncate(ts.num_samples) # Update the tables with the topology - infer_binary_topology(ts, tables) + if topology is None: + topology = infer_binary_topology(ts) + if topology is not None: + add_tree_to_tables(tables, *topology) + tables.sort() binary_ts = tables.tree_sequence() # Now add on mutations under parsimony diff --git a/tests/test_cli.py b/tests/test_cli.py index b104a3c..19bb1cf 100644 --- a/tests/test_cli.py +++ b/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"], + include_samples=[["SRR14631544"]], ) runner = ct.CliRunner() result = runner.invoke( @@ -446,7 +446,7 @@ def test_include_samples_missing_strain(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"], ["NO_SUCH_STRAIN"]], ) runner = ct.CliRunner() with pytest.raises(ValueError, match="not in dataset"): @@ -456,27 +456,22 @@ def test_include_samples_missing_strain(self, tmp_path, fx_ts_map, fx_dataset): 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. + def test_include_samples_bare_string(self, tmp_path, fx_ts_map, fx_dataset): + # A bare strain ID is the old format and is rejected: TOML entries must + # be arrays of strain IDs. config_file = self.make_config( tmp_path, fx_dataset, exclude_sites=[56, 57, 58, 59, 60], - include_samples=[["SRR14631544", "2020-01-01"]], + include_samples=["SRR14631544"], ) 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 + with pytest.raises(ValueError, match="must be a list of strain IDs"): + runner.invoke( + cli.cli, + f"infer {config_file} --stop 2020-01-02", + catch_exceptions=False, + ) def test_override(self, tmp_path, fx_ts_map, fx_dataset): hmm_cost_threshold = 47 diff --git a/tests/test_inference.py b/tests/test_inference.py index b378b71..3175514 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -482,20 +482,66 @@ def test_high_recomb_mutation(self): self.check_double_mirror(ts) -class TestNormaliseIncludeSamples: +class TestCheckIncludeSamples: + metadata = { + "a": {"date": "2021-01-03"}, + "b": {"date": "2021-01-01"}, + "c": {"date": "2021-01-02"}, + } + def test_empty(self): - assert si.normalise_include_samples([]) == [] + assert si.check_include_samples([], self.metadata) == [] + + def test_groups(self): + result = si.check_include_samples([["a", "b"], ["c"]], self.metadata) + assert result == [("a", "b"), ("c",)] + + def test_tuples_accepted(self): + assert si.check_include_samples([("a",)], self.metadata) == [("a",)] + + def test_bare_string_raises(self): + with pytest.raises(ValueError, match="must be a list of strain IDs"): + si.check_include_samples(["a"], self.metadata) - def test_bare_strings(self): - assert si.normalise_include_samples(["a", "b"]) == [("a", None), ("b", None)] + def test_empty_group_raises(self): + with pytest.raises(ValueError, match="must not be empty"): + si.check_include_samples([[]], self.metadata) - def test_tuples_and_lists(self): - result = si.normalise_include_samples([("a", "2021-01-01"), ["b", None], "c"]) - assert result == [("a", "2021-01-01"), ("b", None), ("c", None)] + def test_non_string_strain_raises(self): + with pytest.raises(ValueError, match="must be strings"): + si.check_include_samples([["a", 7]], self.metadata) - def test_malformed_date_raises(self): - with pytest.raises(ValueError, match="isoformat"): - si.normalise_include_samples([("a", "not-a-date")]) + def test_duplicate_strain_raises(self): + with pytest.raises(ValueError, match="more than one include_samples group"): + si.check_include_samples([["a", "b"], ["a"]], self.metadata) + + def test_duplicate_within_group_raises(self): + with pytest.raises(ValueError, match="more than one include_samples group"): + si.check_include_samples([["a", "a"]], self.metadata) + + def test_missing_strain_raises(self): + with pytest.raises(ValueError, match="not in dataset"): + si.check_include_samples([["a", "nosuchstrain"]], self.metadata) + + +class TestSeedGroupDates: + metadata = { + "a": {"date": "2021-01-03"}, + "b": {"date": "2021-01-01"}, + "c": {"date": "2021-01-02"}, + } + + def test_group_gets_minimum_date(self): + result = si.seed_group_dates([("a", "b")], self.metadata) + assert result == [("2021-01-01", ("a", "b"))] + + def test_singleton_gets_own_date(self): + result = si.seed_group_dates([("a",)], self.metadata) + assert result == [("2021-01-03", ("a",))] + + def test_multiple_groups(self): + result = si.seed_group_dates([("a",), ("b", "c")], self.metadata) + assert result == [("2021-01-03", ("a",)), ("2021-01-01", ("b", "c"))] class TestRealData: @@ -616,29 +662,18 @@ def test_2020_02_02(self, tmp_path, fx_ts_map, fx_dataset, num_threads): assert "SRR11597115" not in ts.metadata["sc2ts"]["samples_strain"] ts.tables.assert_equals(fx_ts_map["2020-02-02"].tables, ignore_provenance=True) - @pytest.mark.parametrize( - "include_samples", - ( - ["SRR11597115"], - # The tuple form with a None match date is equivalent to the - # bare-string form. - [("SRR11597115", None)], - [("SRR11597115", "2020-02-02")], - ), - ) def test_2020_02_02_include_samples( self, tmp_path, fx_ts_map, fx_dataset, - include_samples, ): ts = run_extend( dataset=fx_dataset, base_ts=fx_ts_map["2020-02-01"], date="2020-02-02", match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=include_samples, + include_samples=[["SRR11597115"]], ) assert ts.metadata["sc2ts"]["cumulative_stats"]["exact_matches"]["pango"] == { "A": 2, @@ -658,75 +693,107 @@ def test_2020_02_02_include_samples( assert edges[0].left == 0 assert edges[0].right == ts.sequence_length - def test_seed_early_match_date_negative_time(self, tmp_path, fx_ts_map, fx_dataset): - # SRR11597115's actual date is 2020-02-02; match it in two days early. - strain = "SRR11597115" + def test_seed_group_inserted_on_minimum_date(self, tmp_path, fx_ts_map, fx_dataset): + # SRR11494548 is dated 2020-01-31 and SRR11597115 2020-02-02. As one + # group they go in together on the minimum date, 2020-01-31, so the + # later-dated member sits two days in the future at time -2. + group = ["SRR11494548", "SRR11597115"] ts = run_extend( dataset=fx_dataset, base_ts=fx_ts_map["2020-01-30"], date="2020-01-31", match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=[(strain, "2020-01-31")], + include_samples=[group], ) - assert strain in ts.metadata["sc2ts"]["samples_strain"] - u = ts.samples()[ts.metadata["sc2ts"]["samples_strain"].index(strain)] - assert ts.nodes_flags[u] & sc2ts.NODE_IS_UNCONDITIONALLY_INCLUDED > 0 - # The node retains its actual date, two days in the future relative to - # the match-day time-zero, so its time is -2. - assert ts.nodes_time[u] == -2 - # Matched without recombination: a single full-span parent edge, and - # the parent must be older (larger time) than the future-dated seed. - assert ts.nodes_flags[u] & sc2ts.NODE_IS_RECOMBINANT == 0 - edges = [e for e in ts.edges() if e.child == u] + strains = ts.metadata["sc2ts"]["samples_strain"] + nodes = {} + for strain in group: + assert strain in strains + u = ts.samples()[strains.index(strain)] + nodes[strain] = u + assert ts.nodes_flags[u] & sc2ts.NODE_IS_UNCONDITIONALLY_INCLUDED > 0 + assert ts.nodes_flags[u] & sc2ts.NODE_IS_RECOMBINANT == 0 + assert ts.nodes_time[nodes["SRR11494548"]] == 0 + assert ts.nodes_time[nodes["SRR11597115"]] == -2 + # Both members belong to the same sample group. + group_ids = {ts.node(u).metadata["sc2ts"]["group_id"] for u in nodes.values()} + assert len(group_ids) == 1 + + def test_seed_group_attached_as_one_tree(self, tmp_path, fx_ts_map, fx_dataset): + # The whole point of grouping: members that share an ancestor are hung + # off a single internal node, and it is that node which gets placed + # against the ARG. ERR4206180 (2020-02-09) and ERR4206593 (2020-02-13) + # share five derived alleles, so their inferred ancestor is distinct + # from both of them and from the reference. + group = ["ERR4206180", "ERR4206593"] + base_ts = fx_ts_map["2020-02-08"] + ts = run_extend( + dataset=fx_dataset, + base_ts=base_ts, + date="2020-02-09", + match_db=si.MatchDb.initialise(tmp_path / "match.db"), + include_samples=[group], + ) + strains = ts.metadata["sc2ts"]["samples_strain"] + nodes = [ts.samples()[strains.index(strain)] for strain in group] + group_id = ts.node(nodes[0]).metadata["sc2ts"]["group_id"] + assert ts.nodes_time[nodes[0]] == 0 + assert ts.nodes_time[nodes[1]] == -4 + + tree = ts.first() + parents = {tree.parent(u) for u in nodes} + assert len(parents) == 1 + mrca = parents.pop() + # The shared parent is an internal node belonging to the same group. + assert mrca not in nodes + assert ts.node(mrca).metadata["sc2ts"]["group_id"] == group_id + # It hangs off the pre-existing ARG by a single full-span edge. + edges = [e for e in ts.edges() if e.child == mrca] assert len(edges) == 1 assert edges[0].left == 0 assert edges[0].right == ts.sequence_length - assert ts.nodes_time[edges[0].parent] > ts.nodes_time[u] - - def test_seed_match_date_equals_actual_date(self, tmp_path, fx_ts_map, fx_dataset): - # An override match date equal to the sample's actual date behaves like - # an ordinary seed: the node is present-dated (time 0), not in the - # future. SRR11597115's actual date is 2020-02-02. - strain = "SRR11597115" + assert edges[0].parent < base_ts.num_nodes + # The mutations separating the group's inferred root from its match sit + # on that node, not on the individual members. + assert np.sum(ts.mutations_node == mrca) > 0 + + def test_two_seed_groups_stay_separate(self, tmp_path, fx_ts_map, fx_dataset): + # Two independent single-strain groups on the same date must not be + # merged into one group. ts = run_extend( dataset=fx_dataset, base_ts=fx_ts_map["2020-02-01"], date="2020-02-02", match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=[(strain, "2020-02-02")], + include_samples=[["SRR11597115"], ["SRR11597190"]], ) - assert strain in ts.metadata["sc2ts"]["samples_strain"] - u = ts.samples()[ts.metadata["sc2ts"]["samples_strain"].index(strain)] - assert ts.nodes_flags[u] & sc2ts.NODE_IS_UNCONDITIONALLY_INCLUDED > 0 - # match date == actual date, so the node sits at time zero. - assert ts.nodes_time[u] == 0 - # Still matched without recombination: a single full-span parent edge. - assert ts.nodes_flags[u] & sc2ts.NODE_IS_RECOMBINANT == 0 - edges = [e for e in ts.edges() if e.child == u] - assert len(edges) == 1 - assert edges[0].left == 0 - assert edges[0].right == ts.sequence_length - - def test_seed_early_match_date_evolves_over_days( - self, tmp_path, fx_ts_map, fx_dataset - ): - # Inject the seed two days before its actual date, then keep extending - # past the actual date. Its node time must rise by +1/day and hit 0 on - # the actual date, exercising the low-level matcher against a base_ts - # that contains negative ("future") node times. - strain = "SRR11597115" - include_samples = [(strain, "2020-01-31")] + strains = ts.metadata["sc2ts"]["samples_strain"] + group_ids = set() + for strain in ["SRR11597115", "SRR11597190"]: + assert strain in strains + u = ts.samples()[strains.index(strain)] + assert ts.nodes_flags[u] & sc2ts.NODE_IS_UNCONDITIONALLY_INCLUDED > 0 + group_ids.add(ts.node(u).metadata["sc2ts"]["group_id"]) + assert len(group_ids) == 2 + + def test_seed_group_evolves_over_days(self, tmp_path, fx_ts_map, fx_dataset): + # Insert the group on its minimum date, then keep extending past the + # later member's actual date. That member's time must rise by +1/day and + # hit 0 on its actual date, exercising the low-level matcher against a + # base_ts containing negative ("future") node times. It must also not be + # re-added on its natural date. + group = ["SRR11494548", "SRR11597115"] + include_samples = [group] base_path = tmp_path / "base.ts" fx_ts_map["2020-01-30"].dump(base_path) match_db = si.MatchDb.initialise(tmp_path / "match.db") - dates = ["2020-01-31", "2020-02-01", "2020-02-02", "2020-02-03"] expected_time = { "2020-01-31": -2, "2020-02-01": -1, "2020-02-02": 0, "2020-02-03": 1, } - for date in dates: + for date, expected in expected_time.items(): ts = si.extend( dataset=fx_dataset.path, base_ts=base_path, @@ -736,54 +803,26 @@ def test_seed_early_match_date_evolves_over_days( ) ts.dump(base_path) strains = ts.metadata["sc2ts"]["samples_strain"] - # The seed is added exactly once and never double-processed on its - # natural date. - assert strains.count(strain) == 1 - u = ts.samples()[strains.index(strain)] - assert ts.nodes_time[u] == expected_time[date] - - def test_seed_excluded_on_natural_date(self, tmp_path, fx_ts_map, fx_dataset): - # A seed with an override match date must NOT be processed on its - # actual date. SRR11597115 is naturally sampled on 2020-02-02; with an - # override of 2020-01-31 it should be absent when we extend 2020-02-02. - strain = "SRR11597115" - ts = run_extend( - dataset=fx_dataset, - base_ts=fx_ts_map["2020-02-01"], - date="2020-02-02", - match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=[(strain, "2020-01-31")], - ) - assert strain not in ts.metadata["sc2ts"]["samples_strain"] - - def test_seed_override_date_no_samples(self, tmp_path, fx_ts_map, fx_dataset): - # An override match date with no naturally-sampled strains that day is - # never reached by the driver, so the seed is silently not injected. - # Here we simply confirm extend on such a date doesn't crash and the - # seed isn't added. - strain = "SRR11597115" - ts = run_extend( - dataset=fx_dataset, - base_ts=fx_ts_map["2020-02-01"], - date="2020-02-02", - match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=[(strain, "2020-02-12")], - ) - assert strain not in ts.metadata["sc2ts"]["samples_strain"] + # Each member is added exactly once, and never double-processed on + # its natural date. + for strain in group: + assert strains.count(strain) == 1 + u = ts.samples()[strains.index("SRR11597115")] + assert ts.nodes_time[u] == expected @pytest.mark.parametrize( "include_samples", ( - ["SRR11597115", "NOSUCHSTRAIN"], - [("NOSUCHSTRAIN", "2020-02-02")], - [("NOSUCHSTRAIN", None)], + [["SRR11597115"], ["NOSUCHSTRAIN"]], + [["SRR11597115", "NOSUCHSTRAIN"]], + [["NOSUCHSTRAIN"]], ), ) def test_seed_missing_strain_raises( self, tmp_path, fx_ts_map, fx_dataset, include_samples ): - # A seed strain that isn't in the dataset is an error, whether it's a - # bare strain or carries an override match date. + # A seed strain that isn't in the dataset is an error, whether it's on + # its own or grouped with a strain that does exist. with pytest.raises(ValueError, match="not in dataset"): run_extend( dataset=fx_dataset, @@ -793,6 +832,17 @@ def test_seed_missing_strain_raises( include_samples=include_samples, ) + def test_seed_bare_string_raises(self, tmp_path, fx_ts_map, fx_dataset): + # The old format (a bare strain ID) is no longer accepted. + with pytest.raises(ValueError, match="must be a list of strain IDs"): + run_extend( + dataset=fx_dataset, + base_ts=fx_ts_map["2020-02-01"], + date="2020-02-02", + match_db=si.MatchDb.initialise(tmp_path / "match.db"), + include_samples=["SRR11597115"], + ) + def test_2020_02_02_mutation_overlap( self, tmp_path, diff --git a/tests/test_tree_ops.py b/tests/test_tree_ops.py index 3c934cb..4d9b737 100644 --- a/tests/test_tree_ops.py +++ b/tests/test_tree_ops.py @@ -771,6 +771,70 @@ def test_simulation_root_mutations(self, n, num_mutations): self.check_properties(ts2) +def flat_ts(n, genotypes): + """ + A star tree over n samples with one site per column of genotypes, ancestral + state "A" and a mutation wherever the genotype is not "A". + """ + L = len(genotypes[0]) + 1 + tables = tskit.TableCollection(L) + root = n + for j in range(n): + u = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) + tables.edges.add_row(0, L, root, u) + tables.nodes.add_row(time=1) + for site in range(len(genotypes[0])): + tables.sites.add_row(site, "A") + for j in range(n): + if genotypes[j][site] != "A": + tables.mutations.add_row( + site=site, node=j, derived_state=genotypes[j][site] + ) + tables.sort() + return tables.tree_sequence() + + +class TestInferBinaryTopology: + def test_single_sample(self): + ts = flat_ts(1, ["T"]) + assert tree_ops.infer_binary_topology(ts) is None + + @pytest.mark.parametrize("n", [2, 3, 4, 8]) + def test_outgroup_is_root(self, n): + # Each sample gets its own private mutation. + genotypes = ["A" * j + "T" + "A" * (n - j - 1) for j in range(n)] + ts = flat_ts(n, genotypes) + pi, tau = tree_ops.infer_binary_topology(ts) + assert len(tau) == len(pi) + # The n samples and the outgroup are all leaves, so there must be + # internal nodes on top of them. + assert len(pi) > n + 1 + # Node n is the outgroup, and after rerooting it is the sole root. + assert pi[n] == -1 + assert tau[n] == 1 + assert sum(p == -1 for p in pi) == 1 + # Every sample is a leaf, i.e. nothing points at it. + assert set(pi) & set(range(n)) == set() + + def test_supplied_topology_is_used(self): + # Two well separated pairs, so the topology is unambiguous. + genotypes = ["TTAAAA", "TTAAAT", "AAATTA", "AATTTA"] + ts1 = flat_ts(4, genotypes) + topology = tree_ops.infer_binary_topology(ts1) + expected = tree_ops.infer_binary(ts1, topology) + + # A different set of mutations that would give a different topology if + # it were inferred rather than supplied. + ts2 = flat_ts(4, ["TAAAAA", "AATAAA", "AAAATA", "AAAAAT"]) + inferred = tree_ops.infer_binary(ts2) + supplied = tree_ops.infer_binary(ts2, topology) + + assert supplied.tables.edges == expected.tables.edges + assert inferred.tables.edges != expected.tables.edges + # The mutations still come from ts2, not from ts1. + assert_sequences_equal(ts2, supplied) + + class TestFromBiotite: def check_round_trip(self, tsk_tree): node_labels = {u: f"{u}" for u in tsk_tree.tree_sequence.samples()} From 405a3df98434bcfd9befe0ab68b6c88668fbc1d3 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 30 Jul 2026 10:28:00 +0100 Subject: [PATCH 04/10] Fix retrospective matching on days with no new samples create_mask_table records the samples already in the ARG so that the retrospective query skips them, but it was only called when there were samples to match. On a day with no new samples the query therefore ran against the previous day's table, and could add a second copy of a sample that had just been inserted. The mask table is now rebuilt unconditionally in _extend. Update the retro group tests, which were pinning the old behaviour: the 2020-02-15 run has no samples of its own, so ERR4204459 and ERR4206593 were being re-added despite already being samples in the 2020-02-13 base ARG. Assert that samples_strain has no duplicates so this can't regress silently. --- CHANGELOG.md | 5 +++++ tests/test_inference.py | 14 ++++++++++---- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 78504e7..8769b2b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,11 @@ In development ("in the future") node times. The previous format (bare strain IDs, or `(strain, match_date)` tuples) is no longer supported. +- Fix retrospective matching on a day with no new samples. The table of samples + already in the ARG was only rebuilt when there were samples to match, so on + such a day the retrospective query ran against the previous day's table and + could add a second copy of a sample that had just been inserted. + - 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 diff --git a/tests/test_inference.py b/tests/test_inference.py index 3175514..7f67462 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -1056,7 +1056,8 @@ def test_2020_02_14_all_matches(self, tmp_path, fx_ts_map, fx_dataset, fx_match_ min_different_dates=1, ) retro_groups = ts.metadata["sc2ts"]["retro_groups"] - assert len(retro_groups) == 6 + # Everything in the match DB that isn't already a sample in the base ARG. + assert len(retro_groups) == 4 assert retro_groups[0] == { "dates": ["2020-01-29"], "depth": 1, @@ -1069,6 +1070,11 @@ def test_2020_02_14_all_matches(self, tmp_path, fx_ts_map, fx_dataset, fx_match_ "strains": ["SRR15736313"], "date_added": "2020-02-15", } + # 2020-02-15 has no samples of its own, so the retro query is the only + # thing adding samples here. It must not re-add the samples that are + # already in the base ARG. + strains = ts.metadata["sc2ts"]["samples_strain"] + assert len(set(strains)) == len(strains) def test_2020_02_14_allow_pango_lineages( self, tmp_path, fx_ts_map, fx_dataset, fx_match_db @@ -1087,7 +1093,7 @@ def test_2020_02_14_allow_pango_lineages( max_pango_lineages=2, ) retro_groups = ts.metadata["sc2ts"]["retro_groups"] - assert len(retro_groups) == 6 + assert len(retro_groups) == 4 def test_2020_02_14_skip_pango_lineages( self, @@ -1112,7 +1118,7 @@ def test_2020_02_14_skip_pango_lineages( max_pango_lineages=1, ) retro_groups = ts.metadata["sc2ts"]["retro_groups"] - assert len(retro_groups) == 5 + assert len(retro_groups) == 3 assert all(len(set(g["pango_lineages"])) == 1 for g in retro_groups) assert "Skipping num_pango_lineages=2 exceeds threshold" in caplog.text @@ -1168,7 +1174,7 @@ def test_2020_02_14_skip_max_mutations( 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 + "Skipping mean_mutations_per_sample=4.0 exceeds threshold" in caplog.text ) def test_2020_02_14_skip_root_mutations( From 0bb718402c088abbf63d1078fdda057cdfa015c0 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 30 Jul 2026 11:12:11 +0100 Subject: [PATCH 05/10] Rename include_samples to seed_groups The option is a list of seed groups, not a list of samples, so name it for what it is. Pure rename: check_include_samples becomes check_seed_groups and the error messages follow, but the behaviour is unchanged. Since the CLI splats [extend_parameters] straight into extend(), this renames the TOML config key too. An old config now fails with a TypeError from extend(), which is the same way any other unknown parameter is reported. --- CHANGELOG.md | 24 ++++++++++----------- docs/example_config.toml | 4 ++-- sc2ts/inference.py | 33 ++++++++++++++-------------- tests/test_cli.py | 12 +++++------ tests/test_inference.py | 46 ++++++++++++++++++++-------------------- 5 files changed, 59 insertions(+), 60 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8769b2b..21904b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,18 +4,18 @@ In development -- `include_samples` is now a list of seed *groups*, each a list of strain IDs - that are inserted into the ARG together as a single local tree. A tree is - inferred over the group's haplotypes, the haplotype of the group's inferred - ancestor is matched against the ARG with recombination disallowed, and the - whole group is then attached at that single placement. Seed samples are - usually far diverged from the current ARG, so matching each one individually - gave spurious recombinations or matches to derived samples with large - numbers of reversions; the inferred ancestor is much closer to the - contemporaneous ARG. A group is inserted on the minimum date over its - members, and members with later dates keep their real dates via negative - ("in the future") node times. The previous format (bare strain IDs, or - `(strain, match_date)` tuples) is no longer supported. +- Add `seed_groups`, replacing the previous `include_samples` option. It is a + list of seed *groups*, each a list of strain IDs that are inserted into the + ARG together as a single local tree. A tree is inferred over the group's + haplotypes, the haplotype of the group's inferred ancestor is matched against + the ARG with recombination disallowed, and the whole group is then attached at + that single placement. Seed samples are usually far diverged from the current + ARG, so matching each one individually gave spurious recombinations or matches + to derived samples with large numbers of reversions; the inferred ancestor is + much closer to the contemporaneous ARG. A group is inserted on the minimum + date over its members, and members with later dates keep their real dates via + negative ("in the future") node times. The previous `include_samples` formats + (bare strain IDs, or `(strain, match_date)` tuples) are no longer supported. - Fix retrospective matching on a day with no new samples. The table of samples already in the ARG was only rebuilt when there were samples to match, so on diff --git a/docs/example_config.toml b/docs/example_config.toml index e9aa637..0ec944d 100644 --- a/docs/example_config.toml +++ b/docs/example_config.toml @@ -93,8 +93,8 @@ memory_limit=32 # later dates keep their real dates, and so are given negative ("in the # future") node times until their real date is reached. # -# include_samples=[["BA.1_strain_a", "BA.1_strain_b"], ["BA.2_strain_c"]] -include_samples=[] +# seed_groups=[["BA.1_strain_a", "BA.1_strain_b"], ["BA.2_strain_c"]] +seed_groups=[] # Override specific parameter values over a time period. [[override]] diff --git a/sc2ts/inference.py b/sc2ts/inference.py index 959f14c..e726072 100644 --- a/sc2ts/inference.py +++ b/sc2ts/inference.py @@ -502,9 +502,9 @@ def preprocess( return samples -def check_include_samples(include_samples, metadata): +def check_seed_groups(seed_groups, metadata): """ - Check the ``include_samples`` seed specification and return it as a list of + Check the ``seed_groups`` seed specification and return it as a list of tuples of strain IDs, one tuple per seed group. Each entry must be a non-empty list of strain IDs which are inserted into @@ -513,13 +513,13 @@ def check_include_samples(include_samples, metadata): """ groups = [] seen = {} - for entry in include_samples: + for entry in seed_groups: if isinstance(entry, str) or not isinstance(entry, (list, tuple)): raise ValueError( - f"Each include_samples entry must be a list of strain IDs, not {entry!r}" + f"Each seed_groups entry must be a list of strain IDs, not {entry!r}" ) if len(entry) == 0: - raise ValueError("include_samples groups must not be empty") + raise ValueError("Seed groups must not be empty") for strain in entry: if not isinstance(strain, str): raise ValueError(f"Seed strain IDs must be strings, not {strain!r}") @@ -527,8 +527,7 @@ def check_include_samples(include_samples, metadata): for strain in group: if strain in seen: raise ValueError( - f"Seed sample {strain} appears in more than one " - "include_samples group" + f"Seed sample {strain} appears in more than one seed group" ) seen[strain] = group groups.append(group) @@ -539,7 +538,7 @@ def check_include_samples(include_samples, metadata): return groups -def seed_group_dates(include_samples, metadata): +def seed_group_dates(seed_groups, metadata): """ Return a list of ``(date, strains)`` tuples, one per seed group, where ``date`` is the minimum date over the group's members. The whole group is @@ -548,7 +547,7 @@ def seed_group_dates(include_samples, metadata): """ return [ (min(metadata[strain]["date"] for strain in group), group) - for group in include_samples + for group in seed_groups ] @@ -559,7 +558,7 @@ def extend( base_ts, match_db, date_field="date", - include_samples=None, + seed_groups=None, num_mismatches=None, hmm_cost_threshold=None, min_group_size=None, @@ -581,7 +580,7 @@ def extend( Extend the base tree sequence by one day, matching in the samples for the given date. - ``include_samples`` is an optional list of "seed" groups that are inserted + ``seed_groups`` is an optional list of "seed" groups that are inserted unconditionally and without recombination. Each entry is a list of strain IDs which are inserted together as a single local tree: a tree is inferred over the group's haplotypes, the haplotype of the group's inferred ancestor @@ -619,8 +618,8 @@ def extend( max_missing_sites = np.inf if deletions_as_missing is None: deletions_as_missing = False - if include_samples is None: - include_samples = [] + if seed_groups is None: + seed_groups = [] base_ts = str(base_ts) dataset = str(dataset) match_db = str(match_db) @@ -634,7 +633,7 @@ def extend( base_ts = tszip.load(base_ts) ds = _dataset.Dataset(dataset, date_field=date_field) - include_samples = check_include_samples(include_samples, ds.metadata) + seed_groups = check_seed_groups(seed_groups, ds.metadata) with MatchDb(match_db) as matches: tables = _extend( @@ -642,7 +641,7 @@ def extend( base_ts=base_ts, date=date, match_db=matches, - include_samples=include_samples, + seed_groups=seed_groups, num_mismatches=num_mismatches, hmm_cost_threshold=hmm_cost_threshold, min_group_size=min_group_size, @@ -672,7 +671,7 @@ def _extend( date, base_ts, match_db, - include_samples, + seed_groups, num_mismatches, hmm_cost_threshold, min_group_size, @@ -698,7 +697,7 @@ def _extend( # A seed group is inserted on the minimum date over its members, so a seed # is processed on its group's date rather than on its own date. - group_dates = seed_group_dates(include_samples, dataset.metadata) + group_dates = seed_group_dates(seed_groups, dataset.metadata) seed_date = { strain: group_date for group_date, group in group_dates for strain in group } diff --git a/tests/test_cli.py b/tests/test_cli.py index 19bb1cf..0f36768 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -417,12 +417,12 @@ def test_start(self, tmp_path, fx_ts_map, fx_dataset): out_ts = tskit.load(ts_path) out_ts.tables.assert_equals(fx_ts_map[date].tables, ignore_provenance=True) - def test_include_samples(self, tmp_path, fx_ts_map, fx_dataset): + def test_seed_groups(self, tmp_path, fx_ts_map, fx_dataset): config_file = self.make_config( tmp_path, fx_dataset, exclude_sites=[56, 57, 58, 59, 60], - include_samples=[["SRR14631544"]], + seed_groups=[["SRR14631544"]], ) runner = ct.CliRunner() result = runner.invoke( @@ -440,13 +440,13 @@ 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): + def test_seed_groups_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"]], + seed_groups=[["SRR14631544"], ["NO_SUCH_STRAIN"]], ) runner = ct.CliRunner() with pytest.raises(ValueError, match="not in dataset"): @@ -456,14 +456,14 @@ def test_include_samples_missing_strain(self, tmp_path, fx_ts_map, fx_dataset): catch_exceptions=False, ) - def test_include_samples_bare_string(self, tmp_path, fx_ts_map, fx_dataset): + def test_seed_groups_bare_string(self, tmp_path, fx_ts_map, fx_dataset): # A bare strain ID is the old format and is rejected: TOML entries must # be arrays of strain IDs. config_file = self.make_config( tmp_path, fx_dataset, exclude_sites=[56, 57, 58, 59, 60], - include_samples=["SRR14631544"], + seed_groups=["SRR14631544"], ) runner = ct.CliRunner() with pytest.raises(ValueError, match="must be a list of strain IDs"): diff --git a/tests/test_inference.py b/tests/test_inference.py index 7f67462..2929d95 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -482,7 +482,7 @@ def test_high_recomb_mutation(self): self.check_double_mirror(ts) -class TestCheckIncludeSamples: +class TestCheckSeedGroups: metadata = { "a": {"date": "2021-01-03"}, "b": {"date": "2021-01-01"}, @@ -490,38 +490,38 @@ class TestCheckIncludeSamples: } def test_empty(self): - assert si.check_include_samples([], self.metadata) == [] + assert si.check_seed_groups([], self.metadata) == [] def test_groups(self): - result = si.check_include_samples([["a", "b"], ["c"]], self.metadata) + result = si.check_seed_groups([["a", "b"], ["c"]], self.metadata) assert result == [("a", "b"), ("c",)] def test_tuples_accepted(self): - assert si.check_include_samples([("a",)], self.metadata) == [("a",)] + assert si.check_seed_groups([("a",)], self.metadata) == [("a",)] def test_bare_string_raises(self): with pytest.raises(ValueError, match="must be a list of strain IDs"): - si.check_include_samples(["a"], self.metadata) + si.check_seed_groups(["a"], self.metadata) def test_empty_group_raises(self): with pytest.raises(ValueError, match="must not be empty"): - si.check_include_samples([[]], self.metadata) + si.check_seed_groups([[]], self.metadata) def test_non_string_strain_raises(self): with pytest.raises(ValueError, match="must be strings"): - si.check_include_samples([["a", 7]], self.metadata) + si.check_seed_groups([["a", 7]], self.metadata) def test_duplicate_strain_raises(self): - with pytest.raises(ValueError, match="more than one include_samples group"): - si.check_include_samples([["a", "b"], ["a"]], self.metadata) + with pytest.raises(ValueError, match="more than one seed group"): + si.check_seed_groups([["a", "b"], ["a"]], self.metadata) def test_duplicate_within_group_raises(self): - with pytest.raises(ValueError, match="more than one include_samples group"): - si.check_include_samples([["a", "a"]], self.metadata) + with pytest.raises(ValueError, match="more than one seed group"): + si.check_seed_groups([["a", "a"]], self.metadata) def test_missing_strain_raises(self): with pytest.raises(ValueError, match="not in dataset"): - si.check_include_samples([["a", "nosuchstrain"]], self.metadata) + si.check_seed_groups([["a", "nosuchstrain"]], self.metadata) class TestSeedGroupDates: @@ -662,7 +662,7 @@ def test_2020_02_02(self, tmp_path, fx_ts_map, fx_dataset, num_threads): assert "SRR11597115" not in ts.metadata["sc2ts"]["samples_strain"] ts.tables.assert_equals(fx_ts_map["2020-02-02"].tables, ignore_provenance=True) - def test_2020_02_02_include_samples( + def test_2020_02_02_seed_group( self, tmp_path, fx_ts_map, @@ -673,7 +673,7 @@ def test_2020_02_02_include_samples( base_ts=fx_ts_map["2020-02-01"], date="2020-02-02", match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=[["SRR11597115"]], + seed_groups=[["SRR11597115"]], ) assert ts.metadata["sc2ts"]["cumulative_stats"]["exact_matches"]["pango"] == { "A": 2, @@ -703,7 +703,7 @@ def test_seed_group_inserted_on_minimum_date(self, tmp_path, fx_ts_map, fx_datas base_ts=fx_ts_map["2020-01-30"], date="2020-01-31", match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=[group], + seed_groups=[group], ) strains = ts.metadata["sc2ts"]["samples_strain"] nodes = {} @@ -732,7 +732,7 @@ def test_seed_group_attached_as_one_tree(self, tmp_path, fx_ts_map, fx_dataset): base_ts=base_ts, date="2020-02-09", match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=[group], + seed_groups=[group], ) strains = ts.metadata["sc2ts"]["samples_strain"] nodes = [ts.samples()[strains.index(strain)] for strain in group] @@ -765,7 +765,7 @@ def test_two_seed_groups_stay_separate(self, tmp_path, fx_ts_map, fx_dataset): base_ts=fx_ts_map["2020-02-01"], date="2020-02-02", match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=[["SRR11597115"], ["SRR11597190"]], + seed_groups=[["SRR11597115"], ["SRR11597190"]], ) strains = ts.metadata["sc2ts"]["samples_strain"] group_ids = set() @@ -783,7 +783,7 @@ def test_seed_group_evolves_over_days(self, tmp_path, fx_ts_map, fx_dataset): # base_ts containing negative ("future") node times. It must also not be # re-added on its natural date. group = ["SRR11494548", "SRR11597115"] - include_samples = [group] + seed_groups = [group] base_path = tmp_path / "base.ts" fx_ts_map["2020-01-30"].dump(base_path) match_db = si.MatchDb.initialise(tmp_path / "match.db") @@ -799,7 +799,7 @@ def test_seed_group_evolves_over_days(self, tmp_path, fx_ts_map, fx_dataset): base_ts=base_path, date=date, match_db=match_db.path, - include_samples=include_samples, + seed_groups=seed_groups, ) ts.dump(base_path) strains = ts.metadata["sc2ts"]["samples_strain"] @@ -811,7 +811,7 @@ def test_seed_group_evolves_over_days(self, tmp_path, fx_ts_map, fx_dataset): assert ts.nodes_time[u] == expected @pytest.mark.parametrize( - "include_samples", + "seed_groups", ( [["SRR11597115"], ["NOSUCHSTRAIN"]], [["SRR11597115", "NOSUCHSTRAIN"]], @@ -819,7 +819,7 @@ def test_seed_group_evolves_over_days(self, tmp_path, fx_ts_map, fx_dataset): ), ) def test_seed_missing_strain_raises( - self, tmp_path, fx_ts_map, fx_dataset, include_samples + self, tmp_path, fx_ts_map, fx_dataset, seed_groups ): # A seed strain that isn't in the dataset is an error, whether it's on # its own or grouped with a strain that does exist. @@ -829,7 +829,7 @@ def test_seed_missing_strain_raises( base_ts=fx_ts_map["2020-02-01"], date="2020-02-02", match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=include_samples, + seed_groups=seed_groups, ) def test_seed_bare_string_raises(self, tmp_path, fx_ts_map, fx_dataset): @@ -840,7 +840,7 @@ def test_seed_bare_string_raises(self, tmp_path, fx_ts_map, fx_dataset): base_ts=fx_ts_map["2020-02-01"], date="2020-02-02", match_db=si.MatchDb.initialise(tmp_path / "match.db"), - include_samples=["SRR11597115"], + seed_groups=["SRR11597115"], ) def test_2020_02_02_mutation_overlap( From 7e0575b2bceafc52e60533a844a0c3ca10875e9e Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 30 Jul 2026 11:16:16 +0100 Subject: [PATCH 06/10] Add a SeedGroup class A seed group was represented as a bare tuple of strain IDs, with its date and per-day working state carried alongside in separate maps and lists. Give it a class instead, analogous to Sample: strains and date are its specification, and samples, topology and root are the state filled in on the day it is inserted. check_seed_groups now returns SeedGroups, absorbing seed_group_dates, which is removed. The group hash moves to a sample_group_id function shared with SampleGroup, so the two provably agree on the group_id written to node metadata. SeedGroup.sample_group() builds the SampleGroup used to insert the group, which will let add_seed_groups construct it with all fields known rather than patching path and immediate_reversions in afterwards. --- sc2ts/inference.py | 98 +++++++++++++++++++++++++++++++---------- tests/test_inference.py | 84 ++++++++++++++++++++++++++--------- 2 files changed, 138 insertions(+), 44 deletions(-) diff --git a/sc2ts/inference.py b/sc2ts/inference.py index e726072..2c74035 100644 --- a/sc2ts/inference.py +++ b/sc2ts/inference.py @@ -505,11 +505,15 @@ def preprocess( def check_seed_groups(seed_groups, metadata): """ Check the ``seed_groups`` seed specification and return it as a list of - tuples of strain IDs, one tuple per seed group. + SeedGroup instances. Each entry must be a non-empty list of strain IDs which are inserted into the ARG together as a single group. Every strain must be present in ``metadata``, and no strain may appear in more than one group. + + A group is inserted on the minimum date over its members; members with + later dates keep their real dates and so get negative ("in the future") + node times. """ groups = [] seen = {} @@ -535,19 +539,13 @@ def check_seed_groups(seed_groups, metadata): missing = sorted(strain for strain in seen if strain not in metadata) if len(missing) > 0: raise ValueError(f"Seed samples not in dataset: {missing}") - return groups - -def seed_group_dates(seed_groups, metadata): - """ - Return a list of ``(date, strains)`` tuples, one per seed group, where - ``date`` is the minimum date over the group's members. The whole group is - inserted into the ARG on that date; members with later dates keep their - real dates and so get negative ("in the future") node times. - """ return [ - (min(metadata[strain]["date"] for strain in group), group) - for group in seed_groups + SeedGroup( + strains=strains, + date=min(metadata[strain]["date"] for strain in strains), + ) + for strains in groups ] @@ -697,12 +695,9 @@ def _extend( # A seed group is inserted on the minimum date over its members, so a seed # is processed on its group's date rather than on its own date. - group_dates = seed_group_dates(seed_groups, dataset.metadata) - seed_date = { - strain: group_date for group_date, group in group_dates for strain in group - } - todays_groups = [group for group_date, group in group_dates if group_date == date] - todays_seeds = {strain for group in todays_groups for strain in group} + seed_date = {strain: group.date for group in seed_groups for strain in group.strains} + todays_groups = [group for group in seed_groups if group.date == date] + todays_seeds = {strain for group in todays_groups for strain in group.strains} metadata_matches = { strain: dataset.metadata[strain] @@ -762,9 +757,11 @@ def _extend( # group left with no members is skipped entirely. todays_seed_groups = [] for group in todays_groups: - group_samples = [seeds[strain] for strain in group if strain in seeds] + group_samples = [seeds[strain] for strain in group.strains if strain in seeds] if len(group_samples) == 0: - logger.warning(f"Skipping seed group with no usable samples: {list(group)}") + logger.warning( + f"Skipping seed group with no usable samples: {list(group.strains)}" + ) else: todays_seed_groups.append(group_samples) @@ -1068,6 +1065,16 @@ def summary(self): ) +def sample_group_id(strains): + """ + Return the ID of the group consisting of the specified strains. + """ + m = hashlib.md5() + for strain in sorted(strains): + m.update(str(strain).encode()) + return m.hexdigest() + + @dataclasses.dataclass class SampleGroup: """ @@ -1086,10 +1093,7 @@ class SampleGroup: topology: tuple = None def __post_init__(self): - m = hashlib.md5() - for strain in sorted(self.strains): - m.update(strain.encode()) - self.sample_hash = m.hexdigest() + self.sample_hash = sample_group_id(self.strains) @property def strains(self): @@ -1136,6 +1140,52 @@ def add_tree_quality_metrics(self, ts, date): return self.tree_quality_metrics +@dataclasses.dataclass +class SeedGroup: + """ + A group of "seed" samples that are inserted into the ARG unconditionally, + and together, as a single local tree. + + The ``strains`` in the group and the ``date`` it is inserted on are its + specification, as returned by check_seed_groups. The remaining fields are + the working state for that day, and are filled in by add_seed_groups. + """ + + strains: tuple + date: str + samples: List = None + topology: tuple = None + root: Sample = None + + @property + def path(self): + """ + The path that the group's inferred ancestor was matched to. + """ + return tuple(self.root.hmm_match.path) + + @property + def sample_hash(self): + return sample_group_id(s.strain for s in self.samples) + + def __len__(self): + return len(self.strains) + + def summary(self): + return f"{self.date} n={len(self.strains)} strains={list(self.strains)}" + + def sample_group(self): + """ + Return the SampleGroup used to insert this group into the ARG. + """ + return SampleGroup( + samples=self.samples, + path=self.path, + immediate_reversions=(), + topology=self.topology, + ) + + def add_matching_results(where_clause, match_db, ts, date, **kwargs): """ Query the MatchDb, group the resulting matches by their path and set of diff --git a/tests/test_inference.py b/tests/test_inference.py index 2929d95..36b858b 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -494,10 +494,26 @@ def test_empty(self): def test_groups(self): result = si.check_seed_groups([["a", "b"], ["c"]], self.metadata) - assert result == [("a", "b"), ("c",)] + assert [g.strains for g in result] == [("a", "b"), ("c",)] def test_tuples_accepted(self): - assert si.check_seed_groups([("a",)], self.metadata) == [("a",)] + result = si.check_seed_groups([("a",)], self.metadata) + assert [g.strains for g in result] == [("a",)] + + def test_group_gets_minimum_date(self): + result = si.check_seed_groups([["a", "b"]], self.metadata) + assert result[0].date == "2021-01-01" + + def test_singleton_gets_own_date(self): + result = si.check_seed_groups([["a"]], self.metadata) + assert result[0].date == "2021-01-03" + + def test_multiple_groups(self): + result = si.check_seed_groups([["a"], ["b", "c"]], self.metadata) + assert [(g.strains, g.date) for g in result] == [ + (("a",), "2021-01-03"), + (("b", "c"), "2021-01-01"), + ] def test_bare_string_raises(self): with pytest.raises(ValueError, match="must be a list of strain IDs"): @@ -524,24 +540,52 @@ def test_missing_strain_raises(self): si.check_seed_groups([["a", "nosuchstrain"]], self.metadata) -class TestSeedGroupDates: - metadata = { - "a": {"date": "2021-01-03"}, - "b": {"date": "2021-01-01"}, - "c": {"date": "2021-01-02"}, - } - - def test_group_gets_minimum_date(self): - result = si.seed_group_dates([("a", "b")], self.metadata) - assert result == [("2021-01-01", ("a", "b"))] - - def test_singleton_gets_own_date(self): - result = si.seed_group_dates([("a",)], self.metadata) - assert result == [("2021-01-03", ("a",))] - - def test_multiple_groups(self): - result = si.seed_group_dates([("a",), ("b", "c")], self.metadata) - assert result == [("2021-01-03", ("a",)), ("2021-01-01", ("b", "c"))] +class TestSeedGroup: + def example(self): + samples = [si.Sample("a"), si.Sample("b")] + group = si.SeedGroup(strains=("a", "b"), date="2021-01-01") + group.samples = samples + group.topology = ([-1, 2, -1], [0, 0, 1]) + group.root = si.Sample("seed_root") + group.root.hmm_match = si.HmmMatch([si.PathSegment(0, 10, 3)], []) + return group + + def test_specification_only(self): + group = si.SeedGroup(strains=("a", "b"), date="2021-01-01") + assert group.strains == ("a", "b") + assert group.date == "2021-01-01" + assert group.samples is None + assert group.topology is None + assert group.root is None + assert len(group) == 2 + + def test_summary(self): + group = si.SeedGroup(strains=("a", "b"), date="2021-01-01") + summary = group.summary() + assert "2021-01-01" in summary + assert "a" in summary + assert "b" in summary + + def test_path_from_root_match(self): + group = self.example() + assert group.path == (si.PathSegment(0, 10, 3),) + + def test_sample_group(self): + group = self.example() + sample_group = group.sample_group() + assert sample_group.samples is group.samples + assert sample_group.path == group.path + assert sample_group.immediate_reversions == () + assert sample_group.topology is group.topology + + def test_sample_hash_matches_sample_group(self): + # The group ID goes into node metadata, so the two classes must agree. + group = self.example() + assert group.sample_hash == si.sample_group_id(["a", "b"]) + assert group.sample_hash == group.sample_group().sample_hash + + def test_sample_hash_independent_of_order(self): + assert si.sample_group_id(["a", "b"]) == si.sample_group_id(["b", "a"]) class TestRealData: From 7e7db12a63c9ad123a5cd562b14c4d4319731cd0 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 30 Jul 2026 11:36:58 +0100 Subject: [PATCH 07/10] Separate seed processing from the standard pipeline Seeds were preprocessed in the same loop as ordinary samples and then diverted into a dict, with a seed_date map deciding which strains the day's metadata query should suppress and which to inject. Every strain in a seed group is now simply excluded from the standard pipeline on every date, and add_seed_groups preprocesses its own samples. This is the same set of samples by construction: a seed never reached the samples list before either. _extend now mentions seeds twice: one set-difference and one call. The per-day state that was carried in parallel lists lives on the SeedGroup, so the pending list of (group, root_sample) tuples goes away. New helpers, both reusing what was there: - prepare_samples wraps preprocess with the metadata decoration and the no-alignment drop, and is used by both pipelines. - infer_seed_group_root is the group ancestor inference lifted out of add_seed_groups, which makes it testable on its own. Seed samples are no longer merged back into the day's samples, so they no longer appear in the daily samples_processed statistics. A forced match is not an HMM match, and since force_match replaced the old cost = 0.5 sentinel a seed matching its parent exactly was being counted as an exact match. --- sc2ts/inference.py | 291 ++++++++++++++++++++++------------------ tests/test_inference.py | 22 +++ 2 files changed, 185 insertions(+), 128 deletions(-) diff --git a/sc2ts/inference.py b/sc2ts/inference.py index 2c74035..41c21d0 100644 --- a/sc2ts/inference.py +++ b/sc2ts/inference.py @@ -37,6 +37,10 @@ # matching, since solve_num_mismatches drives rho down to its 1e-200 floor. NO_RECOMBINATION_NUM_MISMATCHES = 1000 +# FIXME parametrise +PANGO_LINEAGE_KEY = "Viridian_pangolin" +SCORPIO_KEY = "Viridian_scorpio" + def get_provenance_dict(command, args, start_time): document = { @@ -502,6 +506,45 @@ def preprocess( return samples +def prepare_samples( + strains, + dataset, + date, + *, + keep_sites, + progress_title="", + show_progress=False, +): + """ + Return the Sample objects for the specified strains, with their haplotypes, + metadata, pango, scorpio and date filled in. Strains that have no stored + alignment are dropped. + """ + samples = [] + for sample in preprocess( + strains=strains, + dataset=dataset, + keep_sites=keep_sites, + progress_title=progress_title, + show_progress=show_progress, + ): + if sample.haplotype is None: + logger.debug(f"No alignment stored for {sample.strain}") + continue + md = dataset.metadata[sample.strain] + sample.metadata = md + sample.pango = md.get(PANGO_LINEAGE_KEY, "Unknown") + sample.scorpio = md.get(SCORPIO_KEY, "Unknown") + sample.date = date + logger.debug( + f"Encoded {sample.strain} {sample.scorpio} {sample.pango} " + f"missing={sample.num_missing_sites} " + f"deletions={sample.num_deletion_sites}" + ) + samples.append(sample) + return samples + + def check_seed_groups(seed_groups, metadata): """ Check the ``seed_groups`` seed specification and return it as a list of @@ -693,77 +736,35 @@ def _extend( f"mutations={base_ts.num_mutations};date={previous_date}" ) - # A seed group is inserted on the minimum date over its members, so a seed - # is processed on its group's date rather than on its own date. - seed_date = {strain: group.date for group in seed_groups for strain in group.strains} - todays_groups = [group for group in seed_groups if group.date == date] - todays_seeds = {strain for group in todays_groups for strain in group.strains} + # A seed group is inserted as a unit on its own date by add_seed_groups, so + # its members are excluded from the standard pipeline on every date. + seed_strains = {strain for group in seed_groups for strain in group.strains} + todays_seed_groups = [group for group in seed_groups if group.date == date] - metadata_matches = { - strain: dataset.metadata[strain] + strains = [ + strain for strain in dataset.metadata.samples_for_date(date) - # Exclude a seed from any date other than its group's date; it is - # processed only once, on that date. - if seed_date.get(strain, date) == date - } - # Inject today's seeds that aren't naturally sampled today. Missing strains - # are rejected up-front in extend(). - for strain in todays_seeds: - if strain not in metadata_matches: - metadata_matches[strain] = dataset.metadata[strain] - - logger.info(f"Got {len(metadata_matches)} metadata matches") + if strain not in seed_strains + ] + logger.info(f"Got {len(strains)} samples for {date}") - preprocessed_samples = preprocess( - strains=list(metadata_matches.keys()), - dataset=dataset, + samples = [] + for sample in prepare_samples( + strains, + dataset, + date, keep_sites=base_ts.sites_position.astype(int), progress_title=date, show_progress=show_progress, - ) - # FIXME parametrise - pango_lineage_key = "Viridian_pangolin" - scorpio_key = "Viridian_scorpio" - - seeds = {} - samples = [] - for s in preprocessed_samples: - if s.haplotype is None: - logger.debug(f"No alignment stored for {s.strain}") - continue - md = metadata_matches[s.strain] - s.metadata = md - s.pango = md.get(pango_lineage_key, "Unknown") - s.scorpio = md.get(scorpio_key, "Unknown") - s.date = date - num_missing_sites = s.num_missing_sites - num_deletion_sites = s.num_deletion_sites - logger.debug( - f"Encoded {s.strain} {s.scorpio} {s.pango} missing={num_missing_sites} " - f"deletions={num_deletion_sites}" - ) - if s.strain in todays_seeds: - # Seeds bypass the max_missing_sites and max_daily_samples filters. - s.flags |= core.NODE_IS_UNCONDITIONALLY_INCLUDED - seeds[s.strain] = s - elif num_missing_sites <= max_missing_sites: - samples.append(s) + ): + num_missing_sites = sample.num_missing_sites + if num_missing_sites <= max_missing_sites: + samples.append(sample) else: logger.debug( - f"Filter {s.strain}: missing={num_missing_sites} > {max_missing_sites}" - ) - - # Seeds without a stored alignment are dropped from their group, and a - # group left with no members is skipped entirely. - todays_seed_groups = [] - for group in todays_groups: - group_samples = [seeds[strain] for strain in group.strains if strain in seeds] - if len(group_samples) == 0: - logger.warning( - f"Skipping seed group with no usable samples: {list(group.strains)}" + f"Filter {sample.strain}: " + f"missing={num_missing_sites} > {max_missing_sites}" ) - else: - todays_seed_groups.append(group_samples) if max_daily_samples is not None: if max_daily_samples < len(samples): @@ -806,21 +807,21 @@ def _extend( phase="close", ) - if len(todays_seed_groups) > 0: - # Seeds are attached directly rather than via the MatchDb, so that a - # group can't be split up and can't be added a second time. - ts, seed_samples = add_seed_groups( - ts, - base_ts, - todays_seed_groups, - date, - num_mismatches=num_mismatches, - deletions_as_missing=deletions_as_missing, - show_progress=show_progress, - num_threads=num_threads, - memory_limit=memory_limit, - ) - samples = samples + seed_samples + # Seeds are attached directly rather than via the MatchDb, so that a group + # can't be split up and can't be added a second time. They are not part of + # the day's HMM-matched samples, and so play no part in its statistics. + ts = add_seed_groups( + ts, + base_ts, + todays_seed_groups, + date, + dataset=dataset, + num_mismatches=num_mismatches, + deletions_as_missing=deletions_as_missing, + show_progress=show_progress, + num_threads=num_threads, + memory_limit=memory_limit, + ) samples.sort(key=lambda s: s.strain) @@ -1337,12 +1338,51 @@ def add_sample_groups( return ts, added_groups +def infer_seed_group_root(ts, haplotypes, reference, deletions_as_missing=False): + """ + Infer a tree over the specified group of haplotypes and return the + ``(topology, root_haplotype)`` of the group's inferred ancestor. The + topology is None when there is nothing to infer: either a single + haplotype, or a group identical to ``reference`` at every non-missing site. + """ + flat_ts = flat_group_ts( + ts, + haplotypes, + reference, + deletions_as_missing=deletions_as_missing, + ) + topology = None + if flat_ts.num_samples > 1 and flat_ts.num_mutations > 0: + topology = tree_ops.infer_binary_topology(flat_ts) + + if topology is None: + if len(haplotypes) == 1: + root_haplotype = haplotypes[0].copy() + else: + root_haplotype = reference.copy() + else: + binary_ts = tree_ops.infer_binary(flat_ts, topology) + tree = binary_ts.first() + # The single child of the outgroup root is the group's MRCA. + mrca = tree.children(tree.root)[0] + root_haplotype = reference.copy() + site_index = np.searchsorted(ts.sites_position, binary_ts.sites_position) + root_haplotype[site_index] = node_haplotypes(binary_ts, [mrca])[0] + + # Sites at which every member is missing are unknown at the root, so let + # the HMM ignore them rather than asserting the reference allele. + H = np.array(haplotypes) + root_haplotype[np.all(H == MISSING, axis=0)] = MISSING + return topology, root_haplotype + + def add_seed_groups( ts, base_ts, seed_groups, date, *, + dataset, num_mismatches, deletions_as_missing=False, show_progress=False, @@ -1350,8 +1390,7 @@ def add_seed_groups( memory_limit=-1, ): """ - Add the specified groups of seed samples into the ARG, one local tree per - group. + Add the specified SeedGroups into the ARG, one local tree per group. Rather than matching each seed against the ARG individually, we infer a tree over each group's haplotypes, match the haplotype of the group's @@ -1360,60 +1399,60 @@ def add_seed_groups( piece. Seeds never go through the MatchDb: their matches are constructed directly by forcing them down the ancestor's path. - Returns the updated tree sequence and the list of seed samples added. + Returns the updated tree sequence. """ + if len(seed_groups) == 0: + return ts + + samples = prepare_samples( + [strain for group in seed_groups for strain in group.strains], + dataset, + date, + keep_sites=base_ts.sites_position.astype(int), + progress_title=date, + show_progress=show_progress, + ) + for sample in samples: + # Seeds bypass the max_missing_sites and max_daily_samples filters. + sample.flags |= core.NODE_IS_UNCONDITIONALLY_INCLUDED + sample_map = {sample.strain: sample for sample in samples} + reference = reference_haplotype(base_ts) - root_samples = [] - pending = [] - for samples in seed_groups: - flat_ts = flat_group_ts( + groups = [] + for group in seed_groups: + assert group.date == date + # Seeds without a stored alignment are dropped from their group, and a + # group left with no members is skipped entirely. + group.samples = [ + sample_map[strain] for strain in group.strains if strain in sample_map + ] + if len(group.samples) == 0: + logger.warning( + f"Skipping seed group with no usable samples: {list(group.strains)}" + ) + continue + group.topology, root_haplotype = infer_seed_group_root( base_ts, - [s.haplotype for s in samples], + [sample.haplotype for sample in group.samples], reference, deletions_as_missing=deletions_as_missing, ) - topology = None - if flat_ts.num_samples > 1 and flat_ts.num_mutations > 0: - topology = tree_ops.infer_binary_topology(flat_ts) - - if topology is None: - # Either a single sample, or a group that is identical to the - # reference at every non-missing site. - if len(samples) == 1: - root_haplotype = samples[0].haplotype.copy() - else: - root_haplotype = reference.copy() - else: - binary_ts = tree_ops.infer_binary(flat_ts, topology) - tree = binary_ts.first() - # The single child of the outgroup root is the group's MRCA. - mrca = tree.children(tree.root)[0] - root_haplotype = reference.copy() - site_index = np.searchsorted( - base_ts.sites_position, binary_ts.sites_position - ) - root_haplotype[site_index] = node_haplotypes(binary_ts, [mrca])[0] - - # Sites at which every member is missing are unknown at the root, so - # let the HMM ignore them rather than asserting the reference allele. - H = np.array([s.haplotype for s in samples]) - root_haplotype[np.all(H == MISSING, axis=0)] = MISSING - - group = SampleGroup(samples, topology=topology) - root_sample = Sample( + group.root = Sample( strain=f"seed_root_{group.sample_hash}", date=date, haplotype=root_haplotype, ) - root_samples.append(root_sample) - pending.append((group, root_sample)) + groups.append(group) + + if len(groups) == 0: + return ts # Seed samples are usually far diverged from the current ARG, and matching # them with the standard num_mismatches gives spurious recombinations. # Match the inferred ancestors without recombination instead. match_samples( date, - root_samples, + [group.root for group in groups], base_ts=base_ts, num_mismatches=NO_RECOMBINATION_NUM_MISMATCHES, deletions_as_missing=deletions_as_missing, @@ -1422,36 +1461,32 @@ def add_seed_groups( memory_limit=memory_limit, ) - groups = [] seed_samples = [] - for group, root_sample in pending: - path = root_sample.hmm_match.path - parent_haplotype = path_haplotype(base_ts, path) - for sample in group: + for group in groups: + parent_haplotype = path_haplotype(base_ts, group.path) + for sample in group.samples: force_match( sample, - path, + group.path, parent_haplotype, base_ts.sites_position, deletions_as_missing, num_mismatches, ) logger.warning(f"Unconditionally including {sample.summary()}") - group.path = tuple(path) - group.immediate_reversions = () - groups.append(group) seed_samples.extend(group.samples) characterise_recombinants(base_ts, seed_samples) ts, _ = add_sample_groups( ts, - groups, + [group.sample_group() for group in groups], date, min_group_size=1, show_progress=show_progress, phase="seed", ) - return ts, seed_samples + logger.info(f"Added {len(seed_samples)} seed samples in {len(groups)} groups") + return ts def solve_num_mismatches(k, num_alleles=5): diff --git a/tests/test_inference.py b/tests/test_inference.py index 36b858b..ff712d1 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -820,6 +820,28 @@ def test_two_seed_groups_stay_separate(self, tmp_path, fx_ts_map, fx_dataset): group_ids.add(ts.node(u).metadata["sc2ts"]["group_id"]) assert len(group_ids) == 2 + def test_seed_excluded_from_standard_pipeline(self, tmp_path, fx_ts_map, fx_dataset): + # SRR11597190 is dated 2020-02-02 and is normally added on that date, + # but here it belongs to a group whose date is 2020-01-31. The group is + # therefore not inserted today, and the standard pipeline must not pick + # its members up either, so this day adds one sample fewer than + # test_2020_02_02 does. + base_ts = fx_ts_map["2020-02-01"] + ts = run_extend( + dataset=fx_dataset, + base_ts=base_ts, + date="2020-02-02", + match_db=si.MatchDb.initialise(tmp_path / "match.db"), + seed_groups=[["SRR11494548", "SRR11597190"]], + ) + base_strains = base_ts.metadata["sc2ts"]["samples_strain"] + strains = ts.metadata["sc2ts"]["samples_strain"] + added = [strain for strain in strains if strain not in base_strains] + assert "SRR11597190" not in added + assert "SRR11494548" not in added + assert ts.num_samples == 20 + assert np.sum(ts.nodes_time[ts.samples()] == 0) == 3 + def test_seed_group_evolves_over_days(self, tmp_path, fx_ts_map, fx_dataset): # Insert the group on its minimum date, then keep extending past the # later member's actual date. That member's time must rise by +1/day and From 2e24bc5085460ce0aa185f470977b0c315c26562 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 30 Jul 2026 11:48:28 +0100 Subject: [PATCH 08/10] Tidy up around the seed group code Construct SampleGroup by keyword in add_matching_results, so that topology no longer has to be the trailing field. Let tmp_dataset take a date per sample, so that behaviour depending on samples having different dates can be tested on the 30bp genome rather than through the 21-day fixture. Add unit tests for infer_seed_group_root, which is now separable enough to test directly. --- sc2ts/dataset.py | 4 ++- sc2ts/inference.py | 9 +++--- tests/test_dataset.py | 25 +++++++++++++++ tests/test_inference.py | 70 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 102 insertions(+), 6 deletions(-) diff --git a/sc2ts/dataset.py b/sc2ts/dataset.py index 9032d0e..8f8153b 100644 --- a/sc2ts/dataset.py +++ b/sc2ts/dataset.py @@ -646,9 +646,11 @@ def tmp_dataset(path, alignments, date="2020-01-01", contig_id=None): # Minimal hacky thing for testing. Should refactor into something more useful. if contig_id is None: contig_id = core.REFERENCE_STRAIN + if isinstance(date, str): + date = [date] * len(alignments) sequence_length = len(next(iter(alignments.values()))) + 1 Dataset.new(path, sequence_length=sequence_length, contig_id=contig_id) Dataset.append_alignments(path, alignments) - df = pd.DataFrame({"strain": alignments.keys(), "date": [date] * len(alignments)}) + df = pd.DataFrame({"strain": alignments.keys(), "date": date}) Dataset.add_metadata(path, df.set_index("strain")) return Dataset(path) diff --git a/sc2ts/inference.py b/sc2ts/inference.py index 41c21d0..3c62bc2 100644 --- a/sc2ts/inference.py +++ b/sc2ts/inference.py @@ -1089,8 +1089,7 @@ class SampleGroup: sample_hash: str = None tree_quality_metrics: GroupTreeQualityMetrics = None # Optionally, a pre-inferred (pi, tau) topology to use when building the - # group's local tree instead of inferring one. Must be trailing, because - # add_matching_results constructs SampleGroups positionally. + # group's local tree instead of inferring one. topology: tuple = None def __post_init__(self): @@ -1218,9 +1217,9 @@ def add_matching_results(where_clause, match_db, ts, date, **kwargs): groups = [ SampleGroup( - samples, - key[0], - key[1], + samples=samples, + path=key[0], + immediate_reversions=key[1], ) for key, samples in grouped_matches.items() ] diff --git a/tests/test_dataset.py b/tests/test_dataset.py index d605f7d..b1cb527 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -587,6 +587,31 @@ def test_field_descriptors(self, fx_dataset): assert df.shape[0] == fx_dataset.metadata.num_fields +class TestTmpDataset: + alignments = { + "a": np.array([0, 1, 2, 3]), + "b": np.array([3, 2, 1, 0]), + } + + def test_single_date(self, tmp_path): + ds = sc2ts.dataset.tmp_dataset( + tmp_path / "ds.zarr", self.alignments, date="2020-01-01" + ) + assert ds.metadata["a"]["date"] == "2020-01-01" + assert ds.metadata["b"]["date"] == "2020-01-01" + + def test_date_per_sample(self, tmp_path): + ds = sc2ts.dataset.tmp_dataset( + tmp_path / "ds.zarr", + self.alignments, + date=["2020-01-01", "2020-01-02"], + ) + assert ds.metadata["a"]["date"] == "2020-01-01" + assert ds.metadata["b"]["date"] == "2020-01-02" + ds = sc2ts.Dataset(ds.path, date_field="date") + assert list(ds.metadata.samples_for_date("2020-01-02")) == ["b"] + + class TestEncodeAlignment: @pytest.mark.parametrize( ["hap", "expected"], diff --git a/tests/test_inference.py b/tests/test_inference.py index ff712d1..4398bea 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -540,6 +540,76 @@ def test_missing_strain_raises(self): si.check_seed_groups([["a", "nosuchstrain"]], self.metadata) +class TestInferSeedGroupRoot: + # A 30bp reference, so that a group's inferred ancestor can be checked + # against known haplotypes without running the full pipeline. + reference = "AAAACCCCGGGGTTTTAAAACCCCGGGGTT" + + @property + def base_ts(self): + return si.initial_ts( + reference_sequence="X" + self.reference, + reference_id="chr_test", + reference_date="2019-01-01", + ) + + def haplotype(self, mutations=None): + h = list(self.reference) + for pos, base in (mutations or {}).items(): + h[pos] = base + return jit.encode_alleles(np.array(h)) + + def infer(self, haplotypes): + ts = self.base_ts + return si.infer_seed_group_root(ts, haplotypes, si.reference_haplotype(ts)) + + def test_single_haplotype(self): + h = self.haplotype({4: "T"}) + topology, root = self.infer([h]) + # Nothing to infer over one sample: the ancestor is the sample. + assert topology is None + nt.assert_array_equal(root, h) + assert root is not h + + def test_identical_to_reference(self): + reference = si.reference_haplotype(self.base_ts) + topology, root = self.infer([self.haplotype(), self.haplotype()]) + # No mutations, so there is no tree to infer and the ancestor is the + # reference. + assert topology is None + nt.assert_array_equal(root, reference) + + def test_shared_derived_allele(self): + reference = si.reference_haplotype(self.base_ts) + haplotypes = [ + self.haplotype({4: "T", 12: "A"}), + self.haplotype({4: "T", 21: "A"}), + ] + topology, root = self.infer(haplotypes) + assert topology is not None + # The ancestor carries the shared mutation but neither private one. + expected = reference.copy() + expected[4] = jit.encode_alleles(np.array(["T"]))[0] + nt.assert_array_equal(root, expected) + + def test_all_missing_site_is_missing(self): + haplotypes = [] + for pos in [12, 21]: + h = self.haplotype({pos: "A"}) + h[4] = si.MISSING + haplotypes.append(h) + _, root = self.infer(haplotypes) + # Every member is missing at site 4, so the ancestor is unknown there + # rather than asserting the reference allele. + assert root[4] == si.MISSING + + def test_partially_missing_site_is_not_missing(self): + haplotypes = [self.haplotype({4: "T"}), self.haplotype()] + haplotypes[1][4] = si.MISSING + _, root = self.infer(haplotypes) + assert root[4] != si.MISSING + + class TestSeedGroup: def example(self): samples = [si.Sample("a"), si.Sample("b")] From 82ffb20eebcb66499a294984d52a489d44c4d924 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 30 Jul 2026 11:52:53 +0100 Subject: [PATCH 09/10] Ignore FASTA index and editor swap files reference.fasta.fai is generated on demand when the bundled reference is read, and vim swap files turn up next to whatever is being edited. Both were showing up as untracked and are easy to commit by accident. --- .gitignore | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.gitignore b/.gitignore index b70d402..1d67587 100644 --- a/.gitignore +++ b/.gitignore @@ -3,8 +3,14 @@ *.tsz *.nwk *.nex +# Generated on demand alongside a FASTA +*.fai _version.py +# Editor swap files +*.swp +*.swo + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] From 951c46fcbcf94ab6101d878a835cbf4386248208 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 30 Jul 2026 12:56:23 +0100 Subject: [PATCH 10/10] Simplify seed group tree inference The group's local tree was built by a round trip: infer a topology against the reference, carry it on the group, synthesise a per-sample hmm_match against the matched path so that match_path_ts could rebuild the same tree from HMM mutations, and reuse the stored topology to keep the two inferences consistent. None of that is needed. The stored hmm_match is informational, so every member of a seed group now reports the match that placed the group. Once the tree is no longer derived from hmm_match it can be built straight from the samples' haplotypes, and each inference is a single infer_binary call: - infer_binary_topology and infer_binary go back to their previous form; the only remaining tree_ops change on this branch is detach_future_nodes. - flat_group_ts takes samples rather than bare haplotypes, and optionally a group_id, in which case the sample nodes carry the metadata attach_tree needs. It now writes each site's ancestral state from the supplied ancestral haplotype rather than from the ARG reference, which is a no-op for the reference-rooted call and is what makes it usable for the other. - SampleGroup.topology becomes SampleGroup.flat_ts, the star tree to infer from. MatchDb groups leave it None and build it from their HMM matches. - infer_seed_group_root returns just the ancestor's haplotype. - force_match and path_haplotype are gone. The parent haplotype comes from node_haplotypes, since a seed group attaches to a single node. The group's tree is still inferred twice, once against the reference to find the ancestor to match and once against the matched node's haplotype to build the tree that is attached. That second rooting is not optional: attach_tree hangs the group's root children off the matched node, so a tree built against the reference silently gives the seeds the parent's allele wherever the parent differs from the reference and the group does not. TestSeedGroupAttachment pins this down; it fails if the tree is built against the reference. Also drop the unreachable "haplotype is None" handling: preprocess always assigns a haplotype, and Dataset.alignment and Dataset.metadata share one sample_id_map, so a strain that passes check_seed_groups has an alignment. --- CHANGELOG.md | 4 + sc2ts/inference.py | 223 +++++++++++++++------------------------- sc2ts/tree_ops.py | 36 ++----- tests/test_inference.py | 197 ++++++++++++++++++++++++++++++++--- tests/test_tree_ops.py | 64 ------------ 5 files changed, 283 insertions(+), 241 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 21904b4..43ceaa1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,10 @@ In development negative ("in the future") node times. The previous `include_samples` formats (bare strain IDs, or `(strain, match_date)` tuples) are no longer supported. + Note that the `hmm_match` metadata of a seed node records the match that placed + its *group*, not a per-sample match, so its mutation count is not the number of + mutations between that sample and its parent. + - Fix retrospective matching on a day with no new samples. The table of samples already in the ARG was only rebuilt when there were samples to match, so on such a day the retrospective query ran against the previous day's table and diff --git a/sc2ts/inference.py b/sc2ts/inference.py index 3c62bc2..283a20b 100644 --- a/sc2ts/inference.py +++ b/sc2ts/inference.py @@ -517,8 +517,7 @@ def prepare_samples( ): """ Return the Sample objects for the specified strains, with their haplotypes, - metadata, pango, scorpio and date filled in. Strains that have no stored - alignment are dropped. + metadata, pango, scorpio and date filled in. """ samples = [] for sample in preprocess( @@ -528,9 +527,6 @@ def prepare_samples( progress_title=progress_title, show_progress=show_progress, ): - if sample.haplotype is None: - logger.debug(f"No alignment stored for {sample.strain}") - continue md = dataset.metadata[sample.strain] sample.metadata = md sample.pango = md.get(PANGO_LINEAGE_KEY, "Unknown") @@ -816,7 +812,6 @@ def _extend( todays_seed_groups, date, dataset=dataset, - num_mismatches=num_mismatches, deletions_as_missing=deletions_as_missing, show_progress=show_progress, num_threads=num_threads, @@ -958,33 +953,47 @@ def match_path_ts(group, sequence_length): return tables.tree_sequence() -def flat_group_ts(ts, haplotypes, ancestral_haplotype, deletions_as_missing=False): +def flat_group_ts( + ts, samples, ancestral_haplotype, *, group_id=None, deletions_as_missing=False +): """ - Return a "flat" (star) tree sequence in which each of the specified - haplotypes is a sample node hanging off a root carrying - ``ancestral_haplotype``. - - Site positions and ancestral states are taken from ``ts``, and a mutation - is added wherever a haplotype's non-missing allele differs from - ``ancestral_haplotype``. Missing data therefore reads as the ancestral - state, exactly as it does in the trees built by ``match_path_ts``. + Return a "flat" (star) tree sequence in which each of the specified samples + is a sample node hanging off a root carrying ``ancestral_haplotype``. + + Site positions are taken from ``ts``, each site's ancestral state is the + ``ancestral_haplotype`` allele, and a mutation is added wherever a sample's + non-missing allele differs from it. Missing data therefore reads as the + ancestral state, exactly as it does in the trees built by ``match_path_ts``. + + If ``group_id`` is given the sample nodes carry their full metadata, so + that the tree inferred from this one can be attached to the ARG. Otherwise + bare sample nodes are added, which is all that tree building needs, and is + the only option before the group has been matched. """ + # IUPAC_ALLELES[MISSING] is "." by negative indexing, so a missing ancestral + # allele would quietly give the wrong ancestral state. + assert np.all(ancestral_haplotype != MISSING) tables = tskit.TableCollection(ts.sequence_length) tables.nodes.metadata_schema = tskit.MetadataSchema.permissive_json() + tables.mutations.metadata_schema = tskit.MetadataSchema.permissive_json() sites_position = ts.sites_position - sites_ancestral_state = ts.sites_ancestral_state - root = len(haplotypes) + root = len(samples) site_id_map = {} - for h in haplotypes: - node_id = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) + for sample in samples: + if group_id is None: + node_id = tables.nodes.add_row(flags=sample.flags, time=0) + else: + node_id = add_sample_to_tables(sample, tables, group_id=group_id) tables.edges.add_row(0, tables.sequence_length, parent=root, child=node_id) + h = sample.haplotype if deletions_as_missing: h = np.where(h == DELETION, MISSING, h) (sites,) = np.where((h != ancestral_haplotype) & (h != MISSING)) for site_id in sites: if site_id not in site_id_map: site_id_map[site_id] = tables.sites.add_row( - sites_position[site_id], sites_ancestral_state[site_id] + sites_position[site_id], + core.IUPAC_ALLELES[ancestral_haplotype[site_id]], ) tables.mutations.add_row( site=site_id_map[site_id], @@ -1088,9 +1097,10 @@ class SampleGroup: immediate_reversions: List = None sample_hash: str = None tree_quality_metrics: GroupTreeQualityMetrics = None - # Optionally, a pre-inferred (pi, tau) topology to use when building the - # group's local tree instead of inferring one. - topology: tuple = None + # Optionally, the "flat" (star) tree to infer the group's local tree from. + # Groups coming from the MatchDb leave this as None and build it from their + # samples' HMM matches; seed groups supply it directly. + flat_ts: tskit.TreeSequence = None def __post_init__(self): self.sample_hash = sample_group_id(self.strains) @@ -1154,8 +1164,8 @@ class SeedGroup: strains: tuple date: str samples: List = None - topology: tuple = None root: Sample = None + flat_ts: tskit.TreeSequence = None @property def path(self): @@ -1182,7 +1192,7 @@ def sample_group(self): samples=self.samples, path=self.path, immediate_reversions=(), - topology=self.topology, + flat_ts=self.flat_ts, ) @@ -1269,11 +1279,13 @@ def add_sample_groups( f"threshold: {group.summary()}" ) continue - flat_ts = match_path_ts(group, ts.sequence_length) + flat_ts = group.flat_ts + if flat_ts is None: + 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 else: - binary_ts = tree_ops.infer_binary(flat_ts, group.topology) + 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, date) @@ -1337,30 +1349,26 @@ def add_sample_groups( return ts, added_groups -def infer_seed_group_root(ts, haplotypes, reference, deletions_as_missing=False): +def infer_seed_group_root(ts, samples, reference, deletions_as_missing=False): """ - Infer a tree over the specified group of haplotypes and return the - ``(topology, root_haplotype)`` of the group's inferred ancestor. The - topology is None when there is nothing to infer: either a single - haplotype, or a group identical to ``reference`` at every non-missing site. + Infer a tree over the specified group of samples and return the haplotype + of the group's inferred ancestor, in the ``core.IUPAC_ALLELES`` encoding. """ + haplotypes = [sample.haplotype for sample in samples] flat_ts = flat_group_ts( ts, - haplotypes, + samples, reference, deletions_as_missing=deletions_as_missing, ) - topology = None - if flat_ts.num_samples > 1 and flat_ts.num_mutations > 0: - topology = tree_ops.infer_binary_topology(flat_ts) - - if topology is None: - if len(haplotypes) == 1: - root_haplotype = haplotypes[0].copy() - else: - root_haplotype = reference.copy() + if len(samples) == 1: + # There is nothing to infer over one haplotype: it is its own ancestor. + root_haplotype = haplotypes[0].copy() + elif flat_ts.num_mutations == 0: + # The group is identical to the reference at every non-missing site. + root_haplotype = reference.copy() else: - binary_ts = tree_ops.infer_binary(flat_ts, topology) + binary_ts = tree_ops.infer_binary(flat_ts) tree = binary_ts.first() # The single child of the outgroup root is the group's MRCA. mrca = tree.children(tree.root)[0] @@ -1372,7 +1380,7 @@ def infer_seed_group_root(ts, haplotypes, reference, deletions_as_missing=False) # the HMM ignore them rather than asserting the reference allele. H = np.array(haplotypes) root_haplotype[np.all(H == MISSING, axis=0)] = MISSING - return topology, root_haplotype + return root_haplotype def add_seed_groups( @@ -1382,7 +1390,6 @@ def add_seed_groups( date, *, dataset, - num_mismatches, deletions_as_missing=False, show_progress=False, num_threads=0, @@ -1392,11 +1399,11 @@ def add_seed_groups( Add the specified SeedGroups into the ARG, one local tree per group. Rather than matching each seed against the ARG individually, we infer a - tree over each group's haplotypes, match the haplotype of the group's - inferred ancestor (with recombination disallowed), and then rebuild the - group's tree against that path so that the whole group is attached in one - piece. Seeds never go through the MatchDb: their matches are constructed - directly by forcing them down the ancestor's path. + tree over each group's haplotypes and match the haplotype of the group's + inferred ancestor, with recombination disallowed. The group's local tree is + then inferred again, this time against the haplotype of the node the + ancestor matched to, so that the whole group is attached in one piece. + Seeds never go through the MatchDb. Returns the updated tree sequence. """ @@ -1417,41 +1424,26 @@ def add_seed_groups( sample_map = {sample.strain: sample for sample in samples} reference = reference_haplotype(base_ts) - groups = [] for group in seed_groups: assert group.date == date - # Seeds without a stored alignment are dropped from their group, and a - # group left with no members is skipped entirely. - group.samples = [ - sample_map[strain] for strain in group.strains if strain in sample_map - ] - if len(group.samples) == 0: - logger.warning( - f"Skipping seed group with no usable samples: {list(group.strains)}" - ) - continue - group.topology, root_haplotype = infer_seed_group_root( - base_ts, - [sample.haplotype for sample in group.samples], - reference, - deletions_as_missing=deletions_as_missing, - ) + group.samples = [sample_map[strain] for strain in group.strains] group.root = Sample( strain=f"seed_root_{group.sample_hash}", date=date, - haplotype=root_haplotype, + haplotype=infer_seed_group_root( + base_ts, + group.samples, + reference, + deletions_as_missing=deletions_as_missing, + ), ) - groups.append(group) - - if len(groups) == 0: - return ts # Seed samples are usually far diverged from the current ARG, and matching # them with the standard num_mismatches gives spurious recombinations. # Match the inferred ancestors without recombination instead. match_samples( date, - [group.root for group in groups], + [group.root for group in seed_groups], base_ts=base_ts, num_mismatches=NO_RECOMBINATION_NUM_MISMATCHES, deletions_as_missing=deletions_as_missing, @@ -1460,31 +1452,36 @@ def add_seed_groups( memory_limit=memory_limit, ) - seed_samples = [] - for group in groups: - parent_haplotype = path_haplotype(base_ts, group.path) + num_samples = 0 + for group in seed_groups: + # Recombination is disallowed above, so the group attaches to a single + # node. Multiple segments would also mean characterise_recombinants is + # needed here, and that the group is flagged as a recombinant. + assert len(group.path) == 1 + parent_haplotype = node_haplotypes(base_ts, [group.path[0].parent])[0] for sample in group.samples: - force_match( - sample, - group.path, - parent_haplotype, - base_ts.sites_position, - deletions_as_missing, - num_mismatches, - ) + # The stored HMM match is informational only, so every member of + # the group reports the match that placed the group. + sample.hmm_match = group.root.hmm_match logger.warning(f"Unconditionally including {sample.summary()}") - seed_samples.extend(group.samples) + group.flat_ts = flat_group_ts( + base_ts, + group.samples, + parent_haplotype, + group_id=group.sample_hash, + deletions_as_missing=deletions_as_missing, + ) + num_samples += len(group.samples) - characterise_recombinants(base_ts, seed_samples) ts, _ = add_sample_groups( ts, - [group.sample_group() for group in groups], + [group.sample_group() for group in seed_groups], date, min_group_size=1, show_progress=show_progress, phase="seed", ) - logger.info(f"Added {len(seed_samples)} seed samples in {len(groups)} groups") + logger.info(f"Added {num_samples} seed samples in {len(seed_groups)} groups") return ts @@ -2070,56 +2067,6 @@ def reference_haplotype(ts): return jit.encode_alleles(np.asarray(ts.sites_ancestral_state, dtype="U1")) -def path_haplotype(ts, path): - """ - Return the haplotype copied along the specified HMM path, in the - ``core.IUPAC_ALLELES`` encoding. - """ - H = node_haplotypes(ts, [seg.parent for seg in path]) - h = np.zeros(ts.num_sites, dtype=np.int8) - for seg, parent_haplotype in zip(path, H): - left, right = np.searchsorted(ts.sites_position, [seg.left, seg.right]) - h[left:right] = parent_haplotype[left:right] - return h - - -def force_match( - sample, - path, - parent_haplotype, - sites_position, - deletions_as_missing, - num_mismatches, -): - """ - Set the sample's ``hmm_match`` to the match obtained by forcing it down the - specified path, with one mutation at each site where the sample's - non-missing haplotype differs from the haplotype copied along the path. - """ - h = sample.haplotype - if deletions_as_missing: - h = np.where(h == DELETION, MISSING, h) - (sites,) = np.where( - (h != parent_haplotype) & (h != MISSING) & (parent_haplotype != MISSING) - ) - mutations = [ - MatchMutation( - derived_state=core.IUPAC_ALLELES[h[site_id]], - inherited_state=core.IUPAC_ALLELES[parent_haplotype[site_id]], - site_id=int(site_id), - site_position=int(sites_position[site_id]), - # A forced path has no HMM-match artefacts, so there are no - # reversions to mark up for tree building. Parsimony over the full - # haplotypes handles reversions on the group's root branch instead. - is_reversion=False, - is_immediate_reversion=False, - ) - for site_id in sites - ] - sample.hmm_match = HmmMatch(list(path), mutations) - sample.hmm_match.compute_cost(num_mismatches) - - def characterise_recombinants(ts, samples): """ Update the metadata for any recombinants to add interval information to the metadata. diff --git a/sc2ts/tree_ops.py b/sc2ts/tree_ops.py index 749a8ca..63d701c 100644 --- a/sc2ts/tree_ops.py +++ b/sc2ts/tree_ops.py @@ -148,23 +148,12 @@ def add_tree_to_tables(tables, pi, tau): tables.edges.add_row(0, L, parent, u) -def infer_binary_topology(ts): - """ - Infer a strictly binary topology from the variation data in the specified - tree sequence, returning it as an oriented forest ``(pi, tau)``. - - The root of ``ts`` is included as an outgroup in the tree building so that - the unrooted NJ tree can be rooted; node ``ts.num_samples`` in the returned - arrays is that outgroup, and is the root of the returned topology. - - Returns ``None`` if there are fewer than two samples, in which case there - is no topology to infer. - """ +def infer_binary_topology(ts, tables): assert ts.num_trees == 1 assert ts.num_mutations > 0 if ts.num_samples < 2: - return None + return tables.tree_sequence() samples = ts.samples() tree = ts.first() @@ -187,23 +176,20 @@ def infer_binary_topology(ts): # Node n - 1 is the pre-specified root, so force rerooting around that. reroot(pi, n - 1) - assert n == ts.num_samples + 1 + assert n == len(tables.nodes) + 1 tau = max_leaf_distance(pi, n) tau /= max(1, np.max(tau)) - return pi, tau + add_tree_to_tables(tables, pi, tau) + tables.sort() + + return tables.tree_sequence() # TODO rename this to infer_sample_group_tree -def infer_binary(ts, topology=None): +def infer_binary(ts): """ Infer a strictly binary tree from the variation data in the specified tree sequence. - - If ``topology`` is not None it must be an oriented forest ``(pi, tau)`` as - returned by :func:`infer_binary_topology`, and is used directly instead of - inferring the topology from ``ts``. This lets a group's topology be - inferred once and then reused when the mutations are re-mapped against a - different ancestral haplotype. """ assert ts.num_trees == 1 assert list(ts.samples()) == list(range(ts.num_samples)) @@ -216,11 +202,7 @@ def infer_binary(ts, topology=None): tables.nodes.truncate(ts.num_samples) # Update the tables with the topology - if topology is None: - topology = infer_binary_topology(ts) - if topology is not None: - add_tree_to_tables(tables, *topology) - tables.sort() + infer_binary_topology(ts, tables) binary_ts = tables.tree_sequence() # Now add on mutations under parsimony diff --git a/tests/test_inference.py b/tests/test_inference.py index 4398bea..946ba22 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -561,22 +561,21 @@ def haplotype(self, mutations=None): def infer(self, haplotypes): ts = self.base_ts - return si.infer_seed_group_root(ts, haplotypes, si.reference_haplotype(ts)) + samples = [si.Sample(f"s{j}", haplotype=h) for j, h in enumerate(haplotypes)] + return si.infer_seed_group_root(ts, samples, si.reference_haplotype(ts)) def test_single_haplotype(self): h = self.haplotype({4: "T"}) - topology, root = self.infer([h]) + root = self.infer([h]) # Nothing to infer over one sample: the ancestor is the sample. - assert topology is None nt.assert_array_equal(root, h) assert root is not h def test_identical_to_reference(self): reference = si.reference_haplotype(self.base_ts) - topology, root = self.infer([self.haplotype(), self.haplotype()]) + root = self.infer([self.haplotype(), self.haplotype()]) # No mutations, so there is no tree to infer and the ancestor is the # reference. - assert topology is None nt.assert_array_equal(root, reference) def test_shared_derived_allele(self): @@ -585,8 +584,7 @@ def test_shared_derived_allele(self): self.haplotype({4: "T", 12: "A"}), self.haplotype({4: "T", 21: "A"}), ] - topology, root = self.infer(haplotypes) - assert topology is not None + root = self.infer(haplotypes) # The ancestor carries the shared mutation but neither private one. expected = reference.copy() expected[4] = jit.encode_alleles(np.array(["T"]))[0] @@ -598,7 +596,7 @@ def test_all_missing_site_is_missing(self): h = self.haplotype({pos: "A"}) h[4] = si.MISSING haplotypes.append(h) - _, root = self.infer(haplotypes) + root = self.infer(haplotypes) # Every member is missing at site 4, so the ancestor is unknown there # rather than asserting the reference allele. assert root[4] == si.MISSING @@ -606,16 +604,191 @@ def test_all_missing_site_is_missing(self): def test_partially_missing_site_is_not_missing(self): haplotypes = [self.haplotype({4: "T"}), self.haplotype()] haplotypes[1][4] = si.MISSING - _, root = self.infer(haplotypes) + root = self.infer(haplotypes) assert root[4] != si.MISSING +class TestFlatGroupTs: + reference = "AAAACCCCGGGGTTTTAAAACCCCGGGGTT" + + @property + def base_ts(self): + return si.initial_ts( + reference_sequence="X" + self.reference, + reference_id="chr_test", + reference_date="2019-01-01", + ) + + def haplotype(self, mutations=None): + h = list(self.reference) + for pos, base in (mutations or {}).items(): + h[pos] = base + return jit.encode_alleles(np.array(h)) + + def sample(self, strain, mutations=None): + sample = si.Sample(strain, date="2020-01-01", metadata={"date": "2020-01-01"}) + sample.haplotype = self.haplotype(mutations) + sample.alignment_composition = {} + sample.hmm_match = si.HmmMatch([si.PathSegment(0, 31, 1)], []) + return sample + + def test_mutations_relative_to_ancestral_haplotype(self): + ts = self.base_ts + # An ancestral haplotype that is not the reference: the sample matches it + # at position 4 and differs from it at 8. + ancestral = self.haplotype({4: "T", 8: "A"}) + samples = [self.sample("a", {4: "T"})] + flat_ts = si.flat_group_ts(ts, samples, ancestral) + # Only position 8 is variable, and its ancestral state is the ancestral + # haplotype's allele rather than the reference's. + assert flat_ts.num_sites == 1 + site = flat_ts.site(0) + assert site.position == 9 + assert site.ancestral_state == "A" + assert [m.derived_state for m in site.mutations] == ["G"] + + def test_missing_reads_as_ancestral(self): + ts = self.base_ts + h = self.haplotype() + h[4] = si.MISSING + samples = [self.sample("a")] + samples[0].haplotype = h + flat_ts = si.flat_group_ts(ts, samples, self.haplotype({4: "T"})) + assert flat_ts.num_mutations == 0 + + def test_deletions_as_missing(self): + ts = self.base_ts + samples = [self.sample("a", {4: "-"})] + reference = si.reference_haplotype(ts) + assert si.flat_group_ts(ts, samples, reference).num_mutations == 1 + flat_ts = si.flat_group_ts(ts, samples, reference, deletions_as_missing=True) + assert flat_ts.num_mutations == 0 + + def test_bare_nodes_without_group_id(self): + ts = self.base_ts + samples = [self.sample("a", {4: "T"})] + flat_ts = si.flat_group_ts(ts, samples, si.reference_haplotype(ts)) + assert flat_ts.num_samples == 1 + assert flat_ts.node(0).metadata == {} + + def test_full_metadata_with_group_id(self): + ts = self.base_ts + samples = [self.sample("a", {4: "T"})] + flat_ts = si.flat_group_ts( + ts, samples, si.reference_haplotype(ts), group_id="abc" + ) + md = flat_ts.node(0).metadata + assert md["date"] == "2020-01-01" + assert md["sc2ts"]["group_id"] == "abc" + assert md["sc2ts"]["hmm_match"] == samples[0].hmm_match.asdict() + + def test_missing_ancestral_allele_rejected(self): + ts = self.base_ts + ancestral = si.reference_haplotype(ts) + ancestral[4] = si.MISSING + with pytest.raises(AssertionError): + si.flat_group_ts(ts, [self.sample("a")], ancestral) + + +class TestSeedGroupAttachment: + """ + A seed group's local tree must be inferred against the haplotype of the node + its ancestor matched to, not against the reference, because attach_tree hangs + the group's root children directly off that node. + """ + + reference = "AAAACCCCGGGGTTTTAAAACCCCGGGGTT" + # The day-one sample the group's ancestor will match to. It shares three + # mutations with the group, and has a fourth at position 12 where the group + # keeps the reference allele. + parent = {4: "T", 8: "A", 16: "C", 12: "G"} + # Two seeds, sharing the parent's other three mutations plus one private one. + seeds = { + "s0": {4: "T", 8: "A", 16: "C", 20: "A"}, + "s1": {4: "T", 8: "A", 16: "C", 24: "T"}, + } + + def alignment(self, mutations): + h = list(self.reference) + for pos, base in mutations.items(): + assert h[pos] != base + h[pos] = base + return jit.encode_alleles(np.array(h)) + + def build(self, tmp_path): + alignments = {"p": self.alignment(self.parent)} + for name, mutations in self.seeds.items(): + alignments[name] = self.alignment(mutations) + ds = sc2ts.dataset.tmp_dataset( + tmp_path / "ds.zarr", + alignments, + date=["2020-01-01", "2020-01-02", "2020-01-02"], + contig_id="chr_test", + ) + base_ts = si.initial_ts( + reference_sequence="X" + self.reference, + reference_id="chr_test", + reference_date="2019-01-01", + ) + base_path = tmp_path / "base.ts" + base_ts.dump(base_path) + match_db = si.MatchDb.initialise(tmp_path / "match.db") + for date in ["2020-01-01", "2020-01-02"]: + ts = si.extend( + dataset=ds.path, + base_ts=str(base_path), + date=date, + match_db=str(match_db.path), + min_group_size=1, + seed_groups=[list(self.seeds)], + ) + ts.dump(base_path) + return ts, alignments + + def test_seed_haplotypes_are_preserved(self, tmp_path): + ts, alignments = self.build(tmp_path) + strains = ts.metadata["sc2ts"]["samples_strain"] + # Inferring the group's tree against the reference instead would leave + # position 12 out of the flat tree entirely, so both seeds would silently + # inherit the parent's allele there. + for name in self.seeds: + u = ts.samples()[strains.index(name)] + nt.assert_array_equal(si.node_haplotypes(ts, [u])[0], alignments[name]) + + def test_group_matched_to_parent_sample(self, tmp_path): + ts, _ = self.build(tmp_path) + strains = ts.metadata["sc2ts"]["samples_strain"] + # The ancestor is one mutation from p and three from the reference, so + # the group must match p rather than the root. Assert on the match rather + # than the final topology, which push_up_reversions rearranges: the + # group's G13T is an immediate reversion of p's T13G. + p = ts.samples()[strains.index("p")] + for name in self.seeds: + u = ts.samples()[strains.index(name)] + path = ts.node(u).metadata["sc2ts"]["hmm_match"]["path"] + assert path[0]["parent"] == p + + def test_members_share_the_ancestor_hmm_match(self, tmp_path): + ts, _ = self.build(tmp_path) + strains = ts.metadata["sc2ts"]["samples_strain"] + matches = [] + for name in self.seeds: + u = ts.samples()[strains.index(name)] + hmm_match = ts.node(u).metadata["sc2ts"]["hmm_match"] + # Recombination is disallowed for seeds. + assert len(hmm_match["path"]) == 1 + matches.append(hmm_match) + # The stored match is the one that placed the group, so every member + # reports the same thing. + assert matches[0] == matches[1] + + class TestSeedGroup: def example(self): samples = [si.Sample("a"), si.Sample("b")] group = si.SeedGroup(strains=("a", "b"), date="2021-01-01") group.samples = samples - group.topology = ([-1, 2, -1], [0, 0, 1]) + group.flat_ts = tskit.Tree.generate_star(2, span=10).tree_sequence group.root = si.Sample("seed_root") group.root.hmm_match = si.HmmMatch([si.PathSegment(0, 10, 3)], []) return group @@ -625,7 +798,7 @@ def test_specification_only(self): assert group.strains == ("a", "b") assert group.date == "2021-01-01" assert group.samples is None - assert group.topology is None + assert group.flat_ts is None assert group.root is None assert len(group) == 2 @@ -646,7 +819,7 @@ def test_sample_group(self): assert sample_group.samples is group.samples assert sample_group.path == group.path assert sample_group.immediate_reversions == () - assert sample_group.topology is group.topology + assert sample_group.flat_ts is group.flat_ts def test_sample_hash_matches_sample_group(self): # The group ID goes into node metadata, so the two classes must agree. diff --git a/tests/test_tree_ops.py b/tests/test_tree_ops.py index 4d9b737..3c934cb 100644 --- a/tests/test_tree_ops.py +++ b/tests/test_tree_ops.py @@ -771,70 +771,6 @@ def test_simulation_root_mutations(self, n, num_mutations): self.check_properties(ts2) -def flat_ts(n, genotypes): - """ - A star tree over n samples with one site per column of genotypes, ancestral - state "A" and a mutation wherever the genotype is not "A". - """ - L = len(genotypes[0]) + 1 - tables = tskit.TableCollection(L) - root = n - for j in range(n): - u = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) - tables.edges.add_row(0, L, root, u) - tables.nodes.add_row(time=1) - for site in range(len(genotypes[0])): - tables.sites.add_row(site, "A") - for j in range(n): - if genotypes[j][site] != "A": - tables.mutations.add_row( - site=site, node=j, derived_state=genotypes[j][site] - ) - tables.sort() - return tables.tree_sequence() - - -class TestInferBinaryTopology: - def test_single_sample(self): - ts = flat_ts(1, ["T"]) - assert tree_ops.infer_binary_topology(ts) is None - - @pytest.mark.parametrize("n", [2, 3, 4, 8]) - def test_outgroup_is_root(self, n): - # Each sample gets its own private mutation. - genotypes = ["A" * j + "T" + "A" * (n - j - 1) for j in range(n)] - ts = flat_ts(n, genotypes) - pi, tau = tree_ops.infer_binary_topology(ts) - assert len(tau) == len(pi) - # The n samples and the outgroup are all leaves, so there must be - # internal nodes on top of them. - assert len(pi) > n + 1 - # Node n is the outgroup, and after rerooting it is the sole root. - assert pi[n] == -1 - assert tau[n] == 1 - assert sum(p == -1 for p in pi) == 1 - # Every sample is a leaf, i.e. nothing points at it. - assert set(pi) & set(range(n)) == set() - - def test_supplied_topology_is_used(self): - # Two well separated pairs, so the topology is unambiguous. - genotypes = ["TTAAAA", "TTAAAT", "AAATTA", "AATTTA"] - ts1 = flat_ts(4, genotypes) - topology = tree_ops.infer_binary_topology(ts1) - expected = tree_ops.infer_binary(ts1, topology) - - # A different set of mutations that would give a different topology if - # it were inferred rather than supplied. - ts2 = flat_ts(4, ["TAAAAA", "AATAAA", "AAAATA", "AAAAAT"]) - inferred = tree_ops.infer_binary(ts2) - supplied = tree_ops.infer_binary(ts2, topology) - - assert supplied.tables.edges == expected.tables.edges - assert inferred.tables.edges != expected.tables.edges - # The mutations still come from ts2, not from ts1. - assert_sequences_equal(ts2, supplied) - - class TestFromBiotite: def check_round_trip(self, tsk_tree): node_labels = {u: f"{u}" for u in tsk_tree.tree_sequence.samples()}