Warning
Alpha release (v0.1.0a1). aminx is under active development. The API is functional and validated against the LigandMPNN reference, but may change between releases. You may encounter bugs or rough edges — please open an issue if something breaks.
Aminx is a JAX/Equinox reimplementation of the LigandMPNN codebase. It reproduces the PyTorch reference to ≥ 0.999 Pearson correlation across all five decoding paths, and runs 8–61× faster on a single structure (H200) by trading eager dispatch for jit/vmap/scan kernels.
What you get:
- A functional
sample()/score()API — no model objects to wire up, no inference loop to write. - A composable inference layer (
StageSet) for swapping logit transforms, encode paths, and decode variants without touching kernel math. - Numerical parity with upstream LigandMPNN, validated across unconditional, conditional, autoregressive, membrane, and side-chain-packer paths.
- Native batching across temperatures, backbones, and sequence lengths — operations you'd normally write as Python loops are compiled into single JAX kernels, so there's no recompilation penalty and no padding waste.
The batching point is worth unpacking: in vanilla JAX, running N temperatures requires either a Python loop (N separate compiled calls, or worse, N retraces) or manually writing a vmap. aminx has already done that work. Passing temperature=[0.1, 0.3, 0.7] vmaps over the temperature axis in one compiled call; scoring a mixed-length library reuses a compiled kernel per length bucket rather than recompiling per structure. The speedups in the tables below come from applying this pattern consistently to the operations most commonly written as loops in protein design workflows.
Benchmarked on H200 (NVIDIA SXM5), A100 (PCIe), L40s, and Blackwell (SM120). All figures are warm-call medians; see the full benchmark report for every hardware and configuration.
H200 — single-structure latency (seq_len=76)
| Mode | aminx | ColabDesign (JAX) | LigandMPNN (PyTorch) | Speedup vs PyTorch |
|---|---|---|---|---|
| Autoregressive sample | 17 ms | 38 ms | 149 ms | 8.7× |
| Score conditional | 1.5 ms | 7.0 ms | 92 ms | 61× |
These are capabilities ProteinMPNN doesn't expose, measured against the PyTorch baseline doing the equivalent work.
| Capability | Config | aminx | PyTorch | Speedup vs PyTorch |
|---|---|---|---|---|
| DedupGather | K=1 unique / N=32 total | 1.1 ms | 92 ms | 80× |
| DedupGather | K=32 / N=32 (no dedup) | 36 ms | 2958 ms | 82× |
| Mixed-length batch | lengths [76, 150, 300, 500] | 4.1 ms | 2280 ms | 554× |
| Temperature array | M=8 temperatures | 2.2 ms/temp | — | 8× per-temp |
The 8× floor holds across hardware; ceilings reach 84–91× depending on the operation (A100: 8–91×, L40s: 8–85×, Blackwell SM120: 8–84×).
- Composition Guide —
StageSet,InferencePlan, and the five extension points - Parity Validation — numerical parity report vs the LigandMPNN reference
Aminx is validated against the upstream LigandMPNN reference (which includes ProteinMPNN behavior):
| Decoding Path | Tolerance | Status |
|---|---|---|
| Unconditional | atol/rtol 1e-4, corr ≥ 0.999 | Validated |
| Conditional | atol/rtol 1e-4, corr ≥ 0.999 | Validated |
| Autoregressive | atol/rtol 1e-4, corr ≥ 0.999 | Validated |
| Membrane | atol/rtol 1e-4, corr ≥ 0.999 | Validated |
| Side-chain packer | atol 1e-4/1e-3, corr ≥ 0.999 | Validated |
Full parity suite: 30/30 parity_heavy tests pass on the Engaging cluster (job 14203624). 575 fast tests pass locally (575 passed, 6 skipped, 2 xfailed).
Canonical parity docs (source of truth):
Root-level parity stubs are non-canonical; use the links above.
aminx uses uv as its package manager. Install uv first if you don't have it:
curl -LsSf https://astral.sh/uv/install.sh | shThen install aminx with your target hardware extra:
uv sync --extra cuda # GPU (CUDA)
uv sync --extra tpu # TPU
uv sync --extra cpu # CPU-onlyIf you prefer pip:
pip install "aminx[cuda]" # GPU
pip install "aminx[cpu]" # CPU-onlyInstall as a standalone CLI tool (no virtual environment required):
uv tool install aminx
aminx spec validate run_spec.jsonOr invoke without installing at all:
uvx aminx spec validate run_spec.jsonfrom aminx.io.weights import load_model
from aminx.run import sample, score, SamplingSpecification, ScoringSpecification
model = load_model(model_version="v_48_020", model_weights="original")
# --- Sampling ---
spec = SamplingSpecification(
inputs="path/to/structure.pdb",
num_samples=10,
temperature=0.1,
random_seed=42,
)
results = sample(spec)
sequences = results["sequences"] # (num_samples, seq_len)
logits = results["logits"] # (num_samples, seq_len, 21)
# --- Scoring ---
spec = ScoringSpecification(
inputs="path/to/structure.pdb",
sequences_to_score=["MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEK"],
temperature=1.0,
)
results = score(spec)
scores = results["scores"] # negative log-likelihood per sequenceFor control over fusion strategy, encode path, and decode variant without touching kernel math:
import jax.numpy as jnp
from aminx.host.plan import make_inference_plan, InferencePlan, InferenceComponents
from aminx.inference.encode import make_encode_fn
from aminx.inference import driver
from aminx.inference.logits import GeometricMeanLogits, ARLogitFuse
from aminx.run import SamplingSpecification
from aminx.types.stages import StageSet
# Option A: factory (resolves stages from spec automatically)
spec = SamplingSpecification(
inputs="structure.pdb",
num_samples=10,
multi_state_strategy="geometric_mean",
state_weights=[1.0, 0.8, 0.6],
)
plan = make_inference_plan(model, spec)
# Option B: manual assembly for full control
state_weights = jnp.array([1.0, 0.8, 0.6])
stage_set = StageSet(
logit_transform=GeometricMeanLogits(weights=state_weights, temperature=1.2),
ar_logit_transform=ARLogitFuse(),
)
plan = InferencePlan(
model=model,
components=InferenceComponents(
encode_fn=make_encode_fn(model, use_rolling_state=False),
driver=driver.decode,
stage_set=stage_set,
),
)
result = plan.sample(bundle, key, config) # → SampleResult(sequence, logits)
logits = plan.score(bundle, key, config) # → (L, 21)Plain callables and lambdas work in all StageSet slots (the driver uses eqx.filter_jit). Use eqx.Module only when the callable carries JAX array leaves (e.g. weights that need grad).
See the Composition Guide for the five extension points (logit_transform, ar_logit_transform, decode_step, sample_step, tie_group_fuse).
Pass a list of temperatures to run all M in a single JIT-compiled call. The kernel vmaps over the temperature dimension, so per-temperature cost scales near-ideally.
from aminx.run import sample, SamplingSpecification
spec = SamplingSpecification(
inputs="structure.pdb",
temperature=[0.1, 0.3, 0.7, 1.0], # M=4 temperatures, one JIT call
num_samples=10,
)
results = sample(spec)
# One result dict per temperature, each with "sequences" and "logits"
for temp, result in zip(spec.temperature, results):
print(f"T={temp}: {result['sequences'].shape}")On H200 at M=8 (seq_len=76), per-temperature cost is 2.2 ms — the same total wall-clock as M=1 (17 ms) split across 8 temperatures. Measured scaling: 7.97–8.47× across H200/A100/L40s/Blackwell.
When scoring a large ensemble where many sequences share a backbone, score only the K unique structures rather than all N. aminx JIT-caches a compiled kernel per length bucket, so K sequential plan.score() calls stay fast regardless of N.
from aminx.host.plan import make_inference_plan
from aminx.inference.bundle_builder import build_inference_bundle
from aminx.tiling.bucketing import BucketingConfig
import jax.random as random
plan = make_inference_plan(model, spec)
bucket_cfg = BucketingConfig()
key = random.PRNGKey(0)
# Build one bundle per unique backbone (K unique out of N total)
unique_bundles = []
for coords, mask, residue_index, chain_index, sequence in unique_structures:
bundle, config = build_inference_bundle(
coords=coords,
mask=mask,
residue_index=residue_index,
chain_index=chain_index,
sequence=sequence,
ligand_coords=None,
ligand_atom_types=None,
ligand_mask=None,
temperature=1.0,
mode="score_conditional",
inference=True,
bucket_config=bucket_cfg,
)
unique_bundles.append((bundle, config))
# Score K unique structures; scatter the K scores to your N positions
scores = [plan.score(bundle, key, config) for bundle, config in unique_bundles]At K=1/N=32 on H200, latency is 1.1 ms vs 92 ms for PyTorch scoring all 32 structures — an 80× speedup. The speedup is stable across K (80–82× from K=1 through K=32) because PyTorch's cost scales linearly with N while aminx's scales linearly with K.
Score a library of proteins with different sequence lengths without padding waste. aminx rounds each length to the next power-of-2 boundary and reuses the JIT-compiled kernel across structures in the same bucket — one compile per bucket, not one per structure.
from aminx.host.plan import make_inference_plan
from aminx.inference.bundle_builder import build_inference_bundle
from aminx.tiling.bucketing import BucketingConfig
import equinox as eqx
plan = make_inference_plan(model, spec)
bucket_cfg = BucketingConfig()
# Build one bundle per structure — lengths can differ freely
bundles = []
for coords, mask, residue_index, chain_index, sequence in your_library:
bundle, config = build_inference_bundle(
coords=coords,
mask=mask,
residue_index=residue_index,
chain_index=chain_index,
sequence=sequence,
ligand_coords=None,
ligand_atom_types=None,
ligand_mask=None,
temperature=1.0,
mode="score_conditional",
inference=True,
bucket_config=bucket_cfg,
)
bundles.append((bundle, config))
# Each call reuses the compiled XLA kernel for its length bucket
score_one = eqx.filter_jit(plan.score)
scores = [score_one(bundle, key, config) for bundle, config in bundles]For a batch of [76, 150, 300, 500]-residue structures on H200, total latency is 4.1 ms versus 2280 ms for PyTorch (padded sequential) — a 554× speedup. The gain comes from two sources: JAX's compiled kernels (vs PyTorch eager dispatch) and avoiding padding to the longest sequence.
aminx ships a Typer CLI with four command groups: run, campaign, spec, and the spec emit-* family.
# Full sampling pipeline — constructs a RunSpecification and runs it end-to-end
aminx run sample \
--inputs structure.pdb \
--model-version v_48_020 \
--model-weights original \
--num-samples 10 \
--random-seed 42
# Emit the spec JSON instead of running (useful for inspection or handoff)
aminx run sample --inputs structure.pdb --emit-json
aminx run sample --inputs structure.pdb --emit-json --out sample_spec.json
# score, jacobian, and inspect accept the same options;
# pass --emit-json to get the spec — the runner for these paths is not yet wired
aminx run score --inputs structure.pdb --sequences-to-score ACDEFGHIKLMNPQRSTVWY --emit-jsonAll four subcommands (sample, score, jacobian, inspect) share the same base option surface: --inputs, --model-weights, --model-version, --model-family, --batch-size, --backbone-noise, --random-seed, and the full RunSpecification field set. Spec construction failures exit 1; an unwired runner exits 2.
The --inputs flag accepts multiple URI schemes, allowing you to mix local files and remote-fetched structures in a single command. Remote sources are fetched at CLI time to a local cache directory, making the resolved spec safely reproducible on cluster compute nodes (which are offline).
Accepted URI forms:
| Form | Example | Source | Behavior |
|---|---|---|---|
| Bare path | /data/1ubq.pdb, ./structures/*.cif, data/ |
Local | Files and glob patterns expanded as today; directories scanned for .pdb/.cif files. |
file:// |
file:///data/1ubq.pdb, file://localhost/data/1ubq.pdb |
Local | Explicit file URI; host portion (if present) is stripped; remainder is treated as a local path. |
pdb:// |
pdb://1A3A, pdb://1A3A.pdb, pdb://1A3A.cif |
RCSB | Fetch PDB structure by ID. Default format mmcif; suffix .pdb or .cif overrides. |
afdb:// |
afdb://P12345 |
AlphaFold DB | Fetch AlphaFold v4 structure by UniProt ID. |
mdcath:// |
mdcath://1abcA00 |
MD-CATH | Fetch MD-CATH HDF5 structure by ID. |
Example — mixed local and remote inputs:
# Fetch two RCSB structures and one AlphaFold structure; combine with a local directory
aminx run sample \
--inputs pdb://1A3A \
--inputs pdb://2ABC.pdb \
--inputs afdb://P12345 \
--inputs ./local_structures/ \
--model-version v_48_020 \
--num-samples 10All remote sources are fetched to a cache directory (default: ~/.cache/aminx/inputs or $XDG_CACHE_HOME/aminx/inputs). Subsequent runs with the same accession reuse cached files and skip the network call entirely — the resolved spec is fully deterministic and cluster-safe.
Cache control and URI scheme override:
# Override cache directory
aminx run sample --inputs pdb://1A3A --cache-dir /tmp/aminx_cache --num-samples 10
# Force all --inputs entries to be treated as local paths (suppress scheme detection)
aminx run sample --inputs "pdb://1A3A" --input-type file --num-samples 10
# ^ Treats "pdb://1A3A" as a literal local filename, does not fetch
# Apply a default scheme to schemeless tokens
aminx run sample --inputs 1A3A --input-type pdb --num-samples 10
# ^ Treats "1A3A" as pdb://1A3A; fetches from RCSBError handling: Empty accessions (e.g., pdb://), unknown schemes (e.g., s3://, http://), and fetch failures raise clear errors naming the problematic entry. Use --inputs-fail-fast to exit on first error, or omit it to warn and skip individual entries.
# Create a campaign manifest from a base spec
aminx campaign plan \
--inputs structures/target.pdb \
--campaign-id pilot_v1 \
--manifest-path pilot.manifest.json \
--output-root outputs/pilot \
--designs-per-library-type 50 \
--samples-chunk-size 16
# Execute a single manifest row (used by distributed workers)
aminx campaign worker --manifest-path pilot.manifest.json --row-index 0
# Execute all rows; exits 1 if any row fails
aminx campaign run --manifest-path pilot.manifest.json
# Evaluate quality gates; exits 2 if the campaign is not promoted
aminx campaign gates --manifest-path pilot.manifest.json
# Plan a staged scale ramp
aminx campaign ramp-plan \
--inputs structures/target.pdb \
--campaign-id ramp_v1 \
--manifest-dir ramp_manifests/ \
--output-root outputs/ramp \
--stage-designs-per-library-type 10,50,200 \
--samples-chunk-size 16
# Evaluate ramp stage reports; exits 2 if not promoted
aminx campaign ramp-evaluate --report-path stage1.json --report-path stage2.json--lock-backend distributed is not supported from the CLI (raises an error); use local_fs (the default).
# Validate a spec JSON — exits 0 and prints "OK: <SpecClass>", or exits 1 with the parse error
aminx spec validate run_spec.json
# Round-trip a spec through the codec — prints the re-serialized JSON to stdout
aminx spec roundtrip run_spec.json
aminx spec roundtrip run_spec.json --out validated.json # write to file instead
aminx spec roundtrip run_spec.json --compact # single-line JSON
# Round-trip the portable subset (dict-serializable fields only, no JAX arrays)
aminx spec portable-roundtrip portable_spec.json
aminx spec portable-roundtrip portable_spec.json --compactThe emit-sample, emit-score, emit-jacobian, and emit-inspect subcommands are a convenience complement to aminx run <cmd> --emit-json — identical output, no runner invoked.
# Emit a sample spec JSON — equivalent to: aminx run sample ... --emit-json
aminx spec emit-sample \
--inputs structure.pdb \
--model-version v_48_020 \
--model-weights original \
--compact
aminx spec emit-sample --inputs structure.pdb --out sample_spec.jsonSpecs can also be serialized from Python with run_specification_to_json:
from aminx.run import run_specification_to_json
json_str = run_specification_to_json(spec)The aminx.potts module provides a parallel model family integrating PottsMPNN with
Tree-Reweighted belief propagation (TRW) for global sequence design.
| Component | Description |
|---|---|
PottsModel |
Equinox module wrapping PottsMPNN + DifferentiableTRW |
PoeModel |
N-backbone Product-of-Experts ensemble |
MpnnPottsDesigner |
MPNN-seeded Gibbs sampling coordinator |
CLI: aminx potts run, aminx potts emit
Weight recapture: scripts/recapture/pottsmpnn_to_eqx.py
Architecture: Parallel model family — see .praxia/docs/decisions/260605_potts-parallel-not-stageset.md
- Python ≥ 3.12
- JAX + Equinox (GPU/TPU/CPU via extras)
uv sync --extra cpufor CPU-only;--extra cudafor GPU
| Command | Purpose |
|---|---|
uv run pytest |
Fast test suite (excludes parity_heavy) |
uv run ruff check src |
Lint |
uv run ty check |
Type check (ty strict) |
uv run ruff format . |
Auto-format |
All five decoding paths are validated via parity_heavy tests — see Validation Reference below.
aminx.run ← SamplingSpecification, ScoringSpecification, sample(), score()
aminx.host.plan ← InferencePlan, InferenceComponents, make_inference_plan()
aminx.types.stages ← StageSet (the composition interface)
aminx.inference ← driver.decode, logits (LOGIT_STRATEGIES, TIE_GROUP_STRATEGIES)
aminx.model ← LigandMPNN, Packer (Equinox modules, JIT-safe)
aminx.sampling ← sample() kernel
aminx.scoring ← score() kernel
aminx.cli ← aminx run / campaign / spec (Typer entry point)
StageSet is the seam between the host layer and the JAX-traced kernels: everything above it is Python-land, everything below it is traced. See the Composition Guide.
Importing aminx does not set the multiprocessing start method. If your notebook or script spawns worker processes, call configure_multiprocessing() once at startup (see aminx.runtime); the campaign CLI does this for you.
Running the equivalence and parity suite
# Install project dependencies (CPU/dev/tests path)
uv sync --extra cpu --extra dev --extra tests --group dev
source .venv/bin/activate
# Checkout reference implementation (pinned commit used in CI)
git clone https://github.com/dauparas/LigandMPNN.git reference_ligandmpnn_clone
cd reference_ligandmpnn_clone && git checkout 3870631 && cd ..
# Optional strict preflight per parity tier
REFERENCE_PATH=./reference_ligandmpnn_clone \
uv run python scripts/check_parity_prereqs.py --reference-path "$REFERENCE_PATH" --project-root . --tier parity_heavy
REFERENCE_PATH=./reference_ligandmpnn_clone \
uv run python scripts/check_parity_prereqs.py --reference-path "$REFERENCE_PATH" --project-root . --tier parity_audit
# Validate parity asset cache/checksums
uv run python scripts/check_parity_assets.py --tier parity_fast
REFERENCE_PATH=./reference_ligandmpnn_clone \
uv run python scripts/check_parity_assets.py --tier parity_heavy
REFERENCE_PATH=./reference_ligandmpnn_clone \
uv run python scripts/check_parity_assets.py --tier parity_audit
# Run fast deterministic parity checks
uv run pytest tests/parity -m parity_fast -v
# Run reference-backed heavy parity checks
REFERENCE_PATH=./reference_ligandmpnn_clone \
PRXTEIN_PARITY_TIER=parity_heavy \
uv run pytest tests/parity tests/model/test_ligandmpnn_equivalence.py -m parity_heavy -v
# Convert full checkpoint families and run parity_audit checks
REFERENCE_PATH=./reference_ligandmpnn_clone \
uv run python scripts/convert_parity_family_weights.py \
--project-root . \
--reference-path "$REFERENCE_PATH" \
--tier parity_audit \
--skip-existing
REFERENCE_PATH=./reference_ligandmpnn_clone \
PRXTEIN_PARITY_TIER=parity_audit \
uv run pytest tests/parity tests/model/test_ligandmpnn_equivalence.py -m parity_audit -v
# Collect expanded parity evidence (multi-backbone + synthetic random cases)
REFERENCE_PATH=./reference_ligandmpnn_clone \
uv run python scripts/collect_parity_evidence.py \
--project-root . \
--case-corpus tests/parity/parity_case_corpus.json \
--output-dir docs/parity/reports/evidence
# Render Markdown/HTML report and export PDF with embedded plots/tables
uv run python scripts/generate_parity_report.py --project-root . --output-dir docs/parity --pdfAMINX_VERIFY (runtime jaxtyping + beartype): tests under tests/parity/ set AMINX_VERIFY=1 via tests/parity/conftest.py. Elsewhere, opt in with:
AMINX_VERIFY=1 uv run pytest path/to/test.py -vCI tier routing:
- pull_request/main CI excludes
parity_heavyandparity_auditfrom the default pytest matrix. parity.ymlruns heavy reference-backed checks onmainpush and manual dispatch.parity-audit.ymlruns full-family audit checks on weekly schedule and manual dispatch.ligand-tied-positions-and-multi-stateis staged as warn-only inparity_heavyand fail inparity_audit.
MIT License.
Contributions are welcome — see the contributing guidelines.
- Documentation: http://maraxen.github.io/Aminx
- Issues: GitHub Issues
- Discussions: GitHub Discussions
Built in JAX, for anyone who'd rather design proteins than write inference loops.