gridfm-graphkit 0.0.1__py3-none-any.whl

gridfm_graphkit/utils/visualization.py

@@ -0,0 +1,324 @@
+from gridfm_graphkit.utils.loss import PBELoss
+from gridfm_graphkit.datasets.globals import PQ, PV, REF
+
+import torch
+import networkx as nx
+import matplotlib.pyplot as plt
+import numpy as np
+import copy
+
+
+def visualize_error(data_point, model, baseMVA, device):
+    data_point = copy.deepcopy(data_point)
+    active_loss = None
+    loss = PBELoss(visualization=True)
+
+    # Inference of one data point
+    model.eval()
+    with torch.no_grad():
+        data_point = data_point.to(device)
+        mask_value_expanded = model.mask_value.expand(data_point.x.shape[0], -1)
+        data_point.x[:, : data_point.mask.shape[1]][data_point.mask] = (
+            mask_value_expanded[data_point.mask]
+        )
+        output = model(
+            data_point.x,
+            data_point.pe,
+            data_point.edge_index,
+            data_point.edge_attr,
+            torch.zeros(data_point.x.shape[0], dtype=int).to(device),
+        )
+        loss_dict = loss(
+            output,
+            data_point.y,
+            data_point.edge_index,
+            data_point.edge_attr,
+            data_point.mask,
+        )
+        active_loss = loss_dict["Nodal Active Power Loss in p.u."]
+        active_loss = active_loss.cpu() * baseMVA
+
+    # Create a graph
+    G = nx.Graph()
+    edges = [
+        (u, v)
+        for u, v in zip(
+            data_point.edge_index[0].tolist(),
+            data_point.edge_index[1].tolist(),
+        )
+        if u != v
+    ]
+    G.add_edges_from(edges)
+
+    # Assign labels based on node type
+    node_shapes = {"REF": "s", "PV": "H", "PQ": "o"}
+    num_nodes = data_point.x.shape[0]
+    mask_PQ = data_point.x[:, PQ] == 1
+    mask_PV = data_point.x[:, PV] == 1
+    mask_REF = data_point.x[:, REF] == 1
+    node_labels = {}
+    for i in range(num_nodes):
+        if mask_REF[i]:
+            node_labels[i] = "REF"
+        elif mask_PV[i]:
+            node_labels[i] = "PV"
+        elif mask_PQ[i]:
+            node_labels[i] = "PQ"
+
+    # Set node positions
+    pos = nx.spring_layout(G, seed=42)
+
+    # Define colormap
+    cmap = plt.cm.viridis
+    vmin = min(active_loss)
+    vmax = max(active_loss)
+    norm = plt.Normalize(vmin=vmin, vmax=vmax)
+
+    # Create a figure and axis
+    fig, ax = plt.subplots(figsize=(13, 7))
+
+    # Draw nodes with heatmap coloring
+    for node_type, shape in node_shapes.items():
+        nodes = [i for i in node_labels if node_labels[i] == node_type]
+        nx.draw_networkx_nodes(
+            G,
+            pos,
+            nodelist=nodes,
+            node_color=[active_loss[i] for i in nodes],
+            cmap=cmap,
+            node_size=800,
+            ax=ax,
+            vmin=vmin,
+            vmax=vmax,
+            node_shape=shape,
+        )
+
+    # Draw edges
+    nx.draw_networkx_edges(G, pos, edge_color="gray", alpha=0.5, ax=ax)
+
+    # Draw labels (node types)
+    nx.draw_networkx_labels(
+        G,
+        pos,
+        labels=node_labels,
+        font_size=10,
+        font_color="white",
+        font_weight="bold",
+        ax=ax,
+    )
+
+    # Add colorbar
+    cbar = plt.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), ax=ax)
+    cbar.set_label("Active Power Residuals (MW)", fontsize=12)
+    cbar.ax.tick_params(labelsize=12)
+
+    for spine in ax.spines.values():
+        spine.set_linewidth(2)  # Adjust thickness here (e.g., 2 or any value)
+
+    # Show plot
+    plt.title("Nodal Active Power Residuals", fontsize=14, fontweight="bold")
+    plt.show()
+
+
+def visualize_quantity_heatmap(
+    data_point,
+    model,
+    quantity,
+    quantity_name,
+    unit,
+    node_normalizer,
+    cmap,
+    device,
+):
+    """
+    Visualizes a heatmap of a specified quantity (VM, PD, QD, PG, QG, VA) for a given dataset and model.
+
+    Parameters:
+        data_point: Power grid data.
+        model: The trained model used for inference.
+        quantity: The quantity to visualize (e.g., VM, PD, QD, PG, QG, VA).
+    """
+    data_point = copy.deepcopy(data_point)
+    mask_PQ = data_point.x[:, PQ] == 1
+    mask_PV = data_point.x[:, PV] == 1
+    mask_REF = data_point.x[:, REF] == 1
+    gt_values = data_point.y[:, quantity]  # Extract ground truth values
+
+    # Inference of one data point
+    model.eval()
+    with torch.no_grad():
+        data_point = data_point.to(device)
+        mask_value_expanded = model.mask_value.expand(data_point.x.shape[0], -1)
+        data_point.x[:, : data_point.mask.shape[1]][data_point.mask] = (
+            mask_value_expanded[data_point.mask]
+        )
+        output = model(
+            data_point.x,
+            data_point.pe,
+            data_point.edge_index,
+            data_point.edge_attr,
+            torch.zeros(data_point.x.shape[0], dtype=int).to(device),
+        )
+        output = node_normalizer.inverse_transform(output)
+        denormalized_gt = node_normalizer.inverse_transform(data_point.y)
+
+        gt_values = denormalized_gt[:, quantity]
+        predicted_values = output[:, quantity]
+        predicted_values[~data_point.mask[:, quantity]] = denormalized_gt[
+            ~data_point.mask[:, quantity],
+            quantity,
+        ]
+
+    num_nodes = data_point.x.shape[0]
+    predicted_values = predicted_values.cpu()
+    gt_values = gt_values.cpu()
+
+    node_shapes = {"REF": "s", "PV": "H", "PQ": "o"}
+
+    # Create graph
+    G = nx.Graph()
+    edges = [
+        (u, v)
+        for u, v in zip(
+            data_point.edge_index[0].tolist(),
+            data_point.edge_index[1].tolist(),
+        )
+        if u != v
+    ]
+    G.add_edges_from(edges)
+
+    node_labels = {}
+    for i in range(num_nodes):
+        if mask_REF[i]:
+            node_labels[i] = "REF"
+        elif mask_PV[i]:
+            node_labels[i] = "PV"
+        elif mask_PQ[i]:
+            node_labels[i] = "PQ"
+
+    pos = nx.spring_layout(G, seed=42)
+    cmap = cmap
+    vmin = min(predicted_values)
+    vmax = max(predicted_values)
+    norm = plt.Normalize(vmin=vmin, vmax=vmax)
+
+    masked_node_indices = np.where(data_point.mask[:, quantity].cpu())[0]
+
+    # Create subplots for side-by-side layout (3 plots)
+    fig, axes = plt.subplots(1, 3, figsize=(22, 8))
+
+    # First plot (ground truth values)
+    ax = axes[0]
+    for node_type, shape in node_shapes.items():
+        nodes = [i for i in node_labels if node_labels[i] == node_type]
+        node_size = 390 if node_type == "REF" else 600
+        nx.draw_networkx_nodes(
+            G,
+            pos,
+            nodelist=nodes,
+            node_color=[gt_values[i] for i in nodes],
+            cmap=cmap,
+            node_size=node_size,
+            ax=ax,
+            vmin=vmin,
+            vmax=vmax,
+            node_shape=shape,
+        )
+
+    nx.draw_networkx_edges(G, pos, edge_color="gray", alpha=0.5, ax=ax, width=2)
+    nx.draw_networkx_labels(
+        G,
+        pos,
+        labels=node_labels,
+        font_size=10,
+        font_color="white",
+        font_weight="bold",
+        ax=ax,
+    )
+    ax.set_title(f"Input grid {quantity_name}", fontsize=14, fontweight="bold")
+
+    for spine in ax.spines.values():
+        spine.set_linewidth(2)  # Adjust thickness
+
+    # Second plot (with masked nodes in gray)
+    ax = axes[1]
+    for node_type, shape in node_shapes.items():
+        nodes = [i for i in node_labels if node_labels[i] == node_type]
+        node_size = 390 if node_type == "REF" else 600
+        nx.draw_networkx_nodes(
+            G,
+            pos,
+            nodelist=nodes,
+            node_color=[gt_values[i] for i in nodes],
+            cmap=cmap,
+            node_size=node_size,
+            ax=ax,
+            vmin=vmin,
+            vmax=vmax,
+            node_shape=shape,
+        )
+
+    nx.draw_networkx_nodes(
+        G,
+        pos,
+        nodelist=masked_node_indices,
+        node_color="#D3D3D3",
+        node_size=750,
+        ax=ax,
+    )
+    nx.draw_networkx_edges(G, pos, edge_color="gray", alpha=0.5, ax=ax, width=2)
+    nx.draw_networkx_labels(
+        G,
+        pos,
+        labels=node_labels,
+        font_size=10,
+        font_color="white",
+        font_weight="bold",
+        ax=ax,
+    )
+    ax.set_title(f"Masked grid {quantity_name}", fontsize=14, fontweight="bold")
+
+    for spine in ax.spines.values():
+        spine.set_linewidth(2)  # Adjust thickness
+
+    # Third plot (predicted values without masking)
+    ax = axes[2]
+    for node_type, shape in node_shapes.items():
+        nodes = [i for i in node_labels if node_labels[i] == node_type]
+        node_size = 390 if node_type == "REF" else 600
+        nx.draw_networkx_nodes(
+            G,
+            pos,
+            nodelist=nodes,
+            node_color=[predicted_values[i] for i in nodes],
+            cmap=cmap,
+            node_size=node_size,
+            ax=ax,
+            vmin=vmin,
+            vmax=vmax,
+            node_shape=shape,
+        )
+
+    nx.draw_networkx_edges(G, pos, edge_color="gray", alpha=0.5, ax=ax, width=2)
+    nx.draw_networkx_labels(
+        G,
+        pos,
+        labels=node_labels,
+        font_size=10,
+        font_color="white",
+        font_weight="bold",
+        ax=ax,
+    )
+    ax.set_title(f"Reconstructed grid {quantity_name}", fontsize=14, fontweight="bold")
+
+    for spine in ax.spines.values():
+        spine.set_linewidth(2)  # Adjust thickness
+
+    # Colorbar placement
+    cbar_ax = fig.add_axes([0.93, 0.1, 0.02, 0.8])
+    cbar = plt.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), cax=cbar_ax)
+    cbar.set_label(f"{quantity_name} ({unit})", fontsize=12)
+    cbar.ax.tick_params(labelsize=12)
+
+    plt.subplots_adjust(right=0.9)
+    plt.show()

