gridfm-graphkit 0.0.3__tar.gz → 0.0.4__tar.gz

{gridfm_graphkit-0.0.3 → gridfm_graphkit-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gridfm-graphkit
-Version: 0.0.3
+Version: 0.0.4
 Summary: Grid Foundation Model
 Author-email: Matteo Mazzonelli <matteo.mazzonelli1@ibm.com>, Alban Puech <apuech@seas.harvard.edu>, Tamara Govindasamy <tamara.govindasamy@ibm.com>, Mangaliso Mngomezulu <mngomezulum@ibm.com>, Etienne Vos <etienne.vos@ibm.com>, Celia Cintas <celia.cintas@ibm.com>, Jonas Weiss <jwe@zurich.ibm.com>
 Maintainer-email: Matteo Mazzonelli <matteo.mazzonelli1@ibm.com>
@@ -25,6 +25,7 @@ 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
+Requires-Dist: lightning
 Provides-Extra: dev
 Requires-Dist: mkdocs-material; extra == "dev"
 Requires-Dist: mkdocstrings[python]; extra == "dev"
@@ -38,6 +39,9 @@ Dynamic: license-file
 
 # gridfm-graphkit
 [![Docs](https://img.shields.io/badge/docs-available-brightgreen)](https://gridfm.github.io/gridfm-graphkit/)
+![Coverage](https://img.shields.io/badge/coverage-83%25-yellowgreen)
+![Python](https://img.shields.io/badge/python-3.10%20%E2%80%93%203.12-blue)
+![License](https://img.shields.io/badge/license-Apache%202.0-blue)
 
 This library is brought to you by the GridFM team to train, finetune and interact with a foundation model for the electric power grid.
 
@@ -73,7 +77,7 @@ pip install -e .[dev,test]
 ```
 
 
-# gridfm-graphkit CLI
+# CLI commands
 
 An interface to train, fine-tune, and evaluate GridFM models using configurable YAML files and MLflow tracking.
 
@@ -83,9 +87,10 @@ gridfm_graphkit <command> [OPTIONS]
 
 Available commands:
 
-* `train` – Train a new model
-* `predict` – Evaluate an existing model
-* `finetune` – Fine-tune a pre-trained model
+* `train` – Train a new model from scrathc
+* `finetune` – Fine-tune an existing pre-trained model
+* `evaluate` – Evaluate model performance on a dataset
+* `predict` – Run inference and save predictions
 
 ---
 
@@ -99,75 +104,77 @@ gridfm_graphkit train --config path/to/config.yaml
 
 | 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`  |
+| `--config`       | `str`  | **Required**. Path to the training configuration YAML file.    | `None`  |
+| `--exp_name`     | `str`  | **Optional**. MLflow experiment name.                            | `timestamp`  |
+| `--run_name`     | `str`  | **Optional**. MLflow run name.                                   | `run`  |
+| `--log_dir  `    | `str`  | **Optional**. MLflow logging directory.                              | `mlruns`  |
 | `--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"
+gridfm_graphkit train --config examples/config/case30_ieee_base.yaml --data_path examples/data
 ```
 
-**Grid Search Training:**
+---
+
+## Fine-Tuning Models
 
 ```bash
-gridfm_graphkit train --config config/train.yaml --grid config/grid.yaml
+gridfm_graphkit finetune --config path/to/config.yaml --model_path path/to/model.pth
 ```
 
-**Training from Checkpoint:**
+### 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_name`   | `str` | MLflow experiment name.                         | timestamp |
+| `--run_name`   | `str` | MLflow run name.                                | `run`     |
+| `--log_dir`    | `str` | MLflow logging directory.                       | `mlruns`  |
+| `--data_path`  | `str` | Root dataset directory.                         | `data`    |
 
-```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
+gridfm_graphkit evaluate --config path/to/eval.yaml --model_path path/to/model.pth
 ```
 
 ### 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
-```
+| Argument       | Type  | Description                              | Default   |
+| -------------- | ----- | ---------------------------------------- | --------- |
+| `--config`     | `str` | **Required**. Path to evaluation config. | `None`    |
+| `--model_path` | `str` | Path to the trained model file.          | `None`    |
+| `--exp_name`   | `str` | MLflow experiment name.                  | timestamp |
+| `--run_name`   | `str` | MLflow run name.                         | `run`     |
+| `--log_dir`    | `str` | MLflow logging directory.                | `mlruns`  |
+| `--data_path`  | `str` | Dataset directory.                       | `data`    |
 
 ---
 
-## Fine-Tuning Models
+## Running Predictions
 
 ```bash
-gridfm_graphkit finetune --config path/to/config.yaml --model_path path/to/model.pth
+gridfm_graphkit predict --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`  |
+| Argument        | Type  | Description                                   | Default   |
+| --------------- | ----- | --------------------------------------------- | --------- |
+| `--config`      | `str` | **Required**. Path to prediction config file. | `None`    |
+| `--model_path`  | `str` | Path to the trained model file.               | `None`    |
+| `--exp_name`    | `str` | MLflow experiment name.                       | timestamp |
+| `--run_name`    | `str` | MLflow run name.                              | `run`     |
+| `--log_dir`     | `str` | MLflow logging directory.                     | `mlruns`  |
+| `--data_path`   | `str` | Dataset directory.                            | `data`    |
+| `--output_path` | `str` | Directory where predictions are saved.        | `data`    |
+
+---

gridfm_graphkit-0.0.4/README.md

@@ -0,0 +1,141 @@
+# gridfm-graphkit
+[![Docs](https://img.shields.io/badge/docs-available-brightgreen)](https://gridfm.github.io/gridfm-graphkit/)
+![Coverage](https://img.shields.io/badge/coverage-83%25-yellowgreen)
+![Python](https://img.shields.io/badge/python-3.10%20%E2%80%93%203.12-blue)
+![License](https://img.shields.io/badge/license-Apache%202.0-blue)
+
+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="https://raw.githubusercontent.com/gridfm/gridfm-graphkit/refs/heads/main/docs/figs/pre_training.png" alt="GridFM logo"/>
+  <br/>
+</p>
+
+# Installation
+
+You can install `gridfm-graphkit` directly from PyPI:
+
+```bash
+pip install gridfm-graphkit
+```
+
+To contribute or develop locally, clone the repository and install in editable mode:
+
+```bash
+git clone git@github.com:gridfm/gridfm-graphkit.git
+cd gridfm-graphkit
+python -m venv venv
+source venv/bin/activate
+pip install -e .
+```
+
+For documentation generation and unit testing, install with the optional `dev` and `test` extras:
+
+```bash
+pip install -e .[dev,test]
+```
+
+
+# CLI commands
+
+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 from scrathc
+* `finetune` – Fine-tune an existing pre-trained model
+* `evaluate` – Evaluate model performance on a dataset
+* `predict` – Run inference and save predictions
+
+---
+
+## Training Models
+
+```bash
+gridfm_graphkit train --config path/to/config.yaml
+```
+
+### Arguments
+
+| Argument         | Type   | Description                                                      | Default |
+| ---------------- | ------ | ---------------------------------------------------------------- | ------- |
+| `--config`       | `str`  | **Required**. Path to the training configuration YAML file.    | `None`  |
+| `--exp_name`     | `str`  | **Optional**. MLflow experiment name.                            | `timestamp`  |
+| `--run_name`     | `str`  | **Optional**. MLflow run name.                                   | `run`  |
+| `--log_dir  `    | `str`  | **Optional**. MLflow logging directory.                              | `mlruns`  |
+| `--data_path`    | `str`  | **Optional**. Root dataset directory.                            | `data`  |
+
+### Examples
+
+**Standard Training:**
+
+```bash
+gridfm_graphkit train --config examples/config/case30_ieee_base.yaml --data_path examples/data
+```
+
+---
+
+## 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_name`   | `str` | MLflow experiment name.                         | timestamp |
+| `--run_name`   | `str` | MLflow run name.                                | `run`     |
+| `--log_dir`    | `str` | MLflow logging directory.                       | `mlruns`  |
+| `--data_path`  | `str` | Root dataset directory.                         | `data`    |
+
+
+---
+
+## Evaluating Models
+
+```bash
+gridfm_graphkit evaluate --config path/to/eval.yaml --model_path path/to/model.pth
+```
+
+### Arguments
+
+| Argument       | Type  | Description                              | Default   |
+| -------------- | ----- | ---------------------------------------- | --------- |
+| `--config`     | `str` | **Required**. Path to evaluation config. | `None`    |
+| `--model_path` | `str` | Path to the trained model file.          | `None`    |
+| `--exp_name`   | `str` | MLflow experiment name.                  | timestamp |
+| `--run_name`   | `str` | MLflow run name.                         | `run`     |
+| `--log_dir`    | `str` | MLflow logging directory.                | `mlruns`  |
+| `--data_path`  | `str` | Dataset directory.                       | `data`    |
+
+---
+
+## Running Predictions
+
+```bash
+gridfm_graphkit predict --config path/to/config.yaml --model_path path/to/model.pth
+```
+
+### Arguments
+
+| Argument        | Type  | Description                                   | Default   |
+| --------------- | ----- | --------------------------------------------- | --------- |
+| `--config`      | `str` | **Required**. Path to prediction config file. | `None`    |
+| `--model_path`  | `str` | Path to the trained model file.               | `None`    |
+| `--exp_name`    | `str` | MLflow experiment name.                       | timestamp |
+| `--run_name`    | `str` | MLflow run name.                              | `run`     |
+| `--log_dir`     | `str` | MLflow logging directory.                     | `mlruns`  |
+| `--data_path`   | `str` | Dataset directory.                            | `data`    |
+| `--output_path` | `str` | Directory where predictions are saved.        | `data`    |
+
+---

gridfm_graphkit-0.0.4/gridfm_graphkit/__init__.py

@@ -0,0 +1,6 @@
+import gridfm_graphkit.datasets
+import gridfm_graphkit.models
+
+__all__ = [
+    "gridfm_graphkit",
+]

gridfm_graphkit-0.0.4/gridfm_graphkit/__main__.py

@@ -0,0 +1,58 @@
+import argparse
+from datetime import datetime
+from gridfm_graphkit.cli import main_cli
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="gridfm_graphkit",
+        description="gridfm-graphkit CLI",
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+    exp_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+
+    # ---- TRAIN SUBCOMMAND ----
+    train_parser = subparsers.add_parser("train", help="Run training")
+    train_parser.add_argument("--config", type=str, required=True)
+    train_parser.add_argument("--exp_name", type=str, default=exp_name)
+    train_parser.add_argument("--run_name", type=str, default="run")
+    train_parser.add_argument("--log_dir", type=str, default="mlruns")
+    train_parser.add_argument("--data_path", type=str, default="data")
+
+    # ---- FINETUNE SUBCOMMAND ----
+    finetune_parser = subparsers.add_parser("finetune", help="Run fine-tuning")
+    finetune_parser.add_argument("--config", type=str, required=True)
+    finetune_parser.add_argument("--model_path", type=str, required=True)
+    finetune_parser.add_argument("--exp_name", type=str, default=exp_name)
+    finetune_parser.add_argument("--run_name", type=str, default="run")
+    finetune_parser.add_argument("--log_dir", type=str, default="mlruns")
+    finetune_parser.add_argument("--data_path", type=str, default="data")
+
+    # ---- EVALUATE SUBCOMMAND ----
+    evaluate_parser = subparsers.add_parser(
+        "evaluate",
+        help="Evaluate model performance",
+    )
+    evaluate_parser.add_argument("--model_path", type=str, default=None)
+    evaluate_parser.add_argument("--config", type=str, required=True)
+    evaluate_parser.add_argument("--exp_name", type=str, default=exp_name)
+    evaluate_parser.add_argument("--run_name", type=str, default="run")
+    evaluate_parser.add_argument("--log_dir", type=str, default="mlruns")
+    evaluate_parser.add_argument("--data_path", type=str, default="data")
+
+    # ---- PREDICT SUBCOMMAND ----
+    predict_parser = subparsers.add_parser("predict", help="Evaluate model performance")
+    predict_parser.add_argument("--model_path", type=str, required=None)
+    predict_parser.add_argument("--config", type=str, required=True)
+    predict_parser.add_argument("--exp_name", type=str, default=exp_name)
+    predict_parser.add_argument("--run_name", type=str, default="run")
+    predict_parser.add_argument("--log_dir", type=str, default="mlruns")
+    predict_parser.add_argument("--data_path", type=str, default="data")
+    predict_parser.add_argument("--output_path", type=str, default="data")
+
+    args = parser.parse_args()
+    main_cli(args)
+
+
+if __name__ == "__main__":
+    main()

gridfm_graphkit-0.0.4/gridfm_graphkit/cli.py

@@ -0,0 +1,134 @@
+from gridfm_graphkit.datasets.powergrid_datamodule import LitGridDataModule
+from gridfm_graphkit.io.param_handler import NestedNamespace
+from gridfm_graphkit.training.callbacks import SaveBestModelStateDict
+import numpy as np
+import os
+import yaml
+import torch
+import random
+import pandas as pd
+
+from gridfm_graphkit.tasks.feature_reconstruction_task import FeatureReconstructionTask
+from lightning.pytorch.callbacks.early_stopping import EarlyStopping
+from lightning.pytorch.callbacks.model_checkpoint import ModelCheckpoint
+from lightning.pytorch.loggers import MLFlowLogger
+import lightning as L
+
+
+def get_training_callbacks(args):
+    early_stop_callback = EarlyStopping(
+        monitor="Validation loss",
+        min_delta=args.callbacks.tol,
+        patience=args.callbacks.patience,
+        verbose=False,
+        mode="min",
+    )
+
+    save_best_model_callback = SaveBestModelStateDict(
+        monitor="Validation loss",
+        mode="min",
+        filename="best_model_state_dict.pt",
+    )
+
+    checkpoint_callback = ModelCheckpoint(
+        monitor="Validation loss",  # or whichever metric you track
+        mode="min",
+        save_last=True,
+        save_top_k=0,
+    )
+
+    return [early_stop_callback, save_best_model_callback, checkpoint_callback]
+
+
+def main_cli(args):
+    logger = MLFlowLogger(
+        save_dir=args.log_dir,
+        experiment_name=args.exp_name,
+        run_name=args.run_name,
+    )
+
+    with open(args.config, "r") as f:
+        base_config = yaml.safe_load(f)
+
+    config_args = NestedNamespace(**base_config)
+
+    torch.manual_seed(config_args.seed)
+    random.seed(config_args.seed)
+    np.random.seed(config_args.seed)
+
+    litGrid = LitGridDataModule(config_args, args.data_path)
+    model = FeatureReconstructionTask(
+        config_args,
+        litGrid.node_normalizers,
+        litGrid.edge_normalizers,
+    )
+    if args.command != "train":
+        print(f"Loading model weights from {args.model_path}")
+        state_dict = torch.load(args.model_path)
+        model.load_state_dict(state_dict)
+
+    trainer = L.Trainer(
+        logger=logger,
+        accelerator=config_args.training.accelerator,
+        devices=config_args.training.devices,
+        strategy=config_args.training.strategy,
+        log_every_n_steps=1,
+        default_root_dir=args.log_dir,
+        max_epochs=config_args.training.epochs,
+        callbacks=get_training_callbacks(config_args),
+    )
+    if args.command == "train" or args.command == "finetune":
+        trainer.fit(model=model, datamodule=litGrid)
+
+    if args.command != "predict":
+        trainer.test(model=model, datamodule=litGrid)
+
+    if args.command == "predict":
+        predictions = trainer.predict(model=model, datamodule=litGrid)
+        all_outputs = []
+        all_mask_PQ = []
+        all_mask_PV = []
+        all_mask_REF = []
+        all_scenarios = []
+        all_bus_numbers = []
+
+        for batch in predictions:
+            all_outputs.append(batch["output"])
+            all_mask_PQ.append(batch["mask_PQ"])
+            all_mask_PV.append(batch["mask_PV"])
+            all_mask_REF.append(batch["mask_REF"])
+            all_scenarios.append(batch["scenario_id"])
+            all_bus_numbers.append(batch["bus_number"])
+
+        # Concatenate all
+        outputs = np.concatenate(all_outputs, axis=0)  # shape: [num_nodes, 6]
+        mask_PQ = np.concatenate(all_mask_PQ, axis=0)
+        mask_PV = np.concatenate(all_mask_PV, axis=0)
+        mask_REF = np.concatenate(all_mask_REF, axis=0)
+        scenario_ids = np.concatenate(all_scenarios, axis=0)
+        bus_numbers = np.concatenate(all_bus_numbers, axis=0)
+
+        # Build DataFrame
+        df = pd.DataFrame(
+            {
+                "scenario": scenario_ids,
+                "bus": bus_numbers,
+                "PD": outputs[:, 0],
+                "QD": outputs[:, 1],
+                "PG": outputs[:, 2],
+                "QG": outputs[:, 3],
+                "VM": outputs[:, 4],
+                "VA": outputs[:, 5],
+                "PQ": mask_PQ.astype(int),
+                "PV": mask_PV.astype(int),
+                "REF": mask_REF.astype(int),
+            },
+        )
+
+        # Save CSV
+        output_dir = os.path.join(args.output_path)
+        os.makedirs(output_dir, exist_ok=True)
+        csv_path = os.path.join(output_dir, "predictions.csv")
+        df.to_csv(csv_path, index=False)
+
+        print(f"Saved predictions to {csv_path}")

gridfm_graphkit-0.0.4/gridfm_graphkit/datasets/__init__.py

@@ -0,0 +1,23 @@
+from gridfm_graphkit.datasets.transforms import (
+    AddPFMask,
+    AddIdentityMask,
+    AddRandomMask,
+    AddOPFMask,
+)
+from gridfm_graphkit.datasets.normalizers import (
+    Standardizer,
+    MinMaxNormalizer,
+    BaseMVANormalizer,
+    IdentityNormalizer,
+)
+
+__all__ = [
+    "AddPFMask",
+    "AddIdentityMask",
+    "AddRandomMask",
+    "AddOPFMask",
+    "Standardizer",
+    "MinMaxNormalizer",
+    "BaseMVANormalizer",
+    "IdentityNormalizer",
+]

gridfm_graphkit-0.0.3/gridfm_graphkit/datasets/data_normalization.py → gridfm_graphkit-0.0.4/gridfm_graphkit/datasets/normalizers.py

@@ -1,4 +1,5 @@
 from gridfm_graphkit.datasets.globals import PD, QD, PG, QG, VA
+from gridfm_graphkit.io.registries import NORMALIZERS_REGISTRY
 import torch
 from abc import ABC, abstractmethod
 
@@ -53,13 +54,25 @@ class Normalizer(ABC):
             Original tensor.
         """
 
+    @abstractmethod
+    def get_stats(self) -> dict:
+        """
+        Return the stored normalization statistics for logging/inspection.
+        """
+
 
+@NORMALIZERS_REGISTRY.register("minmax")
 class MinMaxNormalizer(Normalizer):
     """
     Scales each feature to the [0, 1] range.
+
+    Args:
+        node_data (bool): Whether data is node-level or edge-level
+        args (NestedNamespace): Parameters
+
     """
 
-    def __init__(self):
+    def __init__(self, node_data: bool, args):
         self.min_val = None
         self.max_val = None
 
@@ -95,13 +108,25 @@ class MinMaxNormalizer(Normalizer):
         diff[diff == 0] = 1
         return (normalized_data * diff) + self.min_val
 
+    def get_stats(self) -> dict:
+        return {
+            "min_value": self.min_val.tolist() if self.min_val is not None else None,
+            "max_value": self.max_val.tolist() if self.max_val is not None else None,
+        }
 
+
+@NORMALIZERS_REGISTRY.register("standard")
 class Standardizer(Normalizer):
     """
     Standardizes each feature to zero mean and unit variance.
+
+    Args:
+        node_data (bool): Whether data is node-level or edge-level
+        args (NestedNamespace): Parameters
+
     """
 
-    def __init__(self):
+    def __init__(self, node_data: bool, args):
         self.mean = None
         self.std = None
 
@@ -137,7 +162,14 @@ class Standardizer(Normalizer):
         std[std == 0] = 1
         return (normalized_data * std) + self.mean
 
+    def get_stats(self) -> dict:
+        return {
+            "mean": self.mean.tolist() if self.mean is not None else None,
+            "std": self.std.tolist() if self.std is not None else None,
+        }
+
 
+@NORMALIZERS_REGISTRY.register("baseMVAnorm")
 class BaseMVANormalizer(Normalizer):
     """
     In power systems, a suitable normalization strategy must preserve the physical properties of
@@ -148,14 +180,17 @@ class BaseMVANormalizer(Normalizer):
     preserving fundamental physical relationships.
     """
 
-    def __init__(self, node_data: bool, baseMVA_orig: float = 100.0):
+    def __init__(self, node_data: bool, args):
         """
         Args:
-            node_data: Whether data is node-level or edge-level (PD, QD, PG, QG, VA).
-            baseMVA_orig: Original baseMVA (e.g. from MATPOWER).
+            node_data: Whether data is node-level or edge-level
+            args (NestedNamespace): Parameters
+
+        Attributes:
+            baseMVA (float): baseMVA found in casefile. From ``args.data.baseMVA``.
         """
         self.node_data = node_data
-        self.baseMVA_orig = baseMVA_orig
+        self.baseMVA_orig = getattr(args.data, "baseMVA", 100)
         self.baseMVA = None
 
     def to(self, device):
@@ -208,12 +243,26 @@ class BaseMVANormalizer(Normalizer):
 
         return normalized_data
 
+    def get_stats(self) -> dict:
+        return {
+            "baseMVA": self.baseMVA,
+            "baseMVA_orig": self.baseMVA_orig,
+        }
+
 
+@NORMALIZERS_REGISTRY.register("identity")
 class IdentityNormalizer(Normalizer):
     """
     No normalization: returns data unchanged.
+
+    Args:
+            node_data: Whether data is node-level or edge-level
+            args (NestedNamespace): Parameters
     """
 
+    def __init__(self, node_data: bool, args):
+        pass
+
     def fit(self, data: torch.Tensor) -> dict:
         return {}
 
@@ -225,3 +274,6 @@ class IdentityNormalizer(Normalizer):
 
     def inverse_transform(self, normalized_data: torch.Tensor) -> torch.Tensor:
         return normalized_data
+
+    def get_stats(self) -> dict:
+        return {"note": "No normalization applied."}