Genie：通过等变扩散定向残基云进行从头蛋白质设计

本仓库提供了我们 ICML 论文 的实现代码。以下是采样过程的示意图。

安装

克隆此仓库并进入根目录。通过运行 pip install -e . 来设置软件包。这将自动安装代码所需的依赖项，包括 tensorboard 和 wandb 等日志记录包。

数据下载

我们提供了用于下载和清理 SCOPe 数据集的脚本。要下载，请运行

chmod +x scripts/install_dataset.sh
./scripts/install_dataset.sh

训练

要训练 Genie，创建一个目录 runs/[RUN_NAME] 并进入该目录。创建一个名为 configuration 的配置文件。配置文件示例在 example_configuration 中提供，完整的可配置参数列表可在 genie/config.py 中找到。请注意，配置文件中的 name 应与 RUN_NAME 匹配，以便日志记录到正确的目录。要开始训练，运行

python genie/train.py -c runs/RUN_NAME/configuration -g0 &

例如，在 GPU 0 上在后台运行。

采样

要使用自己训练的 Genie 采样域，运行

python genie/sample.py -n RUN_NAME -g0

默认情况下，它使用最新版本和轮次的检查点。您也可以使用 -v 和 -e 标志分别指定版本和轮次。这将在 50 到 128 之间的每个序列长度采样 10 个域，采样批量大小为 5。输出存储在目录 runs/[RUN_NAME]/version_[VERSION]/samples/epoch_[EPOCH] 中。

我们还提供了训练好的模型权重，可在 weights 目录下找到，附带相应的配置文件。要加载模型，运行

from genie.config import Config
from genie.diffusion.genie import Genie

config = Config('weights/configuration')
model = Genie.load_from_checkpoint('weights/genie_l_128_epoch=49999.ckpt', config=config)

评估

为了评估生成的样本，我们基于 ProteinMPNN 和 ESMFold 设置了一个评估流程。要设置评估流程，运行

./scripts/setup_evaluation_pipeline.sh

要运行评估流程，运行

python evaluations/pipeline/evaluate.py --input_dir INPUT_DIR --output_dir OUTPUT_DIR

这里，输入目录包含一个名为 coords 的子目录，其中包含由 Genie 生成的 Ca 坐标。输出目录包含评估结果。