Add in pre-training transform functionality

Author: sidnb13

configs/config_alignn.yml

@@ -3,52 +3,53 @@ trainer: property
 
 task:
   # run_mode: train
-  name: "alignn_first_training"
+  identifier: "alignn_train_job"
 
-  reprocess: "False"
+  reprocess: False
 
-  parallel: "True"
+
+  parallel: True
   seed: 0
   #seed=0 means random initalization
 
-  write_output: "True"
-  parallel: "True"
+
+  write_output: True
+  parallel: True
   #Training print out frequency (print per n number of epochs)
   verbosity: 1
 
-  #Ratios for train/val/test split out of a total of 1
-  train_ratio: 0.85
-  val_ratio: 0.05
-  test_ratio: 0.10
 
-model:
-  name: "ALIGNN_GRAPHITE"
-  load_model: "False"
-  save_model: "True"
-  model_path: "/global/cfs/projectdirs/m3641/Sidharth/MatDeepLearn_dev/testing/models/alignn_model_t1.pth"
 
+model:
+  name: ALIGNN
+  load_model: False
+  save_model: True
+  model_path: "my_model.pth"
+  edge_steps: 50
+  self_loop: True
   #model attributes
-  alignn_layers: 4
-  gcn_layers: 4
-  atom_input_features: 114
-  edge_input_features: 50
-  triplet_input_features: 40
-  embedding_features: 32
-  hidden_features: 64
-  output_features: 1
-  # min_edge_distance: 0.0,
-  # max_edge_distance: 8.0,
-  # min_angle: 0.0,
-  # max_angle: torch.acos(torch.zeros(1)).item() * 2,
-  link: "identity"
+  dim1: 100
+  dim2: 150
+  pre_fc_count: 1
+  gc_count: 4
+  post_fc_count: 3
+  pool: "global_mean_pool"
+  pool_order: "early"
+  batch_norm: True
+  batch_track_stats: True
+  act: "relu"
+  dropout_rate: 0.0
 
 optim:
-  max_epochs: 300
+  max_epochs: 250
   lr: 0.001
-  #Loss functions (from pytorch) examples: l1_loss, mse_loss, binary_cross_entropy
-  loss_fn: "mse_loss"
+  #Either custom or from torch.nn.functional library. If from torch, loss_type is TorchLossWrapper
+  loss:
+    loss_type: "TorchLossWrapper"
+    loss_args: {"loss_fn": "mse_loss"}
+
   batch_size: 64
- 
+
   optimizer:
     optimizer_type: "AdamW"
     optimizer_args: {"weight_decay": 0.00001}
@@ -67,16 +68,18 @@ dataset:
   target_path: "/global/cfs/projectdirs/m3641/Shared/Materials_datasets/MP_data_69K/targets.csv"
   #Path to save processed data.pt file (a directory path not filepath)
   pt_path: "/global/cfs/projectdirs/m3641/Sidharth/datasets/MP_data_69K/"
+  otf: False
   transforms:
     - NumNodeTransform
     - LineGraphMod
     - ToFloat
   #Format of data files (limit to those supported by ASE)
   data_format: "json"
-  #Method of obtaining atom dictionary: available:(one-hot)
+  #Method of obtaining atom idctionary: available:(onehot)
   node_representation: "onehot"
+  additional_attributes: []
   #Print out processing info
-  verbose: "True"
+  verbose: True
 
   #Loading dataset params
   #Index of target column in targets.csv
@@ -85,4 +88,9 @@ dataset:
   #graph specific settings
   cutoff_radius : 8.0
   n_neighbors : 12
-  edge_steps : 50
+  edge_steps : 50
+
+  #Ratios for train/val/test split out of a total of 1
+  train_ratio: 0.8
+  val_ratio: 0.05
+  test_ratio: 0.15

matdeeplearn/common/data.py

