This repository contains the code accompanying our paper in XXX. We propose a method that integrates pre-trained discrete diffusion models (e.g., EvoDiff) for protein sequences with reward models (i.e., seq → target property) at test time for computational protein design. Our algorithm effectively optimizes the reward function while retaining sequence naturalness characterzied by pre-trained diffusion models. Unlike existing single-shot guided approaches in diffusion models, our method uses an iterative refinement approach inspired by evolutionary algorithms, alternating between (derivative-free) reward-guided denoising and noising.

Below are examples of trajectories obtained when optimizing structural properties as rewards

We present results on optimizing several fundamental structural rewards, including symmetry, globularity, match_ss, crmsd. We can further optimize ptm,plddt,tm,lddt,hydrophobic,surface_expose. All rewards are defined based on the outputs of a sequence-to-structure model using ESMFold. Below, we visualize examples of the generated sequences, where ESMFold is used to predict their structures.

CUDA_VISIBLE_DEVICES=0 python refinement.py

Run the above command with additional options. Below is an explanation of the available options.

• --decoding: decoding method e.g., SVDD_edit (our proposal in the paper), SVDD (single-shot generation)
• --repeatnum: batch size
• --duplicate: important hyperparameters in decoding methods for each cycle. It reflects how many states we replicate (e.g., 20). This is the width of the tree (Refer to SVDD).
• --metrics_name: reward function. We have tm, globularity, plddt, ptm, hydrophobic, symmetry, globularity, match_ss, crmsd, ptm,plddt,tm,lddt,hydrophobic,surface_expose.
• --metrics_list: how to set weights for the above rewards. For example, --metric_name match_ss,crmsd,plddt --metrics_list 1,1,1 means we optimize match_ss + crmsd + plddt.
• --proteinname: In tasks such as match_ss, crmsd, tm, we can set some target structure in a pdb format in the folder ./datasets/AlphaFold_model_PDBs (e.g., 5KPH, XX:run1_0367_0004, etc.)
• --iteraiton: number of iterations in our proposal (e.g., 20)
• --seq_length: length of proteins we want to design
• --edit_seqlength: How much portion of the sequence we edit (e.g., 0.5$)

Design symmetric proteins with sevenfold symmetry.

CUDA_VISIBLE_DEVICES=3 python refinement.py --decoding SVDD_edit  --repeatnum 10 --duplicate 20 --seq_length 30 --metrics_name symmetry,hydrophobic,plddt --metrics_list 1,1,1 --iteration 20 --num_symmetry 7

Design globular proteins.

CUDA_VISIBLE_DEVICES=2 python refinement.py --decoding SVDD_edit  --repeatnum 10 --duplicate 20  --metrics_name globularity,plddt  --metrics_list 1,1 --iteration 20 --seq_length 150

Design a sequence that folds into a target structure based on cRMSD.

CUDA_VISIBLE_DEVICES=3 python refinement.py --decoding SVDD_edit  --repeatnum 20 --duplicate 20  --metrics_name crmsd  --metrics_list 1 --proteinname 5KPH --iteration 40

Design a sequence that folds into a target secondary structure.

CUDA_VISIBLE_DEVICES=4 python refinement.py --decoding SVDD_edit  --repeatnum 10 --duplicate 20  --metrics_name match_ss  --metrics_list 1 --proteinname r15_96_TrROS_Hall --iteration 30

Design a sequence that folds into a target structure based on the TM-score.

CUDA_VISIBLE_DEVICES=5 python refinement.py --decoding SVDD_edit  --repeatnum 10 --duplicate 20  --metrics_name tm  --metrics_list 1 --proteinname 5KPH --iteration 30

When we want to edit existing sequences, set --seed_design as True, and --initial_seq as a protein sequence. The following is an example when optimizing symmetry.

CUDA_VISIBLE_DEVICES=6 python refinement.py --decoding SVDD_edit  --repeatnum 10 --duplicate 20 --metrics_name crmsd,plddt,hydrophobic --proteinname r15_96_TrROS_Hall --metrics_list 1,1,1 --iteration 20 --seed_design True --initial_seq MMELEIEIKVEGMTEEELRELAERLAAELTPEGWKVVAVRVERVDEEEGVVRVTVVVEPV  

Refer to the notebook medias/evaluate.ipynb. PDB files of batches and important metrics are saved at each iteration.

Install pytroch, pyrosseta. Then, run the following

conda create -n RERD python=3.9 
conda activate RERD
pip3 install torch torchvision torchaudio
pip3 install -r requirements.txt

Also, to optimize match_ss, crmsd, go to the ./datasets folder and download examples of proteins as follows.

python download_model_data.py

This code puts several pdb files into ./datasets/AlphaFoldPDB/. But, technicallly, you can put any pdb files.

Our codebase is heavily based on evodiff, openfold, ESMfold.