Our paper "AbEgDiffuser: Antibody Sequence-Structure Codesign with Equivariant Graph Neural Networks and Diffusion Models". published in J Chem Theory Comput: https://pubs.acs.org/doi/10.1021/acs.jctc.5c00990

AbEgDiffuser

AbEgDiffuser, a deep generative framework that enables the codesign of antibody sequences and structures conditioned on target antigens. It integrates diffusion models with equivariant graph neural networks and further incorporates evolutionary constraints. During forward diffusion, amino acid sequences, Cα atom coordinates, and residue orientations are progressively corrupted toward a prior distribution. In reverse, a bilevel equivariant graph neural network captures both residue- and atom-level interactions to reconstruct functional antibodies. Evolutionary plausibility is enforced by encoding noisy sequences with the pretrained protein language model ESM-2. Extensive experiments on de novo antibody design and optimization tasks demonstrate that AbEgDiffuser generates antibodies with accurate sequences and structures and high binding affinity, outperforming existing design methods.

This repository contains the data and python script in support of the manuscript: AbEgDiffuser: Antibody sequence-structure co-design with equivariant graph neural networks and diffusion models.

---

1. Requirement

1) System requirements

This tool is supported for Linux.

2) Hardware requirements

We ran the demo using the following specs:

• NVIDIA A800
• CUDA 11.8

3) Environment

We have provided the environment.yml, it is recommended to create a Conda environment and install the necessary dependencies by following these steps:

conda env create -f environment.yaml -n Ab
conda activate Ab

4) [Optional] PyRosetta

PyRosetta is required to relax the generated structures and compute binding energy. Please follow the instruction here to install.

2. How to train and use AbEgDiffuser

1) Dataset preparation and pre-trained models

1. We use SAbDab for training and RAbD for testing.

• Protein structures in the SAbDab dataset can be downloaded here. Extract all_structures.zip into the data folder.
• The data folder contains a snapshot of the dataset index (sabdab_summary_all.tsv). You may replace the index with the latest version here.
• The file data/RAbD_test.idx contains the RAbD test data information.

2. Download the ESM2 model weights from here and the contact regressor weights from here, and save these files in the ./trained_models/esm directory.

2) Training the AbEgDiffuser model

If you use 4 GPUs for training, run the following command:

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 train.py ./configs/train/codesign_single.yml

If only one GPU is used for training, run the following command:

CUDA_VISIBLE_DEVICES=0 torchrun --nproc_per_node=1 train.py ./configs/train/codesign_single.yml

3) Co-design of CDRs

To perform co-design of CDRs using the RAbD test dataset, run the following command:

python design_all_test.py 60 --config ./configs/test/codesign_single.yml --seed 2025 --batch_size 16

4) Optimization

To optimize CDRs in the RAbD test dataset, run the following command:

python design_all_test.py 60 --config ./configs/test/abopt_singlecdr.yml --seed 2025 --batch_size 16

5) Design CDRs given antibody-antigen complex

Taking 7DK2_AB_C antibody-antigen complex as an example, to generate CDRs of given antibdody-antigen complexes, run the following command:

python design_pdb.py ./data/examples/7DK2_AB_C.pdb --config ./configs/test/codesign_single.yml

6) Relaxing the designed proteins

To relax the designed proteins using PyRosetta, run the following command:

python -m AbEgDiffuser.tools.relax.run --root './results/codesign_single/'

7) Evaluation

To compute the evaluation metrics: AAR, RMSD, and IMP for all samples, run the following command:

python -m AbEgDiffuser.tools.eval.run --root './results/codesign_single/'

3. Model availability

The trained AbEgDiffuser model is available here (zenodo). To use it, please download the model checkpoint and save it into the ./trained_models folder.

4. Citation

Please cite our paper if you use the code.

5. Contact

If you have any questions, please contact us via email:

• Xiumin Shi