gridfm-graphkit 0.0.1__py3-none-any.whl

gridfm_graphkit/datasets/data_normalization.py

@@ -0,0 +1,227 @@
+from gridfm_graphkit.datasets.globals import PD, QD, PG, QG, VA
+import torch
+from abc import ABC, abstractmethod
+
+
+class Normalizer(ABC):
+    """
+    Abstract base class for all normalization strategies.
+    """
+
+    @abstractmethod
+    def fit(self, data: torch.Tensor) -> dict:
+        """
+        Fit normalization parameters from data.
+
+        Args:
+            data: Input tensor.
+
+        Returns:
+            Dictionary of computed parameters.
+        """
+
+    @abstractmethod
+    def fit_from_dict(self, params: dict):
+        """
+        Set parameters from a precomputed dictionary.
+
+        Args:
+            params: Dictionary of parameters.
+        """
+
+    @abstractmethod
+    def transform(self, data: torch.Tensor) -> torch.Tensor:
+        """
+        Normalize the input data.
+
+        Args:
+            data: Input tensor.
+
+        Returns:
+            Normalized tensor.
+        """
+
+    @abstractmethod
+    def inverse_transform(self, normalized_data: torch.Tensor) -> torch.Tensor:
+        """
+        Undo normalization.
+
+        Args:
+            normalized_data: Normalized tensor.
+
+        Returns:
+            Original tensor.
+        """
+
+
+class MinMaxNormalizer(Normalizer):
+    """
+    Scales each feature to the [0, 1] range.
+    """
+
+    def __init__(self):
+        self.min_val = None
+        self.max_val = None
+
+    def to(self, device):
+        self.min_val = self.min_val.to(device)
+        self.max_val = self.max_val.to(device)
+
+    def fit(self, data: torch.Tensor) -> dict:
+        self.min_val, _ = data.min(axis=0)
+        self.max_val, _ = data.max(axis=0)
+
+        return {"min_value": self.min_val, "max_value": self.max_val}
+
+    def fit_from_dict(self, params: dict):
+        if self.min_val is None:
+            self.min_val = params.get("min_value")
+        if self.max_val is None:
+            self.max_val = params.get("max_value")
+
+    def transform(self, data: torch.Tensor) -> torch.Tensor:
+        if self.min_val is None or self.max_val is None:
+            raise ValueError("fit must be called before transform.")
+
+        diff = self.max_val - self.min_val
+        diff[diff == 0] = 1  # Avoid division by zero for features with zero range
+        return (data - self.min_val) / diff
+
+    def inverse_transform(self, normalized_data: torch.Tensor) -> torch.Tensor:
+        if self.min_val is None or self.max_val is None:
+            raise ValueError("fit must be called before inverse_transform.")
+
+        diff = self.max_val - self.min_val
+        diff[diff == 0] = 1
+        return (normalized_data * diff) + self.min_val
+
+
+class Standardizer(Normalizer):
+    """
+    Standardizes each feature to zero mean and unit variance.
+    """
+
+    def __init__(self):
+        self.mean = None
+        self.std = None
+
+    def to(self, device):
+        self.mean = self.mean.to(device)
+        self.std = self.std.to(device)
+
+    def fit(self, data: torch.Tensor) -> dict:
+        self.mean = data.mean(axis=0)
+        self.std = data.std(axis=0)
+
+        return {"mean_value": self.mean, "std_value": self.std}
+
+    def fit_from_dict(self, params: dict):
+        if self.mean is None:
+            self.mean = params.get("mean_value")
+        if self.std is None:
+            self.std = params.get("std_value")
+
+    def transform(self, data: torch.Tensor) -> torch.Tensor:
+        if self.mean is None or self.std is None:
+            raise ValueError("fit must be called before transform.")
+
+        std = self.std.clone()
+        std[std == 0] = 1  # Avoid division by zero for features with zero std
+        return (data - self.mean) / std
+
+    def inverse_transform(self, normalized_data: torch.Tensor) -> torch.Tensor:
+        if self.mean is None or self.std is None:
+            raise ValueError("fit must be called before inverse_transform.")
+
+        std = self.std.clone()
+        std[std == 0] = 1
+        return (normalized_data * std) + self.mean
+
+
+class BaseMVANormalizer(Normalizer):
+    """
+    In power systems, a suitable normalization strategy must preserve the physical properties of
+    the system. A known method is the conversion to the per-unit (p.u.) system, which expresses
+    electrical quantities such as voltage, current, power, and impedance as fractions of predefined
+    base values. These base values are usually chosen based on system parameters, such as rated
+    voltage. The per-unit conversion ensures that power system equations remain scale-invariant,
+    preserving fundamental physical relationships.
+    """
+
+    def __init__(self, node_data: bool, baseMVA_orig: float = 100.0):
+        """
+        Args:
+            node_data: Whether data is node-level or edge-level (PD, QD, PG, QG, VA).
+            baseMVA_orig: Original baseMVA (e.g. from MATPOWER).
+        """
+        self.node_data = node_data
+        self.baseMVA_orig = baseMVA_orig
+        self.baseMVA = None
+
+    def to(self, device):
+        pass
+
+    def fit(self, data: torch.Tensor, baseMVA: float = None) -> dict:
+        if self.node_data:
+            self.baseMVA = data[:, [PD, QD, PG, QG]].max()
+        else:
+            self.baseMVA = baseMVA
+
+        return {"baseMVA_orig": self.baseMVA_orig, "baseMVA": self.baseMVA}
+
+    def fit_from_dict(self, params: dict):
+        if self.baseMVA is None:
+            self.baseMVA = params.get("baseMVA")
+        if self.baseMVA_orig is None:
+            self.baseMVA_orig = params.get("baseMVA_orig")
+
+    def transform(self, data: torch.Tensor) -> torch.Tensor:
+        if self.baseMVA is None:
+            raise ValueError("BaseMVA is not specified")
+
+        if self.baseMVA == 0:
+            raise ZeroDivisionError("BaseMVA is 0.")
+
+        if self.node_data:
+            data[:, PD] = data[:, PD] / self.baseMVA
+            data[:, QD] = data[:, QD] / self.baseMVA
+            data[:, PG] = data[:, PG] / self.baseMVA
+            data[:, QG] = data[:, QG] / self.baseMVA
+            data[:, VA] = data[:, VA] * torch.pi / 180.0
+        else:
+            data = data * self.baseMVA_orig / self.baseMVA
+
+        return data
+
+    def inverse_transform(self, normalized_data: torch.Tensor) -> torch.Tensor:
+        if self.baseMVA is None:
+            raise ValueError("fit must be called before inverse_transform.")
+
+        if self.node_data:
+            normalized_data[:, PD] = normalized_data[:, PD] * self.baseMVA
+            normalized_data[:, QD] = normalized_data[:, QD] * self.baseMVA
+            normalized_data[:, PG] = normalized_data[:, PG] * self.baseMVA
+            normalized_data[:, QG] = normalized_data[:, QG] * self.baseMVA
+            normalized_data[:, VA] = normalized_data[:, VA] * 180.0 / torch.pi
+        else:
+            normalized_data = normalized_data * self.baseMVA / self.baseMVA_orig
+
+        return normalized_data
+
+
+class IdentityNormalizer(Normalizer):
+    """
+    No normalization: returns data unchanged.
+    """
+
+    def fit(self, data: torch.Tensor) -> dict:
+        return {}
+
+    def fit_from_dict(self, params: dict):
+        pass
+
+    def transform(self, data: torch.Tensor) -> torch.Tensor:
+        return data
+
+    def inverse_transform(self, normalized_data: torch.Tensor) -> torch.Tensor:
+        return normalized_data

