From 79a9704b53031e631af4740cb0dd2361793b5577 Mon Sep 17 00:00:00 2001 From: Gaurav S Deshmukh Date: Fri, 22 Sep 2023 21:56:51 -0400 Subject: [PATCH] Fixed codestyle --- src/models.py | 120 +++++++++++++++++++++++++++----------------------- 1 file changed, 66 insertions(+), 54 deletions(-) diff --git a/src/models.py b/src/models.py index 000d01b..615ddc8 100644 --- a/src/models.py +++ b/src/models.py @@ -4,8 +4,10 @@ import torch.nn as nn import torch_geometric.nn as gnn + class MultiGCN(nn.Module): """Class to customize the graph neural network.""" + def __init__(self, partition_configs): """Initialize the graph neural network. @@ -15,17 +17,17 @@ def __init__(self, partition_configs): 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 + 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 + 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). """ super().__init__() - + # Store hyperparameters self.n_conv = [config["n_conv"] for config in partition_configs] self.n_hidden = [config["n_hidden"] for config in partition_configs] @@ -74,16 +76,20 @@ def __init__(self, partition_configs): self.hidden_layers = [] for i in range(self.n_partitions): self.hidden_layers.append( - nn.Sequential(*([ - 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]), - ]) + nn.Sequential( + *( + [ + 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]), + ] + ) ) ) @@ -91,7 +97,6 @@ def __init__(self, partition_configs): # 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 = [] @@ -105,10 +110,10 @@ def init_conv_layers(self): dim=self.num_edge_features[i], batch_norm=True, ), - nn.LeakyReLU(inplace=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): @@ -130,13 +135,16 @@ def forward(self, data_objects): # For each data object for i, data in enumerate(data_objects): # Apply initial transform - conv_data = self.init_transform[i](data.x) + 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) + conv_data = layer( + x=conv_data, + edge_index=data.edge_index, + edge_attr=data.edge_attr, + ) else: conv_data = layer(conv_data) @@ -157,12 +165,16 @@ def forward(self, data_objects): output = self.final_lin_transform(contributions) return {"output": output, "contributions": contributions} - + + if __name__ == "__main__": + from pathlib import Path + from ase.io import read - from data import AtomsDatapoints + from constants import REPO_PATH - from pathlib import Path + from data import AtomsDatapoints + # Test for one tensor # Create datapoins data_root_path = Path(REPO_PATH) / "data" / "S_calcs" @@ -185,36 +197,36 @@ def forward(self, data_objects): # Get result 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", - } + { + "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) - result_dict = net(data_objects) \ No newline at end of file + result_dict = net(data_objects)