Protein Workshop

<a href="https://github.com/psf/black"><img alt="Code style: black" src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>

Documentation

This repository provides the code for the protein structure representation learning benchmark detailed in the paper Evaluating Representation Learning on the Protein Structure Universe (ICLR 2024).

In the benchmark, we implement numerous featurisation schemes, datasets for self-supervised pre-training and downstream evaluation, pre-training tasks, and auxiliary tasks.

The benchmark can be used as a working template for a protein representation learning research project, a library of drop-in components for use in your projects, or as a CLI tool for quickly running protein representation learning evaluation and pre-training configurations.

Processed datasets and pre-trained weights are made available. Downloading datasets is not required; upon first run all datasets will be downloaded and processed from their respective source.

Configuration files to run the experiments described in the manuscript are provided in the proteinworkshop/config/sweeps/ directory.

Contents

• Protein Workshop
  • Contents
  • Installation
    • From PyPI
    • Building from source
  • Tutorials
  • Quickstart
    • Downloading datasets
    • Training a model
    • Finetuning a model
    • Running a sweep/experiment
    • Embedding a dataset
    • Visualising a dataset's embeddings
    • Performing attribution of a pre-trained model
    • Verifying a config
    • Using proteinworkshop modules functionally
  • Models
    • Invariant Graph Encoders
    • Equivariant Graph Encoders
      • (Vector-type)
      • (Tensor-type)
  • Datasets
    • Structure-based Pre-training Corpuses
    • Supervised Datasets
  • Tasks
    • Self-Supervised Tasks
    • Generic Supervised Tasks
  • Featurisation Schemes
    • Invariant Node Features
    • Equivariant Node Features
    • Edge Construction
    • Invariant Edge Features
    • Equivariant Edge Features
  • For Developers
    • Dependency Management
    • Code Formatting
    • Documentation

Installation

Below, we outline how one may set up a virtual environment for proteinworkshop. Note that these installation instructions currently target Linux-like systems with NVIDIA CUDA support. Note that Windows and macOS are currently not officially supported.

From PyPI

proteinworkshop is available for install from PyPI. This enables training of specific configurations via the CLI or using individual components from the benchmark, such as datasets, featurisers, or transforms, as drop-ins to other projects. Make sure to install PyTorch (specifically version 2.1.2 or newer) using its official pip installation instructions, with CUDA support as desired.

# install `proteinworkshop` from PyPI
pip install proteinworkshop

# install PyTorch Geometric using the (now-installed) CLI
workshop install pyg

# set a custom data directory for file downloads; otherwise, all data will be downloaded to `site-packages`
export DATA_PATH="where/you/want/data/" # e.g., `export DATA_PATH="proteinworkshop/data"`

However, for full exploration we recommend cloning the repository and building from source.

Building from source

With a local virtual environment activated (e.g., one created with conda create -n proteinworkshop python=3.10):

1. Clone and install the project

   git clone https://github.com/a-r-j/ProteinWorkshop
   cd ProteinWorkshop
   pip install -e .
   

2. Install PyTorch (specifically version 2.1.2 or newer) using its official pip installation instructions, with CUDA support as desired

   # e.g., to install PyTorch with CUDA 11.8 support on Linux:
   pip install torch==2.1.2+cu118 torchvision==0.16.2+cu118 torchaudio==2.1.2+cu118 --index-url https://download.pytorch.org/whl/cu118
   

3. Then use the newly-installed proteinworkshop CLI to install PyTorch Geometric

   workshop install pyg
   

4. Configure paths in .env (optional, will override default paths if set). See .env.example for an example.

5. Download PDB data:

   python proteinworkshop/scripts/download_pdb_mmtf.py
   

Tutorials

We provide a five-part tutorial series of Jupyter notebooks to provide users with examples of how to use and extend proteinworkshop, as outlined below.

1. Training a new model
2. Customizing an existing dataset
3. Adding a new dataset
4. Adding a new model
5. Adding a new task

Quickstart

Downloading datasets

Datasets can either be built from the source structures or downloaded from Zenodo. Datasets will be built from source the first time a dataset is used in a run (or by calling the appropriate setup() method in the corresponding datamodule). We provide a CLI tool for downloading datasets:

workshop download <DATASET_NAME>
workshop download pdb
workshop download cath
workshop download afdb_rep_v4
# etc..

If you wish to build datasets from source, we recommend first downloading the entire PDB first (in MMTF format, c. 24 Gb) to reuse shared PDB data as much as possible:

workshop download pdb
# or
python proteinworkshop/scripts/download_pdb_mmtf.py

Training a model

Launching an experiment minimally requires specification of a dataset, structural encoder, and task (devices can be specified with trainer=cpu/gpu):

workshop train dataset=cath encoder=egnn task=inverse_folding trainer=cpu env.paths.data=where/you/want/data/
# or
python proteinworkshop/train.py dataset=cath encoder=egnn task=inverse_folding trainer=cpu # or trainer=gpu

This command uses the default configurations in configs/train.yaml, which can be overwritten by equivalently named options. For instance, you can use a different input featurisation using the features option, or set the display name of your experiment on wandb using the name option:

workshop train dataset=cath encoder=egnn task=inverse_folding features=ca_bb name=MY-EXPT-NAME trainer=cpu env.paths.data=where/you/want/data/
# or
python proteinworkshop/train.py dataset=cath encoder=egnn task=inverse_folding features=ca_bb name=MY-EXPT-NAME trainer=cpu # or trainer=gpu

Finetuning a model

Finetuning a model additionally requires specification of a checkpoint.

workshop finetune dataset=cath encoder=egnn task=inverse_folding ckpt_path=PATH/TO/CHECKPOINT trainer=cpu env.paths.data=where/you/want/data/
# or
python proteinworkshop/finetune.py dataset=cath encoder=egnn task=inverse_folding ckpt_path=PATH/TO/CHECKPOINT trainer=cpu # or trainer=gpu

Running a sweep/experiment

We can make use of the hydra wandb sweeper plugin to configure experiments as sweeps, allowing searches over hyperparameters, arc