PowerFlowData.py

"""
this file defines the class of PowerFlowData, which is used to load the data of Power Flow
"""
import os
from typing import Callable, Optional, List, Tuple, Union

import torch
import numpy as np
import torch.utils.data as data
import torch_geometric
from torch_geometric.data import Data, InMemoryDataset
from torch_geometric.utils import from_scipy_sparse_matrix, dense_to_sparse
import matplotlib.pyplot as plt

from torch_geometric.datasets import Planetoid

ori_feature_names_x = [
    'index',                # - removed
    'type',                 # --- one-hot encoded,  0, 1, 2, 3
    'voltage magnitude',    # --- this matters,     4, 
    'voltage angle',        # --- this matters,     5,
    'Pd',                   # --- this matters, preprocessed as Pd-Pg   6,
    'Qd',                   # --- this matters                          7,
    'Gs',                   # - equivalent to Pd, Qd                    8,
    'Bs',                   # - equivalent to Pd, Qd                    9,
    'Pg'                    # - removed
]

feature_names_x = [
    'type_0',               # --- one-hot encoded,  0,
    'type_1',               # --- one-hot encoded,  1,
    'type_2',               # --- one-hot encoded,  2,
    'type_3',               # --- one-hot encoded,  3,
    'voltage magnitude',    # --- this matters,     4, 
    'voltage angle',        # --- this matters,     5,
    'Pd',                   # --- this matters, preprocessed as Pd-Pg   6,
    'Qd',                   # --- this matters                          7,
    'Gs',                   # - equivalent to Pd, Qd                    8,
    'Bs',                   # - equivalent to Pd, Qd                    9,
    'to_predict_voltage_magnitude', #               10,
    'to_predict_voltage_angle',     #               11,
    'to_predict_Pd',                #               12,
    'to_predict_Qd'                 #               13,
    'to_predict_Gs',                #               14,
    'to_predict_Bs'                 #               15,
]

ori_feature_names_y = [
    'index',                # - removed
    'type',                 # - removed
    'voltage magnitude',    # --- we care about this
    'voltage angle',        # --- we care about this
    'active power',         # --- we care about this
    'reactive power',       # --- we care about this
    'Gs',                   # -
    'Bs'                    # -
]

feature_names_y = [
    'voltage magnitude',    # --- we care about this
    'voltage angle',        # --- we care about this
    'active power',         # --- we care about this
    'reactive power',       # --- we care about this
    'Gs',                   # -
    'Bs'         
]

edge_feature_names = [
    'r',                    # --- this matters, resistance, pu
    'x',                    # --- this matters, reactance,  pu
    'b',
    'tau',
    'angle'
]
        

