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