gridfm-graphkit 0.0.1__py3-none-any.whl

gridfm_graphkit/datasets/utils.py

@@ -0,0 +1,65 @@
+import numpy as np
+import pandas as pd
+import os
+from torch.utils.data import Subset
+from typing import Tuple
+
+
+def split_dataset(
+    dataset,
+    log_dir: str,
+    val_ratio: float = 0.1,
+    test_ratio: float = 0.1,
+) -> Tuple[Subset, Subset, Subset]:
+    """
+    Splits a dataset into training, validation, and test sets, and logs the indices for each split to CSV files for further analysis
+
+    Args:
+        dataset (torch_geometric.dataDataset): The dataset to split.
+        log_dir (str): Directory where CSV files containing the indices for each split will be saved.
+        val_ratio (float, optional): Proportion of the dataset to include in the validation set.
+        test_ratio (float, optional): Proportion of the dataset to include in the test set.
+
+    Raises:
+        ValueError: If `val_ratio + test_ratio >= 1`, which would leave no data for the training set.
+
+    Returns:
+        tuple: A tuple containing:
+            - train_dataset (torch.utils.data.Subset): The training subset of the dataset.
+            - val_dataset (torch.utils.data.Subset): The validation subset of the dataset.
+            - test_dataset (torch.utils.data.Subset): The test subset of the dataset.
+    """
+
+    if val_ratio + test_ratio >= 1:
+        raise ValueError("The sum of val_ratio and test_ratio must be less than 1.")
+
+    val_size = int(val_ratio * len(dataset))
+    test_size = int(test_ratio * len(dataset))
+    train_size = len(dataset) - val_size - test_size
+
+    # Generate shuffled indices and split manually
+    indices = np.random.permutation(len(dataset))
+    train_indices = indices[:train_size]
+    val_indices = indices[train_size : train_size + val_size]
+    test_indices = indices[train_size + val_size :]
+
+    # Save indices to CSV files
+    pd.DataFrame(train_indices, columns=["index"]).to_csv(
+        os.path.join(log_dir, "train_indices.csv"),
+        index=False,
+    )
+    pd.DataFrame(val_indices, columns=["index"]).to_csv(
+        os.path.join(log_dir, "val_indices.csv"),
+        index=False,
+    )
+    pd.DataFrame(test_indices, columns=["index"]).to_csv(
+        os.path.join(log_dir, "test_indices.csv"),
+        index=False,
+    )
+
+    # Create subsets
+    train_dataset = Subset(dataset, train_indices)
+    val_dataset = Subset(dataset, val_indices)
+    test_dataset = Subset(dataset, test_indices)
+
+    return train_dataset, val_dataset, test_dataset

gridfm_graphkit/io/param_handler.py

