MAXWELL

Source code of MAXWELL: Calibrating the probabilistic outputs of protein language models to the mutation-induced stability change landscape

If you find this work useful, please star the repo! (Click the star button in the top right corner of the page)

Thank you! (づ｡◕‿‿◕｡)づ

Recommended environment

It is highly recommended to start a new conda environment from scratch due to potential CUDA compatibility issues between PyTorch and PyTorch Geometric. The current dependency files target Python 3.12 on Linux, with CUDA-enabled PyTorch.

conda create -n maxwell python=3.12
conda activate maxwell
pip install uv

Clone the repository

git clone https://github.com/ai4protein/Venus-MAXWELL.git
cd Venus-MAXWELL

# Install the project dependencies into the active conda environment
uv sync --active

The dependency versions are specified in pyproject.toml and uv.lock. Run the following commands from the repository root with the maxwell environment activated. A GPU is recommended for training and full-dataset evaluation; the inference example below uses CPU.

Generate mutation-stability scores from a PDB file with Venus-MAXWELL (ProteinMPNN)

The MAXWELL checkpoint and the ProteinMPNN initialization weights are included in the repository:

weights/maxwell_proteinmpnn.ckpt
weights/proteinmpnn/v_48_020.pt

Score the substitutions listed in the example CSV:

python -m maxwell.infer.proteinmpnn \
  --mutant example_data/fireprotdb_1AG2_ddG.csv \
  --pdb example_data/fireprotdb_1AG2_ddG.pdb \
  --checkpoint_path weights/maxwell_proteinmpnn.ckpt \
  --chain A \
  --device cpu \
  --pred_column predicted_stability_score \
  --save_to outputs/example

Predictions are saved to outputs/example/fireprotdb_1AG2_ddG.csv. The input CSV must contain a mutant column with one-based substitutions such as A42G. The existing relative_ddg column in the example file is legacy output and is not used as an experimental target by this command.

The output score is log P(mutant amino acid) - log P(wild-type amino acid). Larger scores indicate greater predicted stability; scores are not absolute ΔΔG or ΔTm estimates. Set --device cuda to use a GPU.

