<jupyter_start><jupyter_text>In this file, we will define our model and parameters Instructions: Finally, you should have at least one notebook where you define and train your notebook.<jupyter_code>!pip install nbimporter
import pandas as pd
import numpy as np
import torch
from dgl.data import DGLDataset
from dgl.data.utils import save_graphs, load_graphs

import nbimporter
import graphize_data as graphize<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>import dgl
import torch
import torch.nn as nn
import torch.nn.functional as F

from dgl.nn.pytorch import NNConv
from dgl.nn.pytorch import GatedGraphConv<jupyter_output><empty_output><jupyter_text>RJ model<jupyter_code># Things to experiment with: SumPooling layer? Kind of like a conv, can sum up neighbors and find important structural properties of molecule
# from dgl.nn import SumPooling

# Also a good read here (https://arxiv.org/pdf/1710.10370.pdf) on Topology Adaptive Graph Convolutional layer might be useful.
# From my takeaways, we can use this "topology" to better track/weight the important parts of the structure
# (maybe the center of the molecule is important, or specific edge portions of the molecule) TAG can find this

# SAGEConv

# Ultimately, the NNConv layer, as used in the MPNN model, is the best one for this situation.
# We have come to this conclusion after trying various other convolution modules
# and finally reading up on this paper which relates closely to what we are doing (https://arxiv.org/pdf/1704.01212.pdf)
# (Quantum Chemistry computations)

# Note they used GatedGraphConv first and that worked okay but not as well as NNConv... we could look to backtrack to GGC
class Electron_MPNN(nn.Module):
    def __init__(self, node_in_feats, edge_in_feats, node_out_feats=200,
                 edge_hidden_feats=300, num_message_passing=3):
        super(Electron_MPNN, self).__init__()

        # Projection
        self.fc1 = nn.Sequential(
            nn.Linear(node_in_feats, node_out_feats),
            nn.ReLU()
        )
    
        self.num_message_passing = num_message_passing
        
        # Multi-iteration Convolution
        edge_func = nn.Sequential(
            nn.Linear(edge_in_feats, edge_hidden_feats*2),
            nn.ReLU(),
            nn.Linear(edge_hidden_feats*2, edge_hidden_feats),
            nn.ReLU(),
            nn.Linear(edge_hidden_feats, node_out_feats * node_out_feats)
        )
        
        self.gnn1 = NNConv(
            in_feats=node_out_feats,
            out_feats=node_out_feats,
            edge_func=edge_func,
            aggregator_type='sum'
        )
        
        # GRU used in MPNN to get rid of vanishing gradient on Conv layer
        self.gru = nn.GRU(node_out_feats, node_out_feats)


    def forward(self, g, node_feats, edge_feats):
        node_feats = self.fc1(node_feats)
        
        hidden_feats = node_feats.unsqueeze(0)
        
        for _ in range(self.num_message_passing):
            # Conv 1
            node_feats = self.gnn1(g, node_feats, edge_feats)
            
            # GRU in between each convolution layer to retain gradient
            node_feats, hidden_feats = self.gru(node_feats.unsqueeze(0), hidden_feats)
            node_feats = node_feats.squeeze(0)
    
        return node_feats
class Electron_Predictor(nn.Module):
    def __init__(self,
                 node_in_feats,
                 edge_in_feats,
                 node_out_feats=128,
                 edge_hidden_feats=256,
                 num_preds=5,
                 num_message_passing=3
                ):
        super(Electron_Predictor, self).__init__()

        self.gnn = Electron_MPNN(node_in_feats=node_in_feats,
                                 node_out_feats=node_out_feats,
                                 edge_in_feats=edge_in_feats,
                                 edge_hidden_feats=edge_hidden_feats,
                                 num_message_passing=num_message_passing
                                )
        
        
        self.predict = nn.Sequential(
            nn.Linear(node_out_feats, 2 * node_out_feats),
            nn.ReLU(),
            nn.Linear(2 * node_out_feats, num_preds)
        )
    
    def forward(self, g, node_feats=None, edge_feats=None):

        # If trying to make a prediction based on smiles object (for evaluation), convert to graph
        if isinstance(g, str):

            g = graphize.smiles_to_bigraph(g,
                                       node_featurizer = graphize.featurize_atoms,
                                       edge_featurizer = graphize.featurize_bonds,
                                       explicit_hydrogens = True
                                 )
            edge_feats = g.edata['bond_feats'].float()
            node_feats = g.ndata['atom_feats'].float()
        

        node_feats = self.gnn(g, node_feats, edge_feats)

        # For our model, we take the mean of 5 different predictions
        return [self.predict(node_feats).mean()]<jupyter_output><empty_output><jupyter_text>Load the graphs to get dimensions for our model<jupyter_code>graphs, _ = load_graphs("./DataGraphs/combined_graph.bin")<jupyter_output><empty_output><jupyter_text>Construct the model<jupyter_code># All graphs in the list have the same scheme size, so pull the dimensions from the first
