CLIPZyme

Implementation of the paper CLIPZyme: Reaction-Conditioned Virtual Screening of Enzymes

Table of contents

• CLIPZyme
• Table of contents
• Installation:
• Checkpoints and Data Files:
• Screening with CLIPZyme
  • Using CLIPZyme's screening set
  • Using your own screening set
    • Interactive (slow)
    • Batched (fast)
• Reproducing published results
  • Data processing
  • Training and evaluation
  • Downloading Batched AlphaFold Database Structures
• CLIPZyme+
• Citations

Installation:

1. Clone the repository:

git clone https://github.com/pgmikhael/CLIPZyme.git

2. Install the dependencies:

cd clipzyme
conda env create -f environment.yml
conda activate clipzyme
python -m pip install clipzyme

3. Download ESM-2 checkpoint esm2_t33_650M_UR50D. The esm_dir argument should point to this directory. The following command will download the checkpoint directly:

wget https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt

Checkpoints and Data Files:

The model checkpoint and data are available on Zenodo here:

• clipzyme_data.zip:

  • The following commands will download the checkpoint directly:

  wget https://zenodo.org/records/11187747/files/clipzyme_data.zip
  unzip clipzyme_data.zip -d files
  

  • Note that the data files should be extracted into the files/ directory.
    • enzymemap.json: contains the EnzymeMap dataset.
    • cached_enzymemap.p: contains the processed EnzymeMap dataset.
    • clipzyme_screening_set.p: contains the screening set as dict of UniProt IDs and precomputed protein embeddings.
    • uniprot2sequence.p: contains the mapping form sequence ID to amino acids.

• clipzyme_model.zip:

  • The following command will download the checkpoint directly:

  wget https://zenodo.org/records/11187747/files/clipzyme_model.zip
  unzip clipzyme_model.zip -d files
  

  • clipzyme_model.ckpt: the trained model checkpoint.

Screening with CLIPZyme

Using CLIPZyme's screening set

First, download the screening set and extract the files into files/.

import pickle
from clipzyme import CLIPZyme

## Load the screening set
##-----------------------
screenset = pickle.load(open("files/clipzyme_screening_set.p", 'rb'))
screen_hiddens = screenset["hiddens"] # hidden representations (261907, 1280)
screen_unis = screenset["uniprots"] # uniprot ids (261907,)

## Load the model and obtain the hidden representations of a reaction
##-------------------------------------------------------------------
model = CLIPZyme(checkpoint_path="files/clipzyme_model.ckpt")
reaction = "[CH3:1][N+:2]([CH3:3])([CH3:4])[CH2:5][CH:6]=[O:7].[O:9]=[O:10].[OH2:8]>>[CH3:1][N+:2]([CH3:3])([CH3:4])[CH2:5][C:6](=[O:7])[OH:8].[OH:9][OH:10]"
reaction_embedding = model.extract_reaction_features(reaction=reaction) # (1,1280)

enzyme_scores = screen_hiddens @ reaction_embedding.T # (261907, 1)

Using your own screening set

Prepare your data as a CSV in the following format, and save it as files/new_data.csv. For the cases where we wish only to obtain the hidden representations of the sequences, the reaction column can be left empty (and vice versa).

reaction	sequence	protein_id	cif
[CH3:1]N+:2([CH3:4])[CH2:5][CH:6]=[O:7].[O:9]=[O:10].[OH2:8]>>[CH3:1]N+:2([CH3:4])[CH2:5]C:6[OH:8].[OH:9][OH:10]	MGLSDGEWQLVLNVWGKVEAD IPGHGQEVLIRLFKGHPETLE KFDKFKHLKSEDEMKASEDLK KHGATVLTALGGILKKKGHHE AELKPLAQSHATKHKIPIKYL EFISEAIIHVLHSRHPGDFGA DAQGAMNKALELFRKDIAAKY KELGYQG	P69905	1a0s.cif

Interactive (slow)

from torch.utils.data import DataLoader
from clipzyme import CLIPZyme
from clipzyme import ReactionDataset
from clipzyme.utils.loading import ignore_None_collate

## Create reaction dataset
#-------------------------
reaction_dataset = DataLoader(
  ReactionDataset(
    dataset_file_path = "files/new_data.csv",
    esm_dir = "/path/to/esm2_dir",
    protein_cache_dir = "/path/to/protein_cache", # optional, where to cache processed protein graphs
  ),
  batch_size=1,
  collate_fn=ignore_None_collate,
)


## Load the model
#----------------
model = CLIPZyme(checkpoint_path="files/clipzyme_model.ckpt")
model = model.eval() # optional 

## For reaction-enzyme pair
#--------------------------
for batch in reaction_dataset:
  output = model(batch) 
  enzyme_scores = output.scores
  protein_hiddens = output.protein_hiddens
  reaction_hiddens = output.reaction_hiddens

## For sequences only
#--------------------
for batch in reaction_dataset:
  protein_hiddens = model.extract_protein_features(batch) 
  
## For reactions only
#--------------------
for batch in reaction_dataset:
  reaction_hiddens = model.extract_reaction_features(batch)

Batched (fast)

1. Update the screening config configs/screening.json with the path to your data and indicate what you want to save and where:

