escalante-bio/mosaic

mosaic – 基于 JAX 的多目标蛋白质设计框架

是什么

一个开源 Python 库（托管于 github.com/escalante-bio/mosaic），允许研究人员将多种可微分的蛋白质属性预测器和结构预测模型组合成一个基于梯度的优化循环。
基于 JAX 构建，所有模型共享一个共同的 JIT 编译后端，可无须复杂的容器编排进行链式调用。

核心理念

统一的损失接口 – 每个预测器（如 Boltz‑1、AlphaFold‑Multimer、ESM‑2、ProteinMPNN、稳定性模型等）都暴露一个 build_loss 方法，返回一个 JAX 兼容的损失项。
连续序列松弛 – 设计表示为 20 种氨基酸上的软最大概率矩阵，从而可在序列空间上进行梯度下降。
可插拔优化器 – 库自带一个简单的投影梯度优化器（simplex_APGM），并提供自定义优化器示例（如带盒约束的 RSO）。

包含内容

模型：Boltz‑1/2、BoltzGen、AlphaFold‑2/Multimer、OpenFold‑3、OpenDDE、ESMFold‑2、Protenix（mini/tiny/base/2025）、ProteinMPNN（standard/soluble/AbMPNN）、ESM‑2/C、稳定性预测器、AbLang/AbLang2、三元语法语言模型，以及若干演示模型。
损失库：结构预测损失（结合体-靶标接触、回转半径、PLDDT 等）、使用 ProteinMPNN 的逆折叠损失、稳定性、n-gram 似然、用户自定义损失。
示例与笔记本：端到端的迷你结合体设计笔记本、竞赛获胜脚本、添加新损失项的教程。
工具：dataclass TargetChain 用于指定序列、MSA 和可选模板；特征提取、模型无关预测和 PSSM 处理的辅助函数。

典型工作流

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

# 加载模型
boltz = Boltz1()
mpnn  = ProteinMPNN.from_pretrained()

# 定义目标和结合体长度
target_seq = "..."
binder_len = 80

# 提取复合物和单体的特征
cplx_feat, _ = boltz.binder_features(binder_len, [TargetChain(sequence=target_seq)])
mono_feat, _ = boltz.binder_features(binder_len, [])

# 构建复合损失（结合性、几何、溶解性、逆折叠等）
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()
)

# 初始化软最大序列矩阵
x0 = jax.nn.softmax(
    0.5 * jax.random.gumbel(jax.random.key(0), (binder_len, 20))
)

# 运行基于梯度的设计
_, pssm = simplex_APGM(loss_function=loss, x=x0, n_steps=150, stepsize=0.1, momentum=0.9)

生成的 pssm 可通过采样获得可用于实验测试的具体氨基酸序列。

为何重要

蛋白质工程通常需要在结合亲和力、特异性、表达性、溶解性、稳定性等多个相互竞争的目标之间取得平衡。mosaic 允许将所有这些目标表示为可微分项并联合优化，这在单独脚本或黑盒管道中难以实现。
由于所有内容都在 JAX 中运行，相同的代码在 CPU、GPU 或 TPU 上只需极少修改即可工作，JIT 编译使首次调用后的内层循环非常快速。

限制 / 注意事项

该库是一个研究框架，而非即插即用的设计工具。用户必须调整学习率、优化器超参数和损失权重；选择不当可能导致设计无法通过基本的计算机模拟检查。
首次 JAX 编译较慢；建议使用交互式笔记本进行快速迭代。
某些结构预测模型（如 AF3 类型）在设计过程中会省略侧链信息，可能影响下游过滤。
重型模型（AlphaFold、OpenFold、Protenix 等）需要 GPU/TPU 兼容的 JAX 构建。

快速入门

克隆仓库并使用 uv（推荐）或 pip 安装：

git clone https://github.com/escalante-bio/mosaic.git
cd mosaic
uv sync --group jax-cuda   # 或手动添加 jax[cuda12]

运行示例笔记本：

uv run marimo edit examples/example_notebook.py

参考 README 中链接的 迷你结合体设计 博客文章，重现针对 PD‑L1 和 IL7Ra 的实验验证结合体。

引用如果在发表工作中使用 mosaic，请引用：

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

escalante-bio/mosaic

mosaic – 基于 JAX 的多目标蛋白质设计框架

相关