escalante-bio/mosaic

mosaic – 多目的タンパク質設計のためのJAXベースフレームワーク

何であるか

• github.com/escalante-bio/mosaic でホストされているオープンソースのPythonライブラリで、研究者がさまざまな微分可能なタンパク質特性予測モデルや構造予測モデルを1つの勾配ベース最適化ループに統合できる。
• JAX上で構築されており、すべてのモデルが共通のJITコンパイルバックエンドを共有し、重いコンテナオーケストレーションなしにチェーン可能である。

コアアイデア

1. 統一された損失インターフェース – 各予測モデル（例：Boltz‑1、AlphaFold‑Multimer、ESM‑2、ProteinMPNN、安定性モデルなど）は、JAX互換の損失項を返す build_loss メソッドを公開する。
2. 連続的な配列リラクゼーション – デザインは20種類のアミノ酸に対するソフトマックス確率行列として表現され、配列空間における勾配降下が可能になる。
3. プラグイン可能な最適化アルゴリズム – ライブラリにはシンプルな投影勾配最適化アルゴリズム（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、3-gram言語モデル、およびいくつかのデモモデル。
• 損失ライブラリ: 構造予測損失（結合体-ターゲット接触、半径の慣性モーメント、PLDDTなど）、ProteinMPNNを用いた逆フォールディング損失、安定性、n-gram尤度、カスタムユーザー定義損失。
• 例とノートブック: ミニバインダー設計のエンドツーエンドノートブック、コンペ勝利スクリプト、新しい損失項の追加チュートリアル。
• ユーティリティ: シーケンス、MSA、オプションのテンプレートを指定するためのdataclass TargetChain；特徴抽出、モデル非依存予測、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ビルドが必要。

導入方法

1. リポジトリをクローンし、uv（推奨）またはpipでインストール：

   git clone https://github.com/escalante-bio/mosaic.git
   cd mosaic
   uv sync --group jax-cuda   # または手動で jax[cuda12] を追加
   

2. 例のノートブックを実行：

   uv run marimo edit examples/example_notebook.py
   

3. READMEにリンクされた「ミニバインダー設計」ブログ記事のウォークスルーに従い、PD‑L1およびIL7Raに対する実験的に検証されたバインダーを再現する。

引用 mosaic を出版物で使用する場合、次のように引用してください：

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