Geometric Vector Perceptron --- a rotation-equivariant GNN for learning from biomolecular structure

Related tags

Deep Learninggvp
Overview

Geometric Vector Perceptron

Code to accompany Learning from Protein Structure with Geometric Vector Perceptrons by B Jing, S Eismann, P Suriana, RJL Townshend, and RO Dror.

This repository serves two purposes. If you would like to use the architecture for protein design, we provide the pipeline for our experiments as well as our final trained model. If you are interested in adapting the architecture for other purposes, we provide instructions for general use of the GVP.

UPDATE: A PyTorch Geometric version of the GVP is now available at https://github.com/drorlab/gvp-pytorch, emphasizing ease of use and modularity. All future changes will be in PyTorch and pushed to this new repository.

Requirements

  • UNIX environment
  • python==3.7.6
  • numpy==1.18.1
  • scipy==1.4.1
  • pandas==1.0.3
  • tensorflow==2.1.0
  • tqdm==4.42.1

Protein design

Our training pipeline uses the CATH 4.2 dataset curated by Ingraham, et al, NeurIPS 2019. We provide code to train, validate, and test the model on this dataset. We also provide a pretrained model in models/cath_pretrained. If you want to test a trained model on new structures, see the section "Using the CPD model" below.

Fetching the datasets

Run getCATH.sh in data/ to fetch the CATH 4.2 dataset. If you are interested in testing on the TS 50 test set, also run grep -Fv -f ts50remove.txt chain_set.jsonl > chain_set_ts50.jsonl to produce a training set without overlap with the TS 50 test set.

Training the CPD model

Run python3 train_cpd.py [dataset] in src/ where [dataset] is the complete CATH 4.2 dataset, ../data/chain_set.jsonl, or the CATH 4.2 with overlap with TS50 removed, ../data/chain_set_ts50.jsonl. Model checkpoints are saved to models/ identified by the timestamp of the start of the run and the epoch number.

Evaluating the CPD model

Perplexity

To evaluate perplexity, run python3 test_cpd_perplexity.py ../models/cath_pretrained in src/.

Command Output
python3 test_cpd_perplexity.py ../models/cath_pretrained ALL TEST PERPLEXITY 5.29298734664917
SHORT TEST PERPLEXITY 7.0954108238220215
SINGLE CHAIN TEST PERPLEXITY 7.4412713050842285

Recovery

To evaluate recovery, run python3 test_cpd_recovery.py [model] [dataset] [output] in src/. [dataset] should be one of cath, short, sc, ts50. [model] should be ../models/ts50_pretrained if evaluating on the TS50 test set and ../models/cath_pretrained otherwise. Recoveries for each target will be dumped into the file [output]. To get the median recovery, run python3 analyze.py [output].

Because the recovery can take some time to run, we have supplied outputs in outputs/.

Command Output
python3 analyze.py ../outputs/cath.out 0.40187938705576753
python3 analyze.py ../outputs/short.out 0.32149746594868545
python3 analyze.py ../outputs/sc.out 0.319731182795699
python3 analyze.py ../outputs/ts50.out 0.44852965747702583

Using the CPD model

To use the CPD model on your own backbone structures, first convert the structures into a json format as follows:

[
    {
        "seq": "TQDCSFQHSP...",
        "coords": [[[74.46, 58.25, -21.65],...],...]
    }
    ...
]

For each structure, coords should be a num_residues x 4 x 3 nested list of the positions of the backbone N, C-alpha, C, and O atoms of each residue (in that order). If only backbone information is available, you can use a placeholder sequence of the same length. Then, run the below instructions (the function sample is defined in test_cpd_recovery.py)

dataset = datasets.load_dataset(PATH_TO_JSON, batch_size=1, shuffle=False)
for structure, seq, mask in dataset:
    n = 1 # number of sequences to sample
    # model is a pretrained CPD model
    design = sample(model, structure, mask, n)
    design = tf.cast(design, tf.int32).numpy()

The output design now an n x num_residues array of n designs, with amino acids represented as integers according to the encodings used to train the model. The encodings used by our pretrained model are in /src/datasets.py.

General usage

We describe our implementation in several levels of abstraction to make it as easy as possible to adapt the GVP to your uses. If you have any questions, please contact [email protected].

Using the core GVP modules

The core GVP modules are implemented in src/GVP.py. It contains code for the GVP itself, the vector/scalar dropout, and the vector/scalar batch norm, each of which is a tf.keras.layers.Module. These modules are initialized as follows:

gvp = GVP(vi, vo, so)
dropout = GVPDropout(drop_rate, nv)
layernorm = GVPLayerNorm(nv)

In the code and comments, vi, vo, si, so refer to number of vector/scalar channels in/out. nv and ns are the number of scalar/vector channels, and nls and nlv are the scalar/vector nonlinearities. The value si doesn't need to be specified because TensorFlow imputes it at the first forward pass.

Because the modules are designed to easily replace dense layers in a GNN, they are designed to take a single tensor x instead of seperate scalar/vector channel tensors. This is accomplished by assigning the first 3*nv channels in the input tensor to be the nv vector channels and the remaining channels to be the ns scalar channels. We provide utility functions merge and split to convert between seperate tensors where the vector tensor has dims [..., 3, nv] and the scalar tensor has dims [..., ns], and a single tensor with dims [..., 3*nv + ns]. For example:

