feat: local MMseqs2 features for AlphaFold 2, with insertions for both backends

[DimaMolod]

Local MMseqs2 features for AlphaFold 2 — and insertions for both backends

The local MMseqs2 feature path (GPU or CPU search, then CPU finalization) was AlphaFold 3 only. The search stage is backend-neutral — one MSA bundle per chain — so AlphaFold 2 support is a second finalizer that turns a bundle into the MonomericObject pickle AlphaFold 2 inference reads. finalize_batch_features.py --data_pipeline=alphafold2 selects it. --use_mmseqs2 keeps meaning the remote ColabFold API.

A defect this fixes for AlphaFold 3 as well

Every local MMseqs2 alignment had an all-zero deletion matrix. result2msa --msa-format-mode 2, the only format the stage ran, keeps full database headers but drops every column the query does not span, so no insertion ever reached AlphaFold. Measured against a real 586-residue query: 98/120 small BFD hits (81.7%) and 712/824 UniProt hits (86.4%) carry insertions — 13 035 residues on UniProt alone. The code assumed mode 2 delivered insertions as query-gap columns and converted them; the pinned build never emits one, so that path never ran on real data.

Mode 5 keeps insertions but cuts the header to the accession (sp|P83570|GWA_SEPOF … → P83570), and AlphaFold reads the species it pairs chains by from the long form. So each search result is now formatted both ways and joined row by row — and the join is verified, not trusted: stripping insertions from every mode-5 row must reproduce its mode-2 row, and a disagreement fails the request rather than labelling one hit with another's header. Cost: ~16 ms per hit (2 s against a 947 s search on small BFD; 13 s against 3739 s on UniProt).

This also removes the search stage's last AlphaFold 3 dependency (alphafold3.cpp.msa_conversion), which the AlphaFold 2 image lacks: the stage imported there and died on its first result.

Bundle provenance moves, so no bundle written before this is reused.

What the AlphaFold 2 finalizer reproduces from native AlphaFold 2

