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] diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ea3f9b..43ceaa1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,10 +4,27 @@ 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. +- 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. + + 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 + 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 diff --git a/docs/example_config.toml b/docs/example_config.toml index 0ee7863..0ec944d 100644 --- a/docs/example_config.toml +++ b/docs/example_config.toml @@ -82,9 +82,19 @@ num_threads=-1 # limit. memory_limit=32 -# A list of sample IDs (strings) for unconditional inclusion (e.g., to -# help seed major saltation events). -include_samples=[] +# 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. +# +# 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/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 a543388..283a20b 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,88 @@ 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. + """ + samples = [] + for sample in preprocess( + strains=strains, + dataset=dataset, + keep_sites=keep_sites, + progress_title=progress_title, + show_progress=show_progress, + ): + 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 + 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 = {} + for entry in seed_groups: + if isinstance(entry, str) or not isinstance(entry, (list, tuple)): + raise ValueError( + f"Each seed_groups entry must be a list of strain IDs, not {entry!r}" + ) + if len(entry) == 0: + 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}") + group = tuple(entry) + for strain in group: + if strain in seen: + raise ValueError( + f"Seed sample {strain} appears in more than one seed 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 [ + SeedGroup( + strains=strains, + date=min(metadata[strain]["date"] for strain in strains), + ) + for strains in groups + ] + + def extend( *, dataset, @@ -509,7 +595,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, @@ -527,7 +613,26 @@ def extend( num_threads=0, memory_limit=0, ): - + """ + Extend the base tree sequence by one day, matching in the samples for the + given date. + + ``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 + 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 if hmm_cost_threshold is None: @@ -550,8 +655,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) @@ -565,13 +670,15 @@ def extend( base_ts = tszip.load(base_ts) ds = _dataset.Dataset(dataset, date_field=date_field) + seed_groups = check_seed_groups(seed_groups, ds.metadata) + with MatchDb(match_db) as matches: tables = _extend( dataset=ds, 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, @@ -601,7 +708,7 @@ def _extend( date, base_ts, match_db, - include_samples, + seed_groups, num_mismatches, hmm_cost_threshold, min_group_size, @@ -625,50 +732,34 @@ def _extend( f"mutations={base_ts.num_mutations};date={previous_date}" ) - metadata_matches = { - strain: dataset.metadata[strain] - for strain in dataset.metadata.samples_for_date(date) - } + # 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] - logger.info(f"Got {len(metadata_matches)} metadata matches") + strains = [ + strain + for strain in dataset.metadata.samples_for_date(date) + 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" - - 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}") - 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 include_strains: - s.flags |= core.NODE_IS_UNCONDITIONALLY_INCLUDED - unconditional_include_samples.append(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}" + f"Filter {sample.strain}: " + f"missing={num_missing_sites} > {max_missing_sites}" ) if max_daily_samples is not None: @@ -680,48 +771,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) @@ -736,6 +803,23 @@ def _extend( phase="close", ) + # 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, + 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) + logger.info("Looking for retrospective matches") assert min_group_size is not None earliest_date = parse_date(date) - datetime.timedelta(days=retrospective_window) @@ -869,6 +953,60 @@ def match_path_ts(group, sequence_length): return tables.tree_sequence() +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 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 + root = len(samples) + site_id_map = {} + 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], + core.IUPAC_ALLELES[ancestral_haplotype[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}") @@ -937,6 +1075,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: """ @@ -949,12 +1097,13 @@ class SampleGroup: immediate_reversions: List = None sample_hash: str = None tree_quality_metrics: GroupTreeQualityMetrics = 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): - 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): @@ -1001,26 +1150,61 @@ 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, -): +@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 + root: Sample = None + flat_ts: tskit.TreeSequence = 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=(), + flat_ts=self.flat_ts, + ) + + +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) @@ -1043,14 +1227,36 @@ def add_matching_results( groups = [ SampleGroup( - samples, - key[0], - key[1], + samples=samples, + path=key[0], + immediate_reversions=key[1], ) 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 = [] @@ -1073,7 +1279,9 @@ def add_matching_results( 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: @@ -1141,6 +1349,142 @@ def add_matching_results( return ts, added_groups +def infer_seed_group_root(ts, samples, reference, deletions_as_missing=False): + """ + 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, + samples, + reference, + deletions_as_missing=deletions_as_missing, + ) + 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) + 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 root_haplotype + + +def add_seed_groups( + ts, + base_ts, + seed_groups, + date, + *, + dataset, + deletions_as_missing=False, + show_progress=False, + num_threads=0, + memory_limit=-1, +): + """ + 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 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. + """ + 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) + for group in seed_groups: + assert group.date == date + group.samples = [sample_map[strain] for strain in group.strains] + group.root = Sample( + strain=f"seed_root_{group.sample_hash}", + date=date, + haplotype=infer_seed_group_root( + base_ts, + group.samples, + reference, + deletions_as_missing=deletions_as_missing, + ), + ) + + # 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 seed_groups], + 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, + ) + + 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: + # 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()}") + 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) + + ts, _ = add_sample_groups( + ts, + [group.sample_group() for group in seed_groups], + date, + min_group_size=1, + show_progress=show_progress, + phase="seed", + ) + logger.info(f"Added {num_samples} seed samples in {len(seed_groups)} groups") + return ts + + def solve_num_mismatches(k, num_alleles=5): """ Return the low-level LS parameters corresponding to accepting @@ -1455,6 +1799,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 = [] @@ -1692,6 +2039,34 @@ 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 characterise_recombinants(ts, samples): """ Update the metadata for any recombinants to add interval information to the metadata. @@ -1755,7 +2130,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..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", "NO_SUCH_STRAIN"], + seed_groups=[["SRR14631544"]], ) runner = ct.CliRunner() result = runner.invoke( @@ -440,6 +440,39 @@ 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_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], + seed_groups=[["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_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], + seed_groups=["SRR14631544"], + ) + runner = ct.CliRunner() + 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 config_file = self.make_config( 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 5d838bb..946ba22 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,355 @@ def test_high_recomb_mutation(self): self.check_double_mirror(ts) +class TestCheckSeedGroups: + metadata = { + "a": {"date": "2021-01-03"}, + "b": {"date": "2021-01-01"}, + "c": {"date": "2021-01-02"}, + } + + def test_empty(self): + assert si.check_seed_groups([], self.metadata) == [] + + def test_groups(self): + result = si.check_seed_groups([["a", "b"], ["c"]], self.metadata) + assert [g.strains for g in result] == [("a", "b"), ("c",)] + + def test_tuples_accepted(self): + 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"): + si.check_seed_groups(["a"], self.metadata) + + def test_empty_group_raises(self): + with pytest.raises(ValueError, match="must not be empty"): + si.check_seed_groups([[]], self.metadata) + + def test_non_string_strain_raises(self): + with pytest.raises(ValueError, match="must be strings"): + si.check_seed_groups([["a", 7]], self.metadata) + + def test_duplicate_strain_raises(self): + 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 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_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 + 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"}) + root = self.infer([h]) + # Nothing to infer over one sample: the ancestor is the sample. + nt.assert_array_equal(root, h) + assert root is not h + + def test_identical_to_reference(self): + reference = si.reference_haplotype(self.base_ts) + root = self.infer([self.haplotype(), self.haplotype()]) + # No mutations, so there is no tree to infer and the ancestor is the + # reference. + 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"}), + ] + 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] + 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 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.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 + + 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.flat_ts 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.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. + 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: dates = [ "2020-01-01", @@ -550,22 +949,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"], ["SRR11597115", "NOSUCHSTRAIN"]) - ) - def test_2020_02_02_include_samples( + def test_2020_02_02_seed_group( 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, + seed_groups=[["SRR11597115"]], ) assert ts.metadata["sc2ts"]["cumulative_stats"]["exact_matches"]["pango"] == { "A": 2, @@ -585,6 +980,178 @@ def test_2020_02_02_include_samples( assert edges[0].left == 0 assert edges[0].right == ts.sequence_length + 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"), + seed_groups=[group], + ) + 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"), + seed_groups=[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 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"), + seed_groups=[["SRR11597115"], ["SRR11597190"]], + ) + 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_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 + # 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"] + 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") + expected_time = { + "2020-01-31": -2, + "2020-02-01": -1, + "2020-02-02": 0, + "2020-02-03": 1, + } + for date, expected in expected_time.items(): + ts = si.extend( + dataset=fx_dataset.path, + base_ts=base_path, + date=date, + match_db=match_db.path, + seed_groups=seed_groups, + ) + ts.dump(base_path) + strains = 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( + "seed_groups", + ( + [["SRR11597115"], ["NOSUCHSTRAIN"]], + [["SRR11597115", "NOSUCHSTRAIN"]], + [["NOSUCHSTRAIN"]], + ), + ) + def test_seed_missing_strain_raises( + 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. + 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"), + seed_groups=seed_groups, + ) + + 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"), + seed_groups=["SRR11597115"], + ) + def test_2020_02_02_mutation_overlap( self, tmp_path, @@ -798,7 +1365,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, @@ -811,6 +1379,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 @@ -829,7 +1402,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, @@ -854,7 +1427,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 @@ -910,7 +1483,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( 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}