Package: proteinworkshop
Protein WorkshopDocumentationThisrepositoryprovides the code for the protein structure representation learning benchmark detailed in the paperEvaluating Representation Learning on the Protein Structure Universe.In the benchmark, we implement numerousfeaturisationschemes,datasetsforself-supervised pre-traininganddownstream evaluation,pre-trainingtasks, andauxiliary tasks.The benchmark can be used as a working template for a protein representation learning research project,a libraryof drop-in components for use in your projects, or as a CLI tool for quickly runningprotein representation learning evaluationandpre-trainingconfigurations.Processed datasetsandpre-trained weightsare 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 theproteinworkshop/config/sweeps/directory.ContentsProtein WorkshopContentsInstallationFrom PyPIBuilding from sourceTutorialsQuickstartDownloading datasetsTraining a modelFinetuning a modelRunning a sweep/experimentEmbedding a datasetVisualising a dataset's embeddingsPerforming attribution of a pre-trained modelVerifying a configUsingproteinworkshopmodules functionallyModelsInvariant Graph EncodersEquivariant Graph Encoders(Vector-type)(Tensor-type)DatasetsStructure-based Pre-training CorpusesSupervised DatasetsTasksSelf-Supervised TasksGeneric Supervised TasksFeaturisation SchemesInvariant Node FeaturesEquivariant Node FeaturesEdge ConstructionInvariant Edge FeaturesEquivariant Edge FeaturesFor DevelopersDependency ManagementCode FormattingDocumentationInstallationBelow, we outline how one may set up a virtual environment forproteinworkshop. 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 PyPIproteinworkshopis available for installfrom PyPI. This enables training of specific configurations via the CLIorusing individual components from the benchmark, such as datasets, featurisers, or transforms, as drop-ins to other projects. Make sure to installPyTorch(specifically version2.1.2or newer) using its officialpipinstallation instructions, with CUDA support as desired.# install `proteinworkshop` from PyPIpipinstallproteinworkshop# install PyTorch Geometric using the (now-installed) CLIworkshopinstallpyg# set a custom data directory for file downloads; otherwise, all data will be downloaded to `site-packages`exportDATA_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 sourceWith a local virtual environment activated (e.g., one created withconda create -n proteinworkshop python=3.10):Clone and install the projectgitclonehttps://github.com/a-r-j/ProteinWorkshopcdProteinWorkshop
pipinstall-e.InstallPyTorch(specifically version2.1.2or newer) using its officialpipinstallation instructions, with CUDA support as desired# e.g., to install PyTorch with CUDA 11.8 support on Linux:pipinstalltorch==2.1.2+cu118torchvision==0.16.2+cu118torchaudio==2.1.2+cu118--index-urlhttps://download.pytorch.org/whl/cu118Then use the newly-installedproteinworkshopCLI to installPyTorch GeometricworkshopinstallpygConfigure paths in.env(optional, will override default paths if set). See.env.examplefor an example.Download PDB data:pythonproteinworkshop/scripts/download_pdb_mmtf.pyTutorialsWe provide a five-part tutorial series of Jupyter notebooks to provide users with examples
of how to use and extendproteinworkshop, as outlined below.Training a new modelCustomizing an existing datasetAdding a new datasetAdding a new modelAdding a new taskQuickstartDownloading datasetsDatasets can either be built from the source structures or downloaded fromZenodo. Datasets will be built from source the first time a dataset is used in a run (or by calling the appropriatesetup()method in the corresponding datamodule). We provide a CLI tool for downloading datasets:workshopdownload<DATASET_NAME>
workshopdownloadpdb
workshopdownloadcath
workshopdownloadafdb_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:workshopdownloadpdb# orpythonproteinworkshop/scripts/download_pdb_mmtf.pyTraining a modelLaunching an experiment minimally requires specification of a dataset, structural encoder, and task (devices can be specified withtrainer=cpu/gpu):workshoptraindataset=cathencoder=egnntask=inverse_foldingtrainer=cpuenv.paths.data=where/you/want/data/# orpythonproteinworkshop/train.pydataset=cathencoder=egnntask=inverse_foldingtrainer=cpu# or trainer=gpuThis command uses the default configurations inconfigs/train.yaml, which can be overwritten by equivalently named options. For instance, you can use a different input featurisation using thefeaturesoption, or set the display name of your experiment on wandb using thenameoption:workshoptraindataset=cathencoder=egnntask=inverse_foldingfeatures=ca_bbname=MY-EXPT-NAMEtrainer=cpuenv.paths.data=where/you/want/data/# orpythonproteinworkshop/train.pydataset=cathencoder=egnntask=inverse_foldingfeatures=ca_bbname=MY-EXPT-NAMEtrainer=cpu# or trainer=gpuFinetuning a modelFinetuning a model additionally requires specification of a checkpoint.workshopfinetunedataset=cathencoder=egnntask=inverse_foldingckpt_path=PATH/TO/CHECKPOINTtrainer=cpuenv.paths.data=where/you/want/data/# orpythonproteinworkshop/finetune.pydataset=cathencoder=egnntask=inverse_foldingckpt_path=PATH/TO/CHECKPOINTtrainer=cpu# or trainer=gpuRunning a sweep/experimentWe can make use of the hydra wandb sweeper plugin to configure experiments as sweeps, allowing searches over hyperparameters, architectures, pre-training/auxiliary tasks and datasets.Seeproteinworkshop/config/sweeps/for examples.Create the sweep with weights and biaseswandbsweepproteinworkshop/config/sweeps/my_new_sweep_config.yamlLaunch job workersWith wandb:wandbagentmywandbgroup/proteinworkshop/2wwtt7oy--count8Or an example SLURM submission script:#!/bin/bash#SBATCH --nodes 1#SBATCH --ntasks-per-node=1#SBATCH --gres=gpu:1#SBATCH --array=0-32source~/.bashrcsource$(condainfo--base)/envs/proteinworkshop/bin/activate