• per-database caps counting the query (uniref90 10 000, MGnify 501), merge order UniRef90 → BFD → MGnify
• templates searched from UniRef90 alone (bundles now record each database's row span, so this is possible), with hmmsearch or --use_hhsearch
• pairing features from the independent UniProt search, species parsed from its headers — not a copy of the unpaired features, which is what the remote path uses
• accession identifiers filled offline from the headers; nothing on this path queries UniProt

The recipe is AlphaFold 2's reduced_dbs set (no BFD/UniRef30 HHblits arm); the docs say so and say where it cannot be exact.

Cache and search correctness

• Raising num_iterations failed every RNA shard. Iterative profile search is protein-only; on the pinned build --num-iterations 3 on a nucleotide search exits 1. It now reaches protein searches only. RNA provenance never recorded it, so no cached RNA MSA changes identity. The real-binary RNA contract test now runs at 1 and 3 iterations, and fails on the previous code.

• A bundle AlphaFold 2 cannot slice was never rebuilt. A bundle whose per-database row spans are missing, name missing or duplicate databases, or don't add up was reported as a failure and left in place, where the workflow's shard completion kept validating it, so every retry failed the same way. It is now deleted like any other damaged bundle, which schedules a repair.

• HHsearch feature reuse now includes PDB70 identity. --use_hhsearch requires --template_pdb70_database_id instead of the unused seqres identity. A changed PDB70 build rebuilds features while reusing MSA bundles; unchanged IDs reuse features, and unused databases do not invalidate them. The companion workflow PR forwards and keys the same selected identity. Operators must change the immutable ID when rebuilding a database, including at the same path. HMMsearch and AF3 retain their cache identities.

• Invalid bundles are removed at the shared reader. Complete, unique database roles and a protein pairing alignment are checked before a finalizer consumes a bundle. RNA retains its valid empty paired alignment. A downstream template-search failure preserves valid MSA data.

Also

• --mmseqs_db_load_mode (memory only, deliberately outside the cache identity)
• compare_msa_backends.py --artifact_format=af2_pickle
• the template stack is built from explicit settings, and the legacy script delegates to the same builder

Tests

Current gate on 6c53a24a, run in an isolated worktree on a compute node (SLURM 61831784, 2026-09-14):

Suite	Passed	Skipped / deselected
unit, AlphaFold 2 environment	589	10 skipped
unit, AlphaFold 3 environment	641	1 skipped
integration, AlphaFold 2 environment	107	27 deselected
integration, AlphaFold 3 environment	107	27 deselected
real MMseqs2 binary (-m external_tools)	14	0

Unit-only coverage was compared against the PR base: all 38 new executable functions are exercised, including nested helpers. The protocol method declaration has no implementation and is excluded. New tests cover both template factories, resolved settings, metadata selection, AF2/AF3 CLI dispatch and validation, PDB70 invalidation/reuse, and malformed-bundle repair. The three integration tests broken by relocated template imports now patch the actual dependency modules and retain their constructor assertions.

The architecture review retained the existing search → bundle → finalizer separation and persistent shard-schedule module. Bundle usability validation now finishes inside the reader's delete-on-invalid operation.

The real-binary tests cover:

• the command contract for protein and RNA, with the two-pass join checked on real output
• an AlphaFold 2 pickle built end to end from a real search, whose template comes from 3L4Q
• ten edge cases on purpose-built databases: a deep MSA that hits AlphaFold 2's MGnify cap at exactly 551 rows, an orphan with nothing anywhere, no templates, X residues, a UniParc-style header with no species, a peptide, and a homomer, a heteromer and a chopped chain assembled into multimer features

Against the remote ColabFold API. A new opt-in cluster check,
test/cluster/check_alphafold2_predictions.py::TestLocalMmseqsAgainstRemote, builds AF2
features for the same chains both ways -- remotely through --use_mmseqs2, locally
through this path on a GPU -- and compares them with compare_msa_backends. It passed at 881642c0 in
1h21 on a GPU node, most of it the search reading the padded databases in full.

chain	remote depth / Neff	local depth / Neff	local / remote
P61626, human lysozyme	5111 / 3023	3354 / 1886	0.66 / 0.62
O31718	1517 / 634	1264 / 594	0.83 / 0.94
a random 90-mer	1 / 1	1 / 1	neither finds a family

Both sides carry insertions (2425 rows remotely and 1317 locally for lysozyme). Local
pairs from its own UniProt search, with 677 species for lysozyme and 149 for O31718,
where the remote path copies its unpaired rows into the paired features. Local finds 20
templates per chain against the remote path's single placeholder. Because the two search
different databases, the test asserts ranges, not equality: at least 100 rows on both
sides, local within 0.2-5x of remote on depth and Neff, insertions present, species to
pair by, a real template for lysozyme, and at most 20 rows for the orphan.

GPU inference smoke test. Those edge-case pickles, folded with AlphaFold 2 (4 monomers and 4 multimers), all completed: 8/8 on a Blackwell MIG slice. The later H100 invocation reused completed predictions.

Snakemake end to end (companion PR KosinskiLab/AlphaPulldownSnakemake#55): the historical 2026-09-11 run used core 71eef1c4 and workflow dd5b1e7, and completed all 35 steps on the cluster. That covered one search shard for five proteins, five finalizations, nine AF2-multimer folds and their analysis. The folds include three edge cases on the real databases, alone and paired: thioredoxin (a deep MSA), a random 90-residue orphan with no homologs or templates, and a 26-residue peptide. A second invocation found nothing to do. This run used hmmsearch and predates the PDB70 cache correction; current fixes are covered by the gate above.

Benchmark against native AlphaFold 2

The set is 32 monomers spread over the quartiles of their native MSA depth, plus 12 heterodimers released after AF2-multimer's training cutoff. Each was featurized with native jackhmmer (reduced_dbs), with this path using GPU and CPU search, and with the remote ColabFold API, all against the same databases and with templates up to 2021-09-30. Each heterodimer was then folded with the same five AF2-multimer models and seed per feature set, and every model scored by DockQ against the experimental structure.

median, unless stated	native	local, GPU	local, CPU
MSA depth, shallowest / deepest quartile	52 / 12 247	29 / 11 756	28 / 5 181
Neff, all monomers	778	627	572
unpaired hits recovered vs native	—	0.76	0.47
templates per chain (overlap with native, Jaccard)	16	16 (0.67)	15 (0.64)
top-ranked DockQ, mean of 12	0.59	0.56	0.56
acceptable interfaces (DockQ ≥ 0.23)	9 / 12	9 / 12	9 / 12

• Structure quality is close to native. Paired over the 12 heterodimers, local features scored 0.028 lower on average. One interface, 9HMX, lost about 0.3 (0.70 → 0.40, although its best model reached 0.67). The other eleven moved by less than 0.1.
• The MSAs are shallower, most on the shallowest families (recovery 0.50 in the lowest quartile), as already measured for AlphaFold 3. Template overlap with native rises from 0.30 on the shallowest quartile to 0.90 on the deepest, because the templates are searched from the UniRef90 alignment.
• A CPU search finds far fewer distant hits than the GPU prefilter on deep families, at MMseqs2's default sensitivity, yet DockQ is identical here. Raising CPU sensitivity is a possible follow-up; this PR doesn't change it.
• Cost: the GPU search of all 56 chains in one shard took 0.9 h and the CPU search 4.25 h, both peaking at 146 GB. An earlier GPU attempt timed out at 4 h while it shared the NFS server with the native jobs and the CPU search. The GPU search is I/O-bound on network storage: the H200 sat at 0% while the padded databases streamed at ~150 MB/s. AlphaFold 2 finalization is dominated by template featurization, not the MSA: median ~1 GB and 2 min per chain, peak 18.8 GB and 87 min, set by which structures the templates come from. The workflow's defaults now follow these numbers.

🤖 Generated with Claude Code

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ae838f9a30

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Include the PDB70 identity in the AF2 cache key

When --use_hhsearch is selected, the template stack actually searches pdb70_database_path, but this signature still keys reuse only on the PDB seqres and mmCIF identifiers. Thus replacing or updating PDB70 while those two IDs remain unchanged causes _read_matching_artifact to reuse pickles whose template hits came from the old PDB70 database; there is also no separate PDB70 identifier flag to force regeneration.

Useful? React with 👍 / 👎.