<jupyter_start><jupyter_text>PSCDB - Baselines<jupyter_code># NBVAL_SKIP
import pandas as pd
import numpy as np

import torch
import torch.nn as nn
import pytorch_lightning as pl
from tqdm.notebook import tqdm
import networkx as nx
import torch_geometric
from torch_geometric.data import Data
from torch_geometric.utils import from_networkx
from sklearn.preprocessing import LabelBinarizer
from sklearn.metrics import f1_score

import warnings
warnings.filterwarnings("ignore")<jupyter_output>RDKit WARNING: [23:36:27] Enabling RDKit 2019.09.3 jupyter extensions
[23:36:27] Enabling RDKit 2019.09.3 jupyter extensions<jupyter_text>Load dataset<jupyter_code># NBVAL_SKIP
df = pd.read_csv("../datasets/pscdb/structural_rearrangement_data.csv")
pdbs = df["Free PDB"]
y = [torch.argmax(torch.Tensor(lab)).type(torch.LongTensor) for lab in LabelBinarizer().fit_transform(df.motion_type)]<jupyter_output><empty_output><jupyter_text>Transformation from Raw Structure to ML-ready Datasets Construction with Graphein<jupyter_code># NBVAL_SKIP
from graphein.protein.config import ProteinGraphConfig
from graphein.protein.edges.distance import add_hydrogen_bond_interactions, add_peptide_bonds, add_k_nn_edges
from graphein.protein.graphs import construct_graph

from functools import partial

# Override config with constructors
constructors = {
    "edge_construction_functions": [partial(add_k_nn_edges, k=3, long_interaction_threshold=0)],
    #"edge_construction_functions": [add_hydrogen_bond_interactions, add_peptide_bonds],
    #"node_metadata_functions": [add_dssp_feature]
}

config = ProteinGraphConfig(**constructors)
print(config.dict())

# Make graphs
graph_list = []
y_list = []
for idx, pdb in enumerate(tqdm(pdbs)):
    try:
        graph_list.append(
            construct_graph(pdb_code=pdb,
                        config=config
                       )
            )
        y_list.append(y[idx])
    except:
        print(str(idx) + ' processing error...')
        pass
# NBVAL_SKIP
pdbs[274]
#pdbs[266]
#pdbs[677]<jupyter_output><empty_output><jupyter_text>Convert Nx graphs to PyTorch Geometric<jupyter_code># NBVAL_SKIP
from graphein.ml.conversion import GraphFormatConvertor

format_convertor = GraphFormatConvertor('nx', 'pyg', 
                                        verbose = 'gnn', 
                                        columns = None)
# NBVAL_SKIP
pyg_list = [format_convertor(graph) for graph in tqdm(graph_list)]
# NBVAL_SKIP
for idx, g in enumerate(pyg_list):
    g.y = y_list[idx] 
    g.coords = torch.FloatTensor(g.coords[0])
# NBVAL_SKIP
for i in pyg_list:
    if i.coords.shape[0] == len(i.node_id):
        pass
    else:
        print(i)
        pyg_list.remove(i)<jupyter_output>Data(coords=[10112, 3], dist_mat=[1], edge_index=[2, 120], name=[1], node_id=[1264], y=1)
Data(coords=[820, 3], dist_mat=[1], edge_index=[2, 1431], name=[1], node_id=[808], y=2)
Data(coords=[668, 3], dist_mat=[1], edge_index=[2, 1166], name=[1], node_id=[666], y=4)
Data(coords=[2720, 3], dist_mat=[1], edge_index=[2, 3], name=[1], node_id=[340], y=5)<jupyter_text>Model Configuration<jupyter_code># NBVAL_SKIP
config_default = dict(
    n_hid = 8,
    n_out = 8,
    batch_size = 4,
    dropout = 0.5,
    lr = 0.001,
    num_heads = 32,
    num_att_dim = 64,
    model_name = 'GCN'
)

class Struct:
    def __init__(self, **entries):
        self.__dict__.update(entries)
        
config = Struct(**config_default)

global model_name
model_name = config.model_name<jupyter_output><empty_output><jupyter_text>Construct DataLoaders<jupyter_code># NBVAL_SKIP
import numpy as np
np.random.seed(42)
idx_all = np.arange(len(pyg_list))
np.random.shuffle(idx_all)

train_idx, valid_idx, test_idx = np.split(idx_all, [int(.8*len(pyg_list)), int(.9*len(pyg_list))])
train, valid, test = [pyg_list[i] for i in train_idx], [pyg_list[i] for i in valid_idx], [pyg_list[i] for i in test_idx]