{
  "dataset_file_path": ["files/new_data.csv"],
  "inference_dir": ["/where/to/save/embeddings_and_scores"],
  "save_hiddens": [true], # whether to save the hidden representations
  "save_predictions": [true], # whether to save the reaction-enzyme pair scores
  "use_as_protein_encoder": [true], # whether to use the model as a protein encoder only
  "use_as_reaction_encoder": [true], # whether to use the model as a reaction encoder only
  "esm_dir": ["/data/esm/checkpoints"], path to ESM-2 checkpoints
  "gpus": [8], # number of gpus to use,
  "protein_cache_dir": ["/path/to/protein_cache"], # where to save the protein cache [optional]
  ...
}

If you want to use specific GPUs, you can specify them in the available_gpus field. For example, to use GPUs 0, 1, and 2, set available_gpus to ["0,1,2"].

2. Run the dispatcher with the screening config:

python scripts/dispatcher.py -c configs/screening.json -l ./logs/

3. Load the saved embeddings and scores:

from clipzyme import collect_screening_results

screen_hiddens, screen_unis, enzyme_scores = collect_screening_results("configs/screening.json")

---

Reproducing published results

Data processing

We obtain the data from the following sources:

• EnzymeMap: Heid et al. Enzymemap: Curation, validation and data-driven prediction of enzymatic reactions. 2023.
• Terpene Synthases: Samusevich et al. Discovery and characterization of terpene synthases powered by machine learning. 2024.

Our processed data is can be downloaded from here.

Training and evaluation

1. To train the models presented in the tables below, run the following command:

   python scripts/dispatcher.py -c {config_path} -l {log_path}
   

   • {config_path} is the path to the config file in the table below
   • {log_path} is the path in which to save the log file.

   For example, to run the first row in Table 1, run:

   python scripts/dispatcher.py -c configs/train/clip_egnn.json -l ./logs/
   

2. Once you've trained the model, run the eval config to evaluate the model on the test set. For example, to evaluate the first row in Table 1, run:

   python scripts/dispatcher.py -c configs/eval/clip_egnn.json -l ./logs/
   

3. We perform all analysis in the jupyter notebook included Results.ipynb. We first calculate the hidden representations of the screening using the eval configs above and collect them into one matrix (saved as a pickle file). These are loaded into the jupyter notebook as well as the test set. All tables are then generated in the notebook.

Downloading Batched AlphaFold Database Structures

Assuming you have a list of uniprot IDs (called uniprot_ids) you can run the following to create a .txt file with the Google Storage urls for the AF2 structures:

file_paths = [f"gs://public-datasets-deepmind-alphafold-v4/AF-{u}-F1-model_v4.cif" for u in uniprot_ids]
output_file = 'uniprot_cif_paths.txt' 
open(output_file, 'w') as file:
    file.write('\n'.join(file_paths))

Then install gsutils (https://cloud.google.com/storage/docs/gsutil_install) and run the following command:

cat uniprot_cif_paths.txt | gsutil -m cp -I /path/to/output/dir/

---

CLIPZyme+

CLIPZyme+ predicts an enzyme's cofactor from its structure. It reuses the CLIPZyme protein encoder to produce a 1280-dim embedding per enzyme, then runs an ensemble of MLP classifiers over that embedding. CLIPZymePlus is inference-only and is used just like CLIPZyme with a ReactionDataset.

The cofactor ensemble ships as a single checkpoint file. Download it from Zenodo into files/ and point cofactor_checkpoint_path at it:

wget https://zenodo.org/records/20673359/files/clipzyme_plus_cofactor_ensemble.pt -P files/

from torch.utils.data import DataLoader
from clipzyme import CLIPZymePlus, ReactionDataset
from clipzyme.utils.loading import ignore_None_collate

## Create reaction dataset (only structures are needed for cofactor prediction)
#-----------------------------------------------------------------------------
loader = DataLoader(
    ReactionDataset(
        dataset_file_path="files/new_data.csv",
        esm_dir="/path/to/esm2_dir",
        use_as_protein_encoder=True,
    ),
    batch_size=1,
    collate_fn=ignore_None_collate,
)

## Load the model and predict cofactors
#--------------------------------------
model = CLIPZymePlus(
    checkpoint_path="files/clipzyme_model.ckpt",
    cofactor_checkpoint_path="files/clipzyme_plus_cofactor_ensemble.pt",
)
model = model.eval()

for batch in loader:
    output = model(batch)
    print(output.sample_ids)            # protein ids
    print(output.predicted_cofactor)    # top-1 cofactor class per sample
    print(output.predicted_probability) # ensemble-mean probability
    print(model.top_k(output, k=5))     # ranked predictions per sample

You can also predict directly from structure files or precomputed CLIPZyme embeddings (e.g. the screening set hiddens), without a ReactionDataset:

# From structure files
output = model.predict_from_structures(["1a0s.cif"], esm_dir="/path/to/esm2_dir")

# From precomputed (N, 1280) protein embeddings
output = model.predict_from_embeddings(screen_hiddens)

The cofactor ensemble checkpoint is a single self-contained .pt file holding each ensemble member's weights, architecture, and class names — no other files from the cofactor_prediction training code are needed at inference time.

Code to reproduce the results of CLIPZyme+ (training the cofactor ensemble) is available in the cofactor_prediction branch. We also make a notebook available for inference, available here.

Citations

For the CLIPZyme model and the results in the main paper, please cite:

@article{mikhael2024clipzyme,
  title={CLIPZyme: Reaction-Conditioned Virtual Screening of Enzymes},
  author={Mikhael, Peter G and Chinn, Itamar and Barzilay, Regina},
  journal={arXiv preprint arXiv:2402.06748},
  year={2024}
}