Merge pull request #22 from GreeleyGroup/fix/contributions

FIX: Fixed convergence issues by adding SlabGCN

src/models.py

@@ -81,25 +81,26 @@ def __init__(self, partition_configs):
                         [
                             nn.Linear(self.hidden_size[i], self.hidden_size[i]),
                             nn.LeakyReLU(inplace=True),
-                            nn.Dropout(p=self.dropout[i]),
                         ]
                         * (self.n_hidden[i] - 1)
                         + [
                             nn.Linear(self.hidden_size[i], 1),
                             nn.LeakyReLU(inplace=True),
-                            nn.Dropout(p=self.dropout[i]),
                         ]
                     )
                 )
             )
 
         # Final linear layer
         # TODO: replace 1 with multiple outputs
-        self.final_lin_transform = nn.Linear(self.n_partitions, 1, bias=False)
-        with torch.no_grad():
-            self.final_lin_transform.weight.copy_(torch.ones(self.n_partitions))
-        for p in self.final_lin_transform.parameters():
-            p.requires_grad = False
+        self.fin1_transform = nn.Linear(self.n_partitions, 50)
+        self.fin1_act = nn.LeakyReLU()
+        self.final_lin_transform = nn.Linear(50, 1)
+        self.final_lin_act = nn.LeakyReLU()
+        # with torch.no_grad():
+        #    self.final_lin_transform.weight.copy_(torch.ones(self.n_partitions))
+        # for p in self.final_lin_transform.parameters():
+        #    p.requires_grad = False
 
     def init_conv_layers(self):
         """Initialize convolutional layers."""
@@ -166,11 +167,199 @@ def forward(self, data_objects):
 
         # Apply final transformation
         contributions = torch.cat(contributions)
-        output = self.final_lin_transform(contributions.view(-1, 3))
+        output = self.fin1_transform(contributions.view(-1, 3))
+        output = self.fin1_act(output)
+        output = self.final_lin_transform(output)
+        output = self.final_lin_act(output)
 
         return {"output": output, "contributions": contributions}
 
 
