feat(alphafold): custom MSA input for monomer and multimer predictions (#52)

[Elarwei001]

Resolves #52

Summary

gget alphafold: Adds a new msa argument (-msa/--msa on the command line) that lets you supply a custom multiple sequence alignment instead of running the internal jackhmmer search. When an MSA is provided, gget alphafold skips the jackhmmer search entirely, so no genetic databases are downloaded (the internal search otherwise fetches ~2 GB) — useful for folding from a manually curated MSA.

Originally scoped to monomers (issue #52); this PR also extends the feature to MULTIMER predictions (one MSA file per chain).

What it does

• Monomer: pass a single MSA file.
• Multimer: pass one MSA file per chain, in the same order as the input sequences. The first sequence in each chain's MSA must be that chain's query (matching the input sequence, ignoring gaps).
• Accepts a3m (.a3m) and aligned FASTA (.fasta / .fa / .afa). a3m lowercase insertions are parsed into the deletion matrix (matching AlphaFold's own a3m handling); aligned FASTA has no insertions, so its deletion matrix is all zeros.
• The multimer path reuses AlphaFold-Multimer's existing cross-chain MSA pairing — each chain's custom MSA is also used as its all_seq MSA, which Multimer pairs across chains by species. Pairing only finds matches when the MSA headers carry species identifiers; otherwise it degrades gracefully to an unpaired multimer.

Implementation notes

• New validation helper _validate_custom_msa_paths(custom_msa_paths, sequences) enforces, before any prediction work:
  1. file count == number of chains,
  2. each file exists,
  3. each file has a supported extension,
  4. the first sequence in each chain's MSA equals that chain's query (gaps removed, case-insensitive).
• Supporting helpers: detect_msa_format, parse_custom_msa (a3m/FASTA → aligned sequences + deletion matrix + descriptions), and read_custom_msa (returns an alphafold.data.parsers.Msa).
• The jackhmmer branch (closest-source lookup, DB list, binary chmod, per-chain search) is now gated behind msa is None and left otherwise unchanged.
• CLI --msa uses nargs="+" so it accepts one or more files (one per chain); main.py forwards args.msa to alphafold(...).

Testing

Unit / validation tests added in tests/test_alphafold.py (with a3m + aligned-FASTA fixtures in tests/fixtures/), covering format detection, a3m/FASTA parsing and deletion-matrix behavior, the CLI --msa flag, and the monomer + multimer validation cases (count mismatch, query mismatch, missing file). These pass. AlphaFold's heavy runtime deps aren't available in CI, so tests run at the parsing/validation level; a full GPU end-to-end fold is in progress and not yet complete.

Docs

• docs/src/en/alphafold.md — documents the -msa/--msa argument.
• docs/src/en/updates.md — changelog entry.

Resolves issue 52.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 48.24561% with 59 lines in your changes missing coverage. Please review.
✅ Project coverage is 57.16%. Comparing base (e43a804) to head (68324d5).
⚠️ Report is 2 commits behind head on dev.

Files with missing lines	Patch %	Lines
gget/gget_alphafold.py	48.24%	59 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##              dev     #235      +/-   ##
==========================================
+ Coverage   56.70%   57.16%   +0.45%     
==========================================
  Files          29       29              
  Lines        9392     9468      +76     
==========================================
+ Hits         5326     5412      +86     
+ Misses       4066     4056      -10     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Elarwei001]

End-to-end verification of the multimer custom-MSA path

The unit tests here cover the argument / parsing / validation layer (which is all CI can run, since AlphaFold's heavy runtime deps aren't available). To cover the part CI can't, I ran a full GPU fold of the new multimer path on a rented CUDA GPU. It passes end-to-end.

Environment

GPU	NVIDIA RTX 6000 Ada (48 GB), Python 3.12
jax / jaxlib	0.4.26 / 0.4.26 (cuda12), jax.devices() == [cuda(id=0)]
AlphaFold params	Colab set, incl. params_model_{1..5}_multimer_v3.npz

Test case — a real small heterodimer: human insulin A chain (21 aa) + B chain (30 aa). Each chain gets a custom a3m whose first sequence is that chain's query, plus three species-tagged homologs (headers use the UniProt db|ACC|ENTRY_SPECIES form that AlphaFold's msa_identifiers parses for cross-chain pairing).

gget.alphafold(
    ["GIVEQCCTSICSLYQLENYCN", "FVNQHLCGSHLVEALYLVCGERGFFYTPKT"],
    msa=["chainA.a3m", "chainB.a3m"],
    relax=False, plot=False,
)

Run log (key lines)

Using the multimer model with 2 sequences.
Using 2 user-provided custom MSA file(s); skipping the internal jackhmmer search.
3 unique sequences found in the custom MSA for sequence 1.
3 unique sequences found in the custom MSA for sequence 2.
Running model_1_multimer_v3 ... model_5_multimer_v3   (all 5 completed)

Result — a valid two-chain complex + PAE:

• selected_prediction.pdb — 2 chains, 2 TER records, 51 CA atoms = 21 + 30 residues, matching insulin A + B.
• predicted_aligned_error.json — PAE matrix.

So the new path (_validate_custom_msa_paths → per-chain read_custom_msa → uniprot_msa = custom_msa, reusing AlphaFold-Multimer's cross-chain pairing) runs end-to-end and produces a real prediction. Fold quality isn't the point — the MSA is intentionally shallow/partly synthetic to exercise the pipeline cheaply.

Full artifacts (inputs, driver, run log, output PDB + PAE, environment snapshot): https://gist.github.com/Elarwei001/0e342e4e76997c7d07553279f778d5a8

---

One side note from setting this up: getting AlphaFold to actually install/import currently needs a couple of workarounds that are unrelated to this PR (they're in the gget setup alphafold path against the current deepmind/alphafold main). I filed those separately in #264 so they don't muddy this PR.

[Elarwei001]

End-to-end verification — monomer path

Following up on the multimer run above: I also verified the monomer custom-MSA path end-to-end on a GPU (the original issue #52 case — a single sequence with a single MSA, passing msa as a string rather than a list, which exercises the monomer branch specifically).

Environment — RTX A6000 (48 GB), Python 3.12, jax 0.4.26 (cuda12), jax.devices() == [cuda(id=0)].

Test case — human insulin B chain (30 aa) with a single custom a3m (query + 3 homologs):

gget.alphafold(
    "FVNQHLCGSHLVEALYLVCGERGFFYTPKT",
    msa="chainM.a3m",          # single string, not a list
    relax=False, plot=False,
)

Run log (key lines)

Using 1 user-provided custom MSA file(s); skipping the internal jackhmmer search.
4 unique sequences found in the custom MSA for sequence 1.
Running model_1 ... model_5, model_2_ptm   (all completed on GPU)
Running model without relaxation stage.
=== FOLD_RETURNED ===

Result — selected_prediction.pdb: single chain A, 30 CA atoms = 30 residues (matching insulin B), 1 TER, first residue PHE A 1; plus predicted_aligned_error.json. So the monomer branch (msa string → single-chain read_custom_msa → monomer feature pipeline, no cross-chain pairing) runs end-to-end and produces a real single-chain prediction.

Artifacts (monomer inputs, driver, run log, output PDB + PAE, environment snapshot) are in the same gist as the multimer run, under the monomer_* / chainM.a3m / run_mono.py files: https://gist.github.com/Elarwei001/0e342e4e76997c7d07553279f778d5a8

With this, both paths added in this PR are verified end-to-end on a real GPU: multimer (one MSA per chain) and monomer (single MSA).

[Elarwei001]

Hi @lauraluebbert , when you get a chance, could you take a look at #235? Quick summary:

1. What it adds — a msa argument (--msa) to gget alphafold that takes a custom MSA and skips the internal jackhmmer search entirely (so no genetic databases are downloaded). Resolves [gget_alphafold] - Feature Request - Add an option to submit a MSA instead of a protein sequence #52.
2. Monomer + multimer — a monomer takes one MSA file; a multimer takes one per chain (in sequence order). The first sequence in each chain's MSA must be that chain's query. The multimer path reuses AlphaFold-Multimer's own cross-chain pairing.
3. Testing — unit tests cover the parsing/validation layer (all CI can run). I also ran both paths end-to-end on a GPU (insulin monomer, and insulin A+B heterodimer); results and artifacts are in two comments on the PR.
4. Unrelated heads-up — while setting that up I hit two pre-existing breakages in the gget setup alphafold path (not from this PR); filed separately as gget setup alphafold / gget alphafold break against the current deepmind/alphafold main branch #264.

Thanks!