From 39ab2ca686c9d33d9c2c0c130e16c223feffdb0d Mon Sep 17 00:00:00 2001
From: Gaurav S Deshmukh <deshmukg@gilbreth-fe01.rcac.purdue.edu>
Date: Fri, 22 Sep 2023 18:32:54 -0400
Subject: [PATCH] Added model

---
 src/models.py | 153 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 153 insertions(+)
 create mode 100644 src/models.py

diff --git a/src/models.py b/src/models.py
new file mode 100644
index 0000000..a7f2e33
--- /dev/null
+++ b/src/models.py
@@ -0,0 +1,153 @@
+"""Graph neural network models."""
+
+import torch
+import torch.nn as nn
+import torch_geometric.nn as gnn
+
+class MultiGCN(gnn.MessagePassing):
+    """Class to customize the graph neural network."""
+    def __init__(self, partition_configs):
+        """Initialize the graph neural network.
+
+        Parameters
+        ----------
+        partition_configs: List[Dict]
+            List of dictionaries containing parameters for the GNN for each
+            partition. The number of different GNNs are judged based on the
+            size of the list. Each partition config should contain the following
+            keys: n_conv (number of convolutional layers, int), n_hidden (number 
+            of hidden layers, int), conv_size (feature size before convolution, int)
+            hidden_size (nodes per hidden layer node, int), dropout (dropout 
+            probability for hidden layers, float), conv_type (type of convolution 
+            layer, str; currently only "CGConv" is supported), pool_type
+            (type of pooling layer, str; currently "add" and "mean" are supported),
+            num_node_features (number of node features, int), num_edge_features
+            (number of edge features, int).
+        """
+        # Store hyperparameters
+        self.n_conv = [config["n_conv"] for config in partition_configs]
+        self.n_hidden = [config["n_hidden"] for config in partition_configs]
+        self.hidden_size = [config["hidden_size"] for config in partition_configs]
+        self.conv_size = [config["conv_size"] for config in partition_configs]
+        self.conv_type = [config["conv_type"] for config in partition_configs]
+        self.dropout = [config["dropout"] for config in partition_configs]
+        self.num_node_features = [
+            config["num_node_features"] for config in partition_configs
+        ]
+        self.num_edge_features = [
+            config["num_node_features"] for config in partition_configs
+        ]
+        self.n_partitions = len(partition_configs)
+
+        # Initialize layers
+        # Initial transform
+        self.init_transform = []
+        for i in range(self.n_partitions):
+            self.init_transform.append(
+                nn.ModuleList(
+                    nn.Linear(self.num_node_features[i], self.conv_size[i]),
+                    nn.LeakyReLU(inplace=True),
+                )
+            )
+
+        # Convolutional layers
+        self.init_conv_layers()
+
+        # Pooling layers
+        self.pool_layers = []
+        for i in range(self.n_partitions):
+            self.pool_layers.append(gnn.pool.global_addpool())
+
+        # Pool transform
+        self.pool_transform = []
+        for i in range(self.n_partitions):
+            self.pool_transform.append(
+                nn.ModuleList(
+                    nn.Linear(self.conv_size[i], self.hidden_size[i]),
+                    nn.LeakyReLU(inplace=True),
+                )
+            )
+
+        # Hidden layers
+        self.hidden_layers = []
+        for i in range(self.n_partitions):
+            self.hidden_layers.append(
+                nn.ModuleList([
+                    nn.Linear(self.hidden_size[i], self.hidden_size[i]),
+                    nn.LeakyReLU(inplace=True),
+                    nn.Dropout(p=self.dropout),
+                ] * (self.hidden_layers - 1) +
+                [
+                    nn.Linear(self.hidden_size[i], 1),
+                    nn.LeakyReLU(inplace=True),
+                    nn.Dropout(p=self.dropout),
+                ]
+                )
+            )
+
+        # Final linear layer
+        # TODO: replace 1 with multiple outputs
+        self.final_lin_transform = nn.Linear(self.n_partitions, 1)
+
+
+    def init_conv_layers(self):
+        """Initialize convolutional layers."""
+        self.conv_layers = []
+        for i in range(self.n_partitions):
+            part_conv_layers = []
+            for j in range(self.n_conv):
+                conv_layer = [
+                    gnn.CGConv(
+                        channels=self.num_node_features[i],
+                        dim=self.num_edge_features[i],
+                        batch_norm=True
+                    ),
+                    nn.LeakyReLU(inplace=True)
+                ]
+                part_conv_layers.append(conv_layer)
+            
+            self.conv_layers.append(nn.ModuleList(part_conv_layers))
+
+    def forward(self, data_objects):
+        """Foward pass of the network(s).
+
+        Parameters
+        ----------
+        data_objects: list
+            List of data objects, each corresponding to a graph of a partition
+            of an atomic structure.
+
+        Returns
+        ------
+        dict
+            Dictionary containing "output" and "contributions".
+        """
+        # Initialize empty list for contributions
+        contributions = []
+        # For each data object
+        for i, data in enumerate(data_objects):
+            # Apply initial transform
+            conv_data = self.init_transform[i](data) 
+
+            # Apply convolutional layers
+            for layer in self.conv_layers[i]:
+                conv_data = layer(conv_data)
+
+            # Apply pooling layer
+            pooled_data = self.pool_layers[i](conv_data)
+
+            # Apply pool-to-hidden transform
+            hidden_data = self.pool_transform[i](pooled_data)
+
+            # Apply hidden layers
+            for layer in self.hidden_layers[i]:
+                hidden_data = layer(hidden_data)
+
+            # Save contribution
+            contributions.append(hidden_data)
+
+        # Apply final transformation
+        output = self.final_lin_transform(*contributions)
+
+        return {"output": output, "contributions": contributions}
+        
\ No newline at end of file