A flexible framework for multi-objective protein design, escalante-bio/mosaic allows researchers to programmatically combine diverse protein property predictors into a unified system. It addresses the challenge of optimizing multiple constraints simultaneously, such as binding affinity, expression, and solubility, offering a powerful tool for custom protein design applications without the typical dependency hell of BioML projects.

How It Works

Mosaic leverages JAX for efficient, just-in-time (JIT) compiled computation, facilitating the programmatic connection of various machine learning models. Its core approach combines a modular interface for defining custom loss terms and optimization algorithms with gradient-based optimization over a continuous, relaxed sequence space. This allows for the simultaneous optimization of multiple learned objective functions, offering a flexible and efficient alternative to containerized or script-heavy workflows.

Quick Start & Requirements

Installation is recommended via uv, using commands like uv sync --group jax-cuda after cloning. Running example notebooks can be done with uv run marimo edit examples/example_notebook.py. A GPU or TPU-compatible version of JAX (e.g., jax[cuda12]) is required for structure prediction. Initial JAX compilation may introduce a delay on the first function call.

Highlighted Details

• Integrates a comprehensive suite of models including Boltz-1/2, BoltzGen, AlphaFold2, OpenFold3, ESMFold2, Protenix (multiple versions), ProteinMPNN (standard, soluble, AbMPNN), ESM (2, C), stability models, AbLang/AbLang2, trigram, and Proteina-Complexa.
• Successfully applied in protein binder design competitions, yielding experimentally validated binders with high hit rates and affinities against targets like PD-L1, IL7Ra, and Nipah virus.
• Supports defining custom loss terms and swapping various optimization algorithms, such as simplex_APGM and gradient_MCMC.
• Enables programmatic composition of models, avoiding complex dependencies and containerization typical in BioML.

Maintenance & Community

No specific community links (Discord, Slack), roadmap, or notable contributor information were found in the provided README.

Licensing & Compatibility

The README does not specify the project's license or provide compatibility notes for commercial use.

Limitations & Caveats

Mosaic requires significant user intervention for tuning learning rates and loss weights, and its outputs often necessitate standard filtering methods as they may fail simple in-silico tests. It is explicitly described as "not for the faint of heart" and intended as a framework for advanced users implementing custom objectives. Predictions from AF3-style models using design features lack sidechains; reprediction with target-only features is needed for sidechain details.