Genie 2: Designing and Scaffing Proteins at the Scale of the Structural Universe

This repository provides the implementation code for our preprint, including training and inference code, as well as model weights. For the in-silico evaluation pipeline, which is used to assess the designability, diversity and novelty of our generated structures, we provide them in a seperate repository since it is independent of Genie 2 and could be applicable for evaluating other protein structure diffusion models. Below provides an illustration of the Genie 2 sampling process for both unconditional generation and motif scaffolding.

Setup

Set up Genie 2 by running

git clone https://github.com/aqlaboratory/genie2.git
cd genie2
pip install -e .

This would clone the repository, install dependencies and set up Genie 2 as a packge.

Docker

Alternatively, you can run Genie 2 in a Docker container. The included Dockerfile contains the required dependencies to run Genie 2 using the pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime default base image.

## Build Docker Container
docker build -t genie2 .

Also, you can use the PYTORCH_TAG build argument to rebuild with a different base image, thus using different PyTorch, CUDA, and/or cuDNN versions.

For example:

## Use build argument to change versions
docker build --build-arg "PYTORCH_TAG=2.2.1-cuda12.1-cudnn8-devel"  -t genie2 .

Then, run the container with GPUs and mount a local volume using the following command:

## Run Container
docker run -v .:/data --gpus all -it genie2

This will open a bash terminal inside the container where you can run the following training and generation commands as needed.

Training

Dataset

In data/afdbreps_l-256_plddt_80, we provide an index file named index.txt. It contains a list of 588,571 AFDB filenames, each of which correponds to a structure in the AlphaFold Protein Structure Database, specifically in the format https://alphafold.ebi.ac.uk/files/[FILENAME].pdb. These structures are representative structures from the FoldSeek-clustered AFDB dataset (filtered with a maximum sequence length of 256 and a minimum pLDDT of 80) and are used for the training of Genie 2. To set up for this dataset, create a subdirectory named pdbs under data/afdbreps_l-256_plddt_80 and download this set of structures from AFDB into this pdbs subdirectory. Our default training data directory is set to data/afdbreps_l-256_plddt_80/pdbs and this could be changed by using the dataDirectory key in the configuration file.

Training

To train a model, create a directory runs/[RUN_NAME] and create a configuration file with name configuration under this directory. An example is provided in runs/example and a complete list of configurable parameters could be found in genie/config.py. Note that in the configuration file, name should match with RUN_NAME in order to log into the correct directory. To start training, run

python genie/train.py --devices [NUM_DEVICES] --num_nodes [NUM_NODES] --config runs/[RUN_NAME]/configuration

Sampling

Directory Setup

To sample using a model, create a directory results/[MODEL_NAME], which consists of

• the model configuration file named configuration
• a subdirectory named checkpoints, where each file is named epoch=[EPOCH].ckpt and contains the model weight checkpointed at the specified epoch.

An example is provided under results/base, which contains the configuration and checkpoints for our trained model. The checkpoints are available in v1.0.0 release. For results reported in our manuscript, we used the 40-epoch checkpoint for unconditional generation and 30-epoch checkpoint for motif scaffolding.

Unconditional Generation

Perform unconditional sampling by running

python genie/sample_unconditional.py --name [NAME] --epoch [EPOCH] --scale [SCALE] --outdir [OUTDIR]

The list of parameters are summarized in the following table.

Parameter	Type	Description	Required	Default
name	string	Model name	✔️
epoch	int	Model epoch	✔️
scale	float	Sampling noise scale (between 0 and 1, inclusive)	✔️
outdir	string	Output directory	✔️
rootdir	string	Model root directory		results
num_samples	int	Number of samples per length		5
batch_size	int	Batch size		4
min_length	int	Minimum sequence length		50
max_length	int	Maximum sequence length		256
length_step	int	Sequence length step		1
num_devices	int	Number of GPUs		1

To balance the sampling time across GPUs, we shuffle the generation tasks. To avoid this behavior and sample in increasing order of length, add the flag --sequential_order.

To reproduce our unconditional generation, run the command below.

python genie/sample_unconditional.py --name base --epoch 40 --scale 0.6 --outdir results/base/outputs

Motif Scaffolding

We provide the benchmark datasets for single-motif and multi-motif scaffolding under data/design25 and data/multimotifs, respectively. The format of a motif scaffolding problem definition file is described here.

Perform motif scaffolding by running

python genie/sample_scaffold.py --name [NAME] --epoch [EPOCH] --scale [SCALE] --outdir [OUTDIR]

The list of parameters are summarized in the following table.

Parameter	Type	Description	Required	Default
name	string	Model name	✔️
epoch	int	Model epoch	✔️
scale	float	Sampling noise scale (between 0 and 1, inclusive)	✔️
outdir	string	Output directory	✔️
rootdir	string	Model root directory		results
num_samples	int	Number of samples per length		100
batch_size	int	Batch size		4
motif_name	string	Name of motif-scaffolding problem (Automatically evaluate all problems in datadir if unspecified)
datadir	string	Directory of motif-scaffolding problems		data/design25
num_devices	int	Number of GPUs		1

To reproduce our single-motif scaffolding, run

python genie/sample_scaffold.py --name base --epoch 30 --scale 0.4 --outdir results/base/design25 --num_samples 1000

To reproduce our multi-motif scaffolding, run

python genie/sample_scaffold.py --name base --epoch 30 --scale 0.4 --outdir results/base/multimotifs --datadir data/multimotifs --num_samples 1000

Format of Motif Scaffolding Problem Definition File

A motif scaffolding problem definition file consists of two parts: the first part defines motif scaffolding specifications, including arrangements of motif segments, minimum/maximum length of each scaffold segment and minimum/maximum sequence length; the second part defines motif structure in a PDB format. The table below details the format for defining motif scaffolding specifications.

Specification	Column	Data	Justification	Data Type
Motif segment	1-16	"REMARK 999 INPUT"		string
	19	Chain index of motif segment in the PDB file		string
	20-23	Starting residue index of motif segment in the PDB file	right	int
	24-27	Ending residue index of motif segment in the PDB file	right	int
	29	Motif group that the segment belongs to (default to 'A' if unspecified)		string
Scaffold segment	1-16	"REMARK 999 INPUT"		string
	20-23	Minimum length of scaffold segment	right	int
	24-27	Maximum length of scaffold segment	right	int
Minimum sequence length	1-31	"REMARK 999 MINIMUM TOTAL LENGTH"		string
	38-40	Minimum sequence length	left	int
Maximum sequence length	1-31	"REMARK 999 MAXIMUM TOTAL LENGTH"		string
	38-40	Maximum sequence length	left	int

One example of motif scaffolding specifications is provided below.

REMARK 999 NAME   1PRW_two
REMARK 999 PDB    1PRW
REMARK 999 INPUT      5  20
REMARK 999 INPUT  A  16  35 A
REMARK 999 INPUT     10  25
REMARK 999 INPUT  A  52  71 A
REMARK 999 INPUT     10  30
REMARK 999 INPUT  A  89 108 B
REMARK 999 INPUT     10  25
REMARK 999 INPUT  A 125 144 B
REMARK 999 INPUT      5  20
REMARK 999 MINIMUM TOTAL LENGTH      120
REMARK 999 MAXIMUM TOTAL LENGTH      200

There are two limitations with our current format:

• When defining motif structures in a PDB format, the residue order should match with the order in motif scaffolding specification.
• PDB name is not specified when defining motif segment in motif scaffolding specification.

We intend to optimize our code to address these two limitations in the future.