v, s = input_data
x = merge(v, s)
x = gvp(x)
x = dropout(x, training=True)
x = layernorm(x)
v, s = split(x, nv=v.shape[-1])

Use vs_concat(x1, x2, nv1, nv2) to concatenate tensors x1 and x2 with nv1 and nv2 implicit vector channels.

Using the protein GNN

Our protein GNN is defined in src/models.py and is adapted from the protein GNN in Ingraham, et al, NeurIPS 2019. We provide two fully specified networks which take in raw protein representations and output a single global scalar prediction (MQAModel) or a 20-dimensional feature vector at each residue (CPDModel). Note that the CPDModel currently uses sequence information autoregressively. Sample usage:

mqa_model = MQAModel(node_dims, edge_dims, hidden_dims, num_layers)
X, S, mask = input_batch
output = mqa_model(X, S, mask) # dims [batch_size, 1]

The input X is a float tensor with dims [batch_size, num_residues, 4, 3] and has the backbone coordinates of N, C-alpha, C, and O atoms of each residue (in that order). S contains the sequence information as a integer tensor with dims [batch_size, num_residues]. The integer encodings can be arbitrary but the ones used by our pretrained model are defined in src/datasets.py. The mask is a float tensor with dims [batch, num_nodes] that is 1 for residues that exist and 0 for nodes that do not.

The three dims arguments should each be tuples (nv, ns) describing the number of vector and scalar channels to use in each embedding. The protein graph is first built using structural features to produce node embeddings with dims node_dims, then transformed into hidden_dims after adding sequence information. Therefore the two arguments are somewhat redundant. Note that the edge embeddings are static and are generated with only one vector feature, so anything greater than edge_nv = 1 is redundant.

If adapting one of the two provided models is insufficient, next we describe the building blocks of the protein GNN.

Structural features

The StructuralFeatures module converts the tensor X of raw backbone coordinates into a proximity graph with structure-based node and edge embeddings described in the paper.

feature_builder = StructuralFeatures(node_dims, edge_dims, top_k=30) # k nearest neighbors
h_V, h_E, E_idx = feature_builder(X, mask) # mask is as described above

h_V is the node embedding tensor with dims [batch, num_nodes, 3*node_nv+node_ns], h_E is the edge embedding tensor with dims [batch, num_nodes, top_k, 3*edge_nv+edge_ns], and E_idx is the tensor of neighbor node indices with dims [batch, num_nodes, top_k].

Message passing layers

A MPNNLayer is a single message-passing layer that takes in a tensor of incoming messages h_M from edges and neighboring nodes to update node embeddings. The layer is initialized as follows:

mpnn_layer = MPNNLayer(vec_in, hidden_dim)

Here, vec_in is the number of vector channels in the incoming message message (node_nv + edge_nv). The layer is then used as follows:

h_V = mpnn_layer(h_V, h_M, mask=None)

The optional mask is as described above. It is also possible to use an edgewise mask mask_attend with dims [batch, num_nodes, num_nodes] --- in autoregressive sampling, for example.

Note that while we also use the local node embedding as part of the message, the MPNNLayer itself will perform this concatenation, so you should only pass in the edge embeddings concatenated to the neighbor node embeddings. That is, h_M should have dims [batch, num_nodes, top_k, 3*vec_in+node_ns+edge_ns]. This tensor can be formed as:

h_M = cat_neighbors_nodes(h_V, h_E, E_idx, node_nv, edge_nv)

The Encoder module is a stack of MPNNLayers that performs multiple graph propagation steps directly using the node and edge embeddings h_V and h_E:

encoder = Encoder(node_dims, edge_dims, num_layers)
h_V = encoder(h_V, h_E, E_idx, mask=None)

The Decoder module is similar, except it incorporates sequence information autoregressively as described in the paper. If you are doing something other than autoregressive protein design, Decoder will likely be less useful to you.

Data pipeline

While we provide a data pipeline in src/datasets.py, it is specific for the training points/labels in protein design, so you will probably need to write your own for a different application. At minimum, you should modify load_dataset and parse_batch to convert your input representation to the model inputs X, S, and any necessary training labels.

Acknowledgements

The initial implementation of portions of the protein GNN and the input data pipeline were adapted from Ingraham, et al, NeurIPS 2019.

Citation

