aminx weights — v0.1.0a1

JAX/Equinox reimplementation of ProteinMPNN and LigandMPNN.

Weights are converted from the original PyTorch checkpoints and verified to ≥0.999 Pearson correlation with the reference implementation (atol/rtol 1e-4 for all families; side-chain packer: atol 1e-4/rtol 1e-3).

Usage

pip install aminx   # huggingface-hub is a required dep; weights download on first use

from aminx.io.weights import load_model

model = load_model("proteinmpnn_v_48_020")
model = load_model("ligandmpnn_v_32_020_25")
model = load_model("solublempnn_v_48_020")
model = load_model("per_residue_label_membrane_mpnn_v_48_020")
model = load_model("ligandmpnn_sc_v_32_002_16")

Weights are cached to ~/.cache/huggingface/hub/ after first download.

Checkpoint topology

All families share: node/edge/hidden features = 128, encoder layers = 3, decoder layers = 3, vocab size = 21. The suffix digits in checkpoint names encode training noise level (e.g. _020 = 0.20 Å backbone noise, per upstream convention) and atom context count (e.g. _25, _16).

Family	Checkpoints	k_neighbors	num_positional_embeddings	Notes
ProteinMPNN	proteinmpnn_v_48_{002,010,020,030}	48	32	Original backbone
SolubleMPNN	solublempnn_v_48_{002,010,020,030}	48	32	Solubility-biased variant
LigandMPNN	ligandmpnn_v_32_{005,010,020,030}_25	32	32	2 ligand-context layers; atom_context_num=25
Membrane	{global,per_residue}_label_membrane_mpnn_v_48_020	48	32	physics_feature_dim=3
Side-chain packer	ligandmpnn_sc_v_32_002_16	32	16	atom_context_num=16

Deprecation notice

The previous maraxen/prxteinmpnn repository is archived. Use this repo going forward.

Citation

If you use these weights, please cite the original works:

@article{dauparas2022robust,
  title={Robust deep learning-based protein sequence design using ProteinMPNN},
  author={Dauparas, Justas and others},
  journal={Science},
  year={2022}
}

@article{dauparas2025atomic,
  title={Atomic context-conditioned protein sequence design using LigandMPNN},
  author={Dauparas, Justas and others},
  journal={Nature Methods},
  volume={22},
  number={4},
  pages={717--723},
  year={2025}
}