from torch_geometric.data import DataLoader
train_loader = DataLoader(train, batch_size=config.batch_size, shuffle = True, drop_last = True)
valid_loader = DataLoader(valid, batch_size=32)
test_loader = DataLoader(test, batch_size=32)
# NBVAL_SKIP
pyg_list[0]<jupyter_output><empty_output><jupyter_text>Define Model<jupyter_code># NBVAL_SKIP
from torch_geometric.nn import GCNConv, GATConv, SAGEConv, global_add_pool
from torch.nn.functional import mse_loss, nll_loss, relu, softmax, cross_entropy
from torch.nn import functional as F
from pytorch_lightning.metrics.functional import accuracy
# NBVAL_SKIP
class GraphNets(pl.LightningModule):
    def __init__(self):
        super().__init__()
        
        if model_name == 'GCN':
            self.layer1 = GCNConv(in_channels=3, out_channels=config.n_hid)
            self.layer2 = GCNConv(in_channels=config.n_hid, out_channels=config.n_out)

        elif model_name == 'GAT':
            self.layer1 = GATConv(3, config.num_att_dim, heads=config.num_heads, dropout=config.dropout)
            self.layer2 = GATConv(config.num_att_dim * config.num_heads, out_channels = config.n_out, heads=1, concat=False,
                                 dropout=config.dropout)
            
        elif model_name == 'GraphSAGE':
            self.layer1 = SAGEConv(3, config.n_hid)
            self.layer2 = SAGEConv(config.n_hid, config.n_out)  
            
        self.decoder = nn.Linear(config.n_out, 7)
        
    def forward(self, g):
        x = g.coords
        x = F.dropout(x, p=config.dropout, training=self.training)
        x = F.elu(self.layer1(x, g.edge_index))
        x = F.dropout(x, p=config.dropout, training=self.training)
        x = self.layer2(x, g.edge_index)
        x = global_add_pool(x, batch=g.batch)
        x = self.decoder(x)
        return softmax(x)

    def training_step(self, batch, batch_idx):
        x = batch   
        y = x.y
        y_hat = self(x)
        loss = cross_entropy(y_hat, y)
        acc = accuracy(y_hat, y)

        self.log("train_loss", loss)
        self.log("train_acc", acc)
        return loss

    def validation_step(self, batch, batch_idx):
        x = batch   
        y = x.y
        y_hat = self(x)
        loss = cross_entropy(y_hat, y)
        acc = accuracy(y_hat, y)
        self.log("valid_loss", loss)
        self.log("valid_acc", acc)

    def test_step(self, batch, batch_idx):
        x = batch   
        y = x.y
        y_hat = self(x)
        loss = cross_entropy(y_hat, y)
        acc = accuracy(y_hat, y)

        y_pred_softmax = torch.log_softmax(y_hat, dim = 1)
        y_pred_tags = torch.argmax(y_pred_softmax, dim = 1) 
        f1 = f1_score(y.detach().cpu().numpy(), y_pred_tags.detach().cpu().numpy(), average = 'weighted')

        self.log("test_loss", loss)
        self.log("test_acc", acc)
        self.log("test_f1", f1)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=config.lr)
        return optimizer
# NBVAL_SKIP
GraphNets()
# NBVAL_SKIP
from pytorch_lightning.callbacks import ModelCheckpoint
import os

file_path = './graphein_model'
if not os.path.exists(file_path):
    os.mkdir(file_path)

checkpoint_callback = ModelCheckpoint(
    monitor="valid_loss",
    dirpath=file_path,
    filename="model-{epoch:02d}-{val_loss:.2f}",
    save_top_k=1,
    mode="min",
)<jupyter_output><empty_output><jupyter_text>Train!<jupyter_code># NBVAL_SKIP
# Train Model
model = GraphNets()
trainer = pl.Trainer(max_epochs=200, gpus=-1, callbacks=[checkpoint_callback])
trainer.fit(model, train_loader, valid_loader)

# evaluate on the model with the best validation set
best_model = GraphNets.load_from_checkpoint(checkpoint_callback.best_model_path)
out_best_test = trainer.test(best_model, test_loader)[0]<jupyter_output>GPU available: True, used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/atj39/anaconda3/envs/graphein-dev/lib/python3.8/multiprocessing/spawn.py", line 116, in spawn_main
    exitcode = _main(fd, parent_sentinel)
  File "/home/atj39/anaconda3/envs/graphein-dev/lib/python3.8/multiprocessing/spawn.py", line 126, in _main
    self = reduction.pickle.load(from_parent)
AttributeError: Can't get attribute 'GraphNets' on <module '__main__' (built-in)>
{'fraction_non_alphanumeric': 0.11678481316832873, 'fraction_numerical': 0.02066562537546558, 'mean_word_length': 3.7756741250717156, 'pattern_counts': {'":': 3, '<': 33, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '38796', 'license': 'MIT', 'n_tokens_mistral': 3130, 'n_tokens_neox': 2968, 'n_words': 593, 'path': 'notebooks/baseline.ipynb', 'repo_name': 'rg314/graphein', 'size': '8323'}