Repository: FreyrS/dMaSIF
Branch: master
Commit: 0dcc26c3c218
Files: 34
Total size: 9.9 MB
Directory structure:
gitextract_4r95syc5/
├── .gitignore
├── Arguments.py
├── LICENSE.md
├── README.md
├── benchmark_layers.py
├── benchmark_models.py
├── benchmark_scripts/
│ ├── DGCNN_site.sh
│ ├── Pointnet_site.sh
│ ├── dMaSIF_search.sh
│ └── dMaSIF_site.sh
├── data.py
├── data_analysis/
│ ├── analyse_atomnet.ipynb
│ ├── analyse_descriptors.py
│ ├── analyse_descriptors_para.py
│ ├── analyse_output.ipynb
│ ├── analyse_site_outputs.py
│ ├── analyse_site_outputs_graph.ipynb
│ ├── plot_search.ipynb
│ └── profiling_surface.ipynb
├── data_iteration.py
├── data_preprocessing/
│ ├── convert_pdb2npy.py
│ ├── convert_ply2npy.py
│ └── download_pdb.py
├── geometry_processing.py
├── helper.py
├── lists/
│ ├── testing.txt
│ ├── testing_ppi.txt
│ ├── training.txt
│ └── training_ppi.txt
├── main_inference.py
├── main_training.py
├── model.py
├── models/
│ └── dMaSIF_search_3layer_12A_16dim
└── requirements.txt
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitignore
================================================
01-benchmark_surfaces/
01-benchmark_surfaces_npy/
01-benchmark_pdbs_npy/
01-benchmark_pdbs/
01-benchmark_pdbs/
shape_index/
masif_preds/
runs/
venv/
preds/
*.log
NeurIPS_2020_benchmarks/
*.out
figures/
timings/
data_analysis/roc_curves
data_analysis/.ipynb_checkpoints/
.ipynb_checkpoints/
================================================
FILE: Arguments.py
================================================
import argparse
parser = argparse.ArgumentParser(description="Network parameters")
# Main parameters
parser.add_argument(
"--experiment_name", type=str, help="Name of experiment", required=True
)
parser.add_argument(
"--use_mesh", type=bool, default=False, help="Use precomputed surfaces"
)
parser.add_argument(
"--embedding_layer",
type=str,
default="dMaSIF",
choices=["dMaSIF", "DGCNN", "PointNet++"],
help="Which convolutional embedding layer to use",
)
parser.add_argument("--profile", type=bool, default=False, help="Profile code")
# Geometric parameters
parser.add_argument(
"--curvature_scales",
type=list,
default=[1.0, 2.0, 3.0, 5.0, 10.0],
help="Scales at which we compute the geometric features (mean and Gauss curvatures)",
)
parser.add_argument(
"--resolution",
type=float,
default=1.0,
help="Resolution of the generated point cloud",
)
parser.add_argument(
"--distance",
type=float,
default=1.05,
help="Distance parameter in surface generation",
)
parser.add_argument(
"--variance",
type=float,
default=0.1,
help="Variance parameter in surface generation",
)
parser.add_argument(
"--sup_sampling", type=int, default=20, help="Sup-sampling ratio around atoms"
)
# Hyper-parameters for the embedding
parser.add_argument(
"--atom_dims",
type=int,
default=6,
help="Number of atom types and dimension of resulting chemical features",
)
parser.add_argument(
"--emb_dims",
type=int,
default=8,
help="Number of input features (+ 3 xyz coordinates for DGCNNs)",
)
parser.add_argument(
"--in_channels",
type=int,
default=16,
help="Number of embedding dimensions",
)
parser.add_argument(
"--orientation_units",
type=int,
default=16,
help="Number of hidden units for the orientation score MLP",
)
parser.add_argument(
"--unet_hidden_channels",
type=int,
default=8,
help="Number of hidden units for TangentConv UNet",
)
parser.add_argument(
"--post_units",
type=int,
default=8,
help="Number of hidden units for the post-processing MLP",
)
parser.add_argument(
"--n_layers", type=int, default=1, help="Number of convolutional layers"
)
parser.add_argument(
"--radius", type=float, default=9.0, help="Radius to use for the convolution"
)
parser.add_argument(
"--k",
type=int,
default=40,
help="Number of nearset neighbours for DGCNN and PointNet++",
)
parser.add_argument(
"--dropout",
type=float,
default=0.0,
help="Amount of Dropout for the input features",
)
# Training
parser.add_argument(
"--n_epochs", type=int, default=50, help="Number of training epochs"
)
parser.add_argument(
"--batch_size", type=int, default=1, help="Number of proteins in a batch"
)
parser.add_argument(
"--device", type=str, default="cuda:0", help="Which gpu/cpu to train on"
)
parser.add_argument(
"--restart_training",
type=str,
default="",
help="Which model to restart the training from",
)
parser.add_argument(
"--n_rocauc_samples",
type=int,
default=100,
help="Number of samples for the Matching ROC-AUC",
)
parser.add_argument(
"--validation_fraction",
type=float,
default=0.1,
help="Fraction of training dataset to use for validation",
)
parser.add_argument("--seed", type=int, default=42, help="Random seed")
parser.add_argument(
"--random_rotation",
type=bool,
default=False,
help="Move proteins to center and add random rotation",
)
parser.add_argument(
"--single_protein",
type=bool,
default=False,
help="Use single protein in a pair or both",
)
parser.add_argument("--site", type=bool, default=False, help="Predict interaction site")
parser.add_argument(
"--search",
type=bool,
default=False,
help="Predict matching between two partners",
)
parser.add_argument(
"--no_chem", type=bool, default=False, help="Predict without chemical information"
)
parser.add_argument(
"--no_geom", type=bool, default=False, help="Predict without curvature information"
)
parser.add_argument(
"--single_pdb",
type=str,
default="",
help="Which structure to do inference on",
)
parser.add_argument(
"--pdb_list",
type=str,
default="",
help="Which structures to do inference on",
)
================================================
FILE: LICENSE.md
================================================
Attribution-NonCommercial-NoDerivatives 4.0 International
=======================================================================
Creative Commons Corporation ("Creative Commons") is not a law firm and
does not provide legal services or legal advice. Distribution of
Creative Commons public licenses does not create a lawyer-client or
other relationship. Creative Commons makes its licenses and related
information available on an "as-is" basis. Creative Commons gives no
warranties regarding its licenses, any material licensed under their
terms and conditions, or any related information. Creative Commons
disclaims all liability for damages resulting from their use to the
fullest extent possible.
Using Creative Commons Public Licenses
Creative Commons public licenses provide a standard set of terms and
conditions that creators and other rights holders may use to share
original works of authorship and other material subject to copyright
and certain other rights specified in the public license below. The
following considerations are for informational purposes only, are not
exhaustive, and do not form part of our licenses.
Considerations for licensors: Our public licenses are
intended for use by those authorized to give the public
permission to use material in ways otherwise restricted by
copyright and certain other rights. Our licenses are
irrevocable. Licensors should read and understand the terms
and conditions of the license they choose before applying it.
Licensors should also secure all rights necessary before
applying our licenses so that the public can reuse the
material as expected. Licensors should clearly mark any
material not subject to the license. This includes other CC-
licensed material, or material used under an exception or
limitation to copyright. More considerations for licensors:
wiki.creativecommons.org/Considerations_for_licensors
Considerations for the public: By using one of our public
licenses, a licensor grants the public permission to use the
licensed material under specified terms and conditions. If
the licensor's permission is not necessary for any reason--for
example, because of any applicable exception or limitation to
copyright--then that use is not regulated by the license. Our
licenses grant only permissions under copyright and certain
other rights that a licensor has authority to grant. Use of
the licensed material may still be restricted for other
reasons, including because others have copyright or other
rights in the material. A licensor may make special requests,
such as asking that all changes be marked or described.
Although not required by our licenses, you are encouraged to
respect those requests where reasonable. More considerations
for the public:
wiki.creativecommons.org/Considerations_for_licensees
=======================================================================
Creative Commons Attribution-NonCommercial-NoDerivatives 4.0
International Public License
By exercising the Licensed Rights (defined below), You accept and agree
to be bound by the terms and conditions of this Creative Commons
Attribution-NonCommercial-NoDerivatives 4.0 International Public
License ("Public License"). To the extent this Public License may be
interpreted as a contract, You are granted the Licensed Rights in
consideration of Your acceptance of these terms and conditions, and the
Licensor grants You such rights in consideration of benefits the
Licensor receives from making the Licensed Material available under
these terms and conditions.
Section 1 -- Definitions.
a. Adapted Material means material subject to Copyright and Similar
Rights that is derived from or based upon the Licensed Material
and in which the Licensed Material is translated, altered,
arranged, transformed, or otherwise modified in a manner requiring
permission under the Copyright and Similar Rights held by the
Licensor. For purposes of this Public License, where the Licensed
Material is a musical work, performance, or sound recording,
Adapted Material is always produced where the Licensed Material is
synched in timed relation with a moving image.
b. Copyright and Similar Rights means copyright and/or similar rights
closely related to copyright including, without limitation,
performance, broadcast, sound recording, and Sui Generis Database
Rights, without regard to how the rights are labeled or
categorized. For purposes of this Public License, the rights
specified in Section 2(b)(1)-(2) are not Copyright and Similar
Rights.
c. Effective Technological Measures means those measures that, in the
absence of proper authority, may not be circumvented under laws
fulfilling obligations under Article 11 of the WIPO Copyright
Treaty adopted on December 20, 1996, and/or similar international
agreements.
d. Exceptions and Limitations means fair use, fair dealing, and/or
any other exception or limitation to Copyright and Similar Rights
that applies to Your use of the Licensed Material.
e. Licensed Material means the artistic or literary work, database,
or other material to which the Licensor applied this Public
License.
f. Licensed Rights means the rights granted to You subject to the
terms and conditions of this Public License, which are limited to
all Copyright and Similar Rights that apply to Your use of the
Licensed Material and that the Licensor has authority to license.
g. Licensor means the individual(s) or entity(ies) granting rights
under this Public License.
h. NonCommercial means not primarily intended for or directed towards
commercial advantage or monetary compensation. For purposes of
this Public License, the exchange of the Licensed Material for
other material subject to Copyright and Similar Rights by digital
file-sharing or similar means is NonCommercial provided there is
no payment of monetary compensation in connection with the
exchange.
i. Share means to provide material to the public by any means or
process that requires permission under the Licensed Rights, such
as reproduction, public display, public performance, distribution,
dissemination, communication, or importation, and to make material
available to the public including in ways that members of the
public may access the material from a place and at a time
individually chosen by them.
j. Sui Generis Database Rights means rights other than copyright
resulting from Directive 96/9/EC of the European Parliament and of
the Council of 11 March 1996 on the legal protection of databases,
as amended and/or succeeded, as well as other essentially
equivalent rights anywhere in the world.
k. You means the individual or entity exercising the Licensed Rights
under this Public License. Your has a corresponding meaning.
Section 2 -- Scope.
a. License grant.
1. Subject to the terms and conditions of this Public License,
the Licensor hereby grants You a worldwide, royalty-free,
non-sublicensable, non-exclusive, irrevocable license to
exercise the Licensed Rights in the Licensed Material to:
a. reproduce and Share the Licensed Material, in whole or
in part, for NonCommercial purposes only; and
b. produce and reproduce, but not Share, Adapted Material
for NonCommercial purposes only.
2. Exceptions and Limitations. For the avoidance of doubt, where
Exceptions and Limitations apply to Your use, this Public
License does not apply, and You do not need to comply with
its terms and conditions.
3. Term. The term of this Public License is specified in Section
6(a).
4. Media and formats; technical modifications allowed. The
Licensor authorizes You to exercise the Licensed Rights in
all media and formats whether now known or hereafter created,
and to make technical modifications necessary to do so. The
Licensor waives and/or agrees not to assert any right or
authority to forbid You from making technical modifications
necessary to exercise the Licensed Rights, including
technical modifications necessary to circumvent Effective
Technological Measures. For purposes of this Public License,
simply making modifications authorized by this Section 2(a)
(4) never produces Adapted Material.
5. Downstream recipients.
a. Offer from the Licensor -- Licensed Material. Every
recipient of the Licensed Material automatically
receives an offer from the Licensor to exercise the
Licensed Rights under the terms and conditions of this
Public License.
b. No downstream restrictions. You may not offer or impose
any additional or different terms or conditions on, or
apply any Effective Technological Measures to, the
Licensed Material if doing so restricts exercise of the
Licensed Rights by any recipient of the Licensed
Material.
6. No endorsement. Nothing in this Public License constitutes or
may be construed as permission to assert or imply that You
are, or that Your use of the Licensed Material is, connected
with, or sponsored, endorsed, or granted official status by,
the Licensor or others designated to receive attribution as
provided in Section 3(a)(1)(A)(i).
b. Other rights.
1. Moral rights, such as the right of integrity, are not
licensed under this Public License, nor are publicity,
privacy, and/or other similar personality rights; however, to
the extent possible, the Licensor waives and/or agrees not to
assert any such rights held by the Licensor to the limited
extent necessary to allow You to exercise the Licensed
Rights, but not otherwise.
2. Patent and trademark rights are not licensed under this
Public License.
3. To the extent possible, the Licensor waives any right to
collect royalties from You for the exercise of the Licensed
Rights, whether directly or through a collecting society
under any voluntary or waivable statutory or compulsory
licensing scheme. In all other cases the Licensor expressly
reserves any right to collect such royalties, including when
the Licensed Material is used other than for NonCommercial
purposes.
Section 3 -- License Conditions.
Your exercise of the Licensed Rights is expressly made subject to the
following conditions.
a. Attribution.
1. If You Share the Licensed Material, You must:
a. retain the following if it is supplied by the Licensor
with the Licensed Material:
i. identification of the creator(s) of the Licensed
Material and any others designated to receive
attribution, in any reasonable manner requested by
the Licensor (including by pseudonym if
designated);
ii. a copyright notice;
iii. a notice that refers to this Public License;
iv. a notice that refers to the disclaimer of
warranties;
v. a URI or hyperlink to the Licensed Material to the
extent reasonably practicable;
b. indicate if You modified the Licensed Material and
retain an indication of any previous modifications; and
c. indicate the Licensed Material is licensed under this
Public License, and include the text of, or the URI or
hyperlink to, this Public License.
For the avoidance of doubt, You do not have permission under
this Public License to Share Adapted Material.
2. You may satisfy the conditions in Section 3(a)(1) in any
reasonable manner based on the medium, means, and context in
which You Share the Licensed Material. For example, it may be
reasonable to satisfy the conditions by providing a URI or
hyperlink to a resource that includes the required
information.
3. If requested by the Licensor, You must remove any of the
information required by Section 3(a)(1)(A) to the extent
reasonably practicable.
Section 4 -- Sui Generis Database Rights.
Where the Licensed Rights include Sui Generis Database Rights that
apply to Your use of the Licensed Material:
a. for the avoidance of doubt, Section 2(a)(1) grants You the right
to extract, reuse, reproduce, and Share all or a substantial
portion of the contents of the database for NonCommercial purposes
only and provided You do not Share Adapted Material;
b. if You include all or a substantial portion of the database
contents in a database in which You have Sui Generis Database
Rights, then the database in which You have Sui Generis Database
Rights (but not its individual contents) is Adapted Material; and
c. You must comply with the conditions in Section 3(a) if You Share
all or a substantial portion of the contents of the database.
For the avoidance of doubt, this Section 4 supplements and does not
replace Your obligations under this Public License where the Licensed
Rights include other Copyright and Similar Rights.
Section 5 -- Disclaimer of Warranties and Limitation of Liability.
a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
c. The disclaimer of warranties and limitation of liability provided
above shall be interpreted in a manner that, to the extent
possible, most closely approximates an absolute disclaimer and
waiver of all liability.
Section 6 -- Term and Termination.
a. This Public License applies for the term of the Copyright and
Similar Rights licensed here. However, if You fail to comply with
this Public License, then Your rights under this Public License
terminate automatically.
b. Where Your right to use the Licensed Material has terminated under
Section 6(a), it reinstates:
1. automatically as of the date the violation is cured, provided
it is cured within 30 days of Your discovery of the
violation; or
2. upon express reinstatement by the Licensor.
For the avoidance of doubt, this Section 6(b) does not affect any
right the Licensor may have to seek remedies for Your violations
of this Public License.
c. For the avoidance of doubt, the Licensor may also offer the
Licensed Material under separate terms or conditions or stop
distributing the Licensed Material at any time; however, doing so
will not terminate this Public License.
d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
License.
Section 7 -- Other Terms and Conditions.
a. The Licensor shall not be bound by any additional or different
terms or conditions communicated by You unless expressly agreed.
b. Any arrangements, understandings, or agreements regarding the
Licensed Material not stated herein are separate from and
independent of the terms and conditions of this Public License.
Section 8 -- Interpretation.
a. For the avoidance of doubt, this Public License does not, and
shall not be interpreted to, reduce, limit, restrict, or impose
conditions on any use of the Licensed Material that could lawfully
be made without permission under this Public License.
b. To the extent possible, if any provision of this Public License is
deemed unenforceable, it shall be automatically reformed to the
minimum extent necessary to make it enforceable. If the provision
cannot be reformed, it shall be severed from this Public License
without affecting the enforceability of the remaining terms and
conditions.
c. No term or condition of this Public License will be waived and no
failure to comply consented to unless expressly agreed to by the
Licensor.
d. Nothing in this Public License constitutes or may be interpreted
as a limitation upon, or waiver of, any privileges and immunities
that apply to the Licensor or You, including from the legal
processes of any jurisdiction or authority.
=======================================================================
Creative Commons is not a party to its public
licenses. Notwithstanding, Creative Commons may elect to apply one of
its public licenses to material it publishes and in those instances
will be considered the “Licensor.” The text of the Creative Commons
public licenses is dedicated to the public domain under the CC0 Public
Domain Dedication. Except for the limited purpose of indicating that
material is shared under a Creative Commons public license or as
otherwise permitted by the Creative Commons policies published at
creativecommons.org/policies, Creative Commons does not authorize the
use of the trademark "Creative Commons" or any other trademark or logo
of Creative Commons without its prior written consent including,
without limitation, in connection with any unauthorized modifications
to any of its public licenses or any other arrangements,
understandings, or agreements concerning use of licensed material. For
the avoidance of doubt, this paragraph does not form part of the
public licenses.
Creative Commons may be contacted at creativecommons.org.
================================================
FILE: README.md
================================================
## dMaSIF - Fast end-to-end learning on protein surfaces
## Abstract
Proteins’ biological functions are defined by the geometric
and chemical structure of their 3D molecular surfaces.
Recent works have shown that geometric deep learning can
be used on mesh-based representations of proteins to identify
potential functional sites, such as binding targets for
potential drugs. Unfortunately though, the use of meshes as
the underlying representation for protein structure has multiple
drawbacks including the need to pre-compute the input
features and mesh connectivities. This becomes a bottleneck
for many important tasks in protein science.
In this paper, we present a new framework for deep
learning on protein structures that addresses these limitations.
Among the key advantages of our method are the computation
and sampling of the molecular surface on-the-fly
from the underlying atomic point cloud and a novel efficient
geometric convolutional layer. As a result, we are able to
process large collections of proteins in an end-to-end fashion,
taking as the sole input the raw 3D coordinates and
chemical types of their atoms, eliminating the need for any
hand-crafted pre-computed features.
To showcase the performance of our approach, we test it
on two tasks in the field of protein structural bioinformatics:
the identification of interaction sites and the prediction
of protein-protein interactions. On both tasks, we achieve
state-of-the-art performance with much faster run times and
fewer parameters than previous models. These results will
considerably ease the deployment of deep learning methods
in protein science and open the door for end-to-end differentiable
approaches in protein modeling tasks such as function
prediction and design.
## Hardware requirements
Models have been trained on either a single NVIDIA RTX 2080 Ti or a single Tesla V100 GPU. Time and memory benchmarks were performed on a single Tesla V100.
## Software prerequisites
Scripts have been tested using the following two sets of core dependencies:
| Dependency | First Option | Second Option |
| ------------- | ------------- | ------------- |
| GCC | 7.5.0 | 8.4.0 |
| CMAKE | 3.10.2 | 3.16.5 |
| CUDA | 10.0.130 | 10.2.89 |
| cuDNN | 7.6.4.38 | 7.6.5.32 |
| Python | 3.6.9 | 3.7.7 |
| PyTorch | 1.4.0 | 1.6.0 |
| PyKeops | 1.4 | 1.4.1 |
| PyTorch Geometric | 1.5.0 | 1.6.1 |
## Code overview
Usage:
- In order to **train models**, run `main_training.py` with the appropriate flags.
Available flags and their descriptions can be found in `Arguments.py`.
- The command line options needed to reproduce the **benchmarks** can be found in `benchmark_scripts/`.
- To make **inference** on the testing set using pretrained models, use `main_inference.py` with the flags that were used for training the models.
Note that the `--experiment_name flag` should be modified to specify the training epoch to use.
Implementation:
- Our **surface generation** algorithm, **curvature** estimation method and **quasi-geodesic convolutions** are implemented in `geometry_processing.py`.
- The **definition of the neural network** along with surface and input features can be found in `model.py`. The convolutional layers are implemented in `benchmark_models.py`.
- The scripts used to **generate the figures** of the paper can be found in `data_analysis/`.
## License
This work is licensed under a Creative Commons Attribution-NonCommercial-NoDerivatives 4.0 International License.
## Reference
Sverrisson, F., Feydy, J., Correia, B. E., & Bronstein, M. M. (2020). Fast end-to-end learning on protein surfaces. [bioRxiv](https://www.biorxiv.org/content/10.1101/2020.12.28.424589v1).
================================================
FILE: benchmark_layers.py
================================================
import torch
from typing import Optional
from pykeops.torch import LazyTensor
from torch_geometric.nn import EdgeConv, Reshape
from torch_cluster import knn
from math import ceil
from torch_geometric.nn.inits import reset
from torch.nn import ELU, Conv1d
from torch.nn import Sequential as S, Linear as L, BatchNorm1d as BN
def ranges_slices(batch):
"""Helper function for the diagonal ranges function."""
Ns = batch.bincount()
indices = Ns.cumsum(0)
ranges = torch.cat((0 * indices[:1], indices))
ranges = (
torch.stack((ranges[:-1], ranges[1:])).t().int().contiguous().to(batch.device)
)
slices = (1 + torch.arange(len(Ns))).int().to(batch.device)
return ranges, slices
def diagonal_ranges(batch_x=None, batch_y=None):
"""Encodes the block-diagonal structure associated to a batch vector."""
if batch_x is None and batch_y is None:
return None
ranges_x, slices_x = ranges_slices(batch_x)
ranges_y, slices_y = ranges_slices(batch_y)
return ranges_x, slices_x, ranges_y, ranges_y, slices_y, ranges_x
@torch.jit.ignore
def keops_knn(
x: torch.Tensor,
y: torch.Tensor,
k: int,
batch_x: Optional[torch.Tensor] = None,
batch_y: Optional[torch.Tensor] = None,
cosine: bool = False,
) -> torch.Tensor:
r"""Straightforward modification of PyTorch_geometric's knn method."""
x = x.view(-1, 1) if x.dim() == 1 else x
y = y.view(-1, 1) if y.dim() == 1 else y
y_i = LazyTensor(y[:, None, :])
x_j = LazyTensor(x[None, :, :])
if cosine:
D_ij = -(y_i | x_j)
else:
D_ij = ((y_i - x_j) ** 2).sum(-1)
D_ij.ranges = diagonal_ranges(batch_y, batch_x)
idy = D_ij.argKmin(k, dim=1) # (N, K)
rows = torch.arange(k * len(y), device=idy.device) // k
return torch.stack([rows, idy.view(-1)], dim=0)
knns = {"torch": knn, "keops": keops_knn}
@torch.jit.ignore
def knn_graph(
x: torch.Tensor,
k: int,
batch: Optional[torch.Tensor] = None,
loop: bool = False,
flow: str = "source_to_target",
cosine: bool = False,
target: Optional[torch.Tensor] = None,
batch_target: Optional[torch.Tensor] = None,
backend: str = "torch",
) -> torch.Tensor:
r"""Straightforward modification of PyTorch_geometric's knn_graph method to allow for source/targets."""
assert flow in ["source_to_target", "target_to_source"]
if target is None:
target = x
if batch_target is None:
batch_target = batch
row, col = knns[backend](
x, target, k if loop else k + 1, batch, batch_target, cosine=cosine
)
row, col = (col, row) if flow == "source_to_target" else (row, col)
if not loop:
mask = row != col
row, col = row[mask], col[mask]
return torch.stack([row, col], dim=0)
class MyDynamicEdgeConv(EdgeConv):
r"""Straightforward modification of PyTorch_geometric's DynamicEdgeConv layer."""
def __init__(self, nn, k, aggr="max", **kwargs):
super(MyDynamicEdgeConv, self).__init__(nn=nn, aggr=aggr, **kwargs)
self.k = k
def forward(self, x, batch=None):
""""""
edge_index = knn_graph(
x, self.k, batch, loop=False, flow=self.flow, backend="keops"
)
return super(MyDynamicEdgeConv, self).forward(x, edge_index)
def __repr__(self):
return "{}(nn={}, k={})".format(self.__class__.__name__, self.nn, self.k)
class MyXConv(torch.nn.Module):
def __init__(
self,
in_channels=None,
out_channels=None,
dim=None,
kernel_size=None,
hidden_channels=None,
dilation=1,
bias=True,
backend="torch",
):
super(MyXConv, self).__init__()
self.in_channels = in_channels
if hidden_channels is None:
hidden_channels = in_channels // 4
if hidden_channels == 0:
hidden_channels = 1
self.hidden_channels = hidden_channels
self.out_channels = out_channels
self.dim = dim
self.kernel_size = kernel_size
self.dilation = dilation
self.backend = backend
C_in, C_delta, C_out = in_channels, hidden_channels, out_channels
D, K = dim, kernel_size
self.mlp1 = S(
L(dim, C_delta),
ELU(),
BN(C_delta),
L(C_delta, C_delta),
ELU(),
BN(C_delta),
Reshape(-1, K, C_delta),
)
self.mlp2 = S(
L(D * K, K ** 2),
ELU(),
BN(K ** 2),
Reshape(-1, K, K),
Conv1d(K, K ** 2, K, groups=K),
ELU(),
BN(K ** 2),
Reshape(-1, K, K),
Conv1d(K, K ** 2, K, groups=K),
BN(K ** 2),
Reshape(-1, K, K),
)
C_in = C_in + C_delta
depth_multiplier = int(ceil(C_out / C_in))
self.conv = S(
Conv1d(C_in, C_in * depth_multiplier, K, groups=C_in),
Reshape(-1, C_in * depth_multiplier),
L(C_in * depth_multiplier, C_out, bias=bias),
)
self.reset_parameters()
def reset_parameters(self):
reset(self.mlp1)
reset(self.mlp2)
reset(self.conv)
def forward(self, x, source, batch_source, target, batch_target):
""""""
# Load data shapes:
# pos = pos.unsqueeze(-1) if pos.dim() == 1 else pos
(Nin, _), (N, D), K = source.size(), target.size(), self.kernel_size
# Compute K-nn:
row, col = knn_graph(
source,
K * self.dilation,
batch_source,
loop=True,
flow="target_to_source",
target=target,
batch_target=batch_target,
backend=self.backend,
)
# row is a vector of size N*K*dilation that indexes "target"
# col is a vector of size N*K*dilation that indexes "source"
# If needed, sup-sample the K-NN graph:
if self.dilation > 1:
dil = self.dilation
index = torch.randint(
K * dil,
(N, K),
dtype=torch.long,
layout=torch.strided,
device=row.device,
)
arange = torch.arange(N, dtype=torch.long, device=row.device)
arange = arange * (K * dil)
index = (index + arange.view(-1, 1)).view(-1) # (N*K,)
row, col = row[index], col[index]
# assert row.max() < N
# assert col.max() < Nin
# Line 1: local difference vector:
pos = source[col] - target[row] # (N * K, D)
# Line 2: compute F_delta
x_star = self.mlp1(pos.view(N * K, D))
# Line 3: concatenate the features and reshape:
if x is not None:
x = x.unsqueeze(-1) if x.dim() == 1 else x
x = x[col].view(N, K, self.in_channels)
x_star = torch.cat([x_star, x], dim=-1)
x_star = x_star.transpose(1, 2).contiguous()
x_star = x_star.view(N, self.in_channels + self.hidden_channels, K, 1)
# Line 4: Compute the transformation matrix:
transform_matrix = self.mlp2(pos.view(N, K * D))
transform_matrix = transform_matrix.view(N, 1, K, K)
# Line 5: Apply it to the neighborhood:
x_transformed = torch.matmul(transform_matrix, x_star)
x_transformed = x_transformed.view(N, -1, K) # (N, I+H, K)
# Line 6: Apply the convolution filter:
out = self.conv(x_transformed) # (N, Cout)
return out
def __repr__(self):
return "{}({}, {})".format(
self.__class__.__name__, self.in_channels, self.out_channels
)
================================================
FILE: benchmark_models.py
================================================
import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.nn import (
Sequential as Seq,
Dropout,
Linear as Lin,
LeakyReLU,
ReLU,
BatchNorm1d as BN,
)
import torch_geometric.transforms as T
from torch_geometric.data import DataLoader
from torch_geometric.nn import (
DynamicEdgeConv,
PointConv,
XConv,
fps,
radius,
global_max_pool,
knn_interpolate,
)
from pykeops.torch import LazyTensor
from benchmark_layers import MyDynamicEdgeConv, MyXConv
from geometry_processing import dMaSIFConv, mesh_normals_areas, tangent_vectors
from helper import diagonal_ranges
DEConv = {"torch": DynamicEdgeConv, "keops": MyDynamicEdgeConv}
# Dynamic Graph CNNs ===========================================================
# Adapted from the PyTorch_geometric gallery to get a close fit to
# the original paper.
def MLP(channels, batch_norm=True):
"""Multi-layer perceptron, with ReLU non-linearities and batch normalization."""
return Seq(
*[
Seq(
Lin(channels[i - 1], channels[i]),
BN(channels[i]) if batch_norm else nn.Identity(),
LeakyReLU(negative_slope=0.2),
)
for i in range(1, len(channels))
]
)
class DGCNN_seg(torch.nn.Module):
def __init__(
self, in_channels, out_channels, n_layers, k=40, aggr="max", backend="keops"
):
super(DGCNN_seg, self).__init__()
self.name = "DGCNN_seg_" + backend
self.I, self.O = (
in_channels + 3,
out_channels,
) # Add coordinates to input channels
self.n_layers = n_layers
self.transform_1 = DEConv[backend](MLP([2 * 3, 64, 128]), k, aggr)
self.transform_2 = MLP([128, 1024])
self.transform_3 = MLP([1024, 512, 256], batch_norm=False)
self.transform_4 = Lin(256, 3 * 3)
self.conv_layers = nn.ModuleList(
[DEConv[backend](MLP([2 * self.I, self.O, self.O]), k, aggr)]
+ [
DEConv[backend](MLP([2 * self.O, self.O, self.O]), k, aggr)
for i in range(n_layers - 1)
]
)
self.linear_layers = nn.ModuleList(
[
nn.Sequential(
nn.Linear(self.O, self.O), nn.ReLU(), nn.Linear(self.O, self.O)
)
for i in range(n_layers)
]
)
self.linear_transform = nn.ModuleList(
[nn.Linear(self.I, self.O)]
+ [nn.Linear(self.O, self.O) for i in range(n_layers - 1)]
)
def forward(self, positions, features, batch_indices):
# Lab: (B,), Pos: (N, 3), Batch: (N,)
pos, feat, batch = positions, features, batch_indices
# TransformNet:
x = pos # Don't use the normals!
x = self.transform_1(x, batch) # (N, 3) -> (N, 128)
x = self.transform_2(x) # (N, 128) -> (N, 1024)
x = global_max_pool(x, batch) # (B, 1024)
x = self.transform_3(x) # (B, 256)
x = self.transform_4(x) # (B, 3*3)
x = x[batch] # (N, 3*3)
x = x.view(-1, 3, 3) # (N, 3, 3)
# Apply the transform:
x0 = torch.einsum("ni,nij->nj", pos, x) # (N, 3)
# Add features to coordinates
x = torch.cat([x0, feat], dim=-1).contiguous()
for i in range(self.n_layers):
x_i = self.conv_layers[i](x, batch)
x_i = self.linear_layers[i](x_i)
x = self.linear_transform[i](x)
x = x + x_i
return x
# Reference PointNet models, from the PyTorch_geometric gallery =========================
class SAModule(torch.nn.Module):
"""Set abstraction module."""
def __init__(self, ratio, r, nn, max_num_neighbors=64):
super(SAModule, self).__init__()
self.ratio = ratio
self.r = r
self.conv = PointConv(nn)
self.max_num_neighbors = max_num_neighbors
def forward(self, x, pos, batch):
# Subsample with Farthest Point Sampling:
# idx = fps(pos, batch, ratio=self.ratio) # Extract self.ratio indices TURN OFF FOR NOW
idx = torch.arange(0, len(pos), device=pos.device)
# For each "cluster", get the list of (up to 64) neighbors in a ball of radius r:
row, col = radius(
pos,
pos[idx],
self.r,
batch,
batch[idx],
max_num_neighbors=self.max_num_neighbors,
)
# Applies the PointNet++ Conv:
edge_index = torch.stack([col, row], dim=0)
x = self.conv(x, (pos, pos[idx]), edge_index)
# Return the features and sub-sampled point clouds:
pos, batch = pos[idx], batch[idx]
return x, pos, batch
class GlobalSAModule(torch.nn.Module):
def __init__(self, nn):
super(GlobalSAModule, self).__init__()
self.nn = nn
def forward(self, x, pos, batch):
x = self.nn(torch.cat([x, pos], dim=1))
x = global_max_pool(x, batch)
pos = pos.new_zeros((x.size(0), 3))
batch = torch.arange(x.size(0), device=batch.device)
return x, pos, batch
class FPModule(torch.nn.Module):
def __init__(self, k, nn):
super(FPModule, self).__init__()
self.k = k
self.nn = nn
def forward(self, x, pos, batch, x_skip, pos_skip, batch_skip):
x = knn_interpolate(x, pos, pos_skip, batch, batch_skip, k=self.k)
if x_skip is not None:
x = torch.cat([x, x_skip], dim=1)
x = self.nn(x)
return x, pos_skip, batch_skip
class PointNet2_seg(torch.nn.Module):
def __init__(self, args, in_channels, out_channels):
super(PointNet2_seg, self).__init__()
self.name = "PointNet2"
self.I, self.O = in_channels, out_channels
self.radius = args.radius
self.k = 10000 # We don't restrict the number of points in a patch
self.n_layers = args.n_layers
# self.sa1_module = SAModule(1.0, self.radius, MLP([self.I+3, self.O, self.O]),self.k)
self.layers = nn.ModuleList(
[SAModule(1.0, self.radius, MLP([self.I + 3, self.O, self.O]), self.k)]
+ [
SAModule(1.0, self.radius, MLP([self.O + 3, self.O, self.O]), self.k)
for i in range(self.n_layers - 1)
]
)
self.linear_layers = nn.ModuleList(
[
nn.Sequential(
nn.Linear(self.O, self.O), nn.ReLU(), nn.Linear(self.O, self.O)
)
for i in range(self.n_layers)
]
)
self.linear_transform = nn.ModuleList(
[nn.Linear(self.I, self.O)]
+ [nn.Linear(self.O, self.O) for i in range(self.n_layers - 1)]
)
def forward(self, positions, features, batch_indices):
x = (features, positions, batch_indices)
for i, layer in enumerate(self.layers):
x_i, pos, b_ind = layer(*x)
x_i = self.linear_layers[i](x_i)
x = self.linear_transform[i](x[0])
x = x + x_i
x = (x, pos, b_ind)
return x[0]
## TangentConv benchmark segmentation
class dMaSIFConv_seg(torch.nn.Module):
def __init__(self, args, in_channels, out_channels, n_layers, radius=9.0):
super(dMaSIFConv_seg, self).__init__()
self.name = "dMaSIFConv_seg_keops"
self.radius = radius
self.I, self.O = in_channels, out_channels
self.layers = nn.ModuleList(
[dMaSIFConv(self.I, self.O, radius, self.O)]
+ [dMaSIFConv(self.O, self.O, radius, self.O) for i in range(n_layers - 1)]
)
self.linear_layers = nn.ModuleList(
[
nn.Sequential(
nn.Linear(self.O, self.O), nn.ReLU(), nn.Linear(self.O, self.O)
)
for i in range(n_layers)
]
)
self.linear_transform = nn.ModuleList(
[nn.Linear(self.I, self.O)]
+ [nn.Linear(self.O, self.O) for i in range(n_layers - 1)]
)
def forward(self, features):
# Lab: (B,), Pos: (N, 3), Batch: (N,)
points, nuv, ranges = self.points, self.nuv, self.ranges
x = features
for i, layer in enumerate(self.layers):
x_i = layer(points, nuv, x, ranges)
x_i = self.linear_layers[i](x_i)
x = self.linear_transform[i](x)
x = x + x_i
return x
def load_mesh(self, xyz, triangles=None, normals=None, weights=None, batch=None):
"""Loads the geometry of a triangle mesh.
Input arguments:
- xyz, a point cloud encoded as an (N, 3) Tensor.
- triangles, a connectivity matrix encoded as an (N, 3) integer tensor.
- weights, importance weights for the orientation estimation, encoded as an (N, 1) Tensor.
- radius, the scale used to estimate the local normals.
- a batch vector, following PyTorch_Geometric's conventions.
The routine updates the model attributes:
- points, i.e. the point cloud itself,
- nuv, a local oriented basis in R^3 for every point,
- ranges, custom KeOps syntax to implement batch processing.
"""
# 1. Save the vertices for later use in the convolutions ---------------
self.points = xyz
self.batch = batch
self.ranges = diagonal_ranges(
batch
) # KeOps support for heterogeneous batch processing
self.triangles = triangles
self.normals = normals
self.weights = weights
# 2. Estimate the normals and tangent frame ----------------------------
# Normalize the scale:
points = xyz / self.radius
# Normals and local areas:
if normals is None:
normals, areas = mesh_normals_areas(points, triangles, 0.5, batch)
tangent_bases = tangent_vectors(normals) # Tangent basis (N, 2, 3)
# 3. Steer the tangent bases according to the gradient of "weights" ----
# 3.a) Encoding as KeOps LazyTensors:
# Orientation scores:
weights_j = LazyTensor(weights.view(1, -1, 1)) # (1, N, 1)
# Vertices:
x_i = LazyTensor(points[:, None, :]) # (N, 1, 3)
x_j = LazyTensor(points[None, :, :]) # (1, N, 3)
# Normals:
n_i = LazyTensor(normals[:, None, :]) # (N, 1, 3)
n_j = LazyTensor(normals[None, :, :]) # (1, N, 3)
# Tangent basis:
uv_i = LazyTensor(tangent_bases.view(-1, 1, 6)) # (N, 1, 6)
# 3.b) Pseudo-geodesic window:
# Pseudo-geodesic squared distance:
rho2_ij = ((x_j - x_i) ** 2).sum(-1) * ((2 - (n_i | n_j)) ** 2) # (N, N, 1)
# Gaussian window:
window_ij = (-rho2_ij).exp() # (N, N, 1)
# 3.c) Coordinates in the (u, v) basis - not oriented yet:
X_ij = uv_i.matvecmult(x_j - x_i) # (N, N, 2)
# 3.d) Local average in the tangent plane:
orientation_weight_ij = window_ij * weights_j # (N, N, 1)
orientation_vector_ij = orientation_weight_ij * X_ij # (N, N, 2)
# Support for heterogeneous batch processing:
orientation_vector_ij.ranges = self.ranges # Block-diagonal sparsity mask
orientation_vector_i = orientation_vector_ij.sum(dim=1) # (N, 2)
orientation_vector_i = (
orientation_vector_i + 1e-5
) # Just in case someone's alone...
# 3.e) Normalize stuff:
orientation_vector_i = F.normalize(orientation_vector_i, p=2, dim=-1) # (N, 2)
ex_i, ey_i = (
orientation_vector_i[:, 0][:, None],
orientation_vector_i[:, 1][:, None],
) # (N,1)
# 3.f) Re-orient the (u,v) basis:
uv_i = tangent_bases # (N, 2, 3)
u_i, v_i = uv_i[:, 0, :], uv_i[:, 1, :] # (N, 3)
tangent_bases = torch.cat(
(ex_i * u_i + ey_i * v_i, -ey_i * u_i + ex_i * v_i), dim=1
).contiguous() # (N, 6)
# 4. Store the local 3D frame as an attribute --------------------------
self.nuv = torch.cat(
(normals.view(-1, 1, 3), tangent_bases.view(-1, 2, 3)), dim=1
)
================================================
FILE: benchmark_scripts/DGCNN_site.sh
================================================
# Load environment
python -W ignore -u main_training.py --experiment_name DGCNN_site_1layer_k200 --batch_size 64 --embedding_layer DGCNN --site True --single_protein True --device cuda:0 --n_layers 1 --random_rotation True --k 200
python -W ignore -u main_training.py --experiment_name DGCNN_site_1layer_k100 --batch_size 64 --embedding_layer DGCNN --site True --single_protein True --device cuda:0 --n_layers 1 --random_rotation True --k 100
python -W ignore -u main_training.py --experiment_name DGCNN_site_3layer_k200 --batch_size 64 --embedding_layer DGCNN --site True --single_protein True --device cuda:0 --n_layers 3 --random_rotation True --k 200
python -W ignore -u main_training.py --experiment_name DGCNN_site_3layer_k100 --batch_size 64 --embedding_layer DGCNN --site True --single_protein True --device cuda:0 --n_layers 3 --random_rotation True --k 100
================================================
FILE: benchmark_scripts/Pointnet_site.sh
================================================
# Load environment
python -W ignore -u main_training.py --experiment_name PointNet_site_3layer_15A --batch_size 64 --embedding_layer PointNet++ --site True --single_protein True --device cuda:0 --random_rotation True --radius 15.0 --n_layers 3
python -W ignore -u main_training.py --experiment_name PointNet_site_3layer_5A --batch_size 64 --embedding_layer PointNet++ --site True --single_protein True --device cuda:0 --random_rotation True --radius 5.0 --n_layers 3
python -W ignore -u main_training.py --experiment_name PointNet_site_3layer_9A --batch_size 64 --embedding_layer PointNet++ --site True --single_protein True --device cuda:0 --random_rotation True --radius 9.0 --n_layers 3
python -W ignore -u main_training.py --experiment_name PointNet_site_1layer_15A --batch_size 64 --embedding_layer PointNet++ --site True --single_protein True --device cuda:0 --random_rotation True --radius 15.0 --n_layers 1
python -W ignore -u main_training.py --experiment_name PointNet_site_1layer_5A --batch_size 64 --embedding_layer PointNet++ --site True --single_protein True --device cuda:0 --random_rotation True --radius 5.0 --n_layers 1
python -W ignore -u main_training.py --experiment_name PointNet_site_1layer_9A --batch_size 64 --embedding_layer PointNet++ --site True --single_protein True --device cuda:0 --random_rotation True --radius 9.0 --n_layers 1
================================================
FILE: benchmark_scripts/dMaSIF_search.sh
================================================
# Load environment
python -W ignore -u main_training.py --experiment_name dMaSIF_search_1layer_12A --batch_size 64 --embedding_layer dMaSIF --search True --device cuda:0 --random_rotation True --radius 12.0 --n_layers 1
python -W ignore -u main_training.py --experiment_name dMaSIF_search_3layer_12A --batch_size 64 --embedding_layer dMaSIF --search True --device cuda:0 --random_rotation True --radius 12.0 --n_layers 3
================================================
FILE: benchmark_scripts/dMaSIF_site.sh
================================================
# Load environment
python -W ignore -u main_training.py --experiment_name dMaSIF_site_1layer_15A --batch_size 64 --embedding_layer dMaSIF --site True --single_protein True --random_rotation True --radius 15.0 --n_layers 1
python -W ignore -u main_training.py --experiment_name dMaSIF_site_1layer_5A --batch_size 64 --embedding_layer dMaSIF --site True --single_protein True --random_rotation True --radius 5.0 --n_layers 1
python -W ignore -u main_training.py --experiment_name dMaSIF_site_1layer_9A --batch_size 64 --embedding_layer dMaSIF --site True --single_protein True --random_rotation True --radius 9.0 --n_layers 1
python -W ignore -u main_training.py --experiment_name dMaSIF_site_3layer_15A --batch_size 64 --embedding_layer dMaSIF --site True --single_protein True --random_rotation True --radius 15.0 --n_layers 3
python -W ignore -u main_training.py --experiment_name dMaSIF_site_3layer_5A --batch_size 64 --embedding_layer dMaSIF --site True --single_protein True --random_rotation True --radius 5.0 --n_layers 3
python -W ignore -u main_training.py --experiment_name dMaSIF_site_3layer_9A --batch_size 64 --embedding_layer dMaSIF --site True --single_protein True --random_rotation True --radius 9.0 --n_layers 3
================================================
FILE: data.py
================================================
import torch
from torch_geometric.data import InMemoryDataset, Data, DataLoader
from torch_geometric.transforms import Compose
import numpy as np
from scipy.spatial.transform import Rotation
import math
import urllib.request
import tarfile
from pathlib import Path
import requests
from data_preprocessing.convert_pdb2npy import convert_pdbs
from data_preprocessing.convert_ply2npy import convert_plys
tensor = torch.FloatTensor
inttensor = torch.LongTensor
def numpy(x):
return x.detach().cpu().numpy()
def iface_valid_filter(protein_pair):
labels1 = protein_pair.y_p1.reshape(-1)
labels2 = protein_pair.y_p2.reshape(-1)
valid1 = (
(torch.sum(labels1) < 0.75 * len(labels1))
and (torch.sum(labels1) > 30)
and (torch.sum(labels1) > 0.01 * labels2.shape[0])
)
valid2 = (
(torch.sum(labels2) < 0.75 * len(labels2))
and (torch.sum(labels2) > 30)
and (torch.sum(labels2) > 0.01 * labels1.shape[0])
)
return valid1 and valid2
class RandomRotationPairAtoms(object):
r"""Randomly rotate a protein"""
def __call__(self, data):
R1 = tensor(Rotation.random().as_matrix())
R2 = tensor(Rotation.random().as_matrix())
data.atom_coords_p1 = torch.matmul(R1, data.atom_coords_p1.T).T
data.xyz_p1 = torch.matmul(R1, data.xyz_p1.T).T
data.normals_p1 = torch.matmul(R1, data.normals_p1.T).T
data.atom_coords_p2 = torch.matmul(R2, data.atom_coords_p2.T).T
data.xyz_p2 = torch.matmul(R2, data.xyz_p2.T).T
data.normals_p2 = torch.matmul(R2, data.normals_p2.T).T
data.rand_rot1 = R1
data.rand_rot2 = R2
return data
def __repr__(self):
return "{}()".format(self.__class__.__name__)
class CenterPairAtoms(object):
r"""Centers a protein"""
def __call__(self, data):
atom_center1 = data.atom_coords_p1.mean(dim=-2, keepdim=True)
atom_center2 = data.atom_coords_p2.mean(dim=-2, keepdim=True)
data.atom_coords_p1 = data.atom_coords_p1 - atom_center1
data.atom_coords_p2 = data.atom_coords_p2 - atom_center2
data.xyz_p1 = data.xyz_p1 - atom_center1
data.xyz_p2 = data.xyz_p2 - atom_center2
data.atom_center1 = atom_center1
data.atom_center2 = atom_center2
return data
def __repr__(self):
return "{}()".format(self.__class__.__name__)
class NormalizeChemFeatures(object):
r"""Centers a protein"""
def __call__(self, data):
pb_upper = 3.0
pb_lower = -3.0
chem_p1 = data.chemical_features_p1
chem_p2 = data.chemical_features_p2
pb_p1 = chem_p1[:, 0]
pb_p2 = chem_p2[:, 0]
hb_p1 = chem_p1[:, 1]
hb_p2 = chem_p2[:, 1]
hp_p1 = chem_p1[:, 2]
hp_p2 = chem_p2[:, 2]
# Normalize PB
pb_p1 = torch.clamp(pb_p1, pb_lower, pb_upper)
pb_p1 = (pb_p1 - pb_lower) / (pb_upper - pb_lower)
pb_p1 = 2 * pb_p1 - 1
pb_p2 = torch.clamp(pb_p2, pb_lower, pb_upper)
pb_p2 = (pb_p2 - pb_lower) / (pb_upper - pb_lower)
pb_p2 = 2 * pb_p2 - 1
# Normalize HP
hp_p1 = hp_p1 / 4.5
hp_p2 = hp_p2 / 4.5
data.chemical_features_p1 = torch.stack([pb_p1, hb_p1, hp_p1]).T
data.chemical_features_p2 = torch.stack([pb_p2, hb_p2, hp_p2]).T
return data
def __repr__(self):
return "{}()".format(self.__class__.__name__)
def load_protein_npy(pdb_id, data_dir, center=False, single_pdb=False):
"""Loads a protein surface mesh and its features"""
# Load the data, and read the connectivity information:
triangles = (
None
if single_pdb
else inttensor(np.load(data_dir / (pdb_id + "_triangles.npy"))).T
)
# Normalize the point cloud, as specified by the user:
points = None if single_pdb else tensor(np.load(data_dir / (pdb_id + "_xyz.npy")))
center_location = None if single_pdb else torch.mean(points, axis=0, keepdims=True)
atom_coords = tensor(np.load(data_dir / (pdb_id + "_atomxyz.npy")))
atom_types = tensor(np.load(data_dir / (pdb_id + "_atomtypes.npy")))
if center:
points = points - center_location
atom_coords = atom_coords - center_location
# Interface labels
iface_labels = (
None
if single_pdb
else tensor(np.load(data_dir / (pdb_id + "_iface_labels.npy")).reshape((-1, 1)))
)
# Features
chemical_features = (
None if single_pdb else tensor(np.load(data_dir / (pdb_id + "_features.npy")))
)
# Normals
normals = (
None if single_pdb else tensor(np.load(data_dir / (pdb_id + "_normals.npy")))
)
protein_data = Data(
xyz=points,
face=triangles,
chemical_features=chemical_features,
y=iface_labels,
normals=normals,
center_location=center_location,
num_nodes=None if single_pdb else points.shape[0],
atom_coords=atom_coords,
atom_types=atom_types,
)
return protein_data
class PairData(Data):
def __init__(
self,
xyz_p1=None,
xyz_p2=None,
face_p1=None,
face_p2=None,
chemical_features_p1=None,
chemical_features_p2=None,
y_p1=None,
y_p2=None,
normals_p1=None,
normals_p2=None,
center_location_p1=None,
center_location_p2=None,
atom_coords_p1=None,
atom_coords_p2=None,
atom_types_p1=None,
atom_types_p2=None,
atom_center1=None,
atom_center2=None,
rand_rot1=None,
rand_rot2=None,
):
super().__init__()
self.xyz_p1 = xyz_p1
self.xyz_p2 = xyz_p2
self.face_p1 = face_p1
self.face_p2 = face_p2
self.chemical_features_p1 = chemical_features_p1
self.chemical_features_p2 = chemical_features_p2
self.y_p1 = y_p1
self.y_p2 = y_p2
self.normals_p1 = normals_p1
self.normals_p2 = normals_p2
self.center_location_p1 = center_location_p1
self.center_location_p2 = center_location_p2
self.atom_coords_p1 = atom_coords_p1
self.atom_coords_p2 = atom_coords_p2
self.atom_types_p1 = atom_types_p1
self.atom_types_p2 = atom_types_p2
self.atom_center1 = atom_center1
self.atom_center2 = atom_center2
self.rand_rot1 = rand_rot1
self.rand_rot2 = rand_rot2
def __inc__(self, key, value):
if key == "face_p1":
return self.xyz_p1.size(0)
if key == "face_p2":
return self.xyz_p2.size(0)
else:
return super(PairData, self).__inc__(key, value)
def __cat_dim__(self, key, value):
if ("index" in key) or ("face" in key):
return 1
else:
return 0
def load_protein_pair(pdb_id, data_dir,single_pdb=False):
"""Loads a protein surface mesh and its features"""
pspl = pdb_id.split("_")
p1_id = pspl[0] + "_" + pspl[1]
p2_id = pspl[0] + "_" + pspl[2]
p1 = load_protein_npy(p1_id, data_dir, center=False,single_pdb=single_pdb)
p2 = load_protein_npy(p2_id, data_dir, center=False,single_pdb=single_pdb)
# pdist = ((p1['xyz'][:,None,:]-p2['xyz'][None,:,:])**2).sum(-1).sqrt()
# pdist = pdist<2.0
# y_p1 = (pdist.sum(1)>0).to(torch.float).reshape(-1,1)
# y_p2 = (pdist.sum(0)>0).to(torch.float).reshape(-1,1)
y_p1 = p1["y"]
y_p2 = p2["y"]
protein_pair_data = PairData(
xyz_p1=p1["xyz"],
xyz_p2=p2["xyz"],
face_p1=p1["face"],
face_p2=p2["face"],
chemical_features_p1=p1["chemical_features"],
chemical_features_p2=p2["chemical_features"],
y_p1=y_p1,
y_p2=y_p2,
normals_p1=p1["normals"],
normals_p2=p2["normals"],
center_location_p1=p1["center_location"],
center_location_p2=p2["center_location"],
atom_coords_p1=p1["atom_coords"],
atom_coords_p2=p2["atom_coords"],
atom_types_p1=p1["atom_types"],
atom_types_p2=p2["atom_types"],
)
return protein_pair_data
class ProteinPairsSurfaces(InMemoryDataset):
url = ""
def __init__(self, root, ppi=False, train=True, transform=None, pre_transform=None):
self.ppi = ppi
super(ProteinPairsSurfaces, self).__init__(root, transform, pre_transform)
path = self.processed_paths[0] if train else self.processed_paths[1]
self.data, self.slices = torch.load(path)
@property
def raw_file_names(self):
return "masif_site_masif_search_pdbs_and_ply_files.tar.gz"
@property
def processed_file_names(self):
if not self.ppi:
file_names = [
"training_pairs_data.pt",
"testing_pairs_data.pt",
"training_pairs_data_ids.npy",
"testing_pairs_data_ids.npy",
]
else:
file_names = [
"training_pairs_data_ppi.pt",
"testing_pairs_data_ppi.pt",
"training_pairs_data_ids_ppi.npy",
"testing_pairs_data_ids_ppi.npy",
]
return file_names
def download(self):
url = 'https://zenodo.org/record/2625420/files/masif_site_masif_search_pdbs_and_ply_files.tar.gz'
target_path = self.raw_paths[0]
response = requests.get(url, stream=True)
if response.status_code == 200:
with open(target_path, 'wb') as f:
f.write(response.raw.read())
#raise RuntimeError(
# "Dataset not found. Please download {} from {} and move it to {}".format(
# self.raw_file_names, self.url, self.raw_dir
# )
#)
def process(self):
pdb_dir = Path(self.root) / "raw" / "01-benchmark_pdbs"
surf_dir = Path(self.root) / "raw" / "01-benchmark_surfaces"
protein_dir = Path(self.root) / "raw" / "01-benchmark_surfaces_npy"
lists_dir = Path('./lists')
# Untar surface files
if not (pdb_dir.exists() and surf_dir.exists()):
tar = tarfile.open(self.raw_paths[0])
tar.extractall(self.raw_dir)
tar.close()
if not protein_dir.exists():
protein_dir.mkdir(parents=False, exist_ok=False)
convert_plys(surf_dir,protein_dir)
convert_pdbs(pdb_dir,protein_dir)
with open(lists_dir / "training.txt") as f_tr, open(
lists_dir / "testing.txt"
) as f_ts:
training_list = sorted(f_tr.read().splitlines())
testing_list = sorted(f_ts.read().splitlines())
with open(lists_dir / "training_ppi.txt") as f_tr, open(
lists_dir / "testing_ppi.txt"
) as f_ts:
training_pairs_list = sorted(f_tr.read().splitlines())
testing_pairs_list = sorted(f_ts.read().splitlines())
pairs_list = sorted(training_pairs_list + testing_pairs_list)
if not self.ppi:
training_pairs_list = []
for p in pairs_list:
pspl = p.split("_")
p1 = pspl[0] + "_" + pspl[1]
p2 = pspl[0] + "_" + pspl[2]
if p1 in training_list:
training_pairs_list.append(p)
if p2 in training_list:
training_pairs_list.append(pspl[0] + "_" + pspl[2] + "_" + pspl[1])
testing_pairs_list = []
for p in pairs_list:
pspl = p.split("_")
p1 = pspl[0] + "_" + pspl[1]
p2 = pspl[0] + "_" + pspl[2]
if p1 in testing_list:
testing_pairs_list.append(p)
if p2 in testing_list:
testing_pairs_list.append(pspl[0] + "_" + pspl[2] + "_" + pspl[1])
# # Read data into huge `Data` list.
training_pairs_data = []
training_pairs_data_ids = []
for p in training_pairs_list:
try:
protein_pair = load_protein_pair(p, protein_dir)
except FileNotFoundError:
continue
training_pairs_data.append(protein_pair)
training_pairs_data_ids.append(p)
testing_pairs_data = []
testing_pairs_data_ids = []
for p in testing_pairs_list:
try:
protein_pair = load_protein_pair(p, protein_dir)
except FileNotFoundError:
continue
testing_pairs_data.append(protein_pair)
testing_pairs_data_ids.append(p)
if self.pre_filter is not None:
training_pairs_data = [
data for data in training_pairs_data if self.pre_filter(data)
]
testing_pairs_data = [
data for data in testing_pairs_data if self.pre_filter(data)
]
if self.pre_transform is not None:
training_pairs_data = [
self.pre_transform(data) for data in training_pairs_data
]
testing_pairs_data = [
self.pre_transform(data) for data in testing_pairs_data
]
training_pairs_data, training_pairs_slices = self.collate(training_pairs_data)
torch.save(
(training_pairs_data, training_pairs_slices), self.processed_paths[0]
)
np.save(self.processed_paths[2], training_pairs_data_ids)
testing_pairs_data, testing_pairs_slices = self.collate(testing_pairs_data)
torch.save((testing_pairs_data, testing_pairs_slices), self.processed_paths[1])
np.save(self.processed_paths[3], testing_pairs_data_ids)
================================================
FILE: data_analysis/analyse_atomnet.ipynb
================================================
{
"metadata": {
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.7-final"
},
"orig_nbformat": 2,
"kernelspec": {
"name": "python37764bitvenvvenv0f8ab664edda4e43b2faa48ba983dc72",
"display_name": "Python 3.7.7 64-bit ('venv': venv)"
}
},
"nbformat": 4,
"nbformat_minor": 2,
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import seaborn as sns\n",
"from sklearn.metrics import mean_squared_error"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"0.8289113443652403\n0.16479157\n"
]
},
{
"output_type": "display_data",
"data": {
"text/plain": "",
"image/svg+xml": "\n\n\n
