diff --git a/README.md b/README.md index 89dc57ca..2c56dbed 100644 --- a/README.md +++ b/README.md @@ -652,6 +652,7 @@ manually). than regenerating them. - `--compress_features` – compress the generated features to save space: `*.pkl.xz` for the AlphaFold2 pipeline, `*_af3_input.json.xz` for AlphaFold3. Both are read back transparently, so compressed feature sets can be used directly (this is how the [features database](https://alphapulldown.s3.embl.de) ships them). - `--skip_existing` – leave existing feature files untouched (safe for reruns). +- `--keep_msas` – refresh **templates only** in features that already exist in `--output_dir`, keeping their MSAs. Use it when the template database or `--max_template_date` has moved but the alignments are still valid: it costs a template search (minutes) instead of a full MSA run (hours). Works for both pipelines — AlphaFold2 features get their `template_*` block replaced, AlphaFold3 features are re-processed through AF3's "search for templates only" path. Proteins with no stored features are generated normally, and it takes precedence over `--skip_existing`. Cannot be combined with `--use_mmseqs2` (which fetches MSAs and templates together) or `--skip_msa` (no MSAs to keep). - `--seq_index N` – only process the N‑th sequence from the FASTA list. - `--use_hhsearch`, `--re_search_templates_mmseqs2` – toggle template search implementations. - `--path_to_mmt`, `--description_file`, `--multiple_mmts` – enable TrueMultimer CSV-driven feature sets. diff --git a/alphapulldown/objects.py b/alphapulldown/objects.py index 41a9a5d3..d9e61b23 100644 --- a/alphapulldown/objects.py +++ b/alphapulldown/objects.py @@ -21,6 +21,7 @@ from alphapulldown.utils.multimeric_template_utils import (extract_multimeric_template_features_for_single_chain, prepare_multimeric_template_meta_info) from alphapulldown.utils.file_handling import temp_fasta_file +from alphapulldown.utils.msa_integrity import validate_precomputed_msas from alphapulldown.utils.mmseqs_species_identifiers import ( enrich_mmseq_feature_dict_with_identifiers, strip_mmseq_comment_lines, @@ -202,6 +203,15 @@ def make_features( logging.info( "will save msa files in :{}".format(msa_output_dir)) plPath(msa_output_dir).mkdir(parents=True, exist_ok=True) + + # AlphaFold reuses whatever alignments it finds here without checking + # them, so a file left behind by a killed job is read back as if it were + # complete. Drop the unsound ones first: an empty uniref90_hits.sto + # otherwise surfaces as StopIteration deep inside the Stockholm parser, + # and a *partially* written one silently yields shallower features. + if use_precomputed_msa: + validate_precomputed_msas(msa_output_dir, remove_invalid=True) + self.skip_msa = skip_msa if skip_msa: self.feature_dict = self._build_query_only_feature_dict() diff --git a/alphapulldown/scripts/create_individual_features.py b/alphapulldown/scripts/create_individual_features.py index 7d623f51..47aaf6cd 100644 --- a/alphapulldown/scripts/create_individual_features.py +++ b/alphapulldown/scripts/create_individual_features.py @@ -29,13 +29,22 @@ # AlphaPulldown helpers from alphapulldown.utils.create_custom_template_db import create_db from alphapulldown.objects import MonomericObject -from alphapulldown.utils.file_handling import iter_seqs, parse_csv_file, resolve_af3_input +from alphapulldown.utils.file_handling import ( + iter_seqs, + parse_csv_file, + read_maybe_xz, + resolve_af3_input, +) from alphapulldown.utils.modelling_setup import create_uniprot_runner from alphapulldown.utils.multimeric_template_utils import ( extract_multimeric_template_features_for_single_chain, ) from alphapulldown.utils import save_meta_data from alphapulldown.utils.feature_metadata import embed_metadata_in_af3_json +from alphapulldown.utils.template_reuse import ( + msa_for_template_search, + search_templates, +) # Try to import AlphaFold3, but it's optional AF3_IMPORT_ERROR = None @@ -156,6 +165,14 @@ flags.DEFINE_integer("seq_index", None, "") flags.DEFINE_boolean("use_hhsearch", False, "") flags.DEFINE_boolean("compress_features", False, "") +flags.DEFINE_boolean( + "keep_msas", False, + "Update existing features in --output_dir in place: keep their MSAs and " + "re-run only the template search. Use this when the template database or " + "--max_template_date has changed but the alignments are still valid; it " + "costs a template search instead of a full MSA run. Features that do not " + "already exist are generated normally.", +) flags.DEFINE_string("path_to_mmt", None, "") flags.DEFINE_string("description_file", None, "") flags.DEFINE_float("threshold_clashes", 1000, "") @@ -205,6 +222,20 @@ def validate_data_pipeline_flags(): "AlphaFold3 does not support --use_mmseqs2. " "Please provide local databases via --data_dir." ) + if FLAGS.keep_msas and FLAGS.use_mmseqs2: + # MMseqs2 obtains MSAs and templates together from the remote server; + # there is no separable template search to re-run against a local + # database, so the combination cannot do what the flag promises. + raise ValueError( + "--keep_msas cannot be combined with --use_mmseqs2: the MMseqs2 " + "path fetches MSAs and templates together and has no standalone " + "template search to re-run." + ) + if FLAGS.keep_msas and FLAGS.skip_msa: + raise ValueError( + "--keep_msas cannot be combined with --skip_msa: there are no MSAs " + "to keep when MSA generation is skipped." + ) def get_database_path(key): """Return the absolute path for a given database key, depending on pipeline.""" @@ -468,6 +499,12 @@ def create_individual_features(): for seq_idx, (seq, desc) in enumerate(iter_seqs(FLAGS.fasta_paths), 1): if FLAGS.seq_index is None or seq_idx == FLAGS.seq_index: + # --keep_msas updates existing features in place. When none exist + # yet there is nothing to keep, so fall through and generate them. + if FLAGS.keep_msas and _update_templates_keeping_msas( + desc, seq, pipeline + ): + continue monomer = MonomericObject(desc, seq) monomer.uniprot_runner = uniprot_runner create_and_save_monomer_objects(monomer, pipeline) @@ -562,6 +599,70 @@ def _replace_template_features(monomer, template_features): ) +def _existing_monomer_for_update(description, sequence): + """Load existing features for in-place update, or None to regenerate. + + Shared by the two modes that rewrite an existing monomer's templates + (TrueMultimer reuse and ``--keep_msas``). A sequence mismatch means the + stored features describe a different protein under the same name, so they + must not be updated in place. + """ + monomer = _load_existing_monomer_from_output_dir(description) + if monomer is None: + return None + stored = getattr(monomer, "sequence", None) + if stored != sequence: + logging.warning( + "Existing features for %s hold sequence %s but the input FASTA has " + "%s; regenerating from scratch instead of updating templates.", + description, stored, sequence, + ) + return None + return monomer + + +def _update_templates_keeping_msas(description, sequence, pipeline): + """Re-run template search for one existing monomer, keeping its MSAs. + + Returns True when the stored features were updated, False when there is + nothing to update and the caller should generate features normally. + """ + monomer = _existing_monomer_for_update(description, sequence) + if monomer is None: + return False + + template_searcher = getattr(pipeline, "template_searcher", None) + template_featurizer = getattr(pipeline, "template_featurizer", None) + if template_searcher is None or template_featurizer is None: + raise RuntimeError( + "--keep_msas needs a template searcher and featurizer, but the " + "pipeline provides none. Template search cannot be skipped in this " + "mode, since re-running it is the entire point." + ) + + msa_output_dir = Path(FLAGS.output_dir) / description + msa = msa_for_template_search(monomer, msa_output_dir) + template_features = search_templates( + template_searcher, + template_featurizer, + query_sequence=sequence, + stockholm_msa=msa.stockholm, + msa_output_dir=msa_output_dir, + ) + _replace_template_features(monomer, template_features) + # Keep the provenance with the features: which alignment drove the search + # decides whether the hits reproduce a fresh run (see template_reuse). + monomer.template_msa_source = msa.source + _persist_monomer_outputs(monomer) + logging.info( + "Updated templates for %s from %s (MSAs kept, %d template(s)).", + description, msa.source, + int(np.asarray(template_features.get("template_domain_names", [])).shape[0]) + if "template_domain_names" in template_features else 0, + ) + return True + + def _reuse_truemultimer_monomer_features(feat): if FLAGS.use_mmseqs2 or len(feat["templates"]) != 1 or len(feat["chains"]) != 1: return None @@ -569,7 +670,10 @@ def _reuse_truemultimer_monomer_features(feat): source_name = _infer_truemultimer_source_name( feat["protein"], feat["templates"], feat["chains"] ) - monomer = _load_existing_monomer_from_output_dir(source_name) + # Shares the load-and-validate step with --keep_msas: both modes rewrite an + # existing monomer's templates and must refuse when the stored features + # describe a different sequence. + monomer = _existing_monomer_for_update(source_name, feat["sequence"]) if monomer is None: return None cached_skip_msa = getattr(monomer, "skip_msa", False) @@ -584,15 +688,6 @@ def _reuse_truemultimer_monomer_features(feat): requested_mode, ) return None - if monomer.sequence != feat["sequence"]: - logging.warning( - "Existing monomer features for %s use sequence %s, but the current " - "TrueMultimer entry expects %s. Falling back to full feature generation.", - source_name, - monomer.sequence, - feat["sequence"], - ) - return None template_path = feat["templates"][0] chain_id = feat["chains"][0] @@ -747,6 +842,86 @@ def _ovr(attr, key): ) return AF3DataPipeline(config) +def _af3_chain_without_templates(chain): + """Copy a protein chain keeping its MSAs but clearing its templates. + + AlphaFold3's data pipeline dispatches on which of the three fields are set, + and one branch is exactly what ``--keep_msas`` wants:: + + elif has_unpaired_msa and has_paired_msa and not has_templates: + # Has MSA, but doesn't have templates. Search for templates only. + + The distinction that matters: ``templates=[]`` means "no templates, do not + search", while ``templates=None`` means "search". A chain with only some of + the three set is rejected by AF3 outright, so both MSAs must be carried + over. Non-protein chains have no template search and pass through. + """ + if folding_input is None or not isinstance(chain, folding_input.ProteinChain): + return chain + if chain.unpaired_msa is None or chain.paired_msa is None: + # Nothing to preserve - let the normal path search everything. + return chain + return folding_input.ProteinChain( + id=chain.id, + sequence=chain.sequence, + ptms=chain.ptms, + residue_ids=chain.residue_ids, + description=chain.description, + paired_msa=chain.paired_msa, + unpaired_msa=chain.unpaired_msa, + templates=None, + ) + + +def _af3_input_keeping_msas(existing_path, desc, sequence=None): + """Build an AF3 Input from existing features that re-searches templates only. + + Returns None when the stored file cannot be reused, so the caller falls back + to generating features normally. + """ + try: + payload = read_maybe_xz(existing_path) + input_obj = folding_input.Input.from_json(payload) + except Exception as exc: # malformed/partial file, or an AF3 schema change + logging.warning( + "Cannot reuse existing AF3 features %s for %s (%s); regenerating.", + existing_path, desc, exc, + ) + return None + + # Features are matched to the input by description alone, so an edited FASTA + # that kept its name would otherwise have its MSAs reused for a different + # protein and be overwritten with features for the wrong sequence. The AF2 + # path refuses that in _existing_monomer_for_update; do the same here. + if sequence is not None: + stored = [getattr(chain, "sequence", None) for chain in input_obj.chains] + if sequence not in stored: + logging.warning( + "Existing AF3 features %s hold sequence(s) %s but the input FASTA " + "has %s; regenerating from scratch instead of keeping their MSAs.", + existing_path, + ", ".join(s for s in stored if s) or "none", + sequence, + ) + return None + + chains = [_af3_chain_without_templates(chain) for chain in input_obj.chains] + if not any( + getattr(chain, "unpaired_msa", None) is not None for chain in chains + ): + logging.warning( + "Existing AF3 features %s carry no MSAs to keep; regenerating.", + existing_path, + ) + return None + + return folding_input.Input( + name=input_obj.name, + chains=chains, + rng_seeds=input_obj.rng_seeds, + ) + + def create_af3_individual_features(): """Generate AlphaFold3 features, one .json per chain.""" validate_data_pipeline_flags() @@ -770,10 +945,16 @@ def create_af3_individual_features(): # silently regenerated (and vice versa). plain = Path(FLAGS.output_dir) / f"{desc}_af3_input.json" outpath = Path(str(plain) + ".xz") if FLAGS.compress_features else plain - if FLAGS.skip_existing and resolve_af3_input(plain) is not None: + existing = resolve_af3_input(plain) + # --keep_msas rewrites existing features, so it takes precedence + # over --skip_existing: skipping would defeat the point. + reuse_input = None + if FLAGS.keep_msas and existing is not None: + reuse_input = _af3_input_keeping_msas(existing, desc, seq) + if reuse_input is None and FLAGS.skip_existing and existing is not None: logging.info(f"Feature file for {desc} already exists. Skipping...") continue - + # Create AlphaFold3 input object with proper chain structure try: # Generate proper chain ID using AlphaFold3's int_id_to_str_id function @@ -788,14 +969,20 @@ def create_af3_individual_features(): chain_id = chain_id + chain_id chain_kind = get_af3_chain_kind(desc, seq) - chain = create_af3_chain(seq, desc, chain_id) - - input_obj = folding_input.Input( - name=desc, - chains=[chain], - rng_seeds=[42] - ) - + if reuse_input is not None: + input_obj = reuse_input + logging.info( + "Keeping MSAs for %s and re-running template search only.", + desc, + ) + else: + chain = create_af3_chain(seq, desc, chain_id) + input_obj = folding_input.Input( + name=desc, + chains=[chain], + rng_seeds=[42] + ) + features = pipeline.process(input_obj) if hasattr(features, "to_json"): feature_payload = json.loads(features.to_json()) diff --git a/alphapulldown/utils/msa_integrity.py b/alphapulldown/utils/msa_integrity.py new file mode 100644 index 00000000..34ac79cb --- /dev/null +++ b/alphapulldown/utils/msa_integrity.py @@ -0,0 +1,184 @@ +"""Integrity checks for precomputed MSA files. + +``--use_precomputed_msas`` makes AlphaFold reuse whatever alignment files it +finds next to the output, which is what makes restartable feature generation +cheap. The failure mode is that it trusts them: a job killed mid-write, or one +whose MSA tool exited non-zero after creating the output file, leaves a +zero-length or truncated alignment behind. The next run reads it back without +complaint and either + +* dies far from the cause -- an empty ``uniref90_hits.sto`` surfaces as + ``StopIteration`` inside ``parsers.deduplicate_stockholm_msa``, which names + neither the file nor the protein; or +* worse, silently succeeds on a *partial* alignment and produces features whose + MSA depth is quietly wrong. + +The second case is the reason these checks exist: a crash is recoverable, but +degraded features are indistinguishable from good ones downstream. + +The checks here are deliberately structural rather than semantic. They ask +whether a file is a complete alignment of the expected format, not whether its +content is biologically sensible -- an alignment can be legitimately shallow. +""" + +from __future__ import annotations + +import gzip +import lzma +import zlib +from dataclasses import dataclass +from pathlib import Path + +from absl import logging + +# Suffixes AlphaFold/AlphaPulldown write alignments to, mapped to their format. +MSA_SUFFIXES = { + ".sto": "sto", + ".a3m": "a3m", + ".fasta": "a3m", # same structural rules: '>' headers plus residue lines +} + +# Files that are not alignments and must never be judged by these rules. +NON_MSA_SUFFIXES = {".hhr", ".hmm", ".pkl", ".json"} + +# Template-search *results*. hmmsearch writes these in Stockholm format, so they +# would otherwise pass for alignments and be deleted when a search legitimately +# returned no hits. They are outputs, not inputs, and are regenerated anyway. +NON_MSA_STEMS = {"pdb_hits"} + + +@dataclass(frozen=True) +class MsaProblem: + """One structurally invalid alignment file.""" + + path: Path + reason: str + + def __str__(self) -> str: # pragma: no cover - trivial + return f"{self.path}: {self.reason}" + + +def _read_text(path: Path) -> str: + """Read an alignment file, transparently handling .gz and .xz.""" + name = path.name + if name.endswith(".gz"): + with gzip.open(path, "rt", errors="replace") as handle: + return handle.read() + if name.endswith(".xz"): + with lzma.open(path, "rt", errors="replace") as handle: + return handle.read() + return path.read_text(errors="replace") + + +def _effective_suffix(path: Path) -> str | None: + """Alignment suffix of a path, looking through a .gz/.xz wrapper.""" + name = path.name + for wrapper in (".gz", ".xz"): + if name.endswith(wrapper): + name = name[: -len(wrapper)] + break + stem_path = Path(name) + if stem_path.suffix in NON_MSA_SUFFIXES or stem_path.stem in NON_MSA_STEMS: + return None + return MSA_SUFFIXES.get(stem_path.suffix) + + +def check_stockholm(text: str) -> str | None: + """Structural problem with a Stockholm alignment, or None if it is sound. + + jackhmmer writes a ``# STOCKHOLM 1.0`` banner, one or more alignment rows, + and a closing ``//``. The terminator is what makes truncation detectable: + a file cut off mid-write keeps a valid-looking header and some rows. + """ + if not text.strip(): + return "empty Stockholm file" + lines = [ln for ln in text.splitlines() if ln.strip()] + if not lines[0].startswith("# STOCKHOLM"): + return "missing '# STOCKHOLM' header" + if not any(ln.strip() == "//" for ln in lines): + return "missing '//' terminator (file is truncated)" + # Alignment rows are the non-comment, non-terminator lines. + rows = [ + ln for ln in lines + if not ln.startswith("#") and ln.strip() != "//" + ] + if not rows: + return "no alignment rows" + return None + + +def check_a3m(text: str) -> str | None: + """Structural problem with an A3M/FASTA alignment, or None if it is sound.""" + if not text.strip(): + return "empty A3M file" + lines = [ln for ln in text.splitlines() if ln.strip()] + headers = [i for i, ln in enumerate(lines) if ln.startswith(">")] + if not headers: + return "no '>' header lines" + # The final record must actually carry residues; a run killed just after + # writing a header leaves a danging description with no sequence. + if headers[-1] == len(lines) - 1: + return "last record has a header but no sequence (file is truncated)" + return None + + +def check_msa_file(path) -> str | None: + """Structural problem with one alignment file, or None if it is sound. + + Returns None for paths that are not alignments, so callers can hand this + whole directory listings. + """ + path = Path(path) + fmt = _effective_suffix(path) + if fmt is None: + return None + if not path.is_file(): + return "file does not exist" + if path.stat().st_size == 0: + return "file is empty (0 bytes)" + try: + text = _read_text(path) + except (OSError, EOFError, lzma.LZMAError, gzip.BadGzipFile, zlib.error) as exc: + # A compressed alignment truncated mid-stream fails to decompress; that + # is itself the finding, not an error to propagate. + return f"unreadable ({exc.__class__.__name__}: {exc})" + return check_stockholm(text) if fmt == "sto" else check_a3m(text) + + +def validate_precomputed_msas(msa_dir, *, remove_invalid: bool = False): + """Check every alignment in ``msa_dir``; optionally delete the bad ones. + + Args: + msa_dir: directory holding a protein's alignment files. + remove_invalid: when True, unlink each unsound file so the caller's next + MSA run regenerates it instead of reusing it. This is the safe default + for automated reruns: recomputing one alignment costs CPU time, whereas + reusing a truncated one corrupts the features. + + Returns: + A list of :class:`MsaProblem`, empty when everything is sound. + """ + msa_dir = Path(msa_dir) + if not msa_dir.is_dir(): + return [] + + problems: list[MsaProblem] = [] + for entry in sorted(msa_dir.iterdir()): + if not entry.is_file(): + continue + reason = check_msa_file(entry) + if reason is None: + continue + problems.append(MsaProblem(entry, reason)) + if remove_invalid: + try: + entry.unlink() + logging.warning( + "Removed unusable precomputed MSA %s (%s); it will be " + "regenerated.", entry, reason, + ) + except OSError as exc: + logging.error("Could not remove unusable MSA %s: %s", entry, exc) + else: + logging.warning("Unusable precomputed MSA %s (%s)", entry, reason) + return problems diff --git a/alphapulldown/utils/save_meta_data.py b/alphapulldown/utils/save_meta_data.py index 0efdeb0d..2ab32b3c 100644 --- a/alphapulldown/utils/save_meta_data.py +++ b/alphapulldown/utils/save_meta_data.py @@ -80,25 +80,54 @@ def get_metadata_for_binary(k, v): return {name: {"version": get_program_version(v)}} +def resolve_database_path(v): + """Follow symlinks to the file a database flag actually reads. + + Several AF3 database versions are encoded in the *filename*, and AF3's own + ``fetch_databases.sh`` pins ``pdb_seqres_2022_09_28.fasta``. Sites that + refresh a database in place commonly keep the pinned name as a symlink to + the current file, so the configured path and the bytes actually searched + disagree. Recording the configured name would then put a date in the + metadata -- and from there into a methods section -- that is simply wrong. + + Returns ``(resolved_path, was_symlink)``. + """ + try: + path = os.fspath(v) + except TypeError: + return str(v), False + try: + resolved = os.path.realpath(path) + except OSError: + return path, False + return resolved, os.path.abspath(path) != resolved + + def get_metadata_for_database(k, v): name = k.replace("_database_path", "").replace("_dir", "") + # Version strings are parsed out of file names, so follow symlinks first: + # a refreshed database hidden behind a pinned name must not be reported + # under the pinned name's date. + resolved, via_symlink = resolve_database_path(v) if name == "pdb_seqres": af3_seqres_release = re.search( - r"pdb_seqres_(\d{4}_\d{2}_\d{2})", str(v) + r"pdb_seqres_(\d{4}_\d{2}_\d{2})", str(resolved) ) if af3_seqres_release: version = af3_seqres_release.group(1) - return { - "PDB seqres": { - "release_date": version.replace("_", "-"), - "version": version, - "location_url": [ - "https://storage.googleapis.com/alphafold-databases/" - f"v3.0/pdb_seqres_{version}.fasta.zst" - ], - } + entry = { + "release_date": version.replace("_", "-"), + "version": version, + "location_url": [ + "https://storage.googleapis.com/alphafold-databases/" + f"v3.0/pdb_seqres_{version}.fasta.zst" + ], } + if via_symlink: + entry["configured_path"] = str(v) + entry["resolved_path"] = resolved + return {"PDB seqres": entry} if name == "rna_central": official_bundle = _looks_like_official_af3_database_bundle(v) diff --git a/alphapulldown/utils/template_reuse.py b/alphapulldown/utils/template_reuse.py new file mode 100644 index 00000000..7fd5457b --- /dev/null +++ b/alphapulldown/utils/template_reuse.py @@ -0,0 +1,222 @@ +"""Re-run template search against existing features, keeping their MSAs. + +Motivation. Template search is cheap; MSA search is not. When the template +database moves underneath a set of features -- a refreshed ``pdb_seqres``, a +corrected ``max_template_date`` -- the MSAs in those features are still +perfectly good and only the ``template_*`` block is stale. Regenerating from +scratch would redo hours of jackhmmer/HHblits work to change minutes of +hmmsearch output. + +What AlphaFold2 actually searches with. ``DataPipeline.process`` does not build +the template profile from the merged MSA that ends up in the features. It uses +the *uniref90* Stockholm alignment alone:: + + msa_for_templates = jackhmmer_uniref90_result['sto'] + msa_for_templates = parsers.deduplicate_stockholm_msa(msa_for_templates) + msa_for_templates = parsers.remove_empty_columns_from_stockholm_msa(...) + pdb_templates_result = template_searcher.query(msa_for_templates) + +So reproducing a template search faithfully means recovering that file, which +is why :func:`msa_for_template_search` prefers ``uniref90_hits.sto`` on disk and +treats reconstruction from the pickled features as a clearly-labelled fallback. +The fallback is not equivalent: the features' ``msa`` is the *merged* alignment +(uniref90 + mgnify + BFD), so the profile built from it is deeper and can +return different hits. That is a defensible result but not a reproduction of +what a fresh run would do, so callers are expected to surface which source was +used rather than let it pass silently. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from absl import logging + +from alphapulldown.utils.msa_encoding import ids_to_a3m +from alphapulldown.utils.msa_integrity import check_msa_file + +# The alignment AF2 builds its template profile from. +UNIREF90_STO = "uniref90_hits.sto" + +SOURCE_UNIREF90_FILE = "uniref90_hits.sto" +SOURCE_RECONSTRUCTED = "reconstructed-from-features" + + +@dataclass(frozen=True) +class TemplateSearchMsa: + """A Stockholm alignment to drive template search, and where it came from.""" + + stockholm: str + source: str + + @property + def is_reconstructed(self) -> bool: + return self.source == SOURCE_RECONSTRUCTED + + +def _iter_a3m_records(a3m_text: str): + """Yield ``(name, sequence)`` from A3M/FASTA text, dropping insertions. + + A3M marks insertions relative to the query with lowercase letters. Removing + them leaves every row at query length, which is what an alignment handed to + hmmbuild has to be. + """ + name = None + chunks: list[str] = [] + for line in a3m_text.splitlines(): + line = line.strip() + if not line: + continue + if line.startswith(">"): + if name is not None: + yield name, "".join(chunks) + name = line[1:].split()[0] if len(line) > 1 else "seq" + chunks = [] + elif name is not None: + chunks.append("".join(ch for ch in line if not ch.islower())) + if name is not None: + yield name, "".join(chunks) + + +def stockholm_from_a3m(a3m_text: str) -> str: + """Convert insertion-free A3M text into a minimal Stockholm alignment. + + Emits the banner, one ``name sequence`` row per record, a ``#=GC RF`` + reference annotation, and the ``//`` terminator. + + The reference annotation is not decorative. ``remove_empty_columns_from_ + stockholm_msa`` treats ``#=GC RF`` as the end of an alignment chunk and only + then moves the buffered rows into its output; without it every row is + dropped and the function fails with a bare ``KeyError``. Real jackhmmer + output always carries one. Every column here is a match column (insertions + were removed), so the annotation is all ``x``. + """ + rows = [] + seen: dict[str, int] = {} + width = None + for name, seq in _iter_a3m_records(a3m_text): + if not seq: + continue + if width is None: + width = len(seq) + elif len(seq) != width: + # Ragged input cannot be a valid alignment; refuse rather than let + # hmmbuild fail with a less obvious message. + raise ValueError( + f"alignment rows differ in length ({len(seq)} vs {width}); " + "input is not a query-anchored alignment" + ) + name = name.replace(" ", "_") + if name in seen: + seen[name] += 1 + name = f"{name}_{seen[name]}" + else: + seen[name] = 0 + rows.append((name, seq)) + + if not rows: + raise ValueError("no sequences to build a Stockholm alignment from") + + pad = max(len(name) for name, _ in rows + [("#=GC RF", "")]) + 2 + body = "\n".join(f"{name:<{pad}}{seq}" for name, seq in rows) + reference = f"{'#=GC RF':<{pad}}{'x' * width}" + return f"# STOCKHOLM 1.0\n{body}\n{reference}\n//\n" + + +def stockholm_from_feature_dict(feature_dict) -> str: + """Rebuild a Stockholm alignment from AF2 features' integer MSA. + + ``feature_dict['msa']`` is the merged, query-length, integer-encoded + alignment, so it decodes to A3M rows that are already insertion-free. + """ + msa = feature_dict.get("msa") + if msa is None: + raise ValueError("features contain no 'msa' array to reconstruct from") + msa = np.asarray(msa) + if msa.ndim != 2 or msa.size == 0: + raise ValueError(f"features have an unusable 'msa' array of shape {msa.shape}") + return stockholm_from_a3m(ids_to_a3m(msa)) + + +def msa_for_template_search(monomer, msa_output_dir) -> TemplateSearchMsa: + """Recover an alignment to re-run template search for an existing monomer. + + Prefers the on-disk ``uniref90_hits.sto`` because that is what AlphaFold2 + itself searches with; falls back to reconstructing the merged alignment from + the stored features when MSA files were not kept (``--save_msa_files=False`` + deletes them once features are written). + """ + sto_path = Path(msa_output_dir) / UNIREF90_STO + problem = check_msa_file(sto_path) if sto_path.exists() else "not present" + if problem is None: + return TemplateSearchMsa(sto_path.read_text(), SOURCE_UNIREF90_FILE) + + logging.warning( + "Cannot use %s (%s); rebuilding the template-search alignment from the " + "stored features instead. That alignment is the merged MSA rather than " + "uniref90 alone, so hits may differ from a fresh run.", + sto_path, problem, + ) + return TemplateSearchMsa( + stockholm_from_feature_dict(monomer.feature_dict), SOURCE_RECONSTRUCTED + ) + + +def search_templates( + template_searcher, + template_featurizer, + *, + query_sequence: str, + stockholm_msa: str, + msa_output_dir=None, +): + """Run template search + featurisation, mirroring ``DataPipeline.process``. + + Returns the ``template_*`` feature mapping. + """ + from alphafold.data import parsers + + # remove_empty_columns_from_stockholm_msa uses '#=GC RF' to delimit an + # alignment chunk and silently drops every row without one, then dies with a + # bare KeyError from its own bookkeeping. jackhmmer always writes the + # annotation, so its absence means the alignment is not what it claims to + # be - say so here rather than 30 frames down. + if "#=GC RF" not in stockholm_msa: + raise ValueError( + "Stockholm alignment has no '#=GC RF' reference annotation; it " + "cannot be used for template search. Genuine jackhmmer output " + "always contains one, so this file is truncated or not Stockholm." + ) + + msa_for_templates = parsers.deduplicate_stockholm_msa(stockholm_msa) + msa_for_templates = parsers.remove_empty_columns_from_stockholm_msa( + msa_for_templates + ) + + if template_searcher.input_format == "sto": + query = msa_for_templates + elif template_searcher.input_format == "a3m": + query = parsers.convert_stockholm_to_a3m(msa_for_templates) + else: + raise ValueError( + f"Unrecognized template input format: {template_searcher.input_format}" + ) + + pdb_templates_result = template_searcher.query(query) + + if msa_output_dir is not None: + hits_path = Path(msa_output_dir) / ( + f"pdb_hits.{template_searcher.output_format}" + ) + hits_path.parent.mkdir(parents=True, exist_ok=True) + hits_path.write_text(pdb_templates_result) + + hits = template_searcher.get_template_hits( + output_string=pdb_templates_result, input_sequence=query_sequence + ) + result = template_featurizer.get_templates( + query_sequence=query_sequence, hits=hits + ) + return dict(result.features) diff --git a/test/integration/test_keep_msas.py b/test/integration/test_keep_msas.py new file mode 100644 index 00000000..2513dd98 --- /dev/null +++ b/test/integration/test_keep_msas.py @@ -0,0 +1,511 @@ +"""CLI wiring for ``--keep_msas``: rewrite existing features, templates only. + +The mode exists because template search is minutes and MSA search is hours, so +when a template database is refreshed underneath a set of features the MSAs are +still good. These tests pin the behaviour that makes that safe: MSAs must come +through untouched, stale templates must be gone, and anything that cannot be +updated in place must fall back to normal generation rather than half-update. +""" + +import json +import lzma +import pickle +import types +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import pytest +from absl import flags + +import alphapulldown.scripts.create_individual_features as create_features +from alphapulldown.objects import MonomericObject +from alphapulldown.utils.template_reuse import ( + SOURCE_RECONSTRUCTED, + SOURCE_UNIREF90_FILE, +) + +FLAGS = flags.FLAGS + +REAL_STO = ( + "# STOCKHOLM 1.0\n\n" + "query ACDE\n" + "hit1 ACDF\n" + "#=GC RF xxxx\n" + "//\n" +) + +OLD_TEMPLATES = { + "template_aatype": np.zeros((1, 4, 22), dtype=np.float32), + "template_domain_names": np.asarray([b"stale_2022"], dtype=object), + "template_sequence": np.asarray([b"OLD"], dtype=object), + "template_sum_probs": np.asarray([0.5], dtype=np.float32), +} + +NEW_TEMPLATES = { + "template_aatype": np.ones((2, 4, 22), dtype=np.float32), + "template_domain_names": np.asarray([b"fresh_A", b"fresh_B"], dtype=object), + "template_sequence": np.asarray([b"NEWA", b"NEWB"], dtype=object), + "template_sum_probs": np.asarray([0.9, 0.8], dtype=np.float32), +} + +MSA_KEYS = ("msa", "deletion_matrix_int", "num_alignments", "msa_species_identifiers") + + +def _make_monomer(sequence="ACDE"): + monomer = MonomericObject("proteinA", sequence) + monomer.feature_dict = { + "msa": np.asarray([[0, 1, 2, 3], [0, 1, 2, 2]], dtype=np.int32), + "deletion_matrix_int": np.zeros((2, 4), dtype=np.int32), + "num_alignments": np.asarray([2, 2, 2, 2], dtype=np.int32), + "msa_species_identifiers": np.asarray([b"", b"9606"], dtype=object), + **OLD_TEMPLATES, + } + return monomer + + +@pytest.fixture +def flagged(tmp_path): + """A flag state with --keep_msas on and an output dir to work in.""" + FLAGS(["test"]) + FLAGS.output_dir = str(tmp_path) + FLAGS.keep_msas = True + FLAGS.use_mmseqs2 = False + FLAGS.skip_msa = False + FLAGS.skip_existing = False + FLAGS.compress_features = False + yield tmp_path + FLAGS.keep_msas = False + + +def _write_monomer(output_dir, monomer, *, compressed=False): + path = Path(output_dir) / ( + f"{monomer.description}.pkl.xz" if compressed else f"{monomer.description}.pkl" + ) + opener = lzma.open if compressed else open + with opener(path, "wb") as handle: + pickle.dump(monomer, handle) + return path + + +def _fake_pipeline(): + return types.SimpleNamespace( + template_searcher=object(), template_featurizer=object() + ) + + +def _load_written(output_dir, name="proteinA"): + with open(Path(output_dir) / f"{name}.pkl", "rb") as handle: + return pickle.load(handle) + + +# ------------------------------------------------------------------- AF2 + +@pytest.mark.parametrize("compressed", [False, True]) +def test_keep_msas_replaces_templates_and_preserves_msas(flagged, compressed): + original = _make_monomer() + _write_monomer(flagged, original, compressed=compressed) + (flagged / "proteinA").mkdir() + (flagged / "proteinA" / "uniref90_hits.sto").write_text(REAL_STO) + FLAGS.compress_features = compressed + + with patch.object( + create_features, "search_templates", return_value=NEW_TEMPLATES + ) as mock_search, patch( + "alphapulldown.utils.save_meta_data.get_meta_dict", return_value={"s": "t"} + ): + updated = create_features._update_templates_keeping_msas( + "proteinA", "ACDE", _fake_pipeline() + ) + + assert updated is True + mock_search.assert_called_once() + # The alignment on disk is what AF2 itself searches with, so it wins. + assert mock_search.call_args.kwargs["stockholm_msa"] == REAL_STO + + if compressed: + with lzma.open(flagged / "proteinA.pkl.xz", "rb") as handle: + written = pickle.load(handle) + else: + written = _load_written(flagged) + + assert written.feature_dict["template_domain_names"].tolist() == [ + b"fresh_A", b"fresh_B", + ] + assert b"stale_2022" not in written.feature_dict["template_domain_names"].tolist() + for key in MSA_KEYS: + np.testing.assert_array_equal( + written.feature_dict[key], original.feature_dict[key] + ) + assert written.template_msa_source == SOURCE_UNIREF90_FILE + + +def test_keep_msas_reconstructs_the_alignment_when_msa_files_were_deleted(flagged): + """--save_msa_files=False removes alignments once features are written.""" + _write_monomer(flagged, _make_monomer()) + + with patch.object( + create_features, "search_templates", return_value=NEW_TEMPLATES + ) as mock_search, patch( + "alphapulldown.utils.save_meta_data.get_meta_dict", return_value={} + ): + assert create_features._update_templates_keeping_msas( + "proteinA", "ACDE", _fake_pipeline() + ) + + # Rebuilt from the stored features, and labelled as such: the merged MSA is + # not the uniref90 alignment, so hits may legitimately differ. + assert "# STOCKHOLM" in mock_search.call_args.kwargs["stockholm_msa"] + assert _load_written(flagged).template_msa_source == SOURCE_RECONSTRUCTED + + +def test_keep_msas_survives_a_search_that_finds_no_templates(flagged): + """A restrictive --max_template_date can legitimately match nothing. + + The features must stay usable: stale templates gone, MSAs intact, and the + template_* block well-formed rather than absent. + """ + original = _make_monomer() + _write_monomer(flagged, original) + + with patch.object(create_features, "search_templates", return_value={}), \ + patch("alphapulldown.utils.save_meta_data.get_meta_dict", return_value={}): + assert create_features._update_templates_keeping_msas( + "proteinA", "ACDE", _fake_pipeline() + ) + + written = _load_written(flagged) + assert "stale_2022" not in str(written.feature_dict.get("template_domain_names", "")) + for key in MSA_KEYS: + np.testing.assert_array_equal( + written.feature_dict[key], original.feature_dict[key] + ) + # Downstream code indexes these unconditionally. + assert "template_sum_probs" in written.feature_dict + assert "template_confidence_scores" in written.feature_dict + assert "template_release_date" in written.feature_dict + + +def test_keep_msas_declines_when_no_features_exist(flagged): + assert create_features._update_templates_keeping_msas( + "proteinA", "ACDE", _fake_pipeline() + ) is False + + +def test_keep_msas_declines_on_a_sequence_mismatch(flagged): + """Same name, different protein: updating in place would corrupt it.""" + _write_monomer(flagged, _make_monomer(sequence="ACDE")) + + with patch.object(create_features, "search_templates") as mock_search: + result = create_features._update_templates_keeping_msas( + "proteinA", "WWWW", _fake_pipeline() + ) + + assert result is False + mock_search.assert_not_called() + + +def test_keep_msas_requires_a_template_stack(flagged): + _write_monomer(flagged, _make_monomer()) + empty_pipeline = types.SimpleNamespace( + template_searcher=None, template_featurizer=None + ) + with pytest.raises(RuntimeError, match="template searcher"): + create_features._update_templates_keeping_msas( + "proteinA", "ACDE", empty_pipeline + ) + + +def test_create_individual_features_updates_instead_of_regenerating(flagged): + """The AF2 loop must not run MSA generation when an update succeeded.""" + _write_monomer(flagged, _make_monomer()) + fasta = flagged / "in.fasta" + fasta.write_text(">proteinA\nACDE\n") + FLAGS.fasta_paths = [str(fasta)] + + with patch.object(create_features, "create_arguments"), \ + patch.object(create_features, "create_pipeline_af2", + return_value=_fake_pipeline()), \ + patch.object(create_features, "create_uniprot_runner") as mock_runner, \ + patch.object(create_features, "create_and_save_monomer_objects") as mock_make, \ + patch.object(create_features, "search_templates", + return_value=NEW_TEMPLATES), \ + patch("alphapulldown.utils.save_meta_data.get_meta_dict", return_value={}): + create_features.create_individual_features() + + mock_make.assert_not_called() + assert mock_runner.called, "uniprot runner is still built for any misses" + + +def test_create_individual_features_falls_through_for_new_proteins(flagged): + """A protein with no stored features has nothing to keep; generate it.""" + fasta = flagged / "in.fasta" + fasta.write_text(">brand_new\nACDE\n") + FLAGS.fasta_paths = [str(fasta)] + + with patch.object(create_features, "create_arguments"), \ + patch.object(create_features, "create_pipeline_af2", + return_value=_fake_pipeline()), \ + patch.object(create_features, "create_uniprot_runner"), \ + patch.object(create_features, "create_and_save_monomer_objects") as mock_make: + create_features.create_individual_features() + + mock_make.assert_called_once() + + +# ------------------------------------------------------- flag combinations + +@pytest.mark.parametrize( + ("conflicting", "message"), + [("use_mmseqs2", "MMseqs2"), ("skip_msa", "skip_msa")], +) +def test_keep_msas_rejects_incompatible_flags(conflicting, message): + FLAGS(["test"]) + FLAGS.keep_msas = True + FLAGS.data_pipeline = "alphafold2" + setattr(FLAGS, conflicting, True) + try: + with pytest.raises(ValueError, match=message): + create_features.validate_data_pipeline_flags() + finally: + FLAGS.keep_msas = False + setattr(FLAGS, conflicting, False) + + +# ------------------------------------------------------------------- AF3 +# +# AlphaFold3 is optional and its compiled parts are absent in CI, so the AF3 +# API is stubbed the way the rest of this suite stubs it. What is under test is +# AlphaPulldown's own logic -- which fields are carried over and which are +# cleared -- not AF3's JSON parser. + + +class _StubProteinChain: + def __init__(self, id, sequence, ptms=(), residue_ids=None, description=None, + paired_msa=None, unpaired_msa=None, templates=None): + self.id = id + self.sequence = sequence + self.ptms = list(ptms) + self.residue_ids = residue_ids + self.description = description + self.paired_msa = paired_msa + self.unpaired_msa = unpaired_msa + self.templates = templates + + +class _StubRnaChain: + def __init__(self, id, sequence, unpaired_msa=None): + self.id = id + self.sequence = sequence + self.unpaired_msa = unpaired_msa + + +class _StubInput: + def __init__(self, name, chains, rng_seeds): + self.name = name + self.chains = list(chains) + self.rng_seeds = list(rng_seeds) + + @classmethod + def from_json(cls, text): + data = json.loads(text) + if "sequences" not in data or "name" not in data: + raise ValueError("not an AF3 input") + chains = [] + for entry in data["sequences"]: + protein = entry["protein"] + chains.append(_StubProteinChain( + id=protein["id"], + sequence=protein["sequence"], + description=protein.get("description"), + paired_msa=protein.get("pairedMsa"), + unpaired_msa=protein.get("unpairedMsa"), + templates=protein.get("templates"), + )) + return cls(data["name"], chains, data.get("modelSeeds", [42])) + + +def _stub_folding_input(): + module = types.SimpleNamespace( + ProteinChain=_StubProteinChain, + RnaChain=_StubRnaChain, + Input=_StubInput, + ) + return module + + +def _af3_json(*, name="proteinA", unpaired=">q\nACDE\n", paired=">q\nACDE\n", + templates=None): + protein = {"id": "A", "sequence": "ACDE", "description": name} + if unpaired is not None: + protein["unpairedMsa"] = unpaired + if paired is not None: + protein["pairedMsa"] = paired + if templates is not None: + protein["templates"] = templates + return json.dumps( + {"name": name, "modelSeeds": [42], "sequences": [{"protein": protein}]} + ) + + +@pytest.fixture +def af3_stub(): + with patch.object(create_features, "folding_input", _stub_folding_input()): + yield + + +def test_af3_chain_without_templates_keeps_msas_and_clears_templates(af3_stub): + """AF3 searches templates only when both MSAs are set and templates is None. + + templates=[] means "no templates, do not search"; a partially populated + chain is rejected by AF3 outright. So None is the specific signal needed. + """ + chain = _StubProteinChain( + id="A", sequence="ACDE", + unpaired_msa=">q\nACDE\n", paired_msa=">q\nACDE\n", + templates=[{"mmcif": "data_x"}], + ) + + stripped = create_features._af3_chain_without_templates(chain) + + assert stripped.templates is None, "None means 'search'; [] means 'do not'" + assert stripped.unpaired_msa == ">q\nACDE\n" + assert stripped.paired_msa == ">q\nACDE\n" + assert stripped.sequence == "ACDE" + assert stripped.id == "A" + + +def test_af3_chain_carries_over_description_and_residue_ids(af3_stub): + """The description holds AlphaPulldown's metadata envelope; losing it would + discard the feature provenance.""" + chain = _StubProteinChain( + id="B", sequence="ACDE", description="prot\n__META__=1", + residue_ids=[1, 2, 3, 4], + unpaired_msa=">q\nACDE\n", paired_msa=">q\nACDE\n", templates=[], + ) + + stripped = create_features._af3_chain_without_templates(chain) + + assert stripped.description == "prot\n__META__=1" + assert stripped.residue_ids == [1, 2, 3, 4] + + +def test_af3_chain_without_msas_is_left_alone(af3_stub): + """Nothing to keep: let the normal path search MSAs and templates.""" + chain = _StubProteinChain(id="A", sequence="ACDE") + assert create_features._af3_chain_without_templates(chain) is chain + + +def test_af3_non_protein_chains_pass_through(af3_stub): + """RNA/DNA chains have no template search to re-run.""" + chain = _StubRnaChain(id="R", sequence="ACGU") + assert create_features._af3_chain_without_templates(chain) is chain + + +def test_af3_input_keeping_msas_round_trips_existing_features(tmp_path, af3_stub): + path = tmp_path / "proteinA_af3_input.json" + path.write_text(_af3_json(templates=[{"mmcif": "data_stale"}])) + + rebuilt = create_features._af3_input_keeping_msas(str(path), "proteinA") + + assert rebuilt is not None + assert rebuilt.name == "proteinA" + (chain,) = rebuilt.chains + assert chain.templates is None, "stale templates must be dropped and re-searched" + assert chain.unpaired_msa == ">q\nACDE\n", "the expensive MSA must survive" + assert chain.paired_msa == ">q\nACDE\n" + + +def test_af3_input_keeping_msas_declines_on_a_sequence_mismatch(tmp_path, af3_stub): + """Same description, different protein: the AF2 path refuses, so must AF3. + + Features are matched to inputs by description alone. If a FASTA is edited + but keeps its name, reusing the cached chain would re-run template search + against the old sequence and then overwrite the file with features for the + wrong protein. + """ + path = tmp_path / "proteinA_af3_input.json" + path.write_text(_af3_json()) # cached chain is "ACDE" + + assert create_features._af3_input_keeping_msas( + str(path), "proteinA", "WWWWWW" + ) is None + + +def test_af3_input_keeping_msas_accepts_a_matching_sequence(tmp_path, af3_stub): + path = tmp_path / "proteinA_af3_input.json" + path.write_text(_af3_json()) + + rebuilt = create_features._af3_input_keeping_msas( + str(path), "proteinA", "ACDE" + ) + + assert rebuilt is not None + assert rebuilt.chains[0].unpaired_msa == ">q\nACDE\n" + + +def test_af3_input_keeping_msas_skips_the_check_without_a_sequence(tmp_path, af3_stub): + """The sequence is optional so the helper stays usable on its own.""" + path = tmp_path / "proteinA_af3_input.json" + path.write_text(_af3_json()) + assert create_features._af3_input_keeping_msas(str(path), "proteinA") is not None + + +def test_af3_loop_regenerates_rather_than_overwriting_a_renamed_protein( + tmp_path, af3_stub, monkeypatch +): + """End to end: an edited FASTA under an old name must not be overwritten.""" + FLAGS(["test"]) + FLAGS.output_dir = str(tmp_path) + FLAGS.keep_msas = True + FLAGS.skip_existing = False + FLAGS.compress_features = False + FLAGS.data_pipeline = "alphafold3" + fasta = tmp_path / "in.fasta" + fasta.write_text(">proteinA\nWWWWWW\n") # edited: no longer ACDE + FLAGS.fasta_paths = [str(fasta)] + (tmp_path / "proteinA_af3_input.json").write_text(_af3_json()) + + seen = {} + + class _Pipeline: + def process(self, input_obj): + seen["chains"] = [c.sequence for c in input_obj.chains] + seen["msa"] = [getattr(c, "unpaired_msa", None) for c in input_obj.chains] + return {"sequences": []} + + with patch.object(create_features, "create_arguments"), \ + patch.object(create_features, "create_pipeline_af3", return_value=_Pipeline()), \ + patch.object(create_features, "validate_data_pipeline_flags"), \ + patch.object(create_features, "get_af3_feature_metadata", return_value={}): + create_features.create_af3_individual_features() + + assert seen["chains"] == ["WWWWWW"], "must build a fresh chain from the FASTA" + assert seen["msa"] == [None], "stale MSAs must not be carried over" + FLAGS.keep_msas = False + FLAGS.data_pipeline = "alphafold2" + + +def test_af3_input_keeping_msas_reads_a_compressed_file(tmp_path, af3_stub): + """AF3 feature JSONs are published .xz; the reuse path must read them.""" + path = tmp_path / "proteinA_af3_input.json.xz" + with lzma.open(path, "wt", encoding="utf-8") as handle: + handle.write(_af3_json()) + + rebuilt = create_features._af3_input_keeping_msas(str(path), "proteinA") + + assert rebuilt is not None + assert rebuilt.chains[0].templates is None + + +def test_af3_input_keeping_msas_declines_a_malformed_file(tmp_path, af3_stub): + """A partial or foreign JSON must fall back, not abort the whole run.""" + path = tmp_path / "broken_af3_input.json" + path.write_text('{"not": "an af3 input"}') + assert create_features._af3_input_keeping_msas(str(path), "broken") is None + + +def test_af3_input_keeping_msas_declines_when_there_are_no_msas(tmp_path, af3_stub): + path = tmp_path / "proteinA_af3_input.json" + path.write_text(_af3_json(unpaired=None, paired=None)) + assert create_features._af3_input_keeping_msas(str(path), "proteinA") is None diff --git a/test/unit/test_msa_integrity.py b/test/unit/test_msa_integrity.py new file mode 100644 index 00000000..7abc93f0 --- /dev/null +++ b/test/unit/test_msa_integrity.py @@ -0,0 +1,159 @@ +"""Structural checks on precomputed MSA files. + +The cases here are the ones actually observed in a large feature-generation +run: a zero-length Stockholm left by a tool that exited after creating its +output file, and alignments cut off mid-write when a job was killed. +""" + +import gzip +import lzma + +import pytest + +from alphapulldown.utils.msa_integrity import ( + check_a3m, + check_msa_file, + check_stockholm, + validate_precomputed_msas, +) + +GOOD_STO = "# STOCKHOLM 1.0\n\nquery ACDE\nhit1 ACDF\n//\n" +GOOD_A3M = ">query\nACDE\n>hit1\nACDF\n" + + +# --------------------------------------------------------------- Stockholm + +@pytest.mark.parametrize( + ("text", "expected_fragment"), + [ + ("", "empty"), + (" \n\n", "empty"), + ("query ACDE\n//\n", "missing '# STOCKHOLM' header"), + # jackhmmer killed mid-write: header and rows present, no terminator. + ("# STOCKHOLM 1.0\n\nquery ACDE\nhit1 ACDF\n", "truncated"), + ("# STOCKHOLM 1.0\n#=GF ID x\n//\n", "no alignment rows"), + ], +) +def test_check_stockholm_rejects(text, expected_fragment): + problem = check_stockholm(text) + assert problem is not None + assert expected_fragment in problem + + +def test_check_stockholm_accepts_a_complete_alignment(): + assert check_stockholm(GOOD_STO) is None + + +# --------------------------------------------------------------------- A3M + +@pytest.mark.parametrize( + ("text", "expected_fragment"), + [ + ("", "empty"), + ("ACDE\nACDF\n", "no '>' header lines"), + # Killed just after writing a description line. + (">query\nACDE\n>hit1\n", "truncated"), + ], +) +def test_check_a3m_rejects(text, expected_fragment): + problem = check_a3m(text) + assert problem is not None + assert expected_fragment in problem + + +def test_check_a3m_accepts_a_complete_alignment(): + assert check_a3m(GOOD_A3M) is None + + +# -------------------------------------------------------------- file level + +def test_check_msa_file_flags_a_zero_byte_stockholm(tmp_path): + """The exact failure seen in production: jackhmmer left an empty .sto. + + Reused, it surfaces as StopIteration inside deduplicate_stockholm_msa, + which names neither the file nor the protein. + """ + empty = tmp_path / "uniref90_hits.sto" + empty.touch() + assert "empty" in check_msa_file(empty) + + +def test_check_msa_file_ignores_non_alignment_files(tmp_path): + """Search results and profiles are not alignments and must not be judged.""" + # pdb_hits.sto is hmmsearch *output* in Stockholm format: it looks like an + # alignment, but an empty one just means no hits were found. + for name in ("pdb_hits.hhr", "pdb_hits.sto", "query.hmm", "features.pkl"): + path = tmp_path / name + path.write_text("whatever") + assert check_msa_file(path) is None + + +def test_check_msa_file_reports_a_missing_file(tmp_path): + assert check_msa_file(tmp_path / "uniref90_hits.sto") == "file does not exist" + + +@pytest.mark.parametrize("wrapper", ["gz", "xz"]) +def test_check_msa_file_reads_through_compression(tmp_path, wrapper): + """zip_msa_files gzips alignments in place, so checks must see through it.""" + path = tmp_path / f"uniref90_hits.sto.{wrapper}" + opener = gzip.open if wrapper == "gz" else lzma.open + with opener(path, "wt") as handle: + handle.write(GOOD_STO) + assert check_msa_file(path) is None + + +def test_check_msa_file_flags_a_corrupt_compressed_alignment(tmp_path): + path = tmp_path / "uniref90_hits.sto.gz" + path.write_bytes(b"\x1f\x8b\x08\x00 truncated garbage") + assert "unreadable" in check_msa_file(path) + + +# ---------------------------------------------------------------- directory + +def test_validate_precomputed_msas_reports_only_bad_files(tmp_path): + (tmp_path / "uniref90_hits.sto").write_text(GOOD_STO) + (tmp_path / "bfd_uniref_hits.a3m").write_text(">q\nACDE\n>h\n") # truncated + (tmp_path / "mgnify_hits.sto").touch() # empty + (tmp_path / "pdb_hits.hhr").write_text("not an alignment") + (tmp_path / "pdb_hits.sto").touch() # a search that found nothing + + problems = validate_precomputed_msas(tmp_path) + + assert {p.path.name for p in problems} == { + "bfd_uniref_hits.a3m", "mgnify_hits.sto", + } + # A sound alignment must survive an inspection that does not remove. + assert (tmp_path / "uniref90_hits.sto").exists() + + +def test_validate_precomputed_msas_removes_only_the_bad_ones(tmp_path): + good = tmp_path / "uniref90_hits.sto" + good.write_text(GOOD_STO) + bad = tmp_path / "mgnify_hits.sto" + bad.touch() + unrelated = tmp_path / "pdb_hits.hhr" + unrelated.write_text("hits") + + problems = validate_precomputed_msas(tmp_path, remove_invalid=True) + + assert [p.path.name for p in problems] == ["mgnify_hits.sto"] + assert not bad.exists(), "unusable alignment should be gone so it is regenerated" + assert good.exists(), "sound alignment must be kept - regenerating it is expensive" + assert unrelated.exists() + + +def test_validate_precomputed_msas_keeps_an_empty_template_hit_file(tmp_path): + """An empty pdb_hits.sto means "no templates matched", not a broken file.""" + hits = tmp_path / "pdb_hits.sto" + hits.touch() + assert validate_precomputed_msas(tmp_path, remove_invalid=True) == [] + assert hits.exists() + + +def test_validate_precomputed_msas_tolerates_a_missing_directory(tmp_path): + assert validate_precomputed_msas(tmp_path / "nope") == [] + + +def test_validate_precomputed_msas_ignores_subdirectories(tmp_path): + (tmp_path / "nested").mkdir() + assert validate_precomputed_msas(tmp_path) == [] diff --git a/test/unit/test_save_meta_data.py b/test/unit/test_save_meta_data.py index dcada607..b14ba233 100644 --- a/test/unit/test_save_meta_data.py +++ b/test/unit/test_save_meta_data.py @@ -237,6 +237,57 @@ def fail_if_called(_): assert "pdb_seqres_2022_09_28" in metadata["location_url"][0] +def test_af3_pdb_seqres_release_follows_a_symlink(tmp_path, monkeypatch): + """A refreshed database behind AF3's pinned name must report its real date. + + AF3's fetch_databases.sh pins pdb_seqres_2022_09_28.fasta, so sites that + update the database in place keep that name as a symlink to the current + file. Reporting the pinned name's date would put a wrong template cutoff + into the metadata and from there into a paper's methods. + """ + monkeypatch.setattr( + save_meta_data, "get_hash", lambda _: pytest.fail("should not hash") + ) + + real = tmp_path / "pdb_seqres_2026_08_19.fasta" + real.write_text(">1abc_A\nACDE\n") + pinned = tmp_path / "pdb_seqres_2022_09_28.fasta" + pinned.symlink_to(real) + + metadata = save_meta_data.get_metadata_for_database( + "pdb_seqres_database_path", str(pinned) + )["PDB seqres"] + + assert metadata["version"] == "2026_08_19" + assert metadata["release_date"] == "2026-08-19" + # Both ends of the symlink are recorded so the substitution stays auditable. + assert metadata["configured_path"] == str(pinned) + assert metadata["resolved_path"] == str(real) + + +def test_af3_pdb_seqres_plain_path_records_no_symlink_fields(tmp_path, monkeypatch): + monkeypatch.setattr( + save_meta_data, "get_hash", lambda _: pytest.fail("should not hash") + ) + real = tmp_path / "pdb_seqres_2022_09_28.fasta" + real.write_text(">1abc_A\nACDE\n") + + metadata = save_meta_data.get_metadata_for_database( + "pdb_seqres_database_path", str(real) + )["PDB seqres"] + + assert metadata["version"] == "2022_09_28" + assert "configured_path" not in metadata + assert "resolved_path" not in metadata + + +def test_resolve_database_path_handles_missing_path(): + """A configured path that does not exist must not raise during metadata.""" + resolved, via_symlink = save_meta_data.resolve_database_path("/no/such/db.fasta") + assert resolved.endswith("db.fasta") + assert via_symlink is False + + def test_get_metadata_for_database_returns_empty_for_unknown_key(): assert save_meta_data.get_metadata_for_database("custom_path", "/db/custom") == {} diff --git a/test/unit/test_template_reuse.py b/test/unit/test_template_reuse.py new file mode 100644 index 00000000..b2b86da3 --- /dev/null +++ b/test/unit/test_template_reuse.py @@ -0,0 +1,244 @@ +"""Recovering an alignment to re-run template search over existing features.""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +from alphapulldown.utils import template_reuse +from alphapulldown.utils.template_reuse import ( + SOURCE_RECONSTRUCTED, + SOURCE_UNIREF90_FILE, + msa_for_template_search, + search_templates, + stockholm_from_a3m, + stockholm_from_feature_dict, +) + +# Shaped like real jackhmmer output, including the #=GC RF annotation that +# AF2's column-pruning step needs to delimit the alignment chunk. +GOOD_STO = ( + "# STOCKHOLM 1.0\n\n" + "query ACDE\n" + "hit1 ACDF\n" + "#=GC RF xxxx\n" + "//\n" +) + + +# ------------------------------------------------------- Stockholm assembly + +def test_stockholm_from_a3m_emits_a_parseable_alignment(): + sto = stockholm_from_a3m(">query\nACDE\n>hit1\nACDF\n") + assert sto.startswith("# STOCKHOLM 1.0") + assert sto.rstrip().endswith("//") + rows = [ + ln.split() for ln in sto.splitlines() + if ln and not ln.startswith("#") and ln.strip() != "//" + ] + assert rows == [["query", "ACDE"], ["hit1", "ACDF"]] + + +def test_stockholm_from_a3m_drops_lowercase_insertions(): + """A3M lowercase marks insertions relative to the query. + + They must go, otherwise rows are ragged and hmmbuild has no alignment. + """ + sto = stockholm_from_a3m(">query\nACDE\n>hit1\nACfgDE\n") + seqs = [ + ln.split()[1] for ln in sto.splitlines() + if ln and not ln.startswith("#") and ln.strip() != "//" + ] + assert seqs == ["ACDE", "ACDE"] + + +def test_stockholm_from_a3m_disambiguates_repeated_names(): + """Stockholm keys rows by name, so duplicates would silently collapse.""" + sto = stockholm_from_a3m(">dup\nACDE\n>dup\nACDF\n>dup\nACDG\n") + names = [ + ln.split()[0] for ln in sto.splitlines() + if ln and not ln.startswith("#") and ln.strip() != "//" + ] + assert len(names) == len(set(names)) == 3 + + +def test_stockholm_from_a3m_uses_only_the_first_token_of_a_header(): + sto = stockholm_from_a3m(">sp|P12345|NAME some description here\nACDE\n") + names = [ + ln.split()[0] for ln in sto.splitlines() + if ln and not ln.startswith("#") and ln.strip() != "//" + ] + assert names == ["sp|P12345|NAME"] + + +def test_stockholm_from_a3m_rejects_ragged_rows(): + with pytest.raises(ValueError, match="differ in length"): + stockholm_from_a3m(">query\nACDE\n>hit1\nACD\n") + + +def test_stockholm_from_a3m_rejects_empty_input(): + with pytest.raises(ValueError, match="no sequences"): + stockholm_from_a3m("") + + +# --------------------------------------------- reconstruction from features + +def test_stockholm_from_feature_dict_decodes_the_integer_msa(): + # HHblits alphabet is alphabetical: 0=A, 1=C, 2=D, 3=E + features = {"msa": np.asarray([[0, 1, 2, 3], [0, 1, 2, 2]], dtype=np.int32)} + sto = stockholm_from_feature_dict(features) + seqs = [ + ln.split()[1] for ln in sto.splitlines() + if ln and not ln.startswith("#") and ln.strip() != "//" + ] + assert seqs == ["ACDE", "ACDD"] + + +@pytest.mark.parametrize( + "features", + [{}, {"msa": np.zeros((0, 0), dtype=np.int32)}, {"msa": np.zeros(4)}], +) +def test_stockholm_from_feature_dict_rejects_unusable_msas(features): + with pytest.raises(ValueError): + stockholm_from_feature_dict(features) + + +# ------------------------------------------------------------ source choice + +def _monomer_with_msa(): + return SimpleNamespace( + sequence="ACDE", + feature_dict={"msa": np.asarray([[0, 1, 2, 3]], dtype=np.int32)}, + ) + + +def test_msa_for_template_search_prefers_the_uniref90_file(tmp_path): + """AF2 builds its template profile from uniref90 alone, so prefer that file.""" + (tmp_path / "uniref90_hits.sto").write_text(GOOD_STO) + + result = msa_for_template_search(_monomer_with_msa(), tmp_path) + + assert result.source == SOURCE_UNIREF90_FILE + assert not result.is_reconstructed + assert result.stockholm == GOOD_STO + + +def test_msa_for_template_search_falls_back_when_the_file_is_absent(tmp_path): + """--save_msa_files=False deletes alignments once features are written.""" + result = msa_for_template_search(_monomer_with_msa(), tmp_path) + + assert result.source == SOURCE_RECONSTRUCTED + assert result.is_reconstructed + assert "ACDE" in result.stockholm + + +def test_msa_for_template_search_falls_back_when_the_file_is_unusable(tmp_path): + """An empty uniref90_hits.sto must not be preferred over reconstruction.""" + (tmp_path / "uniref90_hits.sto").touch() + + result = msa_for_template_search(_monomer_with_msa(), tmp_path) + + assert result.source == SOURCE_RECONSTRUCTED + + +# --------------------------------------------------------------- the search + +class _FakeSearcher: + def __init__(self, input_format="sto"): + self.input_format = input_format + self.output_format = "sto" + self.queried_with = None + + def query(self, text): + self.queried_with = text + return "RAW_HITS" + + def get_template_hits(self, output_string, input_sequence): + assert output_string == "RAW_HITS" + return ["hit-a", "hit-b"] + + +class _FakeFeaturizer: + def __init__(self): + self.seen_hits = None + + def get_templates(self, query_sequence, hits): + self.seen_hits = hits + return SimpleNamespace( + features={"template_domain_names": np.asarray([b"1abc_A"], dtype=object)} + ) + + +def test_search_templates_runs_the_searcher_and_featurizer(tmp_path): + searcher, featurizer = _FakeSearcher(), _FakeFeaturizer() + + features = search_templates( + searcher, featurizer, + query_sequence="ACDE", + stockholm_msa=GOOD_STO, + msa_output_dir=tmp_path, + ) + + assert features["template_domain_names"].tolist() == [b"1abc_A"] + assert featurizer.seen_hits == ["hit-a", "hit-b"] + # The hit file is written where AF2 would write it, for later inspection. + assert (tmp_path / "pdb_hits.sto").read_text() == "RAW_HITS" + + +def test_search_templates_converts_to_a3m_for_an_a3m_searcher(): + """HHsearch consumes A3M; hmmsearch consumes Stockholm. Honour both.""" + searcher = _FakeSearcher(input_format="a3m") + + search_templates( + searcher, _FakeFeaturizer(), + query_sequence="ACDE", + stockholm_msa=GOOD_STO, + msa_output_dir=None, + ) + + assert searcher.queried_with.startswith(">") + + +def test_search_templates_rejects_an_unknown_input_format(): + searcher = _FakeSearcher(input_format="clustal") + with pytest.raises(ValueError, match="Unrecognized template input format"): + search_templates( + searcher, _FakeFeaturizer(), + query_sequence="ACDE", + stockholm_msa=GOOD_STO, + msa_output_dir=None, + ) + + +def test_search_templates_rejects_a_stockholm_without_reference_annotation(): + """Without '#=GC RF' the AF2 helper fails with a bare KeyError; be explicit.""" + with pytest.raises(ValueError, match="#=GC RF"): + search_templates( + _FakeSearcher(), _FakeFeaturizer(), + query_sequence="ACDE", + stockholm_msa="# STOCKHOLM 1.0\nquery ACDE\n//\n", + msa_output_dir=None, + ) + + +def test_generated_stockholm_survives_the_af2_column_pruning(): + """The generator's output must be consumable by the real AF2 helper.""" + sto = stockholm_from_a3m(">query\nACDE\n>hit1\nACDF\n") + features = search_templates( + _FakeSearcher(), _FakeFeaturizer(), + query_sequence="ACDE", + stockholm_msa=sto, + msa_output_dir=None, + ) + assert "template_domain_names" in features + + +def test_search_templates_can_skip_writing_hits(): + searcher = _FakeSearcher() + features = search_templates( + searcher, _FakeFeaturizer(), + query_sequence="ACDE", + stockholm_msa=GOOD_STO, + msa_output_dir=None, + ) + assert "template_domain_names" in features