escalante-bio/mosaic

mosaic – A JAX‑based framework for multi‑objective protein design

What it is

• An open‑source Python library (hosted at github.com/escalante-bio/mosaic) that lets researchers combine many differentiable protein‑property predictors and structure‑prediction models into a single, gradient‑based optimization loop.
• Built on JAX, so all models share a common, JIT‑compiled backend and can be chained without heavy container orchestration.

Core ideas

1. Unified loss interface – each predictor (e.g., Boltz‑1, AlphaFold‑Multimer, ESM‑2, ProteinMPNN, stability models, etc.) exposes a build_loss method that returns a JAX‑compatible loss term.
2. Continuous sequence relaxation – designs are represented as a soft‑max probability matrix over the 20 amino acids, enabling gradient descent on the sequence space.
3. Pluggable optimizers – the library ships with a simple projected‑gradient optimizer (simplex_APGM) and examples of custom optimizers (e.g., box‑constrained RSO).

What’s included

• Models: Boltz‑1/2, BoltzGen, AlphaFold‑2/Multimer, OpenFold‑3, OpenDDE, ESMFold‑2, Protenix (mini/tiny/base/2025), ProteinMPNN (standard/soluble/AbMPNN), ESM‑2/C, stability predictor, AbLang/AbLang2, a trigram language model, and a few demo models.
• Loss library: structure‑prediction losses (binder‑target contact, radius‑of‑gyration, PLDDT, etc.), inverse‑folding losses using ProteinMPNN, stability, n‑gram likelihood, custom user‑defined losses.
• Examples & notebooks: end‑to‑end notebooks for minibinder design, competition‑winning scripts, and tutorials on adding new loss terms.
• Utilities: dataclass TargetChain for specifying sequences, MSAs, and optional templates; helper functions for feature extraction, model‑agnostic prediction, and PSSM handling.

Typical workflow

from mosaic.models.boltz1 import Boltz1
from mosaic.proteinmpnn.mpnn import ProteinMPNN
from mosaic.losses import structure_prediction as sp
from mosaic.optimizers import simplex_APGM
import jax, numpy as np

# Load models
boltz = Boltz1()
mpnn  = ProteinMPNN.from_pretrained()

# Define target and binder length
target_seq = "..."
binder_len = 80

# Extract features for complex and monomer
cplx_feat, _ = boltz.binder_features(binder_len, [TargetChain(sequence=target_seq)])
mono_feat, _ = boltz.binder_features(binder_len, [])

# Build a composite loss (binding, geometry, solubility, inverse‑folding, etc.)
loss = (
    boltz.build_loss(sp.BinderTargetContact() + sp.DistogramRadiusOfGyration(15.0), cplx_feat)
    + 0.3 * sp.HelixLoss()
    + 5.0 * InverseFoldingSequenceRecovery(mpnn, temp=jax.numpy.array(0.01))
    + 0.5 * sp.StabilityLoss()
)

# Initialise a soft‑max sequence matrix
x0 = jax.nn.softmax(
    0.5 * jax.random.gumbel(jax.random.key(0), (binder_len, 20))
)

# Run gradient‑based design
_, pssm = simplex_APGM(loss_function=loss, x=x0, n_steps=150, stepsize=0.1, momentum=0.9)

The resulting pssm can be sampled to obtain concrete amino‑acid sequences for experimental testing.

Why it matters

• Protein engineering often requires balancing many competing objectives (affinity, specificity, expression, solubility, stability). mosaic lets you express all of them as differentiable terms and optimise them jointly, which is hard to do with separate scripts or black‑box pipelines.
• Because everything runs in JAX, the same code works on CPUs, GPUs, or TPUs with minimal changes, and JIT compilation makes the inner loop fast after the first call.

Limitations / cautions

• The library is a research framework, not a turnkey design tool. Users must tune learning rates, optimizer hyper‑parameters, and loss weights; poorly chosen settings can lead to designs that fail basic in‑silico checks.
• First‑time JAX compilation is slow; interactive notebooks are recommended for rapid iteration.
• Some structure‑prediction models (e.g., AF3‑style) omit side‑chain information during design, which may affect downstream filtering.
• GPU/TPU‑compatible JAX builds are required for the heavy models (AlphaFold, OpenFold, Protenix, etc.).

Getting started

1. Clone the repo and install with uv (recommended) or pip:

   git clone https://github.com/escalante-bio/mosaic.git
   cd mosaic
   uv sync --group jax-cuda   # or add jax[cuda12] manually
   

2. Run the example notebook:

   uv run marimo edit examples/example_notebook.py
   

3. Follow the walkthrough in the Minibinder design blog post (linked in the README) to reproduce experimentally validated binders against PD‑L1 and IL7Ra.

Citation If you use mosaic in published work, cite:

Boyd, N., Guns, S. & Escalante Bio. Mosaic. https://github.com/escalante-bio/mosaic (2025).