@@ -0,0 +1,293 @@
+from gridfm_graphkit.datasets.data_normalization import (
+    IdentityNormalizer,
+    MinMaxNormalizer,
+    Standardizer,
+    BaseMVANormalizer,
+)
+from gridfm_graphkit.datasets.transforms import (
+    AddRandomMask,
+    AddPFMask,
+    AddOPFMask,
+    AddIdentityMask,
+)
+from gridfm_graphkit.utils.loss import (
+    PBELoss,
+    MaskedMSELoss,
+    SCELoss,
+    MixedLoss,
+    MSELoss,
+)
+from gridfm_graphkit.models.graphTransformer import GNN_TransformerConv
+from gridfm_graphkit.models.gps_transformer import GPSTransformer
+
+import argparse
+import itertools
+
+
+class NestedNamespace(argparse.Namespace):
+    """
+    A namespace object that supports nested structures, allowing for
+    easy access and manipulation of hierarchical configurations.
+
+    """
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            if isinstance(value, dict):
+                # Recursively convert dictionaries to NestedNamespace
+                setattr(self, key, NestedNamespace(**value))
+            else:
+                setattr(self, key, value)
+
+    def to_dict(self):
+        # Recursively convert NestedNamespace back to dictionary
+        result = {}
+        for key, value in self.__dict__.items():
+            if isinstance(value, NestedNamespace):
+                result[key] = value.to_dict()
+            else:
+                result[key] = value
+        return result
+
+    def flatten(self, parent_key="", sep="."):
+        # Flatten the dictionary with dot-separated keys
+        items = []
+        for key, value in self.__dict__.items():
+            new_key = f"{parent_key}{sep}{key}" if parent_key else key
+            if isinstance(value, NestedNamespace):
+                items.extend(value.flatten(new_key, sep=sep).items())
+            else:
+                items.append((new_key, value))
+        return dict(items)
+
+
+def flatten_dict(d, parent_key="", sep="."):
+    """
+    Flatten a nested dictionary into a single-level dictionary with dot-separated keys.
+
+    Args:
+        d (dict): The dictionary to flatten.
+        parent_key (str, optional): Prefix for the keys in the flattened dictionary.
+        sep (str, optional): Separator for nested keys. Defaults to '.'.
+
+    Returns:
+        dict: A flattened version of the input dictionary.
+    """
+    items = []
+    for key, value in d.items():
+        new_key = f"{parent_key}{sep}{key}" if parent_key else key
+        if isinstance(value, dict):
+            items.extend(flatten_dict(value, new_key, sep=sep).items())
+        else:
+            items.append((new_key, value))
+    return dict(items)
+
+
+def unflatten_dict(d, sep="."):
+    """
+    Reconstruct a nested dictionary from a flattened dictionary with dot-separated keys.
+
+    Args:
+        d (dict): The flattened dictionary to unflatten.
+        sep (str, optional): Separator used in the flattened keys. Defaults to '.'.
+
+    Returns:
+        dict: A nested dictionary reconstructed from the flattened input.
+    """
+    result = {}
+    for key, value in d.items():
+        parts = key.split(sep)
+        target = result
+        for part in parts[:-1]:
+            target = target.setdefault(part, {})
+        target[parts[-1]] = value
+    return result
+
+
+def merge_dict(base, updates):
+    """
+    Recursively merge updates into a base dictionary, but only if the keys exist in the base.
+
+    Args:
+        base (dict): The original dictionary to be updated.
+        updates (dict): The dictionary containing updates.
+
+    Raises:
+        KeyError: If a key in updates does not exist in base.
+        TypeError: If a key in base is not a dictionary but updates attempt to provide nested values.
+    """
+    for key, value in updates.items():
+        if key not in base:
+            raise KeyError(f"Key '{key}' not found in base configuration.")
+
+        if isinstance(value, dict):
+            if not isinstance(base[key], dict):
+                raise TypeError(
+                    f"Default config expects  {type(base[key])}, but got a dict at key '{key}'",
+                )
+            # Recursively merge dictionaries
+            merge_dict(base[key], value)
+        else:
+            # Update the existing key
+            base[key] = value
+
+
+def param_combination_gen(grid_config):
+    """
+    Generate all combinations of parameters from a nested dictionary
+
+    Args:
+        grid_config (dict): A nested dictionary where keys are parameter names
+            and values are lists of possible values.
+
+    Returns:
+        list: A list of dictionaries representing all possible parameter combinations.
+              Each dictionary corresponds to one combination.
+    """
+
+    # Flatten the grid config for combination generation
+    flat_grid = flatten_dict(grid_config)
+
+    # Separate keys and values for itertools.product
+    keys, values = zip(*flat_grid.items())
+
+    # Generate all combinations of parameters
+    combinations = [dict(zip(keys, v)) for v in itertools.product(*values)]
+
+    # Unflatten the combinations back into nested dictionaries
+    nested_combinations = [unflatten_dict(comb) for comb in combinations]
+    return nested_combinations
+
+
+def load_normalizer(args):
+    """
+    Load the appropriate data normalization methods
+
+    Args:
+        args (NestedNamespace): contains configs.
+
+    Returns:
+        tuple: Node and edge normalizers
+
+    Raises:
+        ValueError: If an unknown normalization method is specified.
+    """
+    method = args.data.normalization
+
+    if method == "minmax":
+        return MinMaxNormalizer(), MinMaxNormalizer()
+    elif method == "standard":
+        return Standardizer(), Standardizer()
+    elif method == "baseMVAnorm":
+        return BaseMVANormalizer(
+            node_data=True,
+            baseMVA_orig=args.data.baseMVA,
+        ), BaseMVANormalizer(node_data=False, baseMVA_orig=args.data.baseMVA)
+    elif method == "identity":
+        return IdentityNormalizer(), IdentityNormalizer()
+    else:
+        raise ValueError(f"Unknown normalization method: {method}")
+
+
+def get_loss_function(args):
+    """
+    Load the appropriate loss function
+
+    Args:
+        args (NestedNamespace): contains configs.
+
+    Returns:
+        nn.Module: Loss function
+
+    Raises:
+        ValueError: If an unknown loss function is specified.
+    """
+    loss_functions = []
+    for loss_name in args.training.losses:
+        if loss_name == "MSE":
+            loss_functions.append(MSELoss())
+        elif loss_name == "MaskedMSE":
+            loss_functions.append(MaskedMSELoss())
+        elif loss_name == "SCE":
+            loss_functions.append(SCELoss())
+        elif loss_name == "PBE":
+            loss_functions.append(PBELoss())
+        else:
+            raise ValueError(f"Unknown loss function: {loss_name}")
+
+    return MixedLoss(loss_functions=loss_functions, weights=args.training.loss_weights)
+
+
+def load_model(args):
+    """
+    Load the appropriate model
+
+    Args:
+        args (NestedNamespace): contains configs.
+
+    Returns:
+        nn.Module: The selected model initialized with the provided configurations.
+
+    Raises:
+        ValueError: If an unknown model type is specified.
+    """
+    model_type = args.model.type
+
+    if model_type == "GNN_TransformerConv":
+        return GNN_TransformerConv(
+            input_dim=args.model.input_dim,
+            hidden_dim=args.model.hidden_size,
+            output_dim=args.model.output_dim,
+            edge_dim=args.model.edge_dim,
+            num_layers=args.model.num_layers,
+            heads=args.model.attention_head,
+            mask_dim=args.data.mask_dim,
+            mask_value=args.data.mask_value,
+            learn_mask=args.data.learn_mask,
+        )
+    elif model_type == "GPSTransformer":
+        return GPSTransformer(
+            input_dim=args.model.input_dim,
+            hidden_dim=args.model.hidden_size,
+            output_dim=args.model.output_dim,
+            edge_dim=args.model.edge_dim,
+            pe_dim=args.model.pe_dim,
+            heads=args.model.attention_head,
+            num_layers=args.model.num_layers,
+            dropout=args.model.dropout,
+            mask_dim=args.data.mask_dim,
+            mask_value=args.data.mask_value,
+            learn_mask=args.data.learn_mask,
+        )
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")
+
+
+def get_transform(args):
+    """
+    Load the appropriate dataset transform
+
+    Args:
+        args (NestedNamespace): contains configs.
+
+    Returns:
+        BaseTransform: Transformation
+
+    Raises:
+        ValueError: If an unknown transform is specified.
+    """
+    mask_type = args.data.mask_type
+
+    if mask_type == "rnd":
+        return AddRandomMask(
+            mask_dim=args.data.mask_dim,
+            mask_ratio=args.data.mask_ratio,
+        )
+    elif mask_type == "pf":
+        return AddPFMask()
+    elif mask_type == "opf":
+        return AddOPFMask()
+    elif mask_type == "none":
+        return AddIdentityMask()
+    else:
+        raise ValueError(f"Unknown transformation: {mask_type}")

