proteinmpnn-web

An ESM/TypeScript library for ProteinMPNN, solubleMPNN and CA-only ProteinMPNN inference in browsers, workers and Node.js. ONNX Runtime Web executes the neural network through WebAssembly or WebGPU. No server, Python runtime, UI framework or DOM is required for inference.

The implementation separates structure adapters, graph construction, model execution, autoregressive sampling and upstream JSON adapters. It ships an extensible InferenceBackend interface, declarations, source maps, checkpoint conversion tools and numerical parity tests.

Install and import

This is a source release; it has not been published to npm. From the extracted project:

npm ci
npm run build
npm pack
# In the consuming application:
npm install /path/to/proteinmpnn-web-0.1.0.tgz

import { ProteinMPNN, fromCoordinates, modelManifestURL, parsePDB } from 'proteinmpnn-web';
// Parser-only imports do not load ONNX Runtime:
import { fromResidues } from 'proteinmpnn-web/structure';

Model weights are deliberately excluded from this Git repository and the npm package. The nine converted ONNX bundles are hosted in MurrellLab/webports with their licenses, provenance, file sizes and SHA-256 hashes. modelManifestURL() resolves a bundled model from an immutable Hub revision; pass a custom manifest URL instead when hosting your own weights.

Design from a PDB

const model = await ProteinMPNN.load(modelManifestURL('proteinmpnn', 'v_48_020'), {
  backend: 'auto',                 // 'wasm' | 'webgpu' | 'auto'
  wasmPaths: '/onnxruntime/',      // serve matching ORT dist .wasm and .mjs files here
  numThreads: 1,
  onFallback: reason => console.warn('Using WASM:', reason),
});

try {
  const protein = await model.preparePDB(pdbText);
  const result = await protein.sample({
    seed: 42,
    temperature: 0.1,
    designChains: ['A'],
    fixedPositions: [{ chain: 'A', number: 27 }],
    omitAminoAcids: 'CX',
  });
  console.log(result.sequence, result.chains, result.score);
} finally {
  await model.dispose();
}

Model downloads use standard fetch, so the browser caches them according to the Hugging Face response headers. In Node.js, pass absolute HTTP URLs or load model bytes with node:fs; see examples/node.mjs. The browser runtime deliberately does not import fs or assume that fetch supports file: URLs. Importing this package in an SSR environment does not initialize inference.

Design from coordinates

const structure = fromCoordinates({
  coordinates: backbone, // Float32Array [N,4,3], atom order N, CA, C, O; angstroms
  sequence: nativeSequence,  // omit for X at every position
  chainIds,                 // string[N]; arbitrary identifiers
  residueNumbers,           // optional original PDB numbers
  insertionCodes,           // optional string[N]
  residueIndices,           // optional explicit positional-encoding indices
  mask,                     // optional 0/1 validity mask
});
const protein = await model.prepare(structure, { backboneNoise: 0.02, seed: 5 });
const result = await protein.sample({ seed: 42 });

For another atomic representation, use fromResidues([{ chain, number, aminoAcid, atoms: { N, CA, C, O } }, …]). Each atom is a three-component tuple. CA-only input uses atomMode: 'ca' and flattened [N,3] coordinates. Load a ca checkpoint for that input. A CA model can also extract CA coordinates from a complete backbone structure.

All numeric position references in the API are zero-based flattened residue indices. Object references use original chain/number/insertion-code identities. Upstream JSON helper positions are one-based within each chain; the legacy adapter converts them explicitly.

Nonfinite or missing required atoms mask the entire residue; they are replaced with zero coordinates for the model. An omitted native sequence becomes X; provide the true sequence for fixed positions and sequence-conditioned scoring. Coordinates are copied on input. Do not mutate a prepared structure's arrays while it is in use.

Sampling and constraints

const results = await protein.sampleMany({
  temperatures: [0.1, 0.2, 0.3],
  numSequences: 4,               // per temperature: 12 outputs here
  batchSize: 2,                 // concurrent independent sampling states
  seed: 123,
  designChains: ['A', 'B'],
  fixedPositions: [0, 1, { chain: 'B', number: 50, insertionCode: 'A' }],
  biasAminoAcids: { A: -1.1, F: 0.7 },
  biasByResidue: [{ position: 12, bias: { W: 1.5 } }],
  omitByResidue: [{ position: 13, aminoAcids: 'CP' }],
  tiedPositions: [
    [12, 112],
    { positions: [13, 113], weights: [0.5, 1.5] },
  ],
});

sampleStream() yields results in temperature/sample order without retaining all completed outputs. designPositions optionally restricts the mutable subset. Temperatures must be greater than zero. Tied logits use the upstream weighted sum divided by group size; weights are not automatically normalized.

For compatibility, the default tiedConstraintMode: 'upstream' applies per-residue bias, omissions and PSSM from the last member of each tied group. 'intersection' combines all member omissions and averages their biases; tied PSSM rows must agree. Fixed members pin a whole group; contradictory fixed identities and missing coordinates in a tied group are errors.

PSSMs use flattened [N,21] arrays in ACDEFGHIKLMNPQRSTVWYX order:

await protein.sample({
  pssm: {
    probabilities, coefficients, // probabilities sum to 1; coefficients in [0,1]
    useBias: true, mix: 0.5,
    logOdds, useLogOdds: true, threshold: 0,
  },
});