class PowerFlowData(InMemoryDataset):
    """PowerFlowData(InMemoryDataset)

    Parameters:
        root (str, optional) – Root directory where the dataset should be saved. (optional: None)
        pre_filter (callable)- A function 

    Comments:
        we actually do not need adjacency matrix, since we can use edge_index to represent the graph from `edge_features`

    Returns:
        class instance of PowerFlowData
    """
    partial_file_names = [
        "adjacency_matrix.npy",
        "edge_features.npy",
        "node_features_x.npy",
        "node_features_y.npy"
    ]
    split_order = {
        "train": 0,
        "val": 1,
        "test": 2
    }
    mixed_cases = [
        '118v2',
        '14v2',
    ]

    def __init__(self, 
                root: str, 
                case: str = '14', 
                split: Optional[List[float]] = None, 
                task: str = "train", 
                transform: Optional[Callable] = None, 
                pre_transform: Optional[Callable] = None, 
                pre_filter: Optional[Callable] = None,
                normalize=True):

        assert len(split) == 3
        assert task in ["train", "val", "test"]
        self.normalize = normalize
        self.case = case  # THIS MUST BE EXECUTED BEFORE super().__init__() since it is used in raw_file_names and processed_file_names
        self.split = split
        self.task = task
        super().__init__(root, transform, pre_transform, pre_filter)
        self.mask = torch.tensor([])
        self.data, self.slices, self.mask = self._normalize_dataset(
            *torch.load(self.processed_paths[0]))  # necessary, do not forget!

    def get_data_dimensions(self):
        return self[0].x.shape[1], self[0].y.shape[1], self[0].edge_attr.shape[1]

    def get_data_means_stds(self):
        assert self.normalize == True
        return self.xymean[:1, :], self.xystd[:1, :], self.edgemean[:1, :], self.edgestd[:1, :]

    def _normalize_dataset(self, data, slices) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if not self.normalize:
            # TODO an actual mask, perhaps never necessary though
            return data, slices, torch.tensor([])

        # selecting the right features
        # for x
        print(data.x.shape)
        print(data)
        # exit()
        data.x[:, 4] = data.x[:, 4] - data.x[:, 8]  # Pd = Pd - Pg
        # + 4 for the one-hot encoding for four node types, -2 because we remove the index and Pg
        template = torch.zeros((data.x.shape[0], data.x.shape[1] + 3 - 2))
        template[:, 0:4] = torch.nn.functional.one_hot(
            data.x[:, 1].type(torch.int64), num_classes=4)
        template[:, 4:10] = data.x[:, 2:8]
        data.x = template
        # for y
        data.y = data.y[:, 2:]

        # SHAPE NOW: torch.Size([14, 10]) torch.Size([14, 6]) for x and y

        # normalizing
        # for node attributes
        xy = torch.concat([data.x[:, 4:], data.y], dim=0)
        # 6 for:
        mean = torch.mean(xy, dim=0).unsqueeze(
            dim=0).expand(data.x.shape[0], 6)
        std = torch.std(xy, dim=0).unsqueeze(dim=0).expand(
            data.x.shape[0], 6)  # Vm, Va, Pd, Qd, Gs, Bs
        self.xymean, self.xystd = mean, std
        # + 0.0000001 to avoid NaN's because of division by zero
        data.x[:, 4:] = (data.x[:, 4:] - mean) / (std + 0.0000001)
        data.y = (data.y - mean) / (std + 0.0000001)
        # for edge attributes
        mean = torch.mean(data.edge_attr, dim=0).unsqueeze(dim=0).expand(
            data.edge_attr.shape[0], data.edge_attr.shape[1])
        std = torch.std(data.edge_attr, dim=0).unsqueeze(dim=0).expand(
            data.edge_attr.shape[0], data.edge_attr.shape[1])
        self.edgemean, self.edgestd = mean, std
        data.edge_attr = (data.edge_attr - mean) / (std + 0.0000001)

        # adding the mask
        # where x and y are unequal, the network must predict
        # 1 where value changed, 0 where it did not change
        unequal = (data.x[:, 4:] != data.y).float()
        data.prediction_mask = unequal
        data.x = torch.concat([data.x, unequal], dim=1)

        return data, slices, unequal

    @property
    def raw_file_names(self) -> List[str]:
        if self.case != 'mixed':
            return ["case"+f"{self.case}"+"_"+name for name in self.partial_file_names]
        else:
            return ["case"+f"{case}"+"_"+name for case in self.mixed_cases for name in self.partial_file_names]

    @property
    def processed_file_names(self) -> List[str]:
        return ["case"+f"{self.case}"+"_processed_data.pt"]

    def len(self):
        return self.slices['x'].shape[0]-1

    # def get(self, idx: int) -> Data: # override
    #     return self.data[idx]

    def process(self):
        # then use from_scipy_sparse_matrix()
        assert len(self.raw_paths) % 4 == 0
        raw_paths_per_case = [[self.raw_paths[i], self.raw_paths[i+1], self.raw_paths[i+2], self.raw_paths[i+3],] for i in range(0, len(self.raw_paths), 4)]
        data_list = []
        for case, raw_paths in enumerate(raw_paths_per_case):
            adj_mat = dense_to_sparse(torch.from_numpy(np.load(raw_paths[0])))
            edge_features = torch.from_numpy(np.load(raw_paths[1])).float()
            node_features_x = torch.from_numpy(np.load(raw_paths[2])).float()
            node_features_y = torch.from_numpy(np.load(raw_paths[3])).float()

            if self.split is not None:
                split_len = [int(len(node_features_x) * i) for i in self.split]
                edge_features = torch.split(edge_features, split_len, dim=0)[
                    self.split_order[self.task]]
                node_features_x = torch.split(node_features_x, split_len, dim=0)[
                    self.split_order[self.task]]
                node_features_y = torch.split(node_features_y, split_len, dim=0)[
                    self.split_order[self.task]]

            per_case_data_list = [
                Data(
                    x=node_features_x[i],
                    y=node_features_y[i],
                    edge_index=edge_features[i, :, 0:2].T.to(torch.long)-1,
                    edge_attr=edge_features[i, :, 2:],
                ) for i in range(len(node_features_x))
            ]
            
            data_list.extend(per_case_data_list)

        if self.pre_filter is not None:  # filter out some data
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])


def main():
    try:
        # shape = (N, n_edges, 7)       (from, to, ...)
        edge_features = np.load("data/raw/case14_edge_features.npy")
        # shape = (N, n_nodes, n_nodes)
        adj_matrix = np.load("data/raw/case14_adjacency_matrix.npy")
        # shape = (N, n_nodes, 9)
        node_features_x = np.load("data/raw/case14_node_features_x.npy")
        # shape = (N, n_nodes, 8)
        node_features_y = np.load("data/raw/case14_node_features_y.npy")
    except FileNotFoundError:
        print("File not found.")

    print(f"edge_features.shape = {edge_features.shape}")
    print(f"adj_matrix.shape = {adj_matrix.shape}")
    print(f"node_features_x.shape = {node_features_x.shape}")
    print(f"node_features_y.shape = {node_features_y.shape}")

    trainset = PowerFlowData(root="data", case=14,
                             split=[.5, .2, .3], task="train")
    train_loader = torch_geometric.loader.DataLoader(
        trainset, batch_size=12, shuffle=True)
    print(len(trainset))
    print(trainset[0])
    print(next(iter(train_loader)))
    pass


if __name__ == "__main__":
    main()