gridfm_graphkit/models/gps_transformer.py

@@ -0,0 +1,143 @@
+from torch_geometric.nn import GPSConv, GINEConv
+from torch import nn
+import torch
+
+
+class GPSTransformer(nn.Module):
+    """
+    A GPS (Graph Transformer) model based on [GPSConv](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.GPSConv.html) and [GINEConv](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.GINEConv.html) layers from Pytorch Geometric.
+
+    This model encodes node features and positional encodings separately,
+    then applies multiple graph convolution layers with batch normalization,
+    and finally decodes to the output dimension.
+
+    Args:
+        input_dim (int): Dimension of input node features.
+        hidden_dim (int): Hidden dimension size for all layers.
+        output_dim (int): Dimension of the output node features.
+        edge_dim (int): Dimension of edge features.
+        pe_dim (int): Dimension of the positional encoding.
+            Must be less than hidden_dim.
+        num_layers (int): Number of GPSConv layers.
+        heads (int, optional): Number of attention heads in GPSConv.
+        dropout (float, optional): Dropout rate in GPSConv.
+        mask_dim (int, optional): Dimension of the mask vector.
+        mask_value (float, optional): Initial value for learnable mask parameters.
+        learn_mask (bool, optional): Whether to learn mask values as parameters.
+
+    Raises:
+        ValueError: If `pe_dim` is not less than `hidden_dim`.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        edge_dim: int,
+        pe_dim: int,
+        num_layers: int,
+        heads: int = 1,
+        dropout: float = 0.0,
+        mask_dim: int = 6,
+        mask_value: float = -1.0,
+        learn_mask: bool = True,
+    ):
+        super(GPSTransformer, self).__init__()
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.edge_dim = edge_dim
+        self.pe_dim = pe_dim
+        self.heads = heads
+        self.dropout = dropout
+        self.mask_dim = mask_dim
+        self.mask_value = mask_value
+        self.learn_mask = learn_mask
+
+        if not pe_dim < hidden_dim:
+            raise ValueError(
+                "positional encoding dimension must be smaller than model hidden dimension",
+            )
+
+        self.layers = nn.ModuleList()
+
+        self.encoder = nn.Sequential(
+            nn.Linear(input_dim, self.hidden_dim - self.pe_dim),
+            nn.LeakyReLU(),
+        )
+        self.input_norm = nn.BatchNorm1d(self.hidden_dim - self.pe_dim)
+        self.pe_norm = nn.BatchNorm1d(self.pe_dim)
+
+        for _ in range(self.num_layers):
+            mlp = nn.Sequential(
+                nn.Linear(in_features=self.hidden_dim, out_features=self.hidden_dim),
+                nn.LeakyReLU(),
+            )
+            self.layers.append(
+                nn.ModuleDict(
+                    {
+                        "conv": GPSConv(
+                            channels=self.hidden_dim,
+                            conv=GINEConv(nn=mlp, edge_dim=self.edge_dim),
+                            heads=self.heads,
+                            dropout=self.dropout,
+                        ),
+                        "norm": nn.BatchNorm1d(
+                            self.hidden_dim,
+                        ),  # BatchNorm after each graph layer
+                    },
+                ),
+            )
+
+        self.pre_decoder_norm = nn.BatchNorm1d(self.hidden_dim)
+        # Fully connected (MLP) layers after the GAT layers
+        self.decoder = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.hidden_dim),
+            nn.LeakyReLU(),
+            nn.Linear(self.hidden_dim, output_dim),
+        )
+
+        if learn_mask:
+            self.mask_value = nn.Parameter(
+                torch.randn(mask_dim) + mask_value,
+                requires_grad=True,
+            )
+        else:
+            self.mask_value = nn.Parameter(
+                torch.zeros(mask_dim) + mask_value,
+                requires_grad=False,
+            )
+
+    def forward(self, x, pe, edge_index, edge_attr, batch):
+        """
+        Forward pass for the GPSTransformer.
+
+        Args:
+            x (Tensor): Input node features of shape [num_nodes, input_dim].
+            pe (Tensor): Positional encoding of shape [num_nodes, pe_dim].
+            edge_index (Tensor): Edge indices for graph convolution.
+            edge_attr (Tensor): Edge feature tensor.
+            batch (Tensor): Batch vector assigning nodes to graphs.
+
+        Returns:
+            output (Tensor): Output node features of shape [num_nodes, output_dim].
+        """
+        x_pe = self.pe_norm(pe)
+
+        x = self.encoder(x)
+        x = self.input_norm(x)
+
+        x = torch.cat((x, x_pe), 1)
+        for layer in self.layers:
+            x = layer["conv"](
+                x=x,
+                edge_index=edge_index,
+                edge_attr=edge_attr,
+                batch=batch,
+            )
+            x = layer["norm"](x)
+
+        x = self.pre_decoder_norm(x)
+        x = self.decoder(x)
+
+        return x

gridfm_graphkit/models/graphTransformer.py

@@ -0,0 +1,96 @@
+from torch_geometric.nn import TransformerConv
+from torch import nn
+import torch
+
+
+class GNN_TransformerConv(nn.Module):
+    """
+    Graph Neural Network using [TransformerConv](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.TransformerConv.html) layers from PyTorch Geometric.
+
+    Args:
+        input_dim (int): Dimensionality of input node features.
+        hidden_dim (int): Hidden dimension size for TransformerConv layers.
+        output_dim (int): Output dimension size.
+        edge_dim (int): Dimensionality of edge features.
+        num_layers (int): Number of TransformerConv layers.
+        heads (int, optional): Number of attention heads.
+        mask_dim (int, optional): Dimension of mask vector.
+        mask_value (float, optional): Initial mask value.
+        learn_mask (bool, optional): Whether mask values are learnable.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        edge_dim: int,
+        num_layers: int,
+        heads: int = 1,
+        mask_dim: int = 6,
+        mask_value: float = -1.0,
+        learn_mask: bool = False,
+    ):
+        super(GNN_TransformerConv, self).__init__()
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.edge_dim = edge_dim
+        self.heads = heads
+        self.mask_dim = mask_dim
+        self.mask_value = mask_value
+        self.learn_mask = learn_mask
+
+        self.layers = nn.ModuleList()
+        current_dim = input_dim  # First layer takes `input_dim` as input
+
+        for _ in range(self.num_layers):
+            self.layers.append(
+                TransformerConv(
+                    current_dim,
+                    self.hidden_dim,
+                    heads=self.heads,
+                    edge_dim=self.edge_dim,
+                    beta=False,
+                ),
+            )
+            # Update the dimension for the next layer
+            current_dim = self.hidden_dim * self.heads
+
+        # Fully connected (MLP) layers after the GAT layers
+        self.mlps = nn.Sequential(
+            nn.Linear(self.hidden_dim * self.heads, self.hidden_dim),
+            nn.LeakyReLU(),
+            nn.Linear(self.hidden_dim, output_dim),
+        )
+
+        if learn_mask:
+            self.mask_value = nn.Parameter(
+                torch.randn(mask_dim) + mask_value,
+                requires_grad=True,
+            )
+        else:
+            self.mask_value = nn.Parameter(
+                torch.zeros(mask_dim) + mask_value,
+                requires_grad=False,
+            )
+
+    def forward(self, x, pe, edge_index, edge_attr, batch):
+        """
+        Forward pass for the GPSTransformer.
+
+        Args:
+            x (Tensor): Input node features of shape [num_nodes, input_dim].
+            pe (Tensor): Positional encoding of shape [num_nodes, pe_dim] (not used).
+            edge_index (Tensor): Edge indices for graph convolution.
+            edge_attr (Tensor): Edge feature tensor.
+            batch (Tensor): Batch vector assigning nodes to graphs (not used).
+
+        Returns:
+            output (Tensor): Output node features of shape [num_nodes, output_dim].
+        """
+        for conv in self.layers:
+            x = conv(x, edge_index, edge_attr)
+            x = nn.LeakyReLU()(x)
+
+        x = self.mlps(x)
+        return x

