<jupyter_start><jupyter_text>Molecule solubility prediction with basic GNN modelWater (or any other kind of) solubility is the property of the whole molecule, therefore we will have a simple binary clasification problem for the whole molecule.   References PyG [examples](https://github.com/pyg-team/pytorch_geometric/tree/master/examples) on GNN Textbook: Deep Learning for the Life Sciences: Applying Deep Learning to Genomics, Microscopy, Drug Discovery, and More<jupyter_code>!python -c "import torch; print(torch.__version__)"
!python -c "import torch; print(torch.version.cuda)"
!python --version

!nvidia-smi
pytorch_version = f"torch-{torch.__version__}.html"
!pip install --no-index torch-scatter -f https://pytorch-geometric.com/whl/torch-1.10.0+cu111.html
!pip install --no-index torch-sparse -f https://pytorch-geometric.com/whl/torch-1.10.0+cu111.html
!pip install --no-index torch-cluster -f https://pytorch-geometric.com/whl/torch-1.10.0+cu111.html
!pip install --no-index torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.10.0+cu111.html
!pip install torch-geometric
!pip -q install rdkit-pypi
import torch
from torch.nn import Linear
import torch.nn.functional as F 
from torch_geometric.nn import GCNConv, TopKPooling, global_mean_pool
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
from torch_geometric.data import DataLoader
from torch_geometric.datasets import MoleculeNet

import pandas as pd 
import warnings
import seaborn as sns
import matplotlib.pyplot as plt

import rdkit
from rdkit import Chem
from rdkit.Chem.Draw import IPythonConsole
from rdkit.Chem.Draw import MolsToGridImage<jupyter_output><empty_output><jupyter_text>ESOL [dataset](https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/delaney-processed.csv) will be used[This dataset is created based on MoleculeNet.org](http://ww82.moleculenet.ai/).<jupyter_code># Load the ESOL dataset
data = MoleculeNet(root=".", name="ESOL")
data
print("Dataset type: ", type(data))
print("Dataset features: ", data.num_features)
print("Dataset target: ", data.num_classes)
print("Dataset length: ", data.len)
print("Dataset sample: ", data[0])
print("Sample  nodes: ", data[0].num_nodes)
print("Sample  edges: ", data[0].num_edges)
# Shape of nodes is made of [num_nodes, num_node_features]
data[0].x
# Shape of edges is [2, num_edges]
data[0].edge_index.t()
data[0].y
data[150]["smiles"]
molecule = Chem.MolFromSmiles(data[150]["smiles"])
molecule
type(molecule)<jupyter_output><empty_output><jupyter_text>We implement GNN in form of basic GCN and do aggregation only from 4 hop neighbourhood<jupyter_code>embedding_size = 64

class GCN(torch.nn.Module):
    def __init__(self):
        # Init parent
        super(GCN, self).__init__()
        torch.manual_seed(42)

        # GCN layers
        self.initial_conv = GCNConv(data.num_features, embedding_size)
        self.conv1 = GCNConv(embedding_size, embedding_size)
        self.conv2 = GCNConv(embedding_size, embedding_size)
        self.conv3 = GCNConv(embedding_size, embedding_size)
        self.conv4 = GCNConv(embedding_size, embedding_size)

        # Output layer
        self.out = Linear(embedding_size*2, 1)

    def forward(self, x, edge_index, batch_index):
        # Starting GCN layer
        hidden = self.initial_conv(x, edge_index)
        hidden = F.tanh(hidden)

        # Other GCN layers 
        hidden = self.conv1(hidden, edge_index)
        hidden = F.tanh(hidden)
        hidden = self.conv2(hidden, edge_index)
        hidden = F.tanh(hidden)
        hidden = self.conv3(hidden, edge_index)
        hidden = F.tanh(hidden)
        hidden = self.conv4(hidden, edge_index)
        hidden = F.tanh(hidden)

        # Global Pooling 
        hidden = torch.cat([gmp(hidden, batch_index), 
                            gap(hidden, batch_index)], dim=1)

        # Linear Layer
        out = self.out(hidden)

        return out, hidden

model = GCN()
print(model)
print("Number of parameters: ", sum(p.numel() for p in model.parameters()))<jupyter_output>GCN(
  (initial_conv): GCNConv(9, 64)
  (conv1): GCNConv(64, 64)
  (conv2): GCNConv(64, 64)
  (conv3): GCNConv(64, 64)
  (conv4): GCNConv(64, 64)
  (out): Linear(in_features=128, out_features=1, bias=True)
)
Number of parameters:  17409<jupyter_text>Training<jupyter_code>warnings.filterwarnings("ignore")

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0007)  

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

data_size = len(data)
NUM_GRAPHS_PER_BATCH = 64
loader = DataLoader(data[:int(data_size * 0.8)], 
                    batch_size=NUM_GRAPHS_PER_BATCH, shuffle=True)
test_loader = DataLoader(data[int(data_size * 0.8):], 
                         batch_size=NUM_GRAPHS_PER_BATCH, shuffle=True)

def train(data):
    for batch in loader:
      batch.to(device)  
      optimizer.zero_grad() 
      pred, embedding = model(batch.x.float(), batch.edge_index, batch.batch)    
      loss = torch.sqrt(loss_fn(pred, batch.y))       
      loss.backward()  
      optimizer.step()   
    return loss, embedding

print("Starting training...")
losses = []
for epoch in range(2000):
    loss, h = train(data)
    losses.append(loss)
    if epoch % 100 == 0:
      print(f"Epoch {epoch} | Train Loss {loss}")
# Plotting training loss
losses_float = [float(loss.cpu().detach().numpy()) for loss in losses] 
indices = [i for i,l in enumerate(losses_float)] 
plt.plot(indices, losses_float, c='c');

# Predict
test_batch = next(iter(test_loader))
with torch.no_grad():
    test_batch.to(device)
    pred, embed = model(test_batch.x.float(), test_batch.edge_index, test_batch.batch) 
    df = pd.DataFrame()
    df["y_real"] = test_batch.y.tolist()
    df["y_pred"] = pred.tolist()
df["y_real"] = df["y_real"].apply(lambda row: row[0])
df["y_pred"] = df["y_pred"].apply(lambda row: row[0])

molecules = [Chem.MolFromSmiles(data[index]["smiles"]) for index in range(len(df["y_real"]))]

legends = [f"true/pred = {df['y_real'][i]:.2f}/{df['y_pred'][i]:.2f}" for i in range(len(df['y_real']))]
MolsToGridImage(molecules, molsPerRow=4, legends=legends)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13028792021875502, 'fraction_numerical': 0.018980215538040857, 'mean_word_length': 4.143093465674111, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '63408', 'license': 'MIT', 'n_tokens_mistral': 2256, 'n_tokens_neox': 2142, 'n_words': 481, 'path': 'Pytorch/Molecule_Solubility.ipynb', 'repo_name': 'VladimerKhasia/ML', 'size': '6217'}