@@ -59,7 +59,7 @@ def dataset_split(
 
 
 def get_dataset(
-    data_path, target_index: int = 0, transform_list=[], large_dataset=False
+    data_path, target_index: int = 0, transform_list=[], otf=False, large_dataset=False
 ):
     """
     get dataset according to data_path
@@ -81,15 +81,16 @@ def get_dataset(
 
     transform_list: transformation function/classes to be applied
     """
-    
+
     transforms = [GetY(index=target_index)]
 
     # set transform method
-    for transform in transform_list:
-        if transform in TRANSFORM_REGISTRY:
-            transforms.append(TRANSFORM_REGISTRY[transform]())
-        else:
-            raise ValueError("No such transform found for {transform}")
+    if otf:
+        for transform in transform_list:
+            if transform in TRANSFORM_REGISTRY:
+                transforms.append(TRANSFORM_REGISTRY[transform]())
+            else:
+                raise ValueError("No such transform found for {transform}")
 
     # check if large dataset is needed
     if large_dataset:
@@ -98,7 +99,7 @@ def get_dataset(
         Dataset = StructureDataset
 
     transform = Compose(transforms)
-
+    
     return Dataset(data_path, processed_data_path="", transform=transform)
 
 

matdeeplearn/models/alignn.py

@@ -7,7 +7,6 @@
 from torch_geometric.transforms import Compose
 from matdeeplearn.common.registry import registry
 from matdeeplearn.models.base_model import BaseModel
-from matdeeplearn.preprocessor.transforms import NumNodeTransform, LineGraphMod, ToFloat
 from typing import Optional, Literal
 import numpy as np
 import contextlib
@@ -362,14 +361,6 @@ def target_attr(self):
         return "y"
 
     def forward(self, g: Data):
-        # Compute OTF transform to generate attributes for L(g)
-        
-        # with prof_ctx():
-        
-        with torch.no_grad():
-            otf = Compose([NumNodeTransform(), LineGraphMod(),  ToFloat()])
-            otf(g)
-
         # initial node features
         node_feats = self.atom_embedding(g.x)
         # initial bond features

matdeeplearn/models/alignn_graphite.py

@@ -116,6 +116,10 @@ def __init__(self, dim=100, num_interactions=6, num_species=3, cutoff=3.0, **kwa
         )
 
         self.reset_parameters()
+        
+    @property
+    def target_attr(self):
+        return "y"
 
     def reset_parameters(self):
         self.embed_atm.reset_parameters()
@@ -127,14 +131,10 @@ def embed_ang(self, x_ang):
         cos_ang = torch.cos(x_ang)
         return gaussian(cos_ang, start=-1, end=1, num_basis=self.dim)
 
-    def forward(self, data: Data):
-        with torch.no_grad():
-            otf = Compose([NumNodeTransform(), LineGraphMod(),  ToFloat()])
-            otf(data)
-        
+    def forward(self, data: Data):        
         edge_index_G = data.edge_index
         edge_index_A = data.edge_index_lg
-        h_atm = self.embed_atm(data.x)
+        h_atm = self.embed_atm(data.x.type(torch.long))
         h_bnd = self.embed_bnd(data.edge_attr)
         h_ang = self.embed_ang(data.edge_attr_lg)
 

matdeeplearn/preprocessor/helpers.py

@@ -279,6 +279,7 @@ def generate_node_features(input_data, n_neighbors, device):
 
     for i, data in enumerate(input_data):
         input_data[i] = one_hot_degree(data, n_neighbors+1)
+        
 
 def generate_edge_features(input_data, edge_steps, r, device):
     distance_gaussian = GaussianSmearing(0, 1, edge_steps, 0.2, device=device)
@@ -333,7 +334,7 @@ def compute_bond_angles(pos: torch.Tensor, offsets: torch.Tensor, edge_index: to
 
     # Calculate triplets
     idx_i, idx_j, idx_k, idx_kj, idx_ji = triplets(
-        edge_index, offsets, num_nodes)
+        edge_index, offsets.to(device=edge_index.device), num_nodes)
 
     # Calculate angles.
     pos_i = pos[idx_i]

matdeeplearn/preprocessor/processor.py

@@ -9,6 +9,7 @@
 from ase import io
 from torch_geometric.data import Data, InMemoryDataset
 from torch_geometric.utils import dense_to_sparse
+from torch_geometric.transforms import Compose
 from tqdm import tqdm
 
 from matdeeplearn.preprocessor.helpers import (
@@ -18,6 +19,8 @@
     get_cutoff_distance_matrix,
 )
 
+from matdeeplearn.preprocessor.transforms import TRANSFORM_REGISTRY
+
 
 def process_data(dataset_config):
     root_path = dataset_config["src"]
@@ -41,6 +44,8 @@ def process_data(dataset_config):
         r=cutoff_radius,
         n_neighbors=n_neighbors,
         edge_steps=edge_steps,
+        otf=dataset_config.get("otf", False),
+        transforms=dataset_config.get("transforms", []),
         data_format=data_format,
         image_selfloop=image_selfloop,
         self_loop=self_loop,
@@ -61,6 +66,8 @@ def __init__(
         r: float,
         n_neighbors: int,
         edge_steps: int,
+        otf: bool = False,
+        transforms: list = [],
         data_format: str = "json",
         image_selfloop: bool = True,
         self_loop: bool = True,
@@ -132,6 +139,9 @@ def __init__(
         self.verbose = verbose
         self.device = device
 
+        self.otf = otf
+        self.transforms = transforms
+
         self.disable_tqdm = logging.root.level > logging.INFO
 
     def src_check(self):
@@ -153,14 +163,17 @@ def ase_wrap(self):
         dict_structures = []
         ase_structures = []
 
-        logging.info("Converting data to standardized form for downstream processing.")
+        logging.info(
+            "Converting data to standardized form for downstream processing.")
         for i, structure_id in enumerate(file_names):
-            p = os.path.join(self.root_path, str(structure_id) + "." + self.data_format)
+            p = os.path.join(self.root_path, str(
+                structure_id) + "." + self.data_format)
             ase_structures.append(ase.io.read(p))
 
         for i, s in enumerate(tqdm(ase_structures, disable=self.disable_tqdm)):
             d = {}
-            pos = torch.tensor(s.get_positions(), device=self.device, dtype=torch.float)
+            pos = torch.tensor(s.get_positions(),
+                               device=self.device, dtype=torch.float)
             cell = torch.tensor(
                 np.array(s.get_cell()), device=self.device, dtype=torch.float
             )
@@ -173,7 +186,8 @@ def ase_wrap(self):
 
             # add additional attributes
             if self.additional_attributes:
-                attributes = self.get_csv_additional_attributes(d["structure_id"])
+                attributes = self.get_csv_additional_attributes(
+                    d["structure_id"])
                 for k, v in attributes.items():
                     d[k] = v
 
@@ -189,9 +203,12 @@ def get_csv_additional_attributes(self, structure_id):
         attributes = {}
 
         for attr in self.additional_attributes:
-            p = os.path.join(self.root_path, structure_id + "_" + attr + ".csv")
-            values = np.genfromtxt(p, delimiter=",", dtype=float, encoding=None)
-            values = torch.tensor(values, device=self.device, dtype=torch.float)
+            p = os.path.join(self.root_path, structure_id +
+                             "_" + attr + ".csv")
+            values = np.genfromtxt(
+                p, delimiter=",", dtype=float, encoding=None)
+            values = torch.tensor(
+                values, device=self.device, dtype=torch.float)
             attributes[attr] = values
 
         return attributes
@@ -212,13 +229,17 @@ def json_wrap(self):
 
         dict_structures = []
         y = []
-        y_dim = len(original_structures[0]["y"]) if isinstance(original_structures[0]["y"], list) else 1
+        y_dim = len(original_structures[0]["y"]) if isinstance(
+            original_structures[0]["y"], list) else 1
 
-        logging.info("Converting data to standardized form for downstream processing.")
+        logging.info(
+            "Converting data to standardized form for downstream processing.")
         for i, s in enumerate(tqdm(original_structures, disable=self.disable_tqdm)):
             d = {}
-            pos = torch.tensor(s["positions"], device=self.device, dtype=torch.float)
-            cell = torch.tensor(s["cell"], device=self.device, dtype=torch.float)
+            pos = torch.tensor(
+                s["positions"], device=self.device, dtype=torch.float)
+            cell = torch.tensor(
+                s["cell"], device=self.device, dtype=torch.float)
             atomic_numbers = torch.LongTensor(s["atomic_numbers"])
 
             d["positions"] = pos
@@ -268,6 +289,7 @@ def get_data_list(self, dict_structures, y):
         data_list = [Data() for _ in range(n_structures)]
 
         logging.info("Getting torch_geometric.data.Data() objects.")
+
         for i, sdict in enumerate(tqdm(dict_structures, disable=self.disable_tqdm)):
             target_val = y[i]
             data = data_list[i]
@@ -312,7 +334,26 @@ def get_data_list(self, dict_structures, y):
         generate_node_features(data_list, self.n_neighbors, device=self.device)
 
         logging.info("Generating edge features...")
-        generate_edge_features(data_list, self.edge_steps, self.r, device=self.device)
+        generate_edge_features(data_list, self.edge_steps,
+                               self.r, device=self.device)
+
+        logging.info("Applying transforms...")
+
+        # saving line graph attributes through transforms
+        transforms_list = []
+
+        if not self.otf:
+            for transform in self.transforms:
+                if transform in TRANSFORM_REGISTRY:
+                    transforms_list.append(TRANSFORM_REGISTRY[transform]())
+                else:
+                    raise ValueError(
+                        "No such transform found for {transform}")
+
+        composition = Compose(transforms_list)
+
+        for data in data_list:
+            composition(data)
 
         clean_up(data_list, ["edge_descriptor"])