gridfm_graphkit/training/callbacks.py

@@ -0,0 +1,47 @@
+import torch
+
+
+class EarlyStopper:
+    def __init__(
+        self,
+        saving_path,
+        patience=5,
+        tol=0,
+        min_validation_loss=float("inf"),
+    ):
+        """
+        Args:
+            patience (int): number of epochs to wait before early stopping
+                -1 means no early stopping
+                0 means stop training the first time the validation loss increases
+            tol (float): tolerance to consider validation loss as worse as the best one so far
+        """
+
+        self.patience = patience
+        self.tol = tol
+        self.counter = 0
+        self.min_validation_loss = min_validation_loss
+        self.saving_path = saving_path
+
+    def early_stop(self, validation_loss, model):
+        if validation_loss < self.min_validation_loss:
+            self.min_validation_loss = (
+                validation_loss  # Update the best validation loss
+            )
+            self.counter = 0
+
+            # Save the best model whenever a new minimum is found
+            torch.save(model, self.saving_path)
+
+        # check if the valid loss is worse than the best one so far, accounting for tolerance
+        elif validation_loss > (self.min_validation_loss + self.tol):
+            self.counter += 1
+
+            if self.patience != -1 and self.counter > self.patience:
+                print(
+                    "Early stopping after {} epochs of no improvement.".format(
+                        self.counter,
+                    ),
+                )
+                return True
+        return False