ProteinMPNN.jl

ProteinMPNN.jl is a Julia port of ProteinMPNN with a strict numerical-parity focus against the ColabDesign JAX implementation.

The package is designed around two goals:

• parity and reproducibility against Python/JAX,
• a Julia-native user API for clean, explicit control over protein design.

Installation

using Pkg
Pkg.add(url = "https://github.com/MurrellLab/ProteinMPNN.jl")

Quick Start

using ProteinMPNN

# Download/load default checkpoint (vanilla, v_48_020).
model = mk_mpnn_model()

# Parse structure and cache prepared inputs on the model.
prep_inputs(model; pdb_filename = "example.pdb", chain = "A")

# Sample 10 sequences.
out = sample_parallel(model; batch = 10, temperature = 0.1)
out.seq

Core User API

1) Model + weights

model = mk_mpnn_model(
    family = :vanilla,     # :vanilla, :soluble, :ca
    model_name = "v_48_020",
)

Weights are fetched from Hugging Face through HuggingFaceApi.jl and cached locally by that package.

2) Structure preparation

prepared = prep_inputs(
    model;
    pdb_filename = "example.pdb",
    chain = "A",
    ignore_missing = true,
)

3) Sampling and scoring

samples = sample_parallel(model; batch = 10, temperature = 0.1)

native = score(model; S = prepared.S)
logits = get_logits(model; seq = "ACDEFGHIKLMNPQRSTVWY")
unconditional = get_unconditional_logits(model)

Custom logits-to-AA control:

selector = (position, logits_af) -> AF_ALPHABET[argmax(view(logits_af, 1:20))]
samples = sample_parallel(model; batch = 10, temperature = 0.1, aa_selector = selector)

You can also pass a richer callback ctx::SamplingStepContext for advanced policies (product-of-experts, top-k/nucleus, external model fusion).

4) Design controls

spec = make_design_spec(
    prepared;
    fixed = "A10,A11,A100A",
    numbering = :pdb,
    tied = [[25, 40], [26, 41]],
    global_bias = fill(0.1f0, 20),
    exclude_aa = "C",
)

out = sample_parallel(model; design = spec, batch = 8, temperature = 0.1)

5) Cross-target tied design

inp1 = prep_inputs(; pdb_filename = "state1.pdb", chain = "A")
inp2 = prep_inputs(; pdb_filename = "state2.pdb", chain = "A")

joint = sample_parallel_joint(
    model;
    inputs = [inp1, inp2],
    labels = ["state1", "state2"],
    tied_groups = [[(1, 35), (2, 52)]],
    batch = 4,
    temperature = 0.1,
)

Numbering Modes

Selectors and range controls support:

• :pdb for PDB residue numbering with insertion codes (A100A),
• :chain for contiguous per-chain indices (A:1-120),
• :global for concatenated positions across selected chains.

Full Documentation

A full multi-page docs set is in docs/src:

• Overview
• Getting Started
• Weights and Checkpoints
• PDB Parsing and prep_inputs
• Sampling and Scoring
• Design Controls
• Joint / Multi-Target Design
• Parity, AD, and GPU Notes
• Architecture and Module Layout
• API Reference
• Troubleshooting

Build docs locally:

julia docs/make.jl

Notes

• Default Hugging Face repo: MurrellLab/ProteinMPNN.jl
• Default revision: main
• Alias :orig / :original maps to :vanilla