+class SlabGCN(nn.Module):
+    """Class to customize the graph neural network."""
+
+    def __init__(self, partition_configs, global_config):
+        """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).
+        global_config: Dict
+            This should contain the following keys: n_hidden (Number of hidden
+            layers for the shared FFN, int), hidden_size (Number of nodes per
+            hidden layer, int), dropout (Probability of dropping a node, float).
+        """
+        super().__init__()
+
+        # Store hyperparameters
+        self.n_conv = [config["n_conv"] for config in partition_configs]
+        self.n_hidden = global_config["n_hidden"]
+        self.hidden_size = global_config["hidden_size"]
+        self.conv_size = [config["conv_size"] for config in partition_configs]
+        self.conv_type = [config["conv_type"] for config in partition_configs]
+        self.dropout = global_config.get("dropout", 0)
+        self.num_node_features = [
+            config["num_node_features"] for config in partition_configs
+        ]
+        self.num_edge_features = [
+            config["num_edge_features"] for config in partition_configs
+        ]
+        self.n_partitions = len(partition_configs)
+
+        # Initialize layers
+        # Initial transform
+        self.init_transform = nn.ModuleList()
+        for i in range(self.n_partitions):
+            self.init_transform.append(
+                nn.Sequential(
+                    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 = nn.ModuleList()
+        for i in range(self.n_partitions):
+            self.pool_layers.append(gnn.aggr.SumAggregation())
+
+        # Pool transform
+        self.pool_transform = nn.Sequential(
+            nn.Linear(sum(self.conv_size), self.hidden_size),
+            nn.LeakyReLU(inplace=True),
+        )
+
+        # Hidden layers
+        self.hidden_layers = nn.Sequential(
+            *(
+                [
+                    nn.Linear(self.hidden_size, self.hidden_size),
+                    nn.LeakyReLU(inplace=True),
+                    nn.Dropout(p=self.dropout),
+                ]
+                * (self.n_hidden - 1)
+                + [
+                    nn.Linear(self.hidden_size, 1),
+                    nn.LeakyReLU(inplace=True),
+                    nn.Dropout(p=self.dropout),
+                ]
+            )
+        )
+
+    def init_conv_layers(self):
+        """Initialize convolutional layers."""
+        self.conv_layers = nn.ModuleList()
+        for i in range(self.n_partitions):
+            part_conv_layers = []
+            for j in range(self.n_conv[i]):
+                # TODO Add possibility of changing convolutional layers
+                conv_layer = [
+                    gnn.CGConv(
+                        channels=self.conv_size[i],
+                        dim=self.num_edge_features[i],
+                        batch_norm=True,
+                    ),
+                    nn.LeakyReLU(inplace=True),
+                ]
+                part_conv_layers.extend(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".
+        """
+        # For each data object
+        embeddings = []
+        for i, data in enumerate(data_objects):
+            # Apply initial transform
+            conv_data = self.init_transform[i](data.x)
+
+            # Apply convolutional layers
+            for layer in self.conv_layers[i]:
+                if isinstance(layer, gnn.MessagePassing):
+                    conv_data = layer(
+                        x=conv_data,
+                        edge_index=data.edge_index,
+                        edge_attr=data.edge_attr,
+                    )
+                else:
+                    conv_data = layer(conv_data)
+
+            # Apply pooling layer
+            pooled_data = self.pool_layers[i](x=conv_data, index=data.batch)
+            embeddings.append(pooled_data)
+
+        # Apply pool-to-hidden transform
+        embedding = torch.cat(embeddings, dim=1)
+        hidden_data = self.pool_transform(embedding)
+
+        # Apply hidden layers
+        output = self.hidden_layers(hidden_data)
+
+        return {"output": output}
+
+    def get_embeddings(self, data_objects):
+        """Get the pooled embeddings of each partition.
+
+        Parameters
+        ----------
+        data_objects: list
+            List of data objects, each corresponding to a graph of a partition
+            of an atomic structure.
+
+        Returns
+        ------
+        embeddings: list
+            List of embedding tensors.
+        """
+        # For each data object
+        embeddings = []
+        for i, data in enumerate(data_objects):
+            # Apply initial transform
+            conv_data = self.init_transform[i](data.x)
+
+            # Apply convolutional layers
+            for layer in self.conv_layers[i]:
+                if isinstance(layer, gnn.MessagePassing):
+                    conv_data = layer(
+                        x=conv_data,
+                        edge_index=data.edge_index,
+                        edge_attr=data.edge_attr,
+                    )
+                else:
+                    conv_data = layer(conv_data)
+
+            # Apply pooling layer
+            pooled_data = self.pool_layers[i](x=conv_data, index=data.batch)
+            embeddings.append(pooled_data)
+
+        return embeddings
+
+
 if __name__ == "__main__":
     from pathlib import Path
 
@@ -185,7 +374,7 @@ def forward(self, data_objects):
     atoms = read(data_root_path / "Pt_3_Rh_9_-7-7-S.cif")
     datapoint = AtomsDatapoints(atoms)
     datapoint.process_data(
-        z_cutoffs=[13.0, 20.0],
+        layer_cutoffs=[3, 6],
         node_features=[
             ["atomic_number", "dband_center"],
             ["atomic_number", "reactivity"],
@@ -203,35 +392,26 @@ def forward(self, data_objects):
     partition_configs = [
         {
             "n_conv": 3,
-            "n_hidden": 3,
-            "hidden_size": 30,
             "conv_size": 40,
-            "dropout": 0.1,
             "num_node_features": data_objects[0].num_node_features,
             "num_edge_features": data_objects[0].num_edge_features,
             "conv_type": "CGConv",
         },
         {
             "n_conv": 3,
-            "n_hidden": 3,
-            "hidden_size": 30,
             "conv_size": 40,
-            "dropout": 0.1,
             "num_node_features": data_objects[1].num_node_features,
             "num_edge_features": data_objects[1].num_edge_features,
             "conv_type": "CGConv",
         },
         {
             "n_conv": 3,
-            "n_hidden": 3,
-            "hidden_size": 30,
             "conv_size": 40,
-            "dropout": 0.1,
             "num_node_features": data_objects[2].num_node_features,
             "num_edge_features": data_objects[2].num_edge_features,
             "conv_type": "CGConv",
         },
     ]
-    net = MultiGCN(partition_configs)
+    net = SlabGCN(partition_configs)
     result_dict = net(data_objects)
     print(result_dict)

src/train.py

@@ -7,7 +7,7 @@
 import torch
 from sklearn.metrics import mean_absolute_error, mean_squared_error
 
-from .models import MultiGCN
+from .models import MultiGCN, SlabGCN
 
 
 class Standardizer:
@@ -142,7 +142,7 @@ def __init__(self, global_config, partition_configs, model_path):
             Path where the model is to be saved
         """
         # Create model
-        self.model = MultiGCN(partition_configs)
+        self.model = SlabGCN(partition_configs, global_config)
 
         # Create model path
         self.make_directory_structure(model_path)
@@ -336,14 +336,13 @@ def predict(self, dataset, indices, return_targets=False):
         Returns
         -------
         prediction_dict: dict
-            Dictionary containing "targets", "predictions", "contributions" and
+            Dictionary containing "targets", "predictions", and
             "indices" (copy of predict_idx).
         """
         # Create arrays
         n_partitions = len(dataset.get(indices[0]))
         targets = np.zeros(len(indices))
         predictions = np.zeros(len(indices))
-        contributions = np.zeros((len(indices), n_partitions))
 
         # Enable eval mode of model
         self.model.eval()
@@ -368,16 +367,11 @@ def predict(self, dataset, indices, return_targets=False):
             predictions[i] = self.standardizer.restore(
                 pred_dict["output"].cpu().detach()
             )
-            conts_std = pred_dict["contributions"].cpu().detach()
-            contributions[i, :] = (
-                self.standardizer.restore_cont(conts_std).numpy().flatten()
-            )
 
         predictions_dict = {
             "targets": targets,
             "predictions": predictions,
             "indices": indices,
-            "contributions": contributions,
         }
 
         return predictions_dict