Stack
When parsing, charges, force fields, inverse folding and language-model likelihoods all live in one Automatic differentiation: gradients flow through the whole composition, not just the last step., objectives that used to live in separate processes can be composed and differentiated together.
proteinsmc optimising ProteinMPNN likelihood, ESM likelihood and codon adaptation index as one fitness is the existence proof — and the reason expaloma had to come off PyTorch and DGL in the first place.
Every entry carries a status. working — in use and behaving. partial — real, incomplete. early — exploratory. paused — not currently developed.
Sampling and design
Proposing sequences, under a schedule and against a composed objective.
proteinsmc
earlySequential Monte Carlo over protein sequence space in JAX. Fitness functions compose arbitrarily, so ProteinMPNN likelihood, ESM likelihood and codon adaptation index can be optimised together — selecting for folding and expression at the same time rather than in sequence.
Learned models
What a structure or a family says about which residues belong where.
aminx
workingA functional JAX/Equinox reimplementation of LigandMPNN. Reproduces the PyTorch reference to ≥0.999 Pearson across all five decoding paths and runs 8–61× faster on a single structure, by trading eager dispatch for jit, vmap and scan. A sample/score API with no model objects to wire up.
Physics
Why some positions can vary and others cannot — energies, motion, and where a fold is physically constrained.
prolix
partialProtein physics and molecular dynamics in JAX — force fields, energies, integration, and a generalised Born implicit solvent, bridged to JAX-MD.
psax
earlyProtein Strain Analysis in JAX: per-site finite strain and deformation gradients across functional transitions, after Sartori and Leibler. Strain says which coordinates physically cannot move.
Parameters
The constants the physics needs, inferred rather than tabulated.
expaloma
workingNative JAX partial-charge inference, ported from espaloma_charge so downstream simulators do not have to carry PyTorch and DGL at runtime. Getting the whole pipeline into one framework is what lets objectives compose.
Parsing and IO
Getting structures and sequences into array form. Unglamorous, and the thing that breaks first.
Infrastructure
Cross-cutting: the training and tooling scaffolding every layer above needs.
xtrax
workingComposable building blocks for JAX and Equinox training loops — sharding, axis tiling, inference-time sparsification, checkpointing. jit freezes the program you gave it; the decisions that matter happen in Python before tracing, and that is the code that gets copy-pasted between projects.
jaxlint
workingA JAX-aware AST linter for JAX-specific concerns and doc quality, with a CLI, an LSP and an MCP server.
At the bench
The physical measurement and selection layer of the same stack. Variation can be narrowed by a temperature schedule or by a selection schedule; only one of them tells you whether the protein actually works.
Praxis
pausedA PyLabRobot-based lab automation framework with tiered-access UI, built for provenance: every run recorded, stored for audit, versioned. It came out of watching hastily assembled data repositories fall over, and out of collaborating with scientists who did not want to write code. Paused — AI coding assistants changed the economics of building that layer.
PRANCE
workingContinuous directed evolution at the bench. Characterising TEV protease variants across a panel of peptide substrates — activity and specificity — and the automated split-luciferase and growth-curve assays that made screening at that scale possible.