@inproceedings{
    jing2021learning,
    title={Learning from Protein Structure with Geometric Vector Perceptrons},
    author={Bowen Jing and Stephan Eismann and Patricia Suriana and Raphael John Lamarre Townshend and Ron Dror},
    booktitle={International Conference on Learning Representations},
    year={2021},
    url={https://openreview.net/forum?id=1YLJDvSx6J4}
}
Owner
Dror Lab
Ron Dror's computational biology laboratory at Stanford University
Dror Lab
Contra is a lightweight, production ready Tensorflow alternative for solving time series prediction challenges with AI

Contra AI Engine A lightweight, production ready Tensorflow alternative developed by Styvio styvio.com » How to Use · Report Bug · Request Feature Tab

styvio 14 May 25, 2022
An original implementation of "MetaICL Learning to Learn In Context" by Sewon Min, Mike Lewis, Luke Zettlemoyer and Hannaneh Hajishirzi

MetaICL: Learning to Learn In Context This includes an original implementation of "MetaICL: Learning to Learn In Context" by Sewon Min, Mike Lewis, Lu

Meta Research 141 Jan 07, 2023
Official implementation of EfficientPose

EfficientPose This is the official implementation of EfficientPose. We based our work on the Keras EfficientDet implementation xuannianz/EfficientDet

2 May 17, 2022
Supervised Contrastive Learning for Product Matching

Contrastive Product Matching This repository contains the code and data download links to reproduce the experiments of the paper "Supervised Contrasti

Web-based Systems Group @ University of Mannheim 18 Dec 10, 2022
A lightweight deep network for fast and accurate optical flow estimation.

FastFlowNet: A Lightweight Network for Fast Optical Flow Estimation The official PyTorch implementation of FastFlowNet (ICRA 2021). Authors: Lingtong

Tone 161 Jan 03, 2023
CVPR '21: In the light of feature distributions: Moment matching for Neural Style Transfer

In the light of feature distributions: Moment matching for Neural Style Transfer (CVPR 2021) This repository provides code to recreate results present

Nikolai Kalischek 49 Oct 13, 2022
Score refinement for confidence-based 3D multi-object tracking

Score refinement for confidence-based 3D multi-object tracking Our video gives a brief explanation of our Method. This is the official code for the pa

Cognitive Systems Research Group 47 Dec 26, 2022
Code for "Retrieving Black-box Optimal Images from External Databases" (WSDM 2022)

Retrieving Black-box Optimal Images from External Databases (WSDM 2022) We propose how a user retreives an optimal image from external databases of we

joisino 5 Apr 13, 2022
Neuralnetwork - Basic Multilayer Perceptron Neural Network for deep learning

Neural Network Just a basic Neural Network module Usage Example Importing Module

andreecy 0 Nov 01, 2022
Towards Understanding Quality Challenges of the Federated Learning: A First Look from the Lens of Robustness

FL Analysis This repository contains the code and results for the paper "Towards Understanding Quality Challenges of the Federated Learning: A First L

3 Oct 17, 2022
[ICLR 2021] Rank the Episodes: A Simple Approach for Exploration in Procedurally-Generated Environments.

[ICLR 2021] RAPID: A Simple Approach for Exploration in Reinforcement Learning This is the Tensorflow implementation of ICLR 2021 paper Rank the Episo

Daochen Zha 48 Nov 21, 2022
Web-interface + rest API for classification and regression (https://jeff1evesque.github.io/machine-learning.docs)

Machine Learning This project provides a web-interface, as well as a programmatic-api for various machine learning algorithms. Supported algorithms: S

Jeff Levesque 252 Dec 11, 2022
Deep learning image registration library for PyTorch

TorchIR: Pytorch Image Registration TorchIR is a image registration library for deep learning image registration (DLIR). I have integrated several ide

Bob de Vos 40 Dec 16, 2022
image scene graph generation benchmark

Scene Graph Benchmark in PyTorch 1.7 This project is based on maskrcnn-benchmark Highlights Upgrad to pytorch 1.7 Multi-GPU training and inference Bat

Microsoft 303 Dec 27, 2022
This program writes christmas wish programmatically. It is using turtle as a pen pointer draw christmas trees and stars.

Introduction This is a simple program is written in python and turtle library. The objective of this program is to wish merry Christmas programmatical

Gunarakulan Gunaretnam 1 Dec 25, 2021
An example showing how to use jax to train resnet50 on multi-node multi-GPU

jax-multi-gpu-resnet50-example This repo shows how to use jax for multi-node multi-GPU training. The example is adapted from the resnet50 example in d

Yangzihao Wang 20 Jul 04, 2022
Finite-temperature variational Monte Carlo calculation of uniform electron gas using neural canonical transformation.

CoulombGas This code implements the neural canonical transformation approach to the thermodynamic properties of uniform electron gas. Building on JAX,

FermiFlow 9 Mar 03, 2022
3D dataset of humans Manipulating Objects in-the-Wild (MOW)

MOW dataset [Website] This repository maintains our 3D dataset of humans Manipulating Objects in-the-Wild (MOW). The dataset contains 512 images in th

Zhe Cao 28 Nov 06, 2022
rliable is an open-source Python library for reliable evaluation, even with a handful of runs, on reinforcement learning and machine learnings benchmarks.

Open-source library for reliable evaluation on reinforcement learning and machine learning benchmarks. See NeurIPS 2021 oral for details.

Google Research 529 Jan 01, 2023
PyTorch Implementation of Realtime Multi-Person Pose Estimation project.

PyTorch Realtime Multi-Person Pose Estimation This is a pytorch version of Realtime_Multi-Person_Pose_Estimation, origin code is here Realtime_Multi-P

Dave Fang 157 Nov 12, 2022