Fix transform bugs

Author: sidnb13

.gitignore

@@ -178,3 +178,5 @@ main.py
 
 test*.py
 test*.ipynb
+
+!testing/*

configs/config_alignn.yml

@@ -0,0 +1,88 @@
+
+trainer: property
+
+task:
+  # run_mode: train
+  name: "alignn_first_training"
+
+  reprocess: "False"
+
+  parallel: "True"
+  seed: 0
+  #seed=0 means random initalization
+
+  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 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"
+
+optim:
+  max_epochs: 300
+  lr: 0.001
+  #Loss functions (from pytorch) examples: l1_loss, mse_loss, binary_cross_entropy
+  loss_fn: "mse_loss"
+  batch_size: 64
+ 
+  optimizer:
+    optimizer_type: "AdamW"
+    optimizer_args: {"weight_decay": 0.00001}
+  scheduler:
+    scheduler_type: "OneCycleLR"
+    # Look further into steps per epoch, for now hardcoded calculation from paper
+    scheduler_args: {"max_lr": 0.001, "epochs": 300, "steps_per_epoch": 1}
+
+dataset:
+  processed: True # if False, need to preprocessor data and generate .pt file
+  # Whether to use "inmemory" or "large" format for pytorch-geometric dataset. Reccomend inmemory unless the dataset is too large
+  # dataset_type: "inmemory"
+  #Path to data files
+  src: "/global/cfs/projectdirs/m3641/Shared/Materials_datasets/MP_data_69K/raw/"
+  #Path to target file within data_path
+  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/"
+  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)
+  node_representation: "onehot"
+  #Print out processing info
+  verbose: "True"
+
+  #Loading dataset params
+  #Index of target column in targets.csv
+  target_index: 0
+
+  #graph specific settings
+  cutoff_radius : 8.0
+  n_neighbors : 12
+  edge_steps : 50

matdeeplearn/common/data.py

@@ -3,9 +3,10 @@
 import torch
 from torch.utils.data import random_split
 from torch_geometric.loader import DataLoader
+from torch_geometric.transforms import Compose
 
 from matdeeplearn.preprocessor.datasets import LargeStructureDataset, StructureDataset
-from matdeeplearn.preprocessor.transforms import GetY
+from matdeeplearn.preprocessor.transforms import TRANSFORM_REGISTRY, GetY
 
 
 # train test split
@@ -58,7 +59,7 @@ def dataset_split(
 
 
 def get_dataset(
-    data_path, target_index: int = 0, transform_type="GetY", large_dataset=False
+    data_path, target_index: int = 0, transform_list=[], large_dataset=False
 ):
     """
     get dataset according to data_path
@@ -78,22 +79,25 @@ def get_dataset(
         particular dataset, thus we need to index one column for
         the current run/experiment
 
-    transform_type: transformation function/class to be applied
+    transform_list: transformation function/classes to be applied
     """
+    
+    transforms = [GetY(index=target_index)]
 
     # set transform method
-    if transform_type == "GetY":
-        T = GetY
-    else:
-        raise ValueError("No such transform found for {transform}")
+    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:
         Dataset = LargeStructureDataset
     else:
         Dataset = StructureDataset
 
-    transform = T(index=target_index)
+    transform = Compose(transforms)
 
     return Dataset(data_path, processed_data_path="", transform=transform)
 

matdeeplearn/models/alignn.py

@@ -356,6 +356,10 @@ def __init__(
             )
         elif link == "logit":
             self.link = torch.sigmoid
+            
+    @property
+    def target_attr(self):
+        return "y"
 
     def forward(self, g: Data):
         # Compute OTF transform to generate attributes for L(g)

matdeeplearn/preprocessor/helpers.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 
 import torch
+from torch_sparse import SparseTensor
 import torch.nn.functional as F
 from torch_geometric.utils import dense_to_sparse, degree, add_self_loops
 from torch_geometric.data.data import Data

matdeeplearn/preprocessor/transforms.py

@@ -6,6 +6,7 @@
 from torch_geometric.utils import remove_self_loops
 from matdeeplearn.preprocessor.helpers import compute_bond_angles, triplets
 from scipy.spatial.distance import cdist
+from contextlib import contextmanager
 
 '''
 here resides the transform classes needed for data processing
@@ -16,6 +17,16 @@
     The data object will be transformed before every access.
 '''
 
+TRANSFORM_REGISTRY = {}
+
+
+def register_transform(transform_name):
+    '''Registers a transform function for bookkeeping.'''
+    def registered_transform(transform):
+        TRANSFORM_REGISTRY[transform_name] = transform
+        return transform
+    return registered_transform
+
 
 class GetY(object):
     def __init__(self, index=0):
@@ -28,6 +39,7 @@ def __call__(self, data):
         return data
 
 
+@register_transform("NumNodeTransform")
 class NumNodeTransform(object):
     '''
     Adds the number of nodes to the data object
@@ -38,6 +50,7 @@ def __call__(self, data):
         return data
 
 
+@register_transform("LineGraphMod")
 class LineGraphMod(object):
     '''
     Adds line graph attributes to the data object
@@ -47,59 +60,37 @@ def __call__(self, data):
         # CODE FROM PYG LINEGRAPH TRANSFORM (DIRECTED)
         N = data.num_nodes
         edge_index, edge_attr = data.edge_index, data.edge_attr
-        (row, col), edge_attr = coalesce(edge_index, edge_attr, N, N)
-
-        i = torch.arange(row.size(0), dtype=torch.long, device=row.device)
-        count = scatter_add(torch.ones_like(row), row, dim=0,
-                            dim_size=data.num_nodes)
-        cumsum = torch.cat([count.new_zeros(1), count.cumsum(0)], dim=0)
-
-        cols = [
-            i[cumsum[col[j]]:cumsum[col[j] + 1]]
-            for j in range(col.size(0))
-        ]
-        rows = [row.new_full((c.numel(), ), j) for j, c in enumerate(cols)]
-
-        row, col = torch.cat(rows, dim=0), torch.cat(cols, dim=0)
-
-        data.edge_index_lg = torch.stack([row, col], dim=0)
-        data.x_lg = data.edge_attr
-        data.num_nodes_lg = edge_index.size(1)
-
-        # CUSTOM CODE FOR CALCULATING EDGE ATTRIBUTES
-        edge_attr_lg = torch.zeros(
-            (data.edge_index_lg.shape[1], 1), device='cuda')
+        _, edge_attr = coalesce(edge_index, edge_attr, N, N)
 
         # compute bond angles
         angles, idx_kj, idx_ji = compute_bond_angles(
             data.pos, data.cell_offsets, data.edge_index, data.num_nodes)
         triplet_pairs = torch.stack([idx_kj, idx_ji], dim=0)
 
-        # move triplets and edges to CPU for sklearn based calculation
-        match_indices = torch.Tensor(
-            np.where(cdist(data.edge_index_lg.T.cpu(), triplet_pairs.T.cpu()) == 0)[
-                0].reshape(-1, 1)
-        ).type(torch.long)
+        data.edge_index_lg = triplet_pairs
+        data.x_lg = data.edge_attr
+        data.num_nodes_lg = edge_index.size(1)
 
         # assign bond angles to edge attributes
-        edge_attr_lg[match_indices.squeeze(-1)] = angles.reshape(-1, 1)
+        data.edge_attr_lg = angles.reshape(-1, 1)
 
-        data.edge_attr_lg = edge_attr_lg
-        
         return data
 
+
+@register_transform("ToFloat")
 class ToFloat(object):
     '''
     Convert non-int attributes to float
     '''
+
     def __call__(self, data):
         data.x = data.x.float()
         data.x_lg = data.x_lg.float()
-        
+
         data.distances = data.distances.float()
         data.pos = data.pos.float()
 
         data.edge_attr = data.edge_attr.float()
         data.edge_attr_lg = data.edge_attr_lg.float()
 
-        return data
+        return data

matdeeplearn/trainers/base_trainer.py

@@ -1,6 +1,7 @@
 import copy
 import csv
 import logging
+import re
 import os
 from abc import ABC, abstractmethod
 from datetime import datetime
@@ -41,7 +42,8 @@ def __init__(
         identifier: str = None,
         verbosity: int = None,
     ):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu")
         self.model = model.to(self.device)
         self.dataset = dataset
         self.optimizer = optimizer
@@ -105,7 +107,8 @@ def from_config(cls, config):
         train_loader, val_loader, test_loader = cls._load_dataloader(
             config["optim"], config["dataset"], dataset, sampler
         )
-        scheduler = cls._load_scheduler(config["optim"]["scheduler"], optimizer)
+        scheduler = cls._load_scheduler(
+            config["optim"]["scheduler"], optimizer)
         loss = cls._load_loss(config["optim"]["loss"])
 
         max_epochs = config["optim"]["max_epochs"]
@@ -133,7 +136,7 @@ def _load_dataset(dataset_config):
         dataset_path = dataset_config["pt_path"]
         target_index = dataset_config.get("target_index", 0)
 
-        dataset = get_dataset(dataset_path, target_index)
+        dataset = get_dataset(dataset_path, target_index, transforms_list=dataset_config["transforms"])
 
         return dataset
 
@@ -180,7 +183,8 @@ def _load_dataloader(optim_config, dataset_config, dataset, sampler):
         train_loader = get_dataloader(
             train_dataset, batch_size=batch_size, sampler=sampler
         )
-        val_loader = get_dataloader(val_dataset, batch_size=batch_size, sampler=sampler)
+        val_loader = get_dataloader(
+            val_dataset, batch_size=batch_size, sampler=sampler)
         test_loader = get_dataloader(
             test_dataset, batch_size=batch_size, sampler=sampler
         )
@@ -222,7 +226,8 @@ def predict(self):
 
     def update_best_model(self, val_metrics):
         """Updates the best val metric and model, saves the best model, and saves the best model predictions"""
-        self.best_val_metric = val_metrics[type(self.loss_fn).__name__]["metric"]
+        self.best_val_metric = val_metrics[type(
+            self.loss_fn).__name__]["metric"]
         self.best_model_state = copy.deepcopy(self.model.state_dict())
 
         self.save_model("best_checkpoint.pt", val_metrics, False)
@@ -247,7 +252,8 @@ def save_model(self, checkpoint_file, val_metrics=None, training_state=True):
                 "best_val_metric": self.best_val_metric,
             }
         else:
-            state = {"state_dict": self.model.state_dict(), "val_metrics": val_metrics}
+            state = {"state_dict": self.model.state_dict(),
+                     "val_metrics": val_metrics}
 
         checkpoint_dir = os.path.join(
             self.run_dir, "results", self.timestamp_id, "checkpoint"
@@ -268,7 +274,8 @@ def save_results(self, output, filename, node_level_predictions=False):
         if node_level_predictions:
             id_headers += ["node_id"]
         num_cols = (shape[1] - len(id_headers)) // 2
-        headers = id_headers + ["target"] * num_cols + ["prediction"] * num_cols
+        headers = id_headers + ["target"] * \
+            num_cols + ["prediction"] * num_cols
 
         with open(filename, "w") as f:
             csvwriter = csv.writer(f)