node_dim = graphs[0].ndata['atom_feats'].shape[1]
edge_dim = graphs[0].edata['bond_feats'].shape[1]
print("Dimensions:", node_dim, "(node),", edge_dim, "(edge)")
model = Electron_Predictor(node_dim, edge_dim)
sample_data = graphs[3]
edges = sample_data.edata['bond_feats'].float()
nodes = sample_data.ndata['atom_feats'].float()
model(sample_data, nodes, edges)
# Things to experiment with: SumPooling layer? Kind of like a conv, can sum up neighbors and find important structural properties of molecule
# from dgl.nn import SumPooling

# Also a good read here (https://arxiv.org/pdf/1710.10370.pdf) on Topology Adaptive Graph Convolutional layer might be useful.
# From my takeaways, we can use this "topology" to better track/weight the important parts of the structure
# (maybe the center of the molecule is important, or specific edge portions of the molecule) TAG can find this

# SAGEConv

# Ultimately, the NNConv layer, as used in the MPNN model, is the best one for this situation.
# We have come to this conclusion after trying various other convolution modules
# and finally reading up on this paper which relates closely to what we are doing (https://arxiv.org/pdf/1704.01212.pdf)
# (Quantum Chemistry computations)

# Note they used GatedGraphConv first and that worked okay but not as well as NNConv... we could look to backtrack to GGC
class Electron_MPNN_old(nn.Module):
    def __init__(self, node_in_feats, edge_in_feats, node_out_feats=1,
                 edge_hidden_feats=128):
        super(Electron_MPNN_old, self).__init__()

        # Projection
        self.fc1 = nn.Sequential(
            nn.Linear(node_in_feats, node_out_feats),
            nn.ReLU()
        )
    
        # Multi-iteration Convolution
#         self.message_passing_steps = message_passing_steps
        edge_func = nn.Sequential(
            nn.Linear(edge_in_feats, edge_hidden_feats),
            nn.ReLU(),
            nn.Linear(edge_hidden_feats, node_out_feats * node_out_feats)
        )
        
        self.gnn_layer = NNConv(
            in_feats=node_out_feats,
            out_feats=node_out_feats,
            edge_func=edge_func,
            aggregator_type='sum'
        )
        
        # GRU used in MPNN to get rid of vanishing gradient on Conv layer
#         self.gru = nn.GRU(node_out_feats, node_out_feats)
        
#         self.fc2 = nn.Linear(out_dim, out_dim)


    def forward(self, g, node_feats=None, edge_feats=None):
        # If trying to make a prediction based on smiles object (for evaluation), convert to graph
        if isinstance(g, str):
            g = graphize.smiles_to_bigraph(g,
                                       node_featurizer = graphize.featurize_atoms,
                                       edge_featurizer = graphize.featurize_bonds,
                                       explicit_hydrogens = True
                                 )
            edge_feats = g.edata['bond_feats'].float()
            node_feats = g.ndata['atom_feats'].float()
        
        node_feats = self.fc1(node_feats) # (V, node_out_feats)
        
        hidden_feats = node_feats.unsqueeze(0)           # (1, V, node_out_feats)
        
        node_feats = self.gnn_layer(g, node_feats, edge_feats)
#         node_feats, hidden_feats = self.gru(node_feats.unsqueeze(0), hidden_feats)
#         node_feats = node_feats.squeeze(0)
    
        return [node_feats.mean()]