To score a complete dataset, pass its root directory to --mutant, for example --mutant data/test12k, and omit --pdb. The directory must contain matching mutant/*.csv and pdb/*.pdb files. One scored CSV is saved per protein. Use --pred_column predicted_stability_score to preserve experimental labels stored in a column named score.

Train and Test Venus-MAXWELL (ProteinMPNN) on the stability datasets

The released datasets are:

• data/train226k.zip: Train226K, containing 226,915 records across 255 proteins.
• data/test12k.zip: Test12K, containing 12,496 raw records across 308 proteins. Excluding 394 wild-type-to-wild-type entries leaves 12,102 evaluated substitutions.
• data/full_dtm/: the full ΔTm dataset, containing 2,838 mutations across 86 proteins.

Extract the two archives:

python -m zipfile -e data/train226k.zip data
python -m zipfile -e data/test12k.zip data

Each dataset has matching fasta/, mutant/ and pdb/ subdirectories. FASTA sequences, chain A structures and one-based mutation positions must agree. Check the inputs with:

python scripts/check_dataset.py --dataset_path data/train226k
python scripts/check_dataset.py --dataset_path data/test12k
python scripts/check_dataset.py --dataset_path data/full_dtm

Train226K and Test12K mutation CSVs contain mutant,ddG; the training loader reverses the sign of ddG so that larger targets indicate greater stabilization. Full-DTM CSVs contain mutant,score, where score is ΔTm and larger values already indicate greater stabilization.

To train a new model, first create splits/train_ids.txt and splits/valid_ids.txt, each containing one protein basename per line, without a file extension. Select disjoint training and validation protein sets from Train226K and keep Test12K for testing. These split files are not currently bundled; the following is a training template for a user-defined split.

python -m maxwell.train \
  --model_type proteinmpnn \
  --model_path weights/proteinmpnn/v_48_020.pt \
  --mpnn_score_mode autoregressive \
  --data_root data/train226k \
  --train_data_path data/train226k \
  --valid_data_path data/train226k \
  --test_data_path data/test12k \
  --train_list_txt splits/train_ids.txt \
  --valid_list_txt splits/valid_ids.txt \
  --score_column ddG \
  --valid_score_column ddG \
  --test_score_column ddG \
  --max_seq_length 0 \
  --loss_type pearson \
  --batch_size 16 \
  --learning_rate 5e-4 \
  --max_epochs 100 \
  --check_val_every_n_epoch 1 \
  --save_top_k 1 \
  --monitor val_SC \
  --monitor_mode max \
  --checkpoint_dir checkpoints \
  --experiment_name maxwell_proteinmpnn \
  --accelerator gpu \
  --devices 1 \
  --precision 32

The command saves the best validation checkpoint under checkpoints/ and TensorBoard logs under tb_logs/. --max_seq_length 0 disables length-based exclusion. The existing example_datasets/ directory is a legacy dataset layout and is not used by these instructions.

Evaluate the released checkpoint on Test12K and Full-DTM

The checkpoint can be evaluated directly without training. For Test12K, compute per-protein correlations with:

python scripts/evaluate.py \
  --checkpoint weights/maxwell_proteinmpnn.ckpt \
  --data-dir data/test12k \
  --target-column ddG \
  --target-sign -1 \
  --dataset-name Test12K \
  --device cuda \
  --output-dir outputs/test12k_evaluation

Classification and screening metrics can additionally be enabled with --positive-threshold and --positive-rule. For example, --positive-threshold 0.3 --positive-rule gt defines positives as experimental -ΔΔG > 0.3. Choose the threshold appropriate to the analysis being reproduced.

For Full-DTM, define stabilizing mutations as ΔTm >= 1 °C:

python scripts/evaluate.py \
  --checkpoint weights/maxwell_proteinmpnn.ckpt \
  --data-dir data/full_dtm \
  --target-column score \
  --target-sign 1 \
  --dataset-name Full-DTM \
  --positive-threshold 1 \
  --positive-rule ge \
  --device cuda \
  --output-dir outputs/full_dtm_evaluation

Evaluation uses a fixed checkpoint and does not fit or calibrate scores on the evaluation labels. Self substitutions are excluded by default. Spearman and Pearson require at least three observations and nonconstant targets and predictions. AUROC and AUPRC require both positive and negative mutations; AUPRC is computed as average precision. Enrich@10 requires at least ten candidates and one positive mutation and measures top-10 positive frequency relative to the protein's background positive frequency.

All metrics are computed per protein and summarized as the unweighted mean and sample standard deviation (n - 1) across eligible proteins. The eligible count is reported separately for each metric. Outputs are:

per_protein_predictions/*.csv
combined_predictions.csv
per_protein_metrics.csv
summary.json
runtime.json

Use a new or empty output directory for each run. Set --device cpu for CPU evaluation. Run training and large-scale evaluation on allocated compute nodes when using a cluster.

Citation

If you find this work useful, please cite:

@article {Li2026.07.29.741631,
	author = {Li, Mingchen and Cheng, Xiaoran and Jiang, Fan and Hong, Liang and Yu, Yuanxi},
	title = {MAXWELL: Calibrating the probabilistic outputs of protein language models to the mutation-induced stability change landscape},
	elocation-id = {2026.07.29.741631},
	year = {2026},
	doi = {10.64898/2026.07.29.741631},
	publisher = {Cold Spring Harbor Laboratory},
	URL = {https://www.biorxiv.org/content/early/2026/08/04/2026.07.29.741631},
	eprint = {https://www.biorxiv.org/content/early/2026/08/04/2026.07.29.741631.full.pdf},
	journal = {bioRxiv}
}