From 52552bf8db4d2cc623457895715f2aff4c14d33a Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Sun, 23 Aug 2026 11:17:12 +0200 Subject: [PATCH 1/2] feat: add --keep_msas to refresh templates without redoing MSAs Template search costs minutes; MSA search costs hours. When a template database moves underneath a set of features the alignments are still perfectly good and only the template_* block is stale, but the only way to refresh it was to regenerate everything. This came up at EMBL: AlphaFold3's fetch_databases.sh pins pdb_seqres_2022_09_28.fasta, so every local AF3 run had been searching a 2022 snapshot. The refreshed file is a third larger. Features generated against the old snapshot need new templates, not new MSAs. --keep_msas updates existing features in --output_dir in place: * AlphaFold2 - reloads the monomer, re-runs template search, and swaps the template_* block via the same helper TrueMultimer already used. * AlphaFold3 - reloads the feature JSON and re-processes it with both MSAs set and templates=None, which is AF3's own "search for templates only" branch. (templates=[] means "do not search", and a partially populated chain is rejected outright, so None is the specific signal required.) Proteins with no stored features are generated normally, so the flag is safe to apply to a whole sample sheet. It takes precedence over --skip_existing, and is rejected with --use_mmseqs2 (MSAs and templates arrive together, with no separable search) and with --skip_msa (no MSAs to keep). What AF2 searches with matters. DataPipeline.process builds the template profile from the uniref90 Stockholm alignment alone, not the merged MSA that lands in the features, so the on-disk uniref90_hits.sto is preferred. When --save_msa_files=False deleted it, the alignment is rebuilt from the stored integer MSA - which is the *merged* alignment and so can yield different hits. That is a defensible result but not a reproduction of a fresh run, so the source is logged and recorded on the monomer as template_msa_source rather than passing silently. Two related fixes found in the same round: * Validate precomputed MSAs before reuse. --use_precomputed_msas trusts whatever it finds, so an alignment left by a killed job is read back as if complete. An empty uniref90_hits.sto surfaces as StopIteration inside deduplicate_stockholm_msa naming neither file nor protein; a *partially* written one silently produces shallower features, which is worse because it is indistinguishable from a good result. Unsound files are now removed so they regenerate. Search results (pdb_hits.*) are exempt: an empty one just means no hits. * Follow symlinks when recording database provenance. Sites that refresh a database in place keep AF3's pinned name as a symlink to the current file, and the release date is parsed out of the file name - so metadata reported 2022-09-28 while 2026-08-19 data was searched. That date ends up in methods sections. Both ends of the symlink are now recorded. Verified end to end against the real databases on an existing 74 aa feature set, in both directions, with hmmsearch querying the live pdb_seqres: --max_template_date=1990-01-01 6 templates -> none --max_template_date=2050-01-01 6 templates -> 10 (5 kept, 1 dropped, 5 new) In both runs the MSA came through byte-identical at (11, 74), deletion matrix included. The reshuffle in the second run is the reconstruction caveat above: that feature set had no saved uniref90_hits.sto, so the profile was rebuilt from the deeper merged alignment. 62 new tests across unit and integration (496 pass in total); the AF3 half runs against the stubbed folding_input the suite already uses, so it is covered in environments where AF3 is not built. Co-Authored-By: Claude Opus 5 --- README.md | 1 + alphapulldown/objects.py | 10 + .../scripts/create_individual_features.py | 213 ++++++++- alphapulldown/utils/msa_integrity.py | 184 ++++++++ alphapulldown/utils/save_meta_data.py | 49 +- alphapulldown/utils/template_reuse.py | 222 +++++++++ test/integration/test_keep_msas.py | 441 ++++++++++++++++++ test/unit/test_msa_integrity.py | 159 +++++++ test/unit/test_save_meta_data.py | 51 ++ test/unit/test_template_reuse.py | 244 ++++++++++ 10 files changed, 1543 insertions(+), 31 deletions(-) create mode 100644 alphapulldown/utils/msa_integrity.py create mode 100644 alphapulldown/utils/template_reuse.py create mode 100644 test/integration/test_keep_msas.py create mode 100644 test/unit/test_msa_integrity.py create mode 100644 test/unit/test_template_reuse.py 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..da42be4a 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,70 @@ 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): + """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 + + 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 +929,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) + 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 +953,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..5ff7c3ab --- /dev/null +++ b/test/integration/test_keep_msas.py @@ -0,0 +1,441 @@ +"""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_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 From 6b4b609460cc9f4d6cc1eb02b1a52c2610b00822 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Mon, 24 Aug 2026 09:37:01 +0200 Subject: [PATCH 2/2] fix: validate the cached sequence before reusing AF3 MSAs --keep_msas matched an existing AF3 feature file to the input by description alone. A FASTA edited in place but keeping its name would therefore have the cached chain reused: template search would run against the *old* sequence, and the file would be overwritten with features for the wrong protein. The AF2 path already refuses this in _existing_monomer_for_update, which compares the stored monomer's sequence with the input and regenerates on a mismatch. The AF3 path now does the same, so both fall back to full generation rather than half-updating a different protein. The sequence argument is optional, so the helper stays usable standalone. Reported by codex review on #634. Co-Authored-By: Claude Opus 5 --- .../scripts/create_individual_features.py | 20 +++++- test/integration/test_keep_msas.py | 70 +++++++++++++++++++ 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/alphapulldown/scripts/create_individual_features.py b/alphapulldown/scripts/create_individual_features.py index da42be4a..47aaf6cd 100644 --- a/alphapulldown/scripts/create_individual_features.py +++ b/alphapulldown/scripts/create_individual_features.py @@ -873,7 +873,7 @@ def _af3_chain_without_templates(chain): ) -def _af3_input_keeping_msas(existing_path, desc): +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 @@ -889,6 +889,22 @@ def _af3_input_keeping_msas(existing_path, desc): ) 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 @@ -934,7 +950,7 @@ def create_af3_individual_features(): # 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) + 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 diff --git a/test/integration/test_keep_msas.py b/test/integration/test_keep_msas.py index 5ff7c3ab..2513dd98 100644 --- a/test/integration/test_keep_msas.py +++ b/test/integration/test_keep_msas.py @@ -416,6 +416,76 @@ def test_af3_input_keeping_msas_round_trips_existing_features(tmp_path, af3_stub 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"