torch.cuda.is_available()<jupyter_output><empty_output><jupyter_text>Honorable MentionsThis is the story which describes various models we used and their accuracies. At this point, most of the improvements had already been made, and thus some of the older saved models will not work (though if you really want, you could figure out the architecture from the keys and shapes of weights). With that said, our best model is stashed here while we make changes to it above. It is very similar, but I want to keep this here in case we decide to train this one more.I also understand that the LSTM approach of Set2Set was computationally expensive and perhaps not necessary, but I read that it might be useful from one of the MPNN papers. With our next models, we will cut it altogether.This is ordered chronologically in our trials...***1. 12k_electron.pth / 11.5k_electron_custom.pth w/ lr @ 0.01 (~ 8 epochs on dataset_F completed --> ~ 8k samples trained in total) CPU training***This was actually a stronger start than we anticipated. What I did here was perform a Projection, NNConv, and then a simple regression to choose the best energy (minimizing the loss, of course). What this ended up doing, which I found by analyzing the output, was predicting the same energy for all of the molecules. The regression at the end, the model found, was more insightful than the convolution and projection, and thus the regression simply found the average ground state energy and predicted that value every time. I would predict that with enough epochs, the model would've actually learned about the importance of the convolution and perhaps could reach sub 10k. Regardless, it was our starting point.Note: This was our 16k prediction (wasn't generalizing as well to dataset_E)***2. 15k_electron_mpnn_v1_ReLU.pth / 12k_electron_mpnn_no_ReLU.pth w/ lr @ 0.1 (~ 5 epochs on dataset_F completed --> ~ 5k samples trained in total) CPU training***Next, we look at the prebuilt MPNN. What do they do that makes them so effective? It is quite similar, no shocker here, as the architectures described in the MPNN papers. I'll keep this one short because there's not too much going on here, but I played around with removing ReLUs and adding fully-connected layers because, why not add more complexity, right?! Yeah... not my best decision. But hey, it eventually reached close to the first model, so I guess that's a plus? I took the features mentioned in this paper (https://arxiv.org/pdf/1704.01212.pdf) because they worked well and had scientific research backing it. Not my strong suit, so I let that be the one assumption in this wild journey.One thing to note, the training time took FOREVER (overnight and into the day). I had wrongly assumed that these environments were setup with the BCCS gpus. I found the next day that it was going very very slowly (~ 1 sample every 5-10 seconds) and would take around 3.5 hours per epoch, sooo this wasn't going to work, sadly.***3. rj_12k_mpnn.pth / 3.3k_rj_electron.pth w/ lr @ 0.0001 (20 epochs on the combined dataset completed --> ~ 13k * 20 samples trained in total) GPU training*** I finally remembered about the BCCS GPU cluster that I had access to w/ Professor Wei! I set this up, hastily put together an environment, and the difference was incredible as expected. 15ish samples/sec instead of 0.2 sample/sec, almost 10x the speedAlright, skip through a good amount of studying architectures from various papers related to this subject (all links can be found in a resources section in the train file... not sure why I put it there, but code organization has been a real struggle with this project. Sorry!). I constructed a similar structure outlined in this paper (https://arxiv.org/pdf/1806.03146.pdf) with a little bit of inspiration from the other papers I have cited as well (nothing major from any one, just some ideas on how they deal with convolution). It is also worth mentioning that many components of this model still have the basic structure of DGL's MPNN with some modifications, naturally. I will again spare many details about it, but if you are interested in more about how it works or why I chose certain layers in certain places, I would be happy to talk about it / receive critiques!- brought back multiple message passing iterations- trained to give 5 different predictions (1 prediction with 5 dimensions)- final prediction is the average***4. ???.???k_rj_electron.pth***Okay, this time I will actually keep it short.Same model, remove the readout altogether. Reduce number of passes, increase height of model (taller, not wider). Also, add the energy data (sum of the individual energies of the atoms). We came up with a clever way to implement it efficiently after discussing together, but I will let Ronan take the lead on that one and share! With that said, the sum of the energies of the atoms in the molecule would be a highly useful indicator for the model.- sadly, we ran out of time to fully train this model, though we started and got it down to about 5k- used energy of each atom in feature set***5. ???***Let's talk about the future of what I would do with this if I had more time. I should make this a big title so you can see it when perusing through the file. Here we go  Future Changes If I was continuing, and would make another model, I wanted to expiriment more with another layer of convolution known as GatedGraphConvolution (GGC). I was trying to wrap my head around what makes it special (that and TAG, because I felt that analyzing the topology of a structure and making decisions based on those scans could be super insightful to the model), but ultimately when I tried to put it into the model, I ran into many errors and decided to scrap it in favor of pursuing pre-existing inspirations of architectures from the papers above. I would still like to implement GGC however, and originally tried to do an NNConv, GGC, and then average the two before the GRU and continuing to the next pass. My intention was, of course, to cut down the number of message passing steps by a factor of two. GGC was mentioned in that first paper (https://arxiv.org/pdf/1704.01212.pdf) before the authors switched to NNConv for faster results, but I'd imagine that there are ways to do some sort of dimensionality reduction before the GGC and therefore keep the time complexity down. This is my thought process at least, and I would love to take more time to try to implement it. I would be willing to bet it would be factors better than our 3.3k model!<jupyter_code># No info is past here, these are just archived models 
class Electron_MPNN_3k(nn.Module):
    def __init__(self, node_in_feats, edge_in_feats, node_out_feats=200,
                 edge_hidden_feats=300, num_message_passing=6):
        super(Electron_MPNN_3k, self).__init__()

        # Projection
        self.fc1 = nn.Sequential(
            nn.Linear(node_in_feats, node_out_feats),
            nn.ReLU()
        )
    
    
        self.num_message_passing = num_message_passing
        # Multi-iteration Convolution
        edge_func = nn.Sequential(
            nn.Linear(edge_in_feats, edge_hidden_feats*2),
            nn.ReLU(),
            nn.Linear(edge_hidden_feats*2, edge_hidden_feats),
            nn.ReLU(),
            nn.Linear(edge_hidden_feats, node_out_feats * node_out_feats)
        )
        
        self.gnn1 = NNConv(
            in_feats=node_out_feats,
            out_feats=node_out_feats,
            edge_func=edge_func,
            aggregator_type='sum'
        )
        
#         self.gnn2 = GatedGraphConv(
#             in_feats=node_out_feats,
#             out_feats=node_out_feats,
#             n_steps=2,
#             n_etypes=edge_in_feats
#         )
        
        # GRU used in MPNN to get rid of vanishing gradient on Conv layer
        self.gru = nn.GRU(node_out_feats, node_out_feats)


    def forward(self, g, node_feats, edge_feats):
        node_feats = self.fc1(node_feats) # (V, node_out_feats)
        
        hidden_feats = node_feats.unsqueeze(0)           # (1, V, node_out_feats)
        
        # For each step, perform two convolutions and average them
        for _ in range(self.num_message_passing):
            # Conv 1
            node_feats = self.gnn1(g, node_feats, edge_feats)
            
            # GRU requires 3 dimensions
            
#             node_feats, hidden_feats = self.gru(node_feats.unsqueeze(0), hidden_feats)
#             node_feats = node_feats.squeeze(0)
        
            # Conv 2
#             node_feats_1 = node_feats.detach().clone()
#             node_feats = self.gnn2(g, node_feats.contiguous(), edge_feats.contiguous())
            
#             node_feats = (node_feats_1 + node_feats) / 2
            
            # GRU in between each convolution layer to retain gradient
            # GRU already built into GatedGraphConv
            node_feats, hidden_feats = self.gru(node_feats.unsqueeze(0), hidden_feats)
            node_feats = node_feats.squeeze(0)
    
        return node_feats
from dgl.nn.pytorch import Set2Set
class Electron_Predictor_3k(nn.Module):
    def __init__(self,
                 node_in_feats,
                 edge_in_feats,
                 node_out_feats=128,
                 edge_hidden_feats=256,
                 num_preds=5,
                 num_step_set2set=6,
                 num_layer_set2set=3,
                 num_message_passing=6
                ):
        super(Electron_Predictor_3k, self).__init__()

        self.gnn = Electron_MPNN_3k(node_in_feats=node_in_feats,
                                 node_out_feats=node_out_feats,
                                 edge_in_feats=edge_in_feats,
                                 edge_hidden_feats=edge_hidden_feats,
                                 num_message_passing=num_message_passing
                                )
                    
        self.readout = Set2Set(input_dim=node_out_feats,
                               n_iters=num_step_set2set,
                               n_layers=num_layer_set2set)
        
        self.predict = nn.Sequential(
            nn.Linear(2 * node_out_feats, node_out_feats),
            nn.ReLU(),
            nn.Linear(node_out_feats, num_preds)
        )
    
    def forward(self, g, node_feats=None, edge_feats=None):

        # If trying to make a prediction based on smiles object (for evaluation), convert to graph
        if isinstance(g, str):

            g = graphize.smiles_to_bigraph(g,
                                       node_featurizer = graphize.featurize_atoms,
                                       edge_featurizer = graphize.featurize_bonds,
                                       explicit_hydrogens = True
                                 )
            edge_feats = g.edata['bond_feats'].float()
            node_feats = g.ndata['atom_feats'].float()
        

        node_feats = self.gnn(g, node_feats, edge_feats)

        graph_feats = self.readout(g, node_feats)

        # For our model, we take the mean of 5 different predictions
        return [self.predict(graph_feats).mean()]<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.06507090415367596, 'fraction_numerical': 0.011063059438801167, 'mean_word_length': 2.6125340599455043, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 28, 'https://': 7, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 17, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '164396', 'license': 'MIT', 'n_tokens_mistral': 5908, 'n_tokens_neox': 5617, 'n_words': 2010, 'path': 'model.ipynb', 'repo_name': 'RobertSmithers/PredictingSolubility', 'size': '19886'}