gridfm_graphkit/datasets/globals.py

@@ -0,0 +1,19 @@
+# Global variables
+
+# Node features indices
+PD = 0
+QD = 1
+PG = 2
+QG = 3
+VM = 4
+VA = 5
+PQ = 6
+PV = 7
+REF = 8
+
+# Edge features indices
+G = 0
+B = 1
+
+FEATURES_IDX = {"PD": PD, "QD": QD, "PG": PG, "QG": QG, "VM": VM, "VA": VA}
+BUS_TYPES = ["PQ", "PV", "REF"]

gridfm_graphkit/datasets/powergrid.py

@@ -0,0 +1,192 @@
+from gridfm_graphkit.datasets.data_normalization import Normalizer, BaseMVANormalizer
+from gridfm_graphkit.datasets.transforms import (
+    AddEdgeWeights,
+    AddNormalizedRandomWalkPE,
+)
+
+import os.path as osp
+import torch
+from torch_geometric.data import Data, InMemoryDataset
+import pandas as pd
+from tqdm import tqdm
+from typing import Optional, Callable
+
+
+class GridDatasetMem(InMemoryDataset):
+    """
+    A PyTorch Geometric `InMemoryDataset` for power grid data stored in tabular CSV format.
+    This dataset class reads node and edge data from CSV files, applies normalization using
+    user-specified `Normalizer` instances, and builds graph data objects with edge weights and
+    positional encodings.
+
+    - Reads raw node and edge CSV files (`pf_node.csv`, `pf_edge.csv`).
+    - Applies the normalization method specified on both node and edge features
+    - Stores normalization statistics in the `processed` directory for reuse.
+    - Constructs `torch_geometric.data.Data` objects with edge weights and positional encodings (via random walk embeddings).
+
+    Args:
+        root (str): Root directory where the dataset is stored.
+        norm_method (str): Identifier for normalization method (e.g., "minmax", "standard").
+        node_normalizer (Normalizer): Normalizer used for node features.
+        edge_normalizer (Normalizer): Normalizer used for edge features.
+        pe_dim (int): Length of the random walk used for positional encoding.
+        mask_dim (int, optional): Number of features per-node that could be masked. Usually Pd, Qd, Pg, Qg, Vm, Va
+        transform (callable, optional): Transformation applied at runtime.
+        pre_transform (callable, optional): Transformation applied before saving to disk.
+        pre_filter (callable, optional): Filter to determine which graphs to keep.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        norm_method: str,
+        node_normalizer: Normalizer,
+        edge_normalizer: Normalizer,
+        pe_dim: int,
+        mask_dim: int = 6,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+    ):
+        self.norm_method = norm_method
+        self.node_normalizer = node_normalizer
+        self.edge_normalizer = edge_normalizer
+        self.pe_dim = pe_dim
+        self.mask_dim = mask_dim
+        self.original_transform = None
+
+        super().__init__(root, transform, pre_transform, pre_filter)
+
+        node_stats_path = osp.join(
+            self.processed_dir,
+            f"node_stats_{self.norm_method}.pt",
+        )
+        edge_stats_path = osp.join(
+            self.processed_dir,
+            f"edge_stats_{self.norm_method}.pt",
+        )
+        if osp.exists(node_stats_path) and osp.exists(edge_stats_path):
+            self.node_stats = torch.load(node_stats_path, weights_only=False)
+            self.edge_stats = torch.load(edge_stats_path, weights_only=False)
+            self.node_normalizer.fit_from_dict(self.node_stats)
+            self.edge_normalizer.fit_from_dict(self.edge_stats)
+        self.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self):
+        # No raw files needed for random graphs
+        return ["pf_node.csv", "pf_edge.csv"]
+
+    @property
+    def processed_file_names(self):
+        return [f"data_full_{self.norm_method}.pt"]
+
+    def download(self):
+        pass
+
+    def process(self):
+        node_df = pd.read_csv(osp.join(self.raw_dir, "pf_node.csv"))
+        edge_df = pd.read_csv(osp.join(self.raw_dir, "pf_edge.csv"))
+
+        # Check the unique scenarios available
+        scenarios = node_df["scenario"].unique()
+        # Ensure node and edge data match
+        if not (scenarios == edge_df["scenario"].unique()).all():
+            raise ValueError("Mismatch between node and edge scenario values.")
+
+        # normalize node attributes
+        cols_to_normalize = ["Pd", "Qd", "Pg", "Qg", "Vm", "Va"]
+        to_normalize = torch.tensor(
+            node_df[cols_to_normalize].values,
+            dtype=torch.float,
+        )
+        self.node_stats = self.node_normalizer.fit(to_normalize)
+        node_df[cols_to_normalize] = self.node_normalizer.transform(
+            to_normalize,
+        ).numpy()
+
+        # normalize edge attributes
+        cols_to_normalize = ["G", "B"]
+        to_normalize = torch.tensor(
+            edge_df[cols_to_normalize].values,
+            dtype=torch.float,
+        )
+        if isinstance(self.node_normalizer, BaseMVANormalizer):
+            self.edge_stats = self.edge_normalizer.fit(
+                to_normalize,
+                self.node_normalizer.baseMVA,
+            )
+        else:
+            self.edge_stats = self.edge_normalizer.fit(to_normalize)
+        edge_df[cols_to_normalize] = self.edge_normalizer.transform(
+            to_normalize,
+        ).numpy()
+
+        # save stats
+        node_stats_path = osp.join(
+            self.processed_dir,
+            f"node_stats_{self.norm_method}.pt",
+        )
+        edge_stats_path = osp.join(
+            self.processed_dir,
+            f"edge_stats_{self.norm_method}.pt",
+        )
+        torch.save(self.node_stats, node_stats_path)
+        torch.save(self.edge_stats, edge_stats_path)
+
+        # Create groupby objects for scenarios
+        node_groups = node_df.groupby("scenario")
+        edge_groups = edge_df.groupby("scenario")
+
+        data_list = []
+        for scenario_idx in tqdm(scenarios):
+            # NODE DATA
+            node_data = node_groups.get_group(scenario_idx)
+            x = torch.tensor(
+                node_data[
+                    ["Pd", "Qd", "Pg", "Qg", "Vm", "Va", "PQ", "PV", "REF"]
+                ].values,
+                dtype=torch.float,
+            )
+            y = x[:, : self.mask_dim]
+
+            # EDGE DATA
+            edge_data = edge_groups.get_group(scenario_idx)
+            edge_attr = torch.tensor(edge_data[["G", "B"]].values, dtype=torch.float)
+            edge_index = torch.tensor(
+                edge_data[["index1", "index2"]].values.T,
+                dtype=torch.long,
+            )
+
+            # Create the Data object
+            graph_data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)
+            pe_pre_transform = AddEdgeWeights()
+            graph_data = pe_pre_transform(graph_data)
+            pe_transform = AddNormalizedRandomWalkPE(
+                walk_length=self.pe_dim,
+                attr_name="pe",
+            )
+            graph_data = pe_transform(graph_data)
+            data_list.append(graph_data)
+
+        self.save(data_list, self.processed_paths[0])
+
+    def change_transform(self, new_transform):
+        """
+        Temporarily switch to a new transform function, used when evaluating different tasks.
+
+        Args:
+            new_transform (Callable): The new transform to use.
+        """
+        self.original_transform = self.transform
+        self.transform = new_transform
+
+    def reset_transform(self):
+        """
+        Reverts the transform to the original one set during initialization, usually called after the evaluation step.
+        """
+        if self.original_transform is None:
+            raise ValueError(
+                "The original transform is None or the function change_transform needs to be called before",
+            )
+        self.transform = self.original_transform

gridfm_graphkit/datasets/transforms.py

@@ -0,0 +1,223 @@
+from gridfm_graphkit.datasets.globals import PQ, PV, REF, PG, QG, VM, VA, G, B
+
+import torch
+from torch import Tensor
+from torch_geometric.transforms import BaseTransform
+from typing import Optional
+import torch_geometric.typing
+from torch_geometric.data import Data
+from torch_geometric.utils import (
+    get_self_loop_attr,
+    is_torch_sparse_tensor,
+    to_edge_index,
+    to_torch_coo_tensor,
+    to_torch_csr_tensor,
+)
+
+
+class AddNormalizedRandomWalkPE(BaseTransform):
+    r"""Adds the random walk positional encoding from the
+    [Graph Neural Networks with Learnable Structural and Positional Representations](https://arxiv.org/abs/2110.07875)
+    paper to the given graph. This is an adaptation from the original Pytorch Geometric implementation.
+
+    Args:
+        walk_length (int): The number of random walk steps.
+        attr_name (str, optional): The attribute name of the data object to add
+            positional encodings to. If set to :obj:`None`, will be
+            concatenated to :obj:`data.x`.
+            (default: :obj:`"random_walk_pe"`)
+    """
+
+    def __init__(
+        self,
+        walk_length: int,
+        attr_name: Optional[str] = "random_walk_pe",
+    ) -> None:
+        self.walk_length = walk_length
+        self.attr_name = attr_name
+
+    def forward(self, data: Data) -> Data:
+        if data.edge_index is None:
+            raise ValueError("Expected data.edge_index to be not None")
+        row, col = data.edge_index
+        N = data.num_nodes
+        if N is None:
+            raise ValueError("Expected data.num_nodes to be not None")
+
+        if N <= 2_000:  # Dense code path for faster computation:
+            adj = torch.zeros((N, N), device=row.device)
+            adj[row, col] = data.edge_weight
+            loop_index = torch.arange(N, device=row.device)
+        elif torch_geometric.typing.WITH_WINDOWS:
+            adj = to_torch_coo_tensor(
+                data.edge_index,
+                data.edge_weight,
+                size=data.size(),
+            )
+        else:
+            adj = to_torch_csr_tensor(
+                data.edge_index,
+                data.edge_weight,
+                size=data.size(),
+            )
+
+        row_sums = adj.sum(dim=1, keepdim=True)  # Sum along rows
+        row_sums = row_sums.clamp(min=1e-8)  # Prevent division by zero
+
+        adj = adj / row_sums  # Normalize each row to sum to 1
+
+        def get_pe(out: Tensor) -> Tensor:
+            if is_torch_sparse_tensor(out):
+                return get_self_loop_attr(*to_edge_index(out), num_nodes=N)
+            return out[loop_index, loop_index]
+
+        out = adj
+        pe_list = [get_pe(out)]
+        for _ in range(self.walk_length - 1):
+            out = out @ adj
+            pe_list.append(get_pe(out))
+
+        pe = torch.stack(pe_list, dim=-1)
+        data[self.attr_name] = pe
+
+        return data
+
+
+class AddEdgeWeights(BaseTransform):
+    """
+    Computes and adds edge weight as the magnitude of complex admittance.
+
+    The magnitude is computed from the G and B components in `data.edge_attr` and stored in `data.edge_weight`.
+    """
+
+    def forward(self, data):
+        if not hasattr(data, "edge_attr"):
+            raise AttributeError("Data must have 'edge_attr'.")
+
+        # Extract real and imaginary parts of admittance
+        real = data.edge_attr[:, G]
+        imag = data.edge_attr[:, B]
+
+        # Compute the magnitude of the complex admittance
+        edge_weight = torch.sqrt(real**2 + imag**2)
+
+        # Add the computed edge weights to the data object
+        data.edge_weight = edge_weight
+
+        return data
+
+
+class AddIdentityMask(BaseTransform):
+    """Creates an identity mask, and adds it as a `mask` attribute.
+
+    The mask is generated such that every entry is False, so no masking is actually applied
+    """
+
+    def forward(self, data):
+        if not hasattr(data, "y"):
+            raise AttributeError("Data must have ground truth 'y'.")
+
+        # Generate an identity mask
+        mask = torch.zeros_like(data.y, dtype=torch.bool)
+
+        # Add the mask to the data object
+        data.mask = mask
+
+        return data
+
+
+class AddRandomMask(BaseTransform):
+    """Creates a random mask, and adds it as a `mask` attribute.
+
+    The mask is generated such that each entry is `True` with probability
+    `mask_ratio` and `False` otherwise.
+    """
+
+    def __init__(self, mask_dim, mask_ratio):
+        super().__init__()
+        self.mask_dim = mask_dim
+        self.mask_ratio = mask_ratio
+
+    def forward(self, data):
+        if not hasattr(data, "x"):
+            raise AttributeError("Data must have node features 'x'.")
+
+        # Generate a random mask
+        mask = torch.rand(data.x.size(0), self.mask_dim) < self.mask_ratio
+
+        # Add the mask to the data object
+        data.mask = mask
+
+        return data
+
+
+class AddPFMask(BaseTransform):
+    """Creates a mask according to the power flow problem and assigns it as a `mask` attribute."""
+
+    def forward(self, data):
+        # Ensure the data object has the required attributes
+        if not hasattr(data, "y"):
+            raise AttributeError("Data must have ground truth 'y'.")
+
+        if not hasattr(data, "x"):
+            raise AttributeError("Data must have node features 'x'.")
+
+        # Generate masks for each type of node
+        mask_PQ = data.x[:, PQ] == 1  # PQ buses
+        mask_PV = data.x[:, PV] == 1  # PV buses
+        mask_REF = data.x[:, REF] == 1  # Reference buses
+
+        # Initialize the mask tensor with False values
+        mask = torch.zeros_like(data.y, dtype=torch.bool)
+
+        mask[mask_PQ, VM] = True  # Mask Vm for PQ buses
+        mask[mask_PQ, VA] = True  # Mask Va for PQ buses
+
+        mask[mask_PV, QG] = True  # Mask Qg for PV buses
+        mask[mask_PV, VA] = True  # Mask Va for PV buses
+
+        mask[mask_REF, PG] = True  # Mask Pg for REF buses
+        mask[mask_REF, QG] = True  # Mask Qg for REF buses
+
+        # Attach the mask to the data object
+        data.mask = mask
+
+        return data
+
+
+class AddOPFMask(BaseTransform):
+    """Creates a mask according to the optimal power flow problem and assigns it as a `mask` attribute."""
+
+    def forward(self, data):
+        # Ensure the data object has the required attributes
+        if not hasattr(data, "y"):
+            raise AttributeError("Data must have ground truth 'y'.")
+
+        if not hasattr(data, "x"):
+            raise AttributeError("Data must have node features 'x'.")
+
+        # Generate masks for each type of node
+        mask_PQ = data.x[:, PQ] == 1  # PQ buses
+        mask_PV = data.x[:, PV] == 1  # PV buses
+        mask_REF = data.x[:, REF] == 1  # Reference buses
+
+        # Initialize the mask tensor with False values
+        mask = torch.zeros_like(data.y, dtype=torch.bool)
+
+        mask[mask_PQ, VM] = True  # Mask Vm for PQ
+        mask[mask_PQ, VA] = True  # Mask Va for PQ
+
+        mask[mask_PV, PG] = True  # Mask Pg for PV
+        mask[mask_PV, QG] = True  # Mask Qg for PV
+        mask[mask_PV, VM] = True  # Mask Vm for PV
+        mask[mask_PV, VA] = True  # Mask Va for PV
+
+        mask[mask_REF, PG] = True  # Mask Pg for REF
+        mask[mask_REF, QG] = True  # Mask Qg for REF
+        mask[mask_REF, VM] = True  # Mask Vm for REF
+        mask[mask_REF, VA] = True  # Mask Va for REF
+
+        # Attach the mask to the data object
+        data.mask = mask
+
+        return data