StoiAgent predicts the stoichiometry of a protein complex, meaning how many copies of each subunit it contains, from the subunit sequences in a FASTA file. It searches for structural templates, uses an LLM to propose and rank candidate stoichiometries from the template and literature evidence, models the top candidates with AlphaFold 3 (AF3), and reranks them. The result is saved as final_prediction.json.
- Linux with one NVIDIA GPU with 80 GB of memory, such as an A100 80GB or H100 80GB
- Docker with the NVIDIA Container Toolkit (or Singularity), to run AF3
- 64 GB RAM, about 1 TB of free disk, and internet access
See System requirements for details.
1. Install the environment. You need mamba. If you don't have it, see Installing Miniforge.
git clone https://github.com/BioinfoMachineLearning/StoiAgent.git
cd StoiAgent
mamba env create -f environment.yml
mamba activate stoiagent2. Download the databases and AF3 model parameters (about 700 GB; rerun the command if it is interrupted):
scripts/setup_databases.shThe AF3 model parameters are subject to Google DeepMind's terms of use; review them before running the script.
3. Build the AF3 Docker image:
docker build -t stoiagent-alphafold3 -f tools/alphafold3/docker/Dockerfile tools/alphafold3A FASTA file with one record per distinct protein subunit. StoiAgent names the subunits A, B, C, … in the order of the records. Identical sequences are merged, so a homomer can be given as a single record.
>subunit_1
MSVTSQTSSSGSAAVSDCHRGIIDISGPVPGYEWEPSMTTEPVRGRVW...
>subunit_2
MTGRQRTTVVAPDRPVQDATISQLTTRVWTVAIDGYRTIVVEGETGIV...
stoiagent \
--input_fasta complex.fasta \
--output_dir results/complex \
--description "Protein complex name" \
--organism "Organism name"| Option | Default | Meaning |
|---|---|---|
--input_fasta |
required | Input FASTA |
--output_dir |
required | Folder for all results. Use a separate folder for each complex. |
--description |
FASTA file name | Short description of the complex, given to the LLM |
--organism |
unspecified | Source organism, given to the LLM |
--model |
Qwen3.6-27B |
LLM backbone (see LLM weights) |
--k |
5 | Number of candidates modeled with AF3 and reranked |
--num_models |
100 | AF3 models per candidate (a multiple of 5) |
--config |
config.json |
Configuration file (see Configuration) |
The first run also downloads the LLM weights (about 63 GB).
final_prediction.json in the output folder holds the prediction. For the CASP16 example H0208 (true stoichiometry A1B1):
{
"input_fasta": ".../examples/casp16/H0208.fasta",
"top1_stoichiometry": "A1B1",
"ranked_stoichiometries": ["A1B1", "A2B1", "A2B2", "A1B2", "A4B4"],
"stage1_candidates": ["A1B1", "A2B2", "A1B2", "A2B1", "A4B4"],
"af3_scores_csv": ".../AlphaFold3/h0208/stoichiometry_results.csv",
"stage2_result": ".../S2_fasta_k5_af3/H0208_subunit_phase2_result.json"
}A2B1 means two copies of subunit A and one copy of subunit B. top1_stoichiometry is the prediction, ranked_stoichiometries is the final Stage-2 order of the shortlist, and stage1_candidates is the Stage-1 shortlist that AF3 modeled, in Stage-1 order.
All files of a run go to --output_dir. The target is named after the input FASTA (H0208.fasta → H0208); FASTA names that are not an H or T followed by four digits are named H0001 instead. The examples in this section come from this run of the CASP16 target H0208, with the default --model Qwen3.6-27B and --k 5:
stoiagent \
--input_fasta examples/casp16/H0208.fasta \
--output_dir results/H0208 \
--description "dahAB" \
--num_models 25dahAB is the description from the CASP16 sequence release (examples/casp16_stoichiometry.csv). With 25 AF3 models per candidate instead of the default 100, the run took about 2.5 hours on one A100 80GB: about 15 minutes of template search, 15 minutes for Stage 1, just under 2 hours of AF3 (an hour of it for the eight-chain A4B4) and a few minutes for Stage 2. It produces this folder:
results/H0208/
final_prediction.json final answer (see above)
H0208.fasta subunits renamed A, B, ... (duplicates merged)
target_sequences.fasta sequences with the description and organism given to the LLM
H0208.json template evidence assembled for Stage 1
template_search/ per-subunit template search
raw/H0208/stoichiometry/H0208_subunit/
Agent/Qwen3.6-27B/
S1_fasta/ Stage 1: candidate generation
S2_fasta_k5_af3/ Stage 2: AF3-informed reranking
AlphaFold3/h0208/ AF3 inputs, models and scores
stage2_summary/ one-row summary tables of Stage 2
article_cache/ pdb_json_cache/ pubmed_cache_test/ uniprot_cache_test/
caches of literature and database lookups
For each subunit (A, B, ...) template_search/ keeps the jackhmmer search of UniRef90 (.sto, .tblout, .jackhmmer.log), the alignment and profile built from it (.a3m, .hmm), and the hhsearch hits against the PDB template database (.hhr). <subunit>.search.sha256 records the inputs of the search, so a rerun in the same folder reuses it. H0208.json turns the hits into evidence for Stage 1, such as the copy numbers seen in templates shared by several subunits.
Agent/<model>/S1_fasta/ contains:
| File | Contents |
|---|---|
H0208_phase1_topall_k15.out |
All proposed stoichiometries, ranked by their confidence averaged over the 15 samples. The first --k go to AF3. |
H0208_phase1_top5_k15.out |
The top five of the same ranking. |
H0208_sampleNN_phase1.trace |
The full conversation of sample NN: prompt, tool calls to RCSB, UniProt and PubMed, and the answer (see Reasoning traces). |
phase1_cache/ |
Cached Stage-1 samples, reused by a rerun. |
1. A1B1 avg_conf=0.764 n=15/15
2. A2B2 avg_conf=0.253 n=15/15
3. A1B2 avg_conf=0.037 n=15/15
avg_conf is the candidate's confidence averaged over all samples, counting 0 when a sample did not propose it, and n is how many samples proposed it.
AlphaFold3/h0208/ holds one AF3 job per shortlisted candidate, named h0208_<candidate> (for example h0208_a1b1):
| File | Contents |
|---|---|
input_jsons/h0208_<candidate>.json |
AF3 input of each job. |
h0208_<candidate>/h0208_<candidate>_model.cif |
Best of the job's models by AF3 ranking score. |
h0208_<candidate>/h0208_<candidate>_ranking_scores.csv |
ranking_score of every model (seed, sample). |
h0208_<candidate>/h0208_<candidate>_summary_confidences.json, ..._confidences.json |
AF3 confidence metrics (pTM, ipTM, pLDDT, PAE) of the best model. |
h0208_<candidate>/h0208_<candidate>_data.json |
MSAs and templates used by the job. |
h0208_<candidate>/seed-S_sample-N/ |
Each individual model with its confidences. |
stoichiometry_results.csv |
Best and mean ranking score of every candidate; the AF3 evidence given to Stage 2. |
Only the simplest stoichiometry (A1B1, or A2 for a homomer) runs the AF3 MSA and template search; the other candidates reuse its MSAs and templates. Each job makes --num_models models: --num_models / 5 seeds with five samples each. Candidates over AF3's 5,000-residue or 26-chain limits get no job and are missing from stoichiometry_results.csv.
stoic,max_ranking_score,avg_ranking_score,num_models
A1B1,0.966,0.952,25
A2B2,0.861,0.718,25
A1B2,0.756,0.590,25
A2B1,0.863,0.795,25
A4B4,0.837,0.523,25
A ranking score compares models of the same candidate well, but it is structural confidence, not the probability that a stoichiometry is correct. Larger assemblies can still reach high maximum scores, as A4B4 does here; the gap between the maximum and the mean shows how consistent the models are.
Agent/<model>/S2_fasta_k<k>_af3/ contains:
| File | Contents |
|---|---|
H0208_subunit_phase2_result.json |
The final answer (final_answer), the ranking (ranking) and the LLM's rationale (rationale), plus the Stage-1 ranking and AF3 scores it was given. |
H0208_subunit_phase2_topall_k1.out |
The ranking with each candidate's Stage-1 rank and support and AF3 scores, followed by the rationale. Easiest file to read. |
H0208_subunit_phase2_top5_k1.out |
The top five of the same ranking. |
H0208_subunit_phase2_prompt.json |
The exact Stage-2 prompt. |
H0208_subunit_sample01_phase2.trace |
The full Stage-2 conversation (see Reasoning traces). |
1. A1B1 stage1_rank=1 stage1_support=0.764 af3_max=0.966 af3_avg=0.952 num_models=25
2. A2B1 stage1_rank=4 stage1_support=0.036 af3_max=0.863 af3_avg=0.795 num_models=25
3. A2B2 stage1_rank=2 stage1_support=0.253 af3_max=0.861 af3_avg=0.718 num_models=25
stage2_summary/ repeats the result as one-row tables. af3_score_only_per_target.csv ranks the shortlist by AF3 score alone (af3_max, af3_avg), a useful baseline for the Stage-2 ranking. The truth and hit columns of these tables are for benchmarking and stay empty in normal runs.
Each Stage-1 sample and the Stage-2 reranking leave a plain-text trace of the full LLM conversation (H0208_sampleNN_phase1.trace, H0208_subunit_sample01_phase2.trace):
-
Turns. The trace is a series of
--- Turn N ---blocks, one per model call. Each block repeats the whole conversation up to that call, so the last block is the complete conversation. -
Roles.
SYSTEM:is the agent's instructions. The firstUSER:message is the full prompt: the target description and sequences, the shared-template evidence fromH0208.json, the research tools and their budgets, and the required answer format.ASSISTANT:lines are the model's replies. -
Research. A tool call appears as an
Action:line, such asAction: RCSB.getPDB("2YHE"), followed by theObservation:StoiAgent got back from RCSB, UniProt or PubMed. In the H0208 run the template evidence was decisive and no sample used the tools. -
Feedback from StoiAgent. Later
USER:messages come from StoiAgent's answer checker, not from a person. For example, sample 07 of H0208 listedA1B1twice, so only four of its candidates were valid; it receivedneed 1 more valid unique candidate(s), addedA4B4, and was then asked to choose its final answer. -
What the reasoning is. Qwen3.6-27B runs with its thinking mode switched off, so there is no hidden chain of thought; the trace holds everything the model wrote. In Stage 1 the reasoning is the
Thought:line and thesupport_evidencegiven for each candidate, for example:1) A1B1 support=0.950000 support_evidence=Consistent with DahAB heterodimer identity and multiple high-quality shared templates (2YHE, 6UXU, 5AJL, 4NUR) 2) A2B2 support=0.100000 support_evidence=Possible dimer-of-dimers assembly seen in some homologous aldose reductases, but less supported for specific DahAB -
Stage 2. Turn 1 holds the ranking, and
VALIDATION: okconfirms that it is a complete ordering of the shortlist. Turn 2 asks the model to explain each candidate's position; its answer is markedVALIDATION: explanation only; saved ranking unchanged. The detailed rationale inH0208_subunit_phase2_result.json(rationale, with the first short version kept asoriginal_rationale) is therefore written after the ranking is fixed: it explains the ranking but cannot change it.
Traces show how a candidate got its support, but the model's statements are not checked facts. In H0208, sample 07 called DahAB the lysine-biosynthesis enzymes DapA and DapB while still ranking A1B1 first, so read the traces with the evidence they cite. To list the final answer of every Stage-1 sample (the last Final Answer line of each trace; earlier ones repeat the prompt's answer format and previous turns):
for f in raw/H0208/stoichiometry/H0208_subunit/Agent/Qwen3.6-27B/S1_fasta/*.trace; do
grep "^Final Answer: [A-Z][0-9]" "$f" | tail -n 1
done | sort | uniq -carticle_cache/ stores literature summaries made by Qwen3-4B, and pdb_json_cache/, pubmed_cache_test/ and uniprot_cache_test/ store RCSB, PubMed and UniProt responses; they stay empty when Stage 1 makes no such lookups. A rerun with the same --output_dir reuses the template search, the AF3 jobs whose scores exist and the cached Stage-1 samples, so a stopped run can be continued. Use a new folder after changing the databases, models or configuration.
With Docker, AF3 runs as your own user, so all output files belong to you.
examples/ contains ready-to-run inputs for the stoichiometry-blind protein targets of CASP16 and CASP17, along with their true stoichiometries:
| Dataset | Targets | Input FASTA files | True stoichiometries |
|---|---|---|---|
| CASP16 | 28 | examples/casp16/ |
examples/casp16_stoichiometry.csv |
| CASP17 | 56 | examples/casp17/ |
examples/casp17_stoichiometry.csv |
stoiagent --input_fasta examples/casp17/H1312.fasta --output_dir results/H1312See CASP dataset details for how the datasets were built.
| Component | Requirement |
|---|---|
| Operating system | Linux, x86_64 |
| GPU | One NVIDIA GPU with 80 GB of memory and compute capability 8.0 or higher, such as an A100 80GB or H100 80GB |
| GPU driver | NVIDIA driver 525 or newer |
| Container runtime | Docker with the NVIDIA Container Toolkit, or Singularity/Apptainer, to run AF3 |
| CPU | 8 or more cores (template search runs jackhmmer with 8 threads by default) |
| RAM | 64 GB or more |
| Disk | About 1 TB free, preferably SSD (see below) |
| Network | Internet access for setup, and during predictions for the RCSB, UniProt and PubMed queries that Stage 1 makes |
The GPU needs 80 GB because the default backbone, Qwen3.6-27B, takes about 54 GB in 16-bit precision before any working memory, and AF3 needs the full 80 GB for complexes near StoiAgent's 5,000-residue limit. The LLM stages and AF3 run one after another, so a single GPU is enough.
| Disk use | Size |
|---|---|
| AF3 genetic databases | about 630 GB |
| AF3 model parameters | about 1 GB |
| Default LLM weights (Qwen3.6-27B and Qwen3-4B) | about 63 GB |
| Template metadata and HHsearch database | about 57 GB, plus about 30 GB temporarily during setup |
stoiagent environment |
about 6 GB |
| AF3 Docker image | about 10 GB |
Each prediction also writes its template searches, AF3 models and reasoning traces to its output directory.
Miniforge provides mamba:
curl -L -O "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
bash "Miniforge3-$(uname)-$(uname -m).sh"Accept the license, choose the install location (the default is ~/miniforge3), and answer yes when asked whether to initialize the shell. Then open a new terminal, or run source ~/.bashrc, and check that mamba --version works.
environment.yml creates the stoiagent environment. It contains Python, StoiAgent itself as an editable install, its Python dependencies from pyproject.toml, and the template-search programs jackhmmer (HMMER), hhmake and hhsearch (HH-suite). After pulling new changes, run mamba env update -f environment.yml to pick up new dependencies.
python -m stoiagent is equivalent to stoiagent. The optional Mistral backbone also needs python -m pip install -e '.[mistral]' inside the environment.
scripts/setup_databases.sh downloads the HHsearch template database and the template metadata (about 33 GB), clones AF3, and downloads the AF3 model parameters (about 1 GB) and the AF3 databases (about 630 GB). Everything goes to the default locations below, so StoiAgent finds it without a configuration file. Downloads resume where they stopped. Options:
--db-dir /big/disk/stoiagent_dbputs the databases on another disk and records their paths inconfig.jsonin this folder. StoiAgent reads that file automatically.--only templatesor--only af3sets up just one part.
When it finishes, the script lists the steps it cannot do for you, such as building the AF3 Docker image. If a resource is still missing when you run StoiAgent, it stops before the step that needs it and lists the missing paths.
StoiAgent/
tools/
alphafold3/ AF3 source checkout (setup script)
models/af3.bin.zst AF3 model parameters (setup script)
alphafold3.sif AF3 Singularity image, Singularity only (manual)
databases/
af3/ AF3 genetic databases, incl. UniRef90 (setup script)
sort90_hhsuitedb/ HHsearch template database (prefix hhsuitedb3) (setup script)
templates/ template metadata (setup script)
entity_fasta/
entity_metadata/
chain_atoms/
pdb_stoichiometry.csv
Git ignores databases/, tools/ and config.json.
The setup script clones AF3 into tools/alphafold3, downloads the AF3 model parameters to tools/alphafold3/models/af3.bin.zst, and downloads its databases to databases/af3. AF3 reads the compressed parameters directly. The rest follows the official AlphaFold 3 installation guide:
-
Install Docker, the NVIDIA driver, and the NVIDIA Container Toolkit as described in the official guide. Check
nvidia-smiand that a GPU-enabled Docker container can see the GPU. -
Review the terms of use of the AF3 model parameters. The parameters are not included in this repository; the setup script downloads them from Google.
-
Build the official image and run the official sample prediction to verify the installation:
docker build -t stoiagent-alphafold3 -f tools/alphafold3/docker/Dockerfile tools/alphafold3
With Docker, StoiAgent runs AF3 as your own user (docker run --user), so its outputs belong to you, and mounts tools/alphafold3/models at /models, the AF3 databases at /public_databases and the AF3 output folder at /af_output inside the container. If an AF3 output folder still contains files owned by root, for example from an older run, StoiAgent hands them back to you after the next AF3 job. To use Singularity instead, follow the official Singularity instructions to build tools/alphafold3.sif, and set "af3_runtime": "singularity" in config.json.
Template search runs jackhmmer, hhmake and hhsearch, which the stoiagent environment provides. It uses the UniRef90 file from the AF3 databases. The setup script downloads the HHsearch template database and the template metadata, which turns HHsearch hits into stoichiometry evidence:
| File | Contents |
|---|---|
sort90_hhsuitedb/ |
HHsearch database of PDB chain profiles (prefix hhsuitedb3), searched for templates of each subunit. |
entity_fasta/ |
For each PDB entry, one FASTA record per distinct entity. Record IDs are the subunit letters used in that entry's stoichiometry. |
entity_metadata/ |
For each PDB entry, the RCSB metadata: title, keywords, experimental method, resolution, citation and polymer entities. |
chain_atoms/ |
Coordinates of each template chain, named after its HHsearch hit (PDB ID plus chain ID), such as 101MA.atom.gz. They are used to recover each hit's sequence. |
pdb_stoichiometry.csv |
The assembly stoichiometry (final_stoichiometry) and release_date of each PDB entry. |
StoiAgent loads its language models from Hugging Face. The weights download automatically on the first run and are reused afterwards.
| Model | --model |
Role | Download |
|---|---|---|---|
| Qwen3.6-27B | Qwen3.6-27B (default) |
Stage 1 and Stage 2 reasoning | about 55 GB |
| Qwen3-4B-Instruct-2507 | always used | Literature summaries | about 8 GB |
| Gemma-3-27B-it | Gemma-3-27B-it |
Alternative backbone | about 55 GB |
| OLMo-2-32B-Instruct | OLMo-2-32B-Instruct |
Alternative backbone | about 65 GB |
| Mistral-Small-3.2-24B | Mistral-Small-3.2-24B |
Alternative backbone, needs the mistral extra |
about 50 GB |
- Storage: the weights go to the Hugging Face cache,
~/.cache/huggingfaceby default. To keep them on a larger disk, setexport HF_HOME=/big/disk/huggingfacebefore running StoiAgent, for example in~/.bashrc. - Download in advance: to download before the first run, for example on a login node with internet access, run
hf download Qwen/Qwen3.6-27Bandhf download Qwen/Qwen3-4B-Instruct-2507in thestoiagentenvironment. - Gated models: Gemma is gated. Accept its license on its Hugging Face page, then run
hf auth loginonce with a Hugging Face access token. The Qwen models need no login.
StoiAgent needs no configuration when everything is in its default location. To keep a resource outside this folder, for example a shared AF3 database, set only the paths that differ in config.json in this folder, which StoiAgent reads automatically. You can also pass another file with --config, which is then used instead of config.json. Relative paths are resolved against this folder. Set STOIAGENT_HOME to resolve the default locations against another folder instead.
{
"af3_db_path": "/shared/af3_databases",
"uniref90_path": "/shared/af3_databases/uniref90_2022_05.fa",
"af3_runtime": "singularity",
"gpu_device": "1"
}The available keys are uniref90_path, hhdb_prefix, template_entity_fasta_dir, template_entity_metadata_dir, template_chain_atoms_dir, template_stoichiometry_csv, af3_program_path, af3_params_path, af3_db_path, af3_sif_path, af3_runtime (docker or singularity), af3_image (default stoiagent-alphafold3), gpu_device, jackhmmer_binary, hhmake_binary, hhsearch_binary and search_cpu. The defaults are in src/stoiagent/config.py.
- Template search. For each subunit,
jackhmmersearches UniRef90,hhmakebuilds a profile, andhhsearchfinds PDB templates. The template metadata turns the hits into evidence such as the copy numbers seen in related structures. - Stage 1. The LLM reasons over the template evidence, reads relevant literature on demand, and proposes and ranks candidate stoichiometries over 15 stochastic samples.
- AF3. AF3 models the top
--kcandidates that fit within its 5,000-residue and 26-chain limits. Candidates beyond those limits stay in the shortlist without AF3 scores. - Stage 2. The LLM reranks exactly that shortlist, combining the Stage-1 evidence with the AF3 ranking scores.
AF3 ranking scores are structural-confidence evidence, not probabilities that a stoichiometry is biologically correct. The template-search and AF3 results are reused when you rerun in the same output folder, so use a fresh output folder if the databases, models or configuration change.
For the stoichiometry-blind targets, CASP first released only the subunit sequences, with the stoichiometry withheld ("UNK"). In the next phase it released the stoichiometry of the experimental structure, and the datasets use that as the true answer.
| Dataset | Targets | Heteromers (H) |
Homomers (T) |
Target IDs | True stoichiometry from |
|---|---|---|---|---|---|
| CASP16 | 28 | 19 | 9 | Phase 0: H0xxx, T0xxx |
Phase 1: H1xxx, T1xxx |
| CASP17 | 56 | 43 | 13 | Phase 1: H1xxx, T1xxx |
Phase 2: H2xxx, T2xxx |
Heteromers have 2–9 distinct subunits, and complexes range from 2 to 36 chains. The datasets leave out cancelled targets, targets that turned out to be monomers (A1), and targets whose stoichiometry CASP has not released.
examples/
casp16/<target>.fasta input FASTA for each CASP16 target
casp17/<target>.fasta input FASTA for each CASP17 target
casp16_stoichiometry.csv true stoichiometries for CASP16
casp17_stoichiometry.csv true stoichiometries for CASP17
casp16.HT1.seq.txt, casp17.HT1.seq.txt CASP sequence releases
casp16.targetlist.csv, casp17.targetlist.csv Prediction Center target lists
build_casp_examples.py regenerates the FASTA files and tables
Each FASTA has one record per distinct subunit, in the order of the CASP sequence release, so the subunit letters match those in the true stoichiometry. The stoichiometry tables have these columns:
| Column | Meaning |
|---|---|
target |
Target ID, matching the FASTA file name |
stoichiometry |
True stoichiometry, such as A2B2C2 (subunit A × 2, B × 2, C × 2) |
num_subunits |
Number of distinct subunits in the FASTA |
description, organism |
Target description and source organism from the CASP sequence release |
sequence_release_id |
ID of the target in the CASP sequence release |
stoichiometry_source |
CASP target whose released oligomeric state gives the true stoichiometry |
The CASP sequence releases do not always use the blind-phase ID. For example, the CASP16 release lists H0217 as H1217. The H/T letter and the last three digits identify a target, so the phase digit is swapped to get the blind-phase ID and the ID that carries the truth. This is the same mapping the stage modules use. Release entries with a suffix, such as T1208s1 or T1228v1, are single subunits or variants of another target and are not included. To rebuild the datasets after updating the target lists, run python examples/build_casp_examples.py.
Advanced research interfaces remain available with python -m stoiagent.stage1_agent and python -m stoiagent.stage2_rerank. The end-to-end command creates a local sequence catalog from the FASTA, so no CASP17 input file is needed. Its internal target ID is withheld from LLM prompts. For historical benchmark evaluation, supply the original target-release cutoffs and sequence release file to the standalone stage modules.
The adapted parser retains its original DeepMind copyright and Apache 2.0 notice. A copy of that license is in APACHE-2.0-LICENSE.txt.
