aminx

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).

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}
}