The operation order matches upstream: temperature and biases → global omissions → softmax → PSSM probability mixture → soft log-odds mask (mask + 0.001) → per-position omissions → normalization. Consequently, the PSSM mixture can reintroduce a globally omitted amino acid. Use omitByResidue for a hard final exclusion. Empty feasible distributions raise an error.

Scoring and probabilities

const nativeScore = await protein.score();
const candidateScore = await protein.score('ACDE…/FGHI…', { seed: 7 });
const unconditional = await protein.unconditionalProbabilities();
const conditional = await protein.conditionalProbabilities(nativeSequence, {
  positions: [0, 10, 20], // omitted: all valid residues
  seed: 7,
});
const backboneOnly = await protein.conditionalProbabilities(nativeSequence, {
  backboneOnly: true,
});

score is mean negative log likelihood in nats over selected, valid design positions. globalScore covers every valid residue. Empty selected sets return null. perResidue contains NaN at invalid residues. logProbabilities is the raw, untempered [N,21] model distribution; sample().probabilities is the actual constrained sampling distribution, with zero rows at fixed positions. Sample scores are recomputed with the full teacher-forced decoder and the sampled decoding order, as in upstream.

Conditional probabilities evaluate each requested residue with its sequence revealed last; this takes one full decoder call per requested residue. Backbone-only probabilities use the single-pass unconditional graph. Unrequested conditional rows are NaN. Always consult the structure mask when interpreting probability rows.

decodingOrder accepts an explicit complete permutation. orderNoise accepts upstream's randn values to reproduce its ordering. The built-in seeded PRNG is deterministic within this library; it does not reproduce PyTorch's random stream. Float32 differences between devices can change a categorical draw near a boundary. Seed 0 is deterministic here.

Model families and custom weights

Family	Checkpoints	Input
proteinmpnn	v_48_002, v_48_010, v_48_020, v_48_030	N, CA, C, O
solublempnn	v_48_010, v_48_020	N, CA, C, O
ca	v_48_002, v_48_010, v_48_020	CA

Every bundle has dynamic residue and neighbor dimensions, float32 weights, SHA-256 hashes and checkpoint provenance. Models are checked for integrity on load by default. A raw PyTorch .pt file is not a browser-loadable format; export it once:

python -m pip install -r tools/requirements.txt
python tools/export.py --all
python tools/export.py --family solublempnn --model v_48_020
python tools/export.py --checkpoint custom.pt --family proteinmpnn --model custom

The exporter validates state-dict shapes and reuses the vendored upstream network layers. It supports checkpoints with the upstream architecture; it is not a general converter for unrelated MPNN variants. Source and automatic checkpoint downloads are pinned to the commit in vendor/REVISION. Refer to docs/model-format.md for the graph interface and docs/compatibility.md for option mapping and differences.

Browser/worker integration and performance

• Copy the matching .wasm and .mjs files from node_modules/onnxruntime-web/dist/ to your static asset directory. Do not mix ORT JS and WASM versions. The tested dependency is locked in package-lock.json.
• backend: 'auto' selects WebGPU when available, with WASM fallback at session creation. 'webgpu' requires a working WebGPU session; it does not silently switch the whole session to WASM. ORT may place individual unsupported operations on WASM in a WebGPU session. Runtime/device-loss errors are surfaced.
• WebGPU requires a secure context, such as HTTPS or localhost. One WASM thread works without cross-origin isolation. For multiple WASM threads configure COOP/COEP and the runtime assets in your application. Runtime environment settings are global: configure them before the first model load.
• Run long designs in a module worker; see examples/worker.ts. AbortSignal is checked between graph executions, and onProgress reports residue progress. An in-flight kernel cannot be interrupted.
• Graph search is exact on the CPU: O(N² log K) work and O(NK) output memory. Encoder feature construction uses O(NK) rather than dense N² intermediate features. The default maxLength is a validation guard of 10,000, not a promise that a device can process that length.
• The structure encoder runs once per prepare(). Sampling reuses cached decoder states and runs one small decoder graph per residue; it does not rerun the full network at every residue. Decoder state is CPU-resident, with gathered neighbor tensors transferred for each WebGPU step. This favors portability; small designs may be faster on WASM. This release does not claim fully GPU-resident decoding or hardware speedups.
• batchSize schedules independent sampling states and bounds their memory; ORT session calls are serialized. It is not PyTorch's padded tensor batch size. The low-level step graph additionally supports batching independent residue steps for custom schedulers.
• Use try/finally and model.dispose() to release sessions. One model can prepare multiple structures. Prepared objects become unusable after model disposal.

Verification

npm test                         # input/constraint/lifecycle tests; no model downloads
python tools/validate.py          # upstream PyTorch vs ONNX, writes reference fixtures
npm run test:parity               # actual WASM inference vs those references
npx playwright install chromium --only-shell
npm run test:browser              # browser WASM and WebGPU vs upstream

The validation report included with this release describes the checks actually run and their limitations. Numerical agreement is checked against real official weights, including native scoring, unconditional/conditional probabilities, incremental decoding, PSSMs and tied sampling. See docs/validation.md.

MIT licensed. Original ProteinMPNN: dauparas/ProteinMPNN, Dauparas et al., Science 2022. Runtime: ONNX Runtime Web. Upstream notices are preserved in vendor/LICENSE.ProteinMPNN and THIRD_PARTY_NOTICES.md.