gridfm_graphkit-0.0.1.dist-info/METADATA

@@ -0,0 +1,163 @@
+Metadata-Version: 2.4
+Name: gridfm-graphkit
+Version: 0.0.1
+Summary: Grid Foundation Model
+Author-email: Matteo Mazzonelli <Matteo.Mazzonelli1@ibm.com>, Alban Puech <Alban.Puech1@ibm.com>, Jonas Weiss <jwe@zurich.ibm.com>
+Maintainer-email: Matteo Mazzonelli <Matteo.Mazzonelli1@ibm.com>
+License-Expression: Apache-2.0
+Keywords: electric power grid,foundational model,graph neural networks
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.12.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mlflow>=3.1.0
+Requires-Dist: nbformat>=5.10.4
+Requires-Dist: networkx>=3.4.2
+Requires-Dist: numpy>=2.2.6
+Requires-Dist: pandas>=2.3.0
+Requires-Dist: plotly>=6.1.2
+Requires-Dist: pyyaml>=6.0.2
+Requires-Dist: torch>=2.7.1
+Requires-Dist: torch-geometric>=2.6.1
+Requires-Dist: torchaudio>=2.7.1
+Requires-Dist: torchvision>=0.22.1
+Provides-Extra: dev
+Requires-Dist: mkdocs-material; extra == "dev"
+Requires-Dist: mkdocstrings[python]; extra == "dev"
+Requires-Dist: pre-commit>=4.2.0; extra == "dev"
+Requires-Dist: bandit>=1.8.5; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
+Dynamic: license-file
+
+# gridfm-graphkit
+[![Docs](https://img.shields.io/badge/docs-live-brightgreen)](https://gridfm.github.io/gridfm-graphkit/)
+
+This library is brought to you by the GridFM team to train, finetune and interact with a foundation model for the electric power grid.
+
+---
+
+<p align="center">
+  <img src="docs/figs/pre_training.png" alt="GridFM logo"/>
+  <br/>
+</p>
+
+# Installation
+
+Create a python virtual environment and install the requirements
+```bash
+git clone git@github.com:gridfm/gridfm-graphkit.git
+cd gridfm-graphkit
+python -m venv venv
+source venv/bin/activate
+pip install .
+```
+
+Install the package in editable mode during development phase:
+
+```bash
+pip install -e .[dev,test]
+```
+
+# gridfm-graphkit CLI
+
+An interface to train, fine-tune, and evaluate GridFM models using configurable YAML files and MLflow tracking.
+
+```bash
+gridfm_graphkit <command> [OPTIONS]
+```
+
+Available commands:
+
+* `train` – Train a new model
+* `predict` – Evaluate an existing model
+* `finetune` – Fine-tune a pre-trained model
+
+---
+
+## Training Models
+
+```bash
+gridfm_graphkit train --config path/to/config.yaml
+```
+
+### Arguments
+
+| Argument         | Type   | Description                                                      | Default |
+| ---------------- | ------ | ---------------------------------------------------------------- | ------- |
+| `--config`       | `str`  | **Required for standard training**. Path to base config YAML.    | `None`  |
+| `--grid`         | `str`  | **Optional**. Path to grid search YAML. Not supported with `-c`. | `None`  |
+| `--exp`          | `str`  | **Optional**. MLflow experiment name. Defaults to a timestamp.   | `None`  |
+| `--data_path`    | `str`  | **Optional**. Root dataset directory.                            | `data`  |
+| `-c`             | `flag` | **Optional**. Enable checkpoint mode.                            | `False` |
+| `--model_exp_id` | `str`  | **Required if `-c` is used**. MLflow experiment ID.              | `None`  |
+| `--model_run_id` | `str`  | **Required if `-c` is used**. MLflow run ID.                     | `None`  |
+
+### Examples
+
+**Standard Training:**
+
+```bash
+gridfm_graphkit train --config config/train.yaml --exp "run1"
+```
+
+**Grid Search Training:**
+
+```bash
+gridfm_graphkit train --config config/train.yaml --grid config/grid.yaml
+```
+
+**Training from Checkpoint:**
+
+```bash
+gridfm_graphkit train -c --model_exp_id 123 --model_run_id abc
+```
+
+---
+
+## Evaluating Models
+
+```bash
+gridfm_graphkit predict --model_path model.pth --config config/eval.yaml --eval_name run_eval
+```
+
+### Arguments
+
+| Argument         | Type  | Description                                                       | Default      |
+| ---------------- | ----- | ----------------------------------------------------------------- | ------------ |
+| `--model_path`   | `str` | **Optional**. Path to a model file.                               | `None`       |
+| `--model_exp_id` | `str` | **Required if `--model_path` is not used**. MLflow experiment ID. | `None`       |
+| `--model_run_id` | `str` | **Required if `--model_path` is not used**. MLflow run ID.        | `None`       |
+| `--model_name`   | `str` | **Optional**. Filename inside MLflow artifacts.                   | `best_model` |
+| `--config`       | `str` | **Required**. Path to evaluation config.                          | `None`       |
+| `--eval_name`    | `str` | **Required**. Name of the evaluation run in MLflow.               | `None`       |
+| `--data_path`    | `str` | **Optional**. Path to dataset directory.                          | `data`       |
+
+### Examples
+
+**Evaluate a Logged MLflow Model:**
+
+```bash
+gridfm_graphkit predict --config config/eval.yaml --eval_name run_eval --model_exp_id 1 --model_run_id abc
+```
+
+---
+
+## Fine-Tuning Models
+
+```bash
+gridfm_graphkit finetune --config path/to/config.yaml --model_path path/to/model.pth
+```
+
+### Arguments
+
+| Argument       | Type  | Description                                     | Default |
+| -------------- | ----- | ----------------------------------------------- | ------- |
+| `--config`     | `str` | **Required**. Fine-tuning configuration file.   | `None`  |
+| `--model_path` | `str` | **Required**. Path to a pre-trained model file. | `None`  |
+| `--exp`        | `str` | **Optional**. MLflow experiment name.           | `None`  |
+| `--data_path`  | `str` | **Optional**. Root dataset directory.           | `data`  |

gridfm_graphkit-0.0.1.dist-info/RECORD

@@ -0,0 +1,27 @@
+gridfm_graphkit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+gridfm_graphkit/__main__.py,sha256=nPp_l6YbP62qE907xIrJebLX2DodieZjH79kULAxKgA,2453
+gridfm_graphkit/cli.py,sha256=4j9RN8__EVM2p125k8X7kB507lG_RTUPV_vzPF40Ens,17326
+gridfm_graphkit/datasets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+gridfm_graphkit/datasets/data_normalization.py,sha256=0nlh5wwFATqqdjpTFADi0De4b4HgkkqRHjyd5dvmRlk,7186
+gridfm_graphkit/datasets/globals.py,sha256=3f0Pcap-_XDJl-uflTpxjnhr1-1p84KzwJYvSiessEo,254
+gridfm_graphkit/datasets/powergrid.py,sha256=nShTY8r7q9WZxLZQjkti0pm9qQheJ9NdD4lScrVeoPY,7391
+gridfm_graphkit/datasets/transforms.py,sha256=p8cLhnl2Ey1RVtwbo6KLzYgQjf1W_RYMBkwfiPpbgSs,7397
+gridfm_graphkit/datasets/utils.py,sha256=ZaXf_AM8dlqURocLBrtD7dn2e5msyQIOHXjF5Da51TY,2395
+gridfm_graphkit/io/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+gridfm_graphkit/io/param_handler.py,sha256=s2NIA958lG-5tinYiStDrL4Ua5n9-3B3jeyG8eLjQRM,8998
+gridfm_graphkit/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+gridfm_graphkit/models/gps_transformer.py,sha256=XyCt-fisnISYauhuknM5Zlx9N-gpby71NpzWEzehjEI,5108
+gridfm_graphkit/models/graphTransformer.py,sha256=aFx_ZXzBvAdXmgdIi7mwT6yw_2h_3t4j22V_xdNcGzk,3364
+gridfm_graphkit/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+gridfm_graphkit/training/callbacks.py,sha256=5XSEcRRg-Hhuv5vQ8qgvGyC1h_ccLtnoS-fekQzehRw,1550
+gridfm_graphkit/training/plugins.py,sha256=4w0UHlrafYq2WlFIeXbWludRQVUBErqw4dnslKMt8V8,7485
+gridfm_graphkit/training/trainer.py,sha256=q1Y4NOkvMTQT3m75WSU7lTLa3MLiHrA_2AcaqrQ0ce0,5593
+gridfm_graphkit/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+gridfm_graphkit/utils/loss.py,sha256=OyEu5uheMgKuNE7YMoYI-dmquNHopdH0aiF8vSbS4kE,6519
+gridfm_graphkit/utils/visualization.py,sha256=FXAU5nmo6bY_Pwr8X_p42zQSwr-bknJVDABPTeZSnOs,9491
+gridfm_graphkit-0.0.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+gridfm_graphkit-0.0.1.dist-info/METADATA,sha256=a_cgS-Nm-lz2ozouUBRd7dNYV2EByQQk--uxGPzYPhY,5772
+gridfm_graphkit-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+gridfm_graphkit-0.0.1.dist-info/entry_points.txt,sha256=CVYrtG2_4yKucL63S5klXfmCvfP7l7MBjvfyUuryXcE,66
+gridfm_graphkit-0.0.1.dist-info/top_level.txt,sha256=p3OXCMb-zDtflOl9N4s2Lvp-7hJlmRhf32wE0xbFins,16
+gridfm_graphkit-0.0.1.dist-info/RECORD,,

gridfm_graphkit-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

gridfm_graphkit-0.0.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+gridfm_graphkit = gridfm_graphkit.__main__:main