Working with Protein Tensors in Graphein¶
This tutorial demonstrates working with protein tensors consistent with the PyTorch Geometric API which is supported as of Graphein 1.6.0.
[32]:
#!pip install graphein
import graphein
graphein.verbose(False)
[33]:
%load_ext watermark
%watermark
print("Graphein version: ", graphein.__version__)
The watermark extension is already loaded. To reload it, use:
%reload_ext watermark
Last updated: 2023-03-30T23:14:11.142537+01:00
Python implementation: CPython
Python version : 3.9.16
IPython version : 8.11.0
Compiler : GCC 11.3.0
OS : Linux
Release : 5.15.0-67-generic
Machine : x86_64
Processor : x86_64
CPU cores : 64
Architecture: 64bit
Graphein version: 1.6.0
1. Parsing proteins to Tensors¶
In [graphein.protein.tensor.io] there are several functions for parsing Pandas dataframes from BioPandas into tensors:
We’ll use the following example:
[34]:
from biopandas.pdb import PandasPdb
p = PandasPdb().fetch_pdb("4hhb")
df = p.df["ATOM"]
df.head()
[34]:
| record_name | atom_number | blank_1 | atom_name | alt_loc | residue_name | blank_2 | chain_id | residue_number | insertion | ... | x_coord | y_coord | z_coord | occupancy | b_factor | blank_4 | segment_id | element_symbol | charge | line_idx | |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 0 | ATOM | 1 | N | VAL | A | 1 | ... | 19.323 | 29.727 | 42.781 | 1.0 | 49.05 | N | NaN | 883 | ||||||
| 1 | ATOM | 2 | CA | VAL | A | 1 | ... | 20.141 | 30.469 | 42.414 | 1.0 | 43.14 | C | NaN | 884 | ||||||
| 2 | ATOM | 3 | C | VAL | A | 1 | ... | 21.664 | 29.857 | 42.548 | 1.0 | 24.80 | C | NaN | 885 | ||||||
| 3 | ATOM | 4 | O | VAL | A | 1 | ... | 21.985 | 29.541 | 43.704 | 1.0 | 37.68 | O | NaN | 886 | ||||||
| 4 | ATOM | 5 | CB | VAL | A | 1 | ... | 19.887 | 31.918 | 43.524 | 1.0 | 72.12 | C | NaN | 887 |
5 rows × 21 columns
[35]:
import graphein.protein.tensor as gpt
coords = gpt.io.protein_df_to_tensor(df)
print(coords.shape)
gpt.plot.plot_structure(coords)
torch.Size([574, 37, 3])
Data type cannot be displayed: application/vnd.plotly.v1+json
This parses the atomic coordinates into a tensor of shape (Num Residues x Atom Types (default = 37) x 3). The number of the atoms is determined by:
[36]:
from graphein.protein.resi_atoms import ATOM_NUMBERING
ATOM_NUMBERING
[36]:
{'N': 0,
'CA': 1,
'C': 2,
'O': 3,
'CB': 4,
'OG': 5,
'CG': 6,
'CD1': 7,
'CD2': 8,
'CE1': 9,
'CE2': 10,
'CZ': 11,
'OD1': 12,
'ND2': 13,
'CG1': 14,
'CG2': 15,
'CD': 16,
'CE': 17,
'NZ': 18,
'OD2': 19,
'OE1': 20,
'NE2': 21,
'OE2': 22,
'OH': 23,
'NE': 24,
'NH1': 25,
'NH2': 26,
'OG1': 27,
'SD': 28,
'ND1': 29,
'SG': 30,
'NE1': 31,
'CE3': 32,
'CZ2': 33,
'CZ3': 34,
'CH2': 35,
'OXT': 36}
Similarly, we can get a tensor indicating chain membership:
[37]:
chains = gpt.io.protein_df_to_chain_tensor(df)
print(chains.shape)
print(chains)
chains = gpt.io.protein_df_to_chain_tensor(df, one_hot=True)
print(chains)
print(chains.shape)
torch.Size([574])
tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3,
3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3])
tensor([[1, 0, 0, 0],
[1, 0, 0, 0],
[1, 0, 0, 0],
...,
[0, 0, 0, 1],
[0, 0, 0, 1],
[0, 0, 0, 1]])
torch.Size([574, 4])
2. Parsing to PyG Data¶
We can also parse a pdb code, pdb path or uniprot ID (via AF2 API) to a PyTorch Geometric Data object:
[38]:
import graphein.protein.tensor as gpt
data = gpt.io.protein_to_pyg(
pdb_code="4hhb", # Can alternatively pass a path or a uniprot ID (for AF2) with path=... and uniprot_id=...
chain_selection=["A", "B", "C", "D"], # Select all 4 chains
deprotonate=True, # Deprotonate the structure
keep_insertions=False, # Remove insertions
keep_hets=[], # Remove HETATMs
model_index=1, # Select the first model
# Can select a subset of atoms with atom_types=...
)
print("Data: ", data)
print("ID: ", data.id)
print("Residues: ", data.residues[:10])
print("Residue IDs: ", data.residue_id[:10])
Data: Data(coords=[574, 37, 3], residues=[574], id='4hhb_ABCD', residue_id=[574], residue_type=[574], chains=[574])
ID: 4hhb_ABCD
Residues: ['VAL', 'LEU', 'SER', 'PRO', 'ALA', 'ASP', 'LYS', 'THR', 'ASN', 'VAL']
Residue IDs: ['A:VAL:1', 'A:LEU:2', 'A:SER:3', 'A:PRO:4', 'A:ALA:5', 'A:ASP:6', 'A:LYS:7', 'A:THR:8', 'A:ASN:9', 'A:VAL:10']
[39]:
DataBatch(coords=[1722, 37, 3], residues=[3], id=[3], residue_id=[3], residue_type=[1722], chains=[1722])
3. Parsing to Graphein Protein¶
Graphein provides a Protein class which wraps the PyG data classes with useful protein-specific methods. There are several ways to parse a structure into this format:
[40]:
from rich import inspect
inspect(gpt.Protein)
╭──────────────────────────────── <class 'graphein.protein.tensor.data.Protein'> ─────────────────────────────────╮ │ class Protein(atom_list: List[str] = ['N', 'CA', 'C', 'O', 'CB', 'OG', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', │ │ 'OD1', 'ND2', 'CG1', 'CG2', 'CD', 'CE', 'NZ', 'OD2', 'OE1', 'NE2', 'OE2', 'OH', 'NE', 'NH1', 'NH2', 'OG1', │ │ 'SD', 'ND1', 'SG', 'NE1', 'CE3', 'CZ2', 'CZ3', 'CH2', 'OXT'], fill_value: float = 1e-05) -> None: │ │ │ │ "A data object describing a homogeneous graph. ``Protein`` inherits from │ │ :ref:`torch_geometric.data.Data` and implements structural-biology specific │ │ methods in addition. ``Data`` and ``Protein`` objects can be directly │ │ converted between using ``Protein().from_data()`` and │ │ ``Protein().to_data()``. Think of ``Protein`` as an equivalent to ``Data`` │ │ but with some protein-specific trimmings. │ │ │ │ batch = <property object at 0x7f95ce5a0c20> │ │ edge_attr = <property object at 0x7f95ce5a0b30> │ │ edge_index = <property object at 0x7f95ce5a0a90> │ │ edge_stores = <property object at 0x7f95ce5a0590> │ │ edge_weight = <property object at 0x7f95ce5a0ae0> │ │ is_cuda = <property object at 0x7f95ce58ef90> │ │ keys = <property object at 0x7f95ce58eea0> │ │ node_stores = <property object at 0x7f95ce5a0360> │ │ num_edge_features = <property object at 0x7f95ce5a09f0> │ │ num_edges = <property object at 0x7f95ce58ef40> │ │ num_faces = <property object at 0x7f95ce5a0c70> │ │ num_features = <property object at 0x7f95ce5a09a0> │ │ num_node_features = <property object at 0x7f95ce5a08b0> │ │ num_nodes = <property object at 0x7f95ce58eef0> │ │ pos = <property object at 0x7f95ce5a0bd0> │ │ stores = <property object at 0x7f95ce5a0810> │ │ x = <property object at 0x7f95ce5a0a40> │ │ y = <property object at 0x7f95ce5a0b80> │ ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
[41]:
import graphein.protein.tensor as gpt
# 1. From a PDB code
protein = gpt.Protein()
protein.from_pdb_code("4hhb", chain_selection=["A", "B", "C", "D"])
# 2. From a PDB file
protein = gpt.Protein()
#protein.from_pdb_file("./examples/pdbs/3eiy.pdb")
# 3. From a BioPandas DataFrame
p = PandasPdb().fetch_pdb("4hhb")
df = p.df["ATOM"]
protein = gpt.Protein().from_dataframe(df)
# 4. From a PyG Data object
data = gpt.io.protein_to_pyg(
pdb_code="4hhb", # Can alternatively pass a path or a uniprot ID (for AF2) with path=... and uniprot_id=...
chain_selection=["A", "B", "C", "D"], # Select all 4 chains
deprotonate=True, # Deprotonate the structure
keep_insertions=False, # Remove insertions
keep_hets=[], # Remove HETATMs
model_index=1, # Select the first model
# Can select a subset of atoms with atom_types=...
)
protein = gpt.Protein().from_data(data)
print(protein)
Protein(fill_value=1e-05, atom_list=[37], residue_type=[574], id='4hhb_ABCD', residue_id=[574], chains=[574], coords=[574, 37, 3], residues=[574])
Selecting Views and Representations¶
[42]:
protein.alpha_carbon(cache="ca")
protein.full_atom_coords(cache="fa")
protein.backbone_frames(cache="frame")
print("Alpha Carbon coordinates: ", protein.ca.shape)
print("Backbone coordinates: ", protein.backbone().shape) # TODO
print("Full atom coords: ", protein.fa[0].shape)
print("Backbone Frames: (rotations) ", protein.frame[0].shape)
print("Backbone Frames: (translation) ", protein.frame[1].shape)
Alpha Carbon coordinates: torch.Size([574, 3])
Backbone coordinates: torch.Size([574, 4, 3])
Full atom coords: torch.Size([4384, 3])
Backbone Frames: (rotations) torch.Size([574, 3, 3])
Backbone Frames: (translation) torch.Size([574, 1, 3])
Computing Edges¶
[43]:
protein.edges("eps_8", cache="eps")
protein.edges("knn_16", cache="knn")
print("Radius Edges: ", protein.eps.shape)
print("KNN Edges: ", protein.knn.shape)
protein.edges("eps_8", cache="eps", loop=True)
protein.edges("knn_16", cache="knn", loop=True)
print("Radius Edges: ", protein.eps.shape)
print("KNN Edges: ", protein.knn.shape)
Radius Edges: torch.Size([2, 5634])
KNN Edges: torch.Size([2, 9184])
Radius Edges: torch.Size([2, 6208])
KNN Edges: torch.Size([2, 9184])
[44]:
### Computing Edge Distances
protein.edge_distances(protein.ca, protein.knn).shape
[44]:
torch.Size([9184])
Computing Features¶
[45]:
# Angles
print("Dihedral angles:" , protein.dihedrals())
print("Sidechain Torsion angles: ", protein.sidechain_torsion())
print("Kappa: ", protein.kappa(rad=False, embed=False))
print("Alpha", protein.alpha(embed=False))
Dihedral angles: tensor([[ 1.0000, 0.0000, 0.9463, 0.3232, 0.1277, -0.9918],
[ 0.2636, 0.9646, -0.5480, 0.8365, -0.0836, 0.9965],
[-0.4605, 0.8877, 0.9947, -0.1024, -0.9744, 0.2249],
...,
[-0.8991, -0.4378, -0.5160, 0.8566, -0.6029, -0.7978],
[-0.9993, 0.0365, -0.0748, -0.9972, 0.9891, -0.1474],
[-0.6320, 0.7750, 1.0000, 0.0000, 1.0000, 0.0000]])
Sidechain Torsion angles: tensor([[ 0.9987, -0.0519, 1.0000, ..., 0.0000, 1.0000, 0.0000],
[ 0.9996, -0.0280, 1.0000, ..., 0.0000, 1.0000, 0.0000],
[ 0.9997, 0.0230, 1.0000, ..., 0.0000, 1.0000, 0.0000],
...,
[ 1.0000, -0.0071, 0.9994, ..., -0.0541, 0.9991, 0.0427],
[ 0.9997, -0.0246, 0.9994, ..., 0.0000, 1.0000, 0.0000],
[ 0.9999, 0.0168, 0.9995, ..., 0.0000, 1.0000, 0.0000]])
Kappa: tensor([ 0.0000, 0.0000, 29.4452, 124.3686, 102.3605, 105.1655, 107.3162,
111.5782, 113.1725, 113.7545, 113.4106, 114.2164, 116.8045, 104.4386,
104.9213, 102.4106, 84.3757, 71.5867, 122.8545, 66.6200, 76.8520,
110.8692, 113.7741, 109.5457, 109.9186, 110.1803, 107.4435, 106.8849,
111.2042, 110.0257, 109.5954, 109.5770, 116.5330, 121.7168, 129.2709,
67.6034, 70.8772, 86.2147, 81.8529, 75.6061, 90.3254, 109.3141,
58.3846, 74.0334, 95.8026, 80.0438, 27.2018, 36.8015, 75.7786,
112.4792, 83.7933, 48.7484, 108.6910, 110.3211, 113.5658, 112.2140,
111.9837, 106.3771, 105.7953, 108.6027, 109.4749, 112.2587, 105.7999,
106.2173, 109.5791, 108.2811, 108.0186, 109.0980, 107.6228, 110.0023,
121.9857, 73.7189, 73.5920, 120.6167, 66.8503, 68.7163, 106.9677,
118.6948, 116.4287, 83.6626, 74.6209, 113.9280, 108.6968, 107.7958,
112.9154, 113.7857, 121.0273, 125.8149, 134.0292, 124.2631, 74.0040,
7.9420, 35.1083, 30.7914, 103.5490, 87.2995, 86.6578, 98.6479,
106.6020, 111.5976, 114.1221, 106.9308, 111.0820, 109.5989, 108.8881,
110.6327, 110.3807, 111.3263, 111.7163, 108.3039, 120.6460, 115.2387,
56.1405, 78.3323, 118.9786, 89.7774, 54.2490, 59.7364, 123.6486,
111.6888, 109.9301, 109.8677, 112.2128, 113.8482, 108.5114, 112.2281,
109.2737, 110.0861, 108.2194, 110.4529, 109.4557, 111.4844, 109.7848,
109.9008, 111.3323, 115.0153, 90.3321, 80.7043, 82.1504, 86.2720,
87.2934, 84.9535, 69.2108, 39.6504, 30.0937, 127.5169, 113.4369,
112.3019, 110.4729, 107.1005, 110.5840, 111.2768, 108.3251, 110.6971,
108.3146, 93.5624, 102.3668, 89.3321, 63.9790, 22.4254, 85.1417,
114.8967, 118.2151, 113.9653, 109.7958, 114.5515, 112.2329, 113.8213,
114.4161, 111.0322, 107.9713, 109.1654, 117.2920, 117.7615, 129.8250,
75.2088, 71.8777, 81.1670, 78.9547, 67.2388, 86.7595, 112.3475,
55.1522, 72.0690, 87.5263, 96.6408, 67.2954, 53.7085, 38.7398,
74.5141, 77.5260, 118.9177, 111.0021, 112.5657, 111.4128, 116.8910,
84.6676, 57.2858, 95.5644, 108.6854, 110.1830, 112.8179, 109.9710,
107.6189, 107.4156, 106.8071, 112.3577, 114.7365, 105.8917, 107.8107,
110.5926, 107.2135, 107.5485, 105.2241, 90.2548, 94.4620, 106.9309,
65.8641, 72.8784, 122.7250, 71.7637, 74.4414, 113.5534, 115.0589,
120.6483, 87.4814, 83.4399, 114.1236, 106.3206, 107.0970, 109.2812,
108.5372, 115.5591, 121.4693, 134.8265, 128.3885, 81.2117, 31.7297,
38.3107, 23.7137, 99.5002, 82.9109, 94.5869, 108.8382, 110.8891,
111.4456, 108.8123, 109.6006, 112.7954, 109.4395, 106.8459, 110.8349,
111.0942, 107.7398, 111.7465, 110.7079, 118.6557, 122.4717, 82.7955,
73.4489, 122.0826, 80.5036, 77.7440, 41.5167, 124.2869, 111.6792,
109.2399, 106.9950, 106.0186, 111.2928, 106.5566, 105.5396, 114.7271,
112.2458, 111.9217, 111.4784, 110.0751, 111.3186, 114.5204, 111.1401,
111.7100, 115.2878, 88.0681, 87.7466, 79.6632, 78.4938, 63.0532,
57.3877, 40.7848, 29.7910, 126.4864, 103.3695, 108.3819, 108.6733,
110.5776, 112.4978, 111.3139, 112.9073, 111.3377, 117.5661, 111.0254,
103.8846, 107.1842, 79.8694, 67.4254, 122.0520, 69.0583, 77.5331,
111.0134, 113.8138, 106.5949, 108.2526, 110.3287, 107.2449, 110.2363,
111.7854, 110.5423, 112.2503, 109.8368, 116.1135, 122.3657, 127.9525,
68.7193, 71.3466, 82.6621, 82.8563, 78.6296, 88.8570, 109.7955,
56.8601, 76.2219, 94.2776, 80.3246, 24.1925, 35.2117, 72.5125,
121.0505, 84.6559, 33.4994, 101.5312, 108.8929, 115.6663, 110.7305,
112.8359, 108.1926, 103.1844, 107.3720, 111.0626, 111.7052, 109.3754,
106.6239, 110.8514, 109.0949, 109.7291, 108.4103, 109.5334, 108.0790,
122.0732, 65.6744, 72.8142, 119.5030, 70.8820, 66.4881, 107.6618,
112.7306, 118.3036, 83.8566, 76.4978, 112.7944, 109.5748, 110.9455,
112.3806, 112.9551, 119.6298, 125.5823, 134.9212, 125.9777, 74.0451,
12.8238, 35.5617, 24.9626, 104.6106, 90.4906, 85.9967, 97.8541,
106.6222, 111.0453, 111.8815, 108.2535, 111.4015, 107.7559, 109.5380,
109.3212, 110.1219, 112.7612, 111.2595, 110.5335, 121.7879, 115.1693,
54.9851, 86.6916, 116.9939, 92.1268, 57.6106, 61.4862, 126.4773,
112.1987, 108.7125, 108.2527, 112.7414, 112.6068, 110.2034, 114.4493,
113.2548, 110.7587, 112.4181, 108.1768, 112.8213, 111.0445, 109.6948,
110.8729, 110.4764, 118.0310, 88.8105, 83.8834, 81.3157, 85.5275,
85.5625, 84.1886, 68.1180, 31.3709, 33.1545, 121.5646, 111.7370,
114.0073, 112.3546, 107.0779, 115.7818, 115.7658, 111.4645, 112.7774,
108.1568, 95.3530, 103.1948, 92.2902, 62.8486, 14.7445, 86.1638,
112.5176, 122.1774, 111.3493, 110.5394, 114.9414, 111.3867, 116.3332,
112.4701, 109.2141, 105.9539, 107.6427, 117.1809, 118.7265, 130.1242,
74.6725, 71.4855, 82.5081, 79.1204, 67.8084, 88.8452, 108.2482,
58.0370, 72.4352, 90.4025, 94.9953, 66.0537, 57.3020, 35.3184,
68.1088, 86.6186, 118.7876, 115.8533, 111.5778, 119.1641, 115.9867,
85.7971, 56.1209, 96.1703, 106.1886, 110.8123, 110.6988, 115.4252,
107.3056, 106.8170, 107.7311, 112.8368, 114.0857, 108.5903, 108.3581,
112.6980, 105.4932, 101.5349, 99.4108, 88.5853, 93.3055, 111.7297,
68.5167, 69.9066, 113.9631, 58.5684, 71.7351, 109.8313, 116.6967,
120.1699, 88.2144, 84.0390, 115.9588, 109.9733, 107.4564, 111.2211,
103.9748, 120.8010, 123.7707, 137.4949, 131.4860, 80.5471, 33.5931,
37.3247, 18.8009, 100.6905, 84.3534, 95.4535, 106.1850, 110.0926,
110.5705, 107.6609, 110.2489, 112.9340, 106.6505, 107.3167, 109.4810,
110.2208, 109.5875, 110.9688, 113.7061, 118.0853, 124.1593, 84.5591,
68.4251, 121.4505, 78.6792, 79.8214, 43.0172, 124.8438, 109.5220,
110.0105, 105.9807, 109.9260, 111.4670, 108.3692, 104.8462, 113.4009,
112.4410, 113.8966, 112.5260, 110.6384, 111.6085, 112.2110, 111.0163,
111.0780, 111.4477, 85.4452, 82.7611, 81.2657, 0.0000, 0.0000])
Alpha tensor([ 0.0000, -2.4767, -1.8137, 0.9845, 1.0530, 0.9134, 0.8459, 0.8743,
0.7471, 0.9543, 0.7302, 0.8200, 0.7925, 1.0518, 0.9065, 1.0274,
2.5421, -0.0797, 1.2547, 1.9817, 0.7334, 0.8684, 0.8448, 0.9557,
0.8410, 0.8284, 1.0277, 0.8229, 0.8718, 0.9070, 0.8649, 0.8861,
0.6979, 0.6893, 0.4872, 3.1380, 1.3071, 1.2007, 1.2675, 1.4149,
0.9349, 0.8872, 2.2476, 0.9992, 1.5108, -2.1848, -2.9639, 2.2535,
-1.5325, 0.4023, 2.8076, -2.2437, 1.0250, 0.7859, 0.8745, 0.8128,
0.8563, 1.0173, 0.8819, 0.8685, 0.8876, 0.8675, 0.9296, 0.9464,
0.8129, 0.9869, 0.8484, 0.8921, 0.9276, 0.8988, 0.5535, 2.4848,
1.0371, 0.3872, 2.6298, 1.1691, 0.6008, 0.9654, 0.5958, 1.9895,
0.9205, 0.7226, 1.0220, 0.8177, 0.9216, 0.6738, 0.8244, 0.4688,
0.2469, 0.6336, -2.9338, -3.0731, -1.9021, -2.1063, 1.2842, 1.0675,
1.2741, 0.8843, 0.9427, 0.7140, 0.9503, 0.9203, 0.8126, 0.9171,
0.8868, 0.8570, 0.8748, 0.8461, 0.8776, 0.9115, 0.5731, 0.9066,
2.7466, 0.8159, 0.6411, 1.6533, 2.6772, -1.7638, 0.8977, 0.7734,
0.9886, 0.8506, 0.7941, 0.8426, 0.9372, 0.8285, 0.9057, 0.8424,
0.9569, 0.8220, 0.8751, 0.9200, 0.8427, 0.8511, 0.9189, 0.7012,
1.7875, 0.6959, 1.9252, -2.5421, -0.8640, 1.4141, 2.7971, -1.7047,
-1.8703, 0.7833, 0.9281, 0.8164, 0.8828, 0.8918, 0.8985, 0.8558,
0.9284, 0.7310, 1.0688, 1.1470, 0.8735, 1.7871, -2.4753, -2.7288,
0.8465, 0.7635, 0.7689, 0.9352, 0.7778, 0.8900, 0.7745, 0.8667,
0.7957, 0.8481, 1.0474, 0.7395, 0.8118, 0.6633, 0.5597, -3.1325,
1.2299, 1.3197, 1.3817, 1.5566, 0.7882, 0.9803, 1.9878, 1.2283,
1.3423, -2.2988, 0.4446, 2.6727, 1.9214, 0.6146, -2.0513, 0.8181,
0.9210, 0.7991, 0.8512, 0.7204, 2.0591, -2.5587, 0.9295, 0.9179,
0.8016, 0.8717, 0.8954, 0.9498, 0.9047, 0.8463, 0.8356, 0.8375,
0.9849, 0.8456, 0.8777, 0.9996, 0.8272, 1.0393, 1.1962, 1.0500,
0.7470, 2.4272, 0.9708, 0.5438, 2.4546, 0.9886, 0.6457, 0.9665,
0.5650, 1.8601, 0.7157, 0.9071, 1.0081, 0.8627, 0.8475, 0.9144,
0.7045, 0.6209, 0.2665, 0.7096, -2.5762, 2.9030, -2.2988, -2.2315,
1.4952, 1.0621, 1.0510, 0.7648, 0.9207, 0.8453, 0.9342, 0.8241,
0.8325, 0.9728, 0.8598, 0.8750, 0.8993, 0.8441, 0.9058, 0.8048,
0.7765, 0.5704, -3.0346, 0.0335, 1.2383, 1.3833, -3.0827, -1.7459,
0.8982, 0.8058, 0.9437, 0.9452, 0.8664, 0.8599, 0.9927, 0.8707,
0.7997, 0.8538, 0.8473, 0.9004, 0.8576, 0.8496, 0.7928, 0.8984,
0.8463, 0.7809, 1.6586, 0.7535, 2.2789, -2.6393, 3.0752, -2.5848,
-2.4985, -1.8321, 0.9389, 0.9351, 0.9285, 0.8637, 0.8385, 0.8564,
0.8480, 0.8950, 0.7124, 0.8210, 0.9888, 0.8656, 1.0904, 2.4883,
-0.0280, 1.1883, 1.9463, 0.8018, 0.8579, 0.8535, 0.9611, 0.8505,
0.8728, 0.9781, 0.8132, 0.8610, 0.8653, 0.8649, 0.8722, 0.7340,
0.6546, 0.5138, -3.0645, 1.2742, 1.2396, 1.2028, 1.4131, 0.9680,
0.8155, 2.3941, 1.0347, 1.3958, -2.1916, -2.9761, 2.3825, -1.4723,
0.3984, 2.9090, -2.5935, 1.0288, 0.7320, 0.9219, 0.8104, 0.8202,
1.0242, 0.8950, 0.9106, 0.8318, 0.9029, 0.8363, 0.9503, 0.8525,
0.9212, 0.8610, 0.8844, 0.9059, 0.9159, 0.6239, 2.4699, 0.9326,
0.5656, 2.5664, 1.1222, 0.7031, 0.8850, 0.6522, 2.0087, 0.8236,
0.8160, 0.9703, 0.8045, 0.9197, 0.6630, 0.8711, 0.3860, 0.3955,
0.5209, -2.9011, -3.1086, -1.8160, -2.1631, 1.2303, 1.0968, 1.2292,
0.8951, 0.9638, 0.7226, 0.9768, 0.8532, 0.8937, 0.9142, 0.8273,
0.9509, 0.8084, 0.8497, 0.9016, 0.8562, 0.5749, 0.9444, 2.8210,
0.6647, 0.7028, 1.5307, 2.6638, -1.6503, 0.8121, 0.8810, 0.9544,
0.8516, 0.7858, 0.9059, 0.8464, 0.8124, 0.8570, 0.8270, 0.9592,
0.7824, 0.9039, 0.8760, 0.8670, 0.8393, 0.8971, 0.6940, 1.7212,
0.7853, 1.8236, -2.4965, -0.7853, 1.8931, 2.9265, -2.3745, -2.0759,
1.0130, 0.7460, 0.9621, 0.7372, 0.9418, 0.7982, 0.7790, 0.9590,
0.6846, 1.0372, 1.1351, 0.8748, 1.8342, -2.5203, -2.6841, 0.8150,
0.7111, 0.7533, 0.9782, 0.7634, 0.8788, 0.8050, 0.7815, 0.9218,
0.8402, 1.0117, 0.8161, 0.7691, 0.6837, 0.4683, 3.1336, 1.2848,
1.2405, 1.3313, 1.5973, 0.8628, 0.8897, 2.2251, 1.0760, 1.4797,
-1.8754, -0.2620, 2.9249, 2.2411, 0.0370, -1.4542, 0.6660, 0.9846,
0.7101, 0.8467, 0.6792, 2.0849, -2.5900, 1.0109, 0.8875, 0.8602,
0.8669, 0.7245, 1.0654, 0.8056, 1.0066, 0.7141, 0.8364, 1.0122,
0.7631, 0.9302, 1.0189, 0.9493, 1.0090, 1.2321, 1.0946, 0.7533,
2.6030, 0.9697, 0.6877, 2.5798, 1.0007, 0.7571, 0.8855, 0.5384,
1.8301, 0.8320, 0.7005, 1.0919, 0.7942, 0.9179, 0.8546, 0.7828,
0.4611, 0.3893, 0.5756, -2.5794, 2.9308, -2.1747, -2.3238, 1.3766,
1.0887, 1.0639, 0.8196, 0.9156, 0.8057, 1.0209, 0.7666, 0.8630,
0.9791, 0.8680, 0.9260, 0.8150, 0.9118, 0.8510, 0.8066, 0.7456,
0.5757, -2.9436, 0.0156, 1.2893, 1.3047, -3.1369, -1.5998, 0.8619,
0.8617, 0.9449, 0.9133, 0.8429, 0.8396, 1.0184, 0.8336, 0.8828,
0.7521, 0.8720, 0.9158, 0.8022, 0.8534, 0.8598, 0.8915, 0.8547,
0.8510, 1.6121, 0.8990, 2.1613, 0.0000, 0.0000])
Plotting¶
We provide several plotting utilities built with Plotly. This means plots can easily be logged to WandB.
[46]:
Data type cannot be displayed: application/vnd.plotly.v1+json
[47]:
Data type cannot be displayed: application/vnd.plotly.v1+json
[48]:
protein.plot_dihedrals()
Data type cannot be displayed: application/vnd.plotly.v1+json
Testing¶
[49]:
# Check if the protein has a complete backbone
protein.has_complete_backbone()
[49]:
tensor(True)
[50]:
# Checks all expected atoms are present for each residue
protein.is_complete()
[50]:
True
ProteinBatch¶
[51]:
import graphein.protein.tensor as gpt
batch = gpt.data.ProteinBatch().from_pdb_codes(pdb_codes=["3eiy", "4hhb", "1a0q"])
print(batch)
DataProteinBatch(fill_value=[3], atom_list=[3], residue_type=[1159], id=[3], residue_id=[3], chains=[1159], coords=[1159, 37, 3], residues=[3])
[52]:
batch.plot_structure()
Data type cannot be displayed: application/vnd.plotly.v1+json
Unbatching and apply¶
We can quickly unbatch a ProteinBatch into a list of Proteins
[53]:
proteins = batch.to_protein_list()
proteins
[53]:
[Protein(fill_value=[1], atom_list=[37], residue_type=[174], id='3eiy', residue_id=[174], chains=[174], coords=[174, 37, 3], residues=[174]),
Protein(fill_value=[1], atom_list=[37], residue_type=[574], id='4hhb', residue_id=[574], chains=[574], coords=[574, 37, 3], residues=[574]),
Protein(fill_value=[1], atom_list=[37], residue_type=[411], id='1a0q', residue_id=[411], chains=[411], coords=[411, 37, 3], residues=[411])]
We can also define a function that operates on a singe Protein and apply it to all proteins in the ProteinBatch using the protein_apply() method:
[54]:
from graphein.protein.tensor.plot import plot_structure
def single_plot(protein: gpt.Protein):
return plot_structure(protein.coords, lines=False)
plots = batch.protein_apply(single_plot)
plots[2]
Data type cannot be displayed: application/vnd.plotly.v1+json
If we want to apply the function to a single protein, we can use apply_to with a specified index:
[55]:
batch.apply_to(single_plot, 0)
Data type cannot be displayed: application/vnd.plotly.v1+json
[56]:
batch.apply_to(single_plot, 1)
Data type cannot be displayed: application/vnd.plotly.v1+json
Saving Proteins and ProteinBatches¶
Saving as Tensor¶
[57]:
import torch
import graphein.protein.tensor as gpt
# Save a protein
protein = gpt.data.get_random_protein()
torch.save(protein, "test_protein.pt")
protein.save("test_protein_2.pt")
batch = gpt.data.get_random_batch(16)
torch.save(batch, "test_batch.pt")
batch.save("test_batch_2.pt")
[ ]:
print(torch.load("test_protein.pt", weights_only=False))
print(torch.load("test_protein.pt", weights_only=False))
[ ]:
print(torch.load("test_batch.pt", weights_only=False))
print(torch.load("test_batch_2.pt", weights_only=False))
To PDB Files¶
N.B. if you work with PDBs in VSCode you may find this extension useful :+)
[60]:
import torch
import graphein.protein.tensor as gpt
# Save a protein
protein = gpt.data.get_random_protein()
protein.to_pdb("test_protein.pdb")
[61]:
from biopandas.pdb import PandasPdb
p = PandasPdb().read_pdb("test_protein.pdb")
p.df["ATOM"].head()
[61]:
| record_name | atom_number | blank_1 | atom_name | alt_loc | residue_name | blank_2 | chain_id | residue_number | insertion | ... | x_coord | y_coord | z_coord | occupancy | b_factor | blank_4 | segment_id | element_symbol | charge | line_idx | |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 0 | ATOM | 1 | N | GLN | A | 1 | ... | 25.569 | 27.140 | -2.829 | 1.0 | 0.0 | N | 0.0 | 0 | ||||||
| 1 | ATOM | 2 | CA | GLN | A | 1 | ... | 26.022 | 27.178 | -1.417 | 1.0 | 0.0 | C | 0.0 | 1 | ||||||
| 2 | ATOM | 3 | C | GLN | A | 1 | ... | 25.202 | 26.176 | -0.607 | 1.0 | 0.0 | C | 0.0 | 2 | ||||||
| 3 | ATOM | 4 | O | GLN | A | 1 | ... | 24.418 | 25.401 | -1.149 | 1.0 | 0.0 | O | 0.0 | 3 | ||||||
| 4 | ATOM | 5 | CB | GLN | A | 1 | ... | 25.856 | 28.603 | -0.862 | 1.0 | 0.0 | C | 0.0 | 4 |
5 rows × 21 columns
To DataFrame¶
[62]:
import graphein.protein.tensor as gpt
protein = gpt.data.get_random_protein()
df = protein.to_dataframe()
df.head()
[62]:
| record_name | atom_number | blank_1 | atom_name | alt_loc | residue_name | blank_2 | chain_id | residue_number | insertion | ... | x_coord | y_coord | z_coord | occupancy | b_factor | blank_4 | segment_id | element_symbol | charge | line_idx | |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 0 | ATOM | 1 | N | GLN | 0 | 1 | ... | 25.569000 | 27.139999 | -2.829 | 1.0 | 0.0 | N | 0 | 1 | ||||||
| 1 | ATOM | 2 | CA | GLN | 0 | 1 | ... | 26.021999 | 27.177999 | -1.417 | 1.0 | 0.0 | C | 0 | 2 | ||||||
| 2 | ATOM | 3 | C | GLN | 0 | 1 | ... | 25.202000 | 26.176001 | -0.607 | 1.0 | 0.0 | C | 0 | 3 | ||||||
| 3 | ATOM | 4 | O | GLN | 0 | 1 | ... | 24.417999 | 25.400999 | -1.149 | 1.0 | 0.0 | O | 0 | 4 | ||||||
| 4 | ATOM | 5 | CB | GLN | 0 | 1 | ... | 25.856001 | 28.603001 | -0.862 | 1.0 | 0.0 | C | 0 | 5 |
5 rows × 21 columns