wandbagentmywandbgroup/proteinworkshop/2wwtt7oy--count1Reproduce the sweeps performed in the manuscript:# reproduce the baseline tasks sweep (i.e., those performed without pre-training each model)wandbsweepproteinworkshop/config/sweeps/baseline_fold.yaml
wandbagentmywandbgroup/proteinworkshop/2awtt7oy--count8wandbsweepproteinworkshop/config/sweeps/baseline_ppi.yaml
wandbagentmywandbgroup/proteinworkshop/2bwtt7oy--count8wandbsweepproteinworkshop/config/sweeps/baseline_inverse_folding.yaml
wandbagentmywandbgroup/proteinworkshop/2cwtt7oy--count8# reproduce the model pre-training sweepwandbsweepproteinworkshop/config/sweeps/pre_train.yaml
wandbagentmywandbgroup/proteinworkshop/2dwtt7oy--count8# reproduce the pre-trained tasks sweep (i.e., those performed after pre-training each model)wandbsweepproteinworkshop/config/sweeps/pt_fold.yaml
wandbagentmywandbgroup/proteinworkshop/2ewtt7oy--count8wandbsweepproteinworkshop/config/sweeps/pt_ppi.yaml
wandbagentmywandbgroup/proteinworkshop/2fwtt7oy--count8wandbsweepproteinworkshop/config/sweeps/pt_inverse_folding.yaml
wandbagentmywandbgroup/proteinworkshop/2gwtt7oy--count8Embedding a datasetWe provide a utility inproteinworkshop/embed.pyfor embedding a dataset using a pre-trained model.
To run it:pythonproteinworkshop/embed.pyckpt_path=PATH/TO/CHECKPOINTcollection_name=COLLECTION_NAMESee theembedsection ofproteinworkshop/config/embed.yamlfor additional parameters.Visualising pre-trained model embeddings for a given datasetWe provide a utility inproteinworkshop/visualise.pyfor visualising the UMAP embeddings of a pre-trained model for a given dataset.
To run it:pythonproteinworkshop/visualise.pyckpt_path=PATH/TO/CHECKPOINTplot_filepath=VISUALISATION/FILEPATH.pngSee thevisualisesection ofproteinworkshop/config/visualise.yamlfor additional parameters.Performing attribution of a pre-trained modelWe provide a utility inproteinworkshop/explain.pyfor performing attribution of a pre-trained model using integrated gradients.This will write PDB files for all the structures in a dataset for a supervised task with residue-level attributions in theb_factorcolumn. To visualise the attributions, we recommend using theProtein Viewer VSCode extensionand changing the 3D representation to colour byUncertainty/Disorder.To run the attribution:pythonproteinworkshop/explain.pyckpt_path=PATH/TO/CHECKPOINToutput_dir=ATTRIBUTION/DIRECTORYSee theexplainsection ofproteinworkshop/config/explain.yamlfor additional parameters.Verifying a configpythonproteinworkshop/validate_config.pydataset=cathfeatures=full_atomtask=inverse_foldingUsingproteinworkshopmodules functionallyOne may use the modules (e.g., datasets, models, featurisers, and utilities) ofproteinworkshopfunctionally by importing them directly. When installing this package using PyPi, this makes building
on top of the assets ofproteinworkshopstraightforward and convenient.For example, to use any datamodule available inproteinworkshop:fromproteinworkshop.datasets.cathimportCATHDataModuledatamodule=CATHDataModule(path="data/cath/",pdb_dir="data/pdb/",format="mmtf",batch_size=32)datamodule.download()train_dl=datamodule.train_dataloader()To use any model or featuriser available inproteinworkshop:fromproteinworkshop.models.graph_encoders.dimenetppimportDimeNetPPModelfromproteinworkshop.features.factoryimportProteinFeaturiserfromproteinworkshop.datasets.utilsimportcreate_example_batchmodel=DimeNetPPModel(hidden_channels=64,num_layers=3)ca_featuriser=ProteinFeaturiser(representation="CA",scalar_node_features=["amino_acid_one_hot"],vector_node_features=[],edge_types=["knn_16"],scalar_edge_features=["edge_distance"],vector_edge_features=[],)example_batch=create_example_batch()batch=ca_featuriser(example_batch)model_outputs=model(example_batch)Readthe docsfor a full list of modules available inproteinworkshop.ModelsInvariant Graph EncodersNameSourceProtein SpecificGearNetZhang et al.✓DimeNet++Gasteiger et al.✗SchNetSchütt et al.✗CDConvFan et al.✓Equivariant Graph Encoders(Vector-type)NameSourceProtein SpecificGCPNetMorehead et al.✓GVP-GNNJing et al.✓EGNNSatorras et al.✗(Tensor-type)NameSourceProtein SpecificTensor Field NetworkCorso et al.✓Multi-ACEBatatia et al.✗Sequence-based EncodersNameSourceProtein SpecificESM2Lin et al.✓DatasetsTo download a (processed) dataset from Zenodo, you can runworkshopdownload<DATASET_NAME>where<DATASET_NAME>is given the first column in the tables below.Otherwise, simply starting a training run will download and process the data from source.Structure-based Pre-training CorpusesPre-training corpuses (with the exception ofpdb,cath, andastral) are provided in FoldComp database format. This format is highly compressed, resulting in very small disk space requirements despite the large size.pdbis provided as a collection ofMMTFfiles, which are significantly smaller in size than conventional.pdbor.ciffile.NameDescriptionSourceSizeDisk SizeLicenseastralSCOPedomain structuresSCOPe/ASTRAL1 - 2.2 GbPublicly availableafdb_rep_v4Representative structures identified from theAlphaFold databasebyFoldSeekstructural clusteringBarrio-Hernandez et al.2.27M Chains9.6 GbGPL-3.0afdb_rep_dark_v4Dark proteome structures identied by structural clustering of theAlphaFold database.Barrio-Hernandez et al.~800k2.2 GbGPL-3.0afdb_swissprot_v4AlphaFold2 predictions forSwissProt/UniProtKBKim et al.542k Chains2.9 GbGPL-3.0afdb_uniprot_v4AlphaFold2 predictions forUniProtKim et al.214M Chains1 TbGPL-3.0/CC-BY 4.0cathCATH4.2 40% split by CATH topologies.Ingraham et al.~18k chains4.3 GbCC-BY 4.0esmatlasESMAtlaspredictions  (full)Kim et al.1 TbGPL-3.0/CC-BY 4.0esmatlas_v2023_02ESMAtlaspredictions (v2023_02 release)Kim et al.137 GbGPL-3.0/CC-BY 4.0highquality_clust30ESMAtlasHigh Quality predictionsKim et al.37M Chains114 GbGPL-3.0/CC-BY 4.0igfold_paired_oasIGFold Predictions forPaired OASRuffolo et al.104,994 paired Ab chainsCC-BY 4.0igfold_jaffeIGFold predictions forJaffe2022dataRuffolo et al.1,340,180 paired Ab chainsCC-BY 4.0pdbExperimental structures deposited in theRCSB Protein Data BankwwPDB consortium~800k Chains23 GbCC0 1.0Additionally, we provide several species-specific compilations (mostly reference species)NameDescriptionSourceSizea_thalianaArabidopsis thaliana(thale cress) proteomeAlphaFold2c_albicansCandida albicans(a fungus) proteomeAlphaFold2c_elegansCaenorhabditis elegans(roundworm) proteomeAlphaFold2d_discoideumDictyostelium discoideum(slime mold) proteomeAlphaFold2d_melanogasterDrosophila melanogaster(fruit fly) proteomeAlphaFold2d_rerioDanio rerio(zebrafish) proteomeAlphaFold2e_coliEscherichia coli(a bacteria) proteomeAlphaFold2g_maxGlycine max(soy bean) proteomeAlphaFold2h_sapiensHomo sapiens(human) proteomeAlphaFold2m_jannaschiiMethanocaldococcus jannaschii(an archaea) proteomeAlphaFold2m_musculusMus musculus(mouse) proteomeAlphaFold2o_sativaOryza sative(rice) proteomeAlphaFold2r_norvegicusRattus norvegicus(brown rat) proteomeAlphaFold2s_cerevisiaeSaccharomyces cerevisiae(brewer's yeast) proteomeAlphaFold2s_pombeSchizosaccharomyces pombe(a fungus) proteomeAlphaFold2z_maysZea mays(corn) proteomeAlphaFold2Supervised DatasetsNameDescriptionSourceLicenseantibody_developabilityAntibody developability predictionChen et al.CC-BY 3.0atom3d_mspMutation stability predictionTownshend et al.MITatom3d_ppiProtein-protein interaction predictionTownshend et al.MITatom3d_psrProtein structure rankingTownshend et al.MITatom3d_resResidue identity predictionTownshend et al.MITccpdb_ligandsLigand binding residue predictionAgrawal et al.Publicly Availableccpdb_metalMetal ion binding residue predictionAgrawal et al.Publicly Availableccpdb_nucleicNucleic acid binding residue predictionAgrawal et al.Publicly Availableccpdb_nucleotidesNucleotide binding residue predictionAgrawal et al.Publicly Availabledeep_sea_proteinsGene Ontology prediction (Biological Process)Sieg et al.Public domaingo-bpGene Ontology prediction (Biological Process)Gligorijevic et alCC-BY 4.0go-ccGene Ontology (Cellular Component)Gligorijevic et alCC-BY 4.0go-mfGene Ontology (Molecular Function)Gligorijevic et alCC-BY 4.0ec-reactionEnzyme Commission (EC) Number PredictionHermosilla et al.MITfold-foldFold prediction, split at the fold levelHou et al.CC-BY 4.0fold-familyFold prediction, split at the family levelHou et al.CC-BY 4.0fold-superfamilyFold prediction, split at the superfamily levelHou et al.CC-BY 4.0masif-siteProtein-protein interaction site predictionGainza et al.Apache 2.0metal_3dZinc Binding Site PredictionDuerr et al.MITptmPost Translational Modification Side PredictionYan et al.CC-BY 4.0TasksSelf-Supervised TasksNameDescriptionSourceinverse_foldingPredict amino acid sequence given structureresidue_predictionMasked residue type predictiondistance_predictionMasked edge distance predictionZhang et al.angle_predictionMasked triplet angle predictionZhang et al.dihedral_angle_predictionMasked quadruplet dihedral predictionZhang et al.multiview_contrastContrastive learning with multiple crops and InfoNCE lossZhang et al.structural_denoisingDenoising of atomic coordinates with SE(3) decodersGeneric Supervised TasksGeneric supervised tasks can be applied broadly across datasets. The labels are directly extracted from the PDB structures.These are likely to be most frequently used with thepdbdataset class which wraps thePDB Dataset curatorfromGraphein.NameDescriptionRequiresbinding_site_predictionPredict ligand binding residuesHETATM ligands (for training)ppi_site_predictionPredict protein binding residuesgraph_yattribute in data objects specifying the desired chain to select interactions for (for training)Featurisation SchemesPart of the goal of theproteinworkshopbenchmark is to investigate the impact of the degree to which increasing granularity of structural detail affects performance. To achieve this, we provide several featurisation schemes for protein structures.Invariant Node FeaturesN.B. All angular features are provided in [sin, cos] transformed form. E.g.: $\textrm{dihedrals} = [sin(\phi), cos(\phi), sin(\psi), cos(\psi), sin(\omega), \cos(\omega)]$, hence their dimensionality will be double the number of angles.NameDescriptionDimensionalityresidue_typeOne-hot encoding of amino acid type21positional_encodingTransformer-like positional encoding of sequence position16alphaVirtual torsion angle defined by four $C_\alpha$ atoms of residues $I_{-1}, I, I_{+1}, I_{+2}$2kappaVirtual bond angle (bend angle) defined by the three $C_\alpha$ atoms of residues $I_{-2}, I, I_{+2}$2dihedralsBackbone dihedral angles $(\phi, \psi, \omega)$6sidechain_torsionsSidechain torsion angles  $(\chi_{1-4})$8Equivariant Node FeaturesNameDescriptionDimensionalityorientationForward and backward node orientation vectors (unit-normalized)2Edge ConstructionWe predominanty support two types of edges: $k$-NN and $\epsilon$ edges.Edge types can be specified as follows:pythonproteinworkshop/train.py...features.edge_types=[knn_16,knn_32,eps_16]Where the suffix afterknnorepsspecifies $k$ (number of neighbours) or $\epsilon$ (distance threshold in angstroms).Invariant Edge FeaturesNameDescriptionDimensionalityedge_distanceEuclidean distance between source and target nodes1node_featuresConcatenated scalar node features of the source and target nodesNumber of scalar node features $\times 2$edge_typeType annotation for each edge1sequence_distanceSequence-based distance between source and target nodes1pos_embStructured Transformer-inspired positional embedding of $i - j$ for source node $i$ and target node $j$16Equivariant Edge FeaturesNameDescriptionDimensionalityedge_vectorsEdge directional vectors (unit-normalized)1For DevelopersDependency ManagementWe usepoetryto manage the project's underlying dependencies and to push updates to the project's PyPI package. To make changes to the project's dependencies, follow the instructions below to (1) installpoetryon your local machine; (2) customize the dependencies; or (3) (de)activate the project's virtual environment usingpoetry:Installpoetryfor platform-agnostic dependency management using itsinstallation instructionsAfter installingpoetry, to avoid potentialkeyring errors, disable its keyring usage by addingPYTHON_KEYRING_BACKEND=keyring.backends.null.Keyringto your shell's startup configuration and restarting your shell environment (e.g.,echo 'export PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring' >> ~/.bashrc && source ~/.bashrcfor a Bash shell environment and likewise for other shell environments).Install, add, or upgrade project dependenciespoetryinstall# install the latest project dependencies# orpoetryaddXYZ# add dependency `XYZ` to the project# orpoetryshow# list all dependencies currently installed# orpoetrylock# standardize the (now-)installed dependenciesActivate the newly-created virtual environment followingpoetry'susage documentation# activate the environment on a `posix`-like (e.g., macOS or Linux) systemsource$(poetryenvinfo--path)/bin/activate# activate the environment on a `Windows`-like system&((poetryenvinfo--path)+"\Scripts\activate.ps1")# if desired, deactivate the environmentdeactivateCode FormattingTo keep with the code style for theproteinworkshoprepository, using the following lines, please format your commits before opening a pull request:# assuming you are located in the `ProteinWorkshop` top-level directoryisort.
autoflake-r--in-place--remove-unused-variables--remove-all-unused-imports--ignore-init-module-imports.
black--config=pyproject.toml.DocumentationTo build a local version of the project's Sphinx documentation web pages:# assuming you are located in the `ProteinWorkshop` top-level directorypipinstall-rdocs/.docs.requirements# one-time onlyrm-rfdocs/build/&&sphinx-builddocs/source/docs/build/# NOTE: errors can safely be ignored