diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ea3f9b..ae15f7a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,11 @@ In development usually far diverged from the current ARG, which previously led to spurious recombinations. +- Entries in `include_samples` may now be `(strain, match_date)` tuples (in + addition to bare strain IDs). A non-None `match_date` matches the seed in on + that date instead of its actual date; it may be *before* the actual date, in + which case the seed node retains its actual date via a negative node time. + - Add basic support for non-SARS-CoV-2 genomes via an optional reference FASTA. Supply `--reference` to `import-alignments` and a `reference_fasta` key in the inference config; both default to the built-in SARS-CoV-2 reference, so diff --git a/sc2ts/inference.py b/sc2ts/inference.py index a543388..fc86863 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,21 @@ def extend( num_threads=0, memory_limit=0, ): - + """ + Extend the base tree sequence by one day, matching in the samples for the + given date. + + ``include_samples`` is an optional list of "seed" samples that are matched + in unconditionally and without recombination. Each entry is either a bare + strain ID, or a ``(strain, match_date)`` tuple. When ``match_date`` is not + None the seed is matched in on ``match_date`` instead of its actual date; + this may be *before* the actual date, in which case the seed node retains + its actual date and is given a negative ("in the future") node time + relative to the match-day time-zero. + A ``match_date`` must be a date that is otherwise processed by the pipeline + (i.e. has real samples in the dataset); otherwise the seed is never + injected. + """ if num_mismatches is None: num_mismatches = 3 if hmm_cost_threshold is None: @@ -552,6 +587,7 @@ def extend( deletions_as_missing = False if include_samples is None: include_samples = [] + include_samples = normalise_include_samples(include_samples) base_ts = str(base_ts) dataset = str(dataset) match_db = str(match_db) @@ -565,6 +601,20 @@ def extend( base_ts = tszip.load(base_ts) ds = _dataset.Dataset(dataset, date_field=date_field) + missing = sorted( + strain for strain, _ in include_samples if strain not in ds.metadata + ) + if len(missing) > 0: + raise ValueError(f"Seed samples not in dataset: {missing}") + for strain, match_date in include_samples: + if match_date is not None: + actual_date = ds.metadata[strain]["date"] + if match_date > actual_date: + logger.warning( + f"Seed sample {strain} match date {match_date} is after its " + f"actual date {actual_date}; this is unusual for a seed sample" + ) + with MatchDb(match_db) as matches: tables = _extend( dataset=ds, @@ -625,10 +675,27 @@ def _extend( f"mutations={base_ts.num_mutations};date={previous_date}" ) + include_strains = {strain for strain, _ in include_samples} + # Seeds with an explicit match date are matched in on that date instead of + # their actual date, so we override which day they're processed on. + override_dates = { + strain: match_date + for strain, match_date in include_samples + if match_date is not None + } + metadata_matches = { strain: dataset.metadata[strain] for strain in dataset.metadata.samples_for_date(date) + # Exclude a seed with an override match date from its actual date; it + # is processed only on the override date. + if override_dates.get(strain, date) == date } + # Inject seeds whose override match date is today but which aren't + # naturally sampled today. Missing strains are rejected up-front in extend(). + for strain, match_date in override_dates.items(): + if match_date == date and strain not in metadata_matches: + metadata_matches[strain] = dataset.metadata[strain] logger.info(f"Got {len(metadata_matches)} metadata matches") @@ -643,7 +710,6 @@ def _extend( pango_lineage_key = "Viridian_pangolin" scorpio_key = "Viridian_scorpio" - include_strains = set(include_samples) unconditional_include_samples = [] samples = [] for s in preprocessed_samples: @@ -1455,6 +1521,9 @@ def match_tsinfer( num_alleles = 4 if deletions_as_missing else 5 mu, rho = solve_num_mismatches(num_mismatches, num_alleles) + # Detach any future (negative-time) nodes so that samples can't copy from + # them. Node IDs are preserved, so the returned match paths stay valid. + ts = tree_ops.detach_future_nodes(ts) tsb, coord_map = make_tsb(ts, num_alleles, mirror_coordinates) work = [] @@ -1755,7 +1824,11 @@ def attach_tree( node = child_ts.node(u) sample_date = parse_date(node.metadata["date"]) node_time[u] = (current_date - sample_date).days - assert node_time[u] >= 0.0 + # Seed samples can be matched in on a date before their actual + # date, giving a negative ("in the future") node time. + assert node_time[u] >= 0.0 or ( + node.flags & core.NODE_IS_UNCONDITIONALLY_INCLUDED + ) max_sample_time = max(node_time.values()) node_id_map = {} diff --git a/sc2ts/tree_ops.py b/sc2ts/tree_ops.py index 6e22fba..63d701c 100644 --- a/sc2ts/tree_ops.py +++ b/sc2ts/tree_ops.py @@ -657,6 +657,41 @@ def drop_vestigial_root_edge(ts): return tables.tree_sequence() +def detach_future_nodes(ts): + """ + Return a copy of ``ts`` in which every node with a negative ("in the + future") time is fully detached: all edges incident to it are removed, any + mutations over it are dropped, and its ``NODE_IS_SAMPLE`` flag is cleared. + + Such nodes are seed samples that were matched in on a date before their + actual date, and so lie in the future relative to the current time-zero. + Detaching them prevents other samples from copying from them during + matching. Node IDs and times are preserved so that any match paths + referring to the returned tree sequence remain valid against the original. + """ + future = ts.nodes_time < 0 + if not np.any(future): + return ts + tables = ts.dump_tables() + keep_edges = ~(future[tables.edges.parent] | future[tables.edges.child]) + keep_mutations = ~future[tables.mutations.node] + logger.debug( + f"Detaching {int(np.sum(future))} future nodes " + f"({len(tables.edges) - int(np.sum(keep_edges))} edges, " + f"{len(tables.mutations) - int(np.sum(keep_mutations))} mutations removed)" + ) + tables.edges.keep_rows(keep_edges) + tables.mutations.keep_rows(keep_mutations) + # A detached future node must not be treated as a sample to copy from. + flags = tables.nodes.flags + flags[future] &= ~np.uint32(tskit.NODE_IS_SAMPLE) + tables.nodes.flags = flags + tables.sort() + tables.build_index() + tables.compute_mutation_parents() + return tables.tree_sequence() + + def insert_vestigial_root_edge(ts): """ Insert an edge between node 0 and 1 at the end of the edge table, if diff --git a/tests/test_cli.py b/tests/test_cli.py index dd066c5..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,44 @@ def test_include_samples(self, tmp_path, fx_ts_map, fx_dataset): assert np.sum(ts.nodes_time[ts.samples()] == 0) == 1 assert ts.num_samples == 1 + def test_include_samples_missing_strain(self, tmp_path, fx_ts_map, fx_dataset): + # A seed strain not in the dataset makes the run fail. + config_file = self.make_config( + tmp_path, + fx_dataset, + exclude_sites=[56, 57, 58, 59, 60], + include_samples=["SRR14631544", "NO_SUCH_STRAIN"], + ) + runner = ct.CliRunner() + with pytest.raises(ValueError, match="not in dataset"): + runner.invoke( + cli.cli, + f"infer {config_file} --stop 2020-01-02", + catch_exceptions=False, + ) + + def test_include_samples_with_dates(self, tmp_path, fx_ts_map, fx_dataset): + # The (strain, date) tuple form is expressed in TOML as a 2-element + # array, mixed with bare strings. + config_file = self.make_config( + tmp_path, + fx_dataset, + exclude_sites=[56, 57, 58, 59, 60], + include_samples=[["SRR14631544", "2020-01-01"]], + ) + runner = ct.CliRunner() + result = runner.invoke( + cli.cli, + f"infer {config_file} --stop 2020-01-02", + catch_exceptions=False, + ) + assert result.exit_code == 0 + date = "2020-01-01" + ts_path = tmp_path / "results" / "test" / f"test_{date}.ts" + ts = tskit.load(ts_path) + assert "SRR14631544" in ts.metadata["sc2ts"]["samples_strain"] + assert ts.num_samples == 1 + def test_override(self, tmp_path, fx_ts_map, fx_dataset): hmm_cost_threshold = 47 config_file = self.make_config( diff --git a/tests/test_inference.py b/tests/test_inference.py index 5d838bb..b378b71 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,14 @@ 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"], + # 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 +658,141 @@ 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_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 + ): + # 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"] + + @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, tmp_path, diff --git a/tests/test_tree_ops.py b/tests/test_tree_ops.py index c8bb5bd..3c934cb 100644 --- a/tests/test_tree_ops.py +++ b/tests/test_tree_ops.py @@ -957,3 +957,188 @@ 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 _incident_edges(ts, node): + return [(e.left, e.right, e.parent) for e in ts.edges() if 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): + # 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_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) + # Nodes preserved (ids and times), future node now isolated. + assert result.num_nodes == ts.num_nodes + 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) + + result = tree_ops.detach_future_nodes(ts) + assert result.num_edges == 1 + 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) + + result = tree_ops.detach_future_nodes(ts) + # 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}