Add dataset wrapper

.github/workflows/ci-build.yaml

@@ -49,6 +49,8 @@ jobs:
         run: |
           python -m pip install --upgrade pip wheel
           pip install -e ".[test]"
+          TORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])")
+          pip install torch-scatter -f "https://data.pyg.org/whl/torch-${TORCH_VERSION}+cpu.html"
 
       - name: Unit tests
         run: |
@@ -67,11 +69,11 @@ jobs:
           python3 -m pip install --upgrade pip
           python3 -m pip install --upgrade "git+https://github.com/ibm/detect-secrets.git@master#egg=detect-secrets"
           python3 -m pip install boxsdk
-          
+
       - name: Scan repository & write snapshot
         run: |
           mkdir -p security-outputs
-      
+
           # Run detect-secrets while skipping binary files
           detect-secrets scan \
             --exclude-files '.*\.ipynb$|.*\.(png|jpg|jpeg|gif|pdf|onnx|pt|pth|bin|zip)$' \
@@ -166,22 +168,30 @@ jobs:
           pip install -e .[dev,test] || pip install -e .
       - name: Run pip-audit
         uses: pypa/gh-action-pip-audit@v1.1.0
+        with:
+          # CVE-2026-4539: pygments AdlLexer ReDoS, local-only attack vector, no fix released yet
+          ignore-vulns: CVE-2026-4539
 
   trivy_repo:
-    name: Trivy (repo scan)
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - name: Run Trivy filesystem scan
-        uses: aquasecurity/trivy-action@0.33.1
-        with:
-          scan-type: 'fs'
-          scan-ref: '.'
-          format: 'sarif'
-          output: 'trivy-results.sarif'
-          severity: 'HIGH,CRITICAL'
-          ignore-unfixed: true
-      - name: Upload SARIF to Code Scanning
-        uses: github/codeql-action/upload-sarif@v3
-        with:
-          sarif_file: trivy-results.sarif
+      name: Trivy (repo scan)
+      runs-on: ubuntu-latest
+      permissions:
+        security-events: write
+      steps:
+        - uses: actions/checkout@v4
+
+        - name: Run Trivy vulnerability scanner in repo mode
+          # We use the official container-based action to avoid binary install issues
+          uses: aquasecurity/trivy-action@master
+          with:
+            scan-type: 'fs'
+            ignore-unfixed: true
+            format: 'sarif'
+            output: 'trivy-results.sarif'
+            severity: 'HIGH,CRITICAL'
+
+        - name: Upload SARIF to Code Scanning
+          uses: github/codeql-action/upload-sarif@v3
+          if: always() # Upload results even if vulnerabilities are found
+          with:
+            sarif_file: 'trivy-results.sarif'

docs/datasets/data_modules.md

@@ -1,3 +1,3 @@
-# LitGridDataModule
+# LitGridHeteroDataModule
 
-::: gridfm_graphkit.datasets.powergrid_datamodule.LitGridDataModule
+::: gridfm_graphkit.datasets.hetero_powergrid_datamodule.LitGridHeteroDataModule

docs/datasets/data_normalization.md

@@ -3,12 +3,10 @@
 
 
 Normalization improves neural network training by ensuring features are well-scaled, preventing issues like exploding gradients and slow convergence. In power grids, where variables like voltage and power span wide ranges, normalization is essential.
-The `gridfm-graphkit` package offers four methods:
+The `gridfm-graphkit` package offers normalization methods based on the per-unit (p.u.) system:
 
-- [`Min-Max Normalization`](#minmaxnormalizer)
-- [`Standardization (Z-score)`](#standardizer)
-- [`Identity (no normalization)`](#identitynormalizer)
-- [`BaseMVA Normalization`](#basemvanormalizer)
+- [`BaseMVA Normalization`](#heterodatamvanormalizer)
+- [`Per-Sample BaseMVA Normalization`](#heterodatapersamplemvanormalizer)
 
 Each of these strategies implements a unified interface and can be used interchangeably depending on the learning task and data characteristics.
 
@@ -25,27 +23,15 @@ Each of these strategies implements a unified interface and can be used intercha
 
 ---
 
-### `MinMaxNormalizer`
+### `HeteroDataMVANormalizer`
 
-::: gridfm_graphkit.datasets.normalizers.MinMaxNormalizer
+::: gridfm_graphkit.datasets.normalizers.HeteroDataMVANormalizer
 
 ---
 
-### `Standardizer`
+### `HeteroDataPerSampleMVANormalizer`
 
-::: gridfm_graphkit.datasets.normalizers.Standardizer
-
----
-
-### `BaseMVANormalizer`
-
-::: gridfm_graphkit.datasets.normalizers.BaseMVANormalizer
-
----
-
-### `IdentityNormalizer`
-
-::: gridfm_graphkit.datasets.normalizers.IdentityNormalizer
+::: gridfm_graphkit.datasets.normalizers.HeteroDataPerSampleMVANormalizer
 
 ---
 
@@ -54,13 +40,18 @@ Each of these strategies implements a unified interface and can be used intercha
 Example:
 
 ```python
-from gridfm_graphkit.datasets.normalizers import MinMaxNormalizer
-import torch
+from gridfm_graphkit.datasets.normalizers import HeteroDataMVANormalizer
+from torch_geometric.data import HeteroData
+
+# Create normalizer
+normalizer = HeteroDataMVANormalizer(args)
+
+# Fit on training data
+params = normalizer.fit(data_path, scenario_ids)
 
-data = torch.randn(100, 5)  # Example tensor
+# Transform data
+normalizer.transform(hetero_data)
 
-normalizer = MinMaxNormalizer(node_data=True,args=None)
-params = normalizer.fit(data)
-normalized = normalizer.transform(data)
-restored = normalizer.inverse_transform(normalized)
+# Inverse transform to restore original scale
+normalizer.inverse_transform(hetero_data)
 ```

docs/datasets/powergrid.md

@@ -1,3 +1,3 @@
-## `GridDatasetDisk`
+## `HeteroGridDatasetDisk`
 
-::: gridfm_graphkit.datasets.powergrid_dataset.GridDatasetDisk
+::: gridfm_graphkit.datasets.powergrid_hetero_dataset.HeteroGridDatasetDisk

docs/datasets/transforms.md

@@ -2,26 +2,18 @@
 
 > Each transformation class inherits from [`BaseTransform`](https://pytorch-geometric.readthedocs.io/en/latest/modules/transforms.html#torch_geometric.transforms.BaseTransform) provided by [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/).
 
-### `AddNormalizedRandomWalkPE`
+### `RemoveInactiveGenerators`
 
-::: gridfm_graphkit.datasets.transforms.AddNormalizedRandomWalkPE
+::: gridfm_graphkit.datasets.transforms.RemoveInactiveGenerators
 
-### `AddEdgeWeights`
+### `RemoveInactiveBranches`
 
-::: gridfm_graphkit.datasets.transforms.AddEdgeWeights
+::: gridfm_graphkit.datasets.transforms.RemoveInactiveBranches
 
-### `AddIdentityMask`
+### `ApplyMasking`
 
-::: gridfm_graphkit.datasets.transforms.AddIdentityMask
+::: gridfm_graphkit.datasets.transforms.ApplyMasking
 
-### `AddRandomMask`
+### `LoadGridParamsFromPath`
 
-::: gridfm_graphkit.datasets.transforms.AddRandomMask
-
-### `AddPFMask`
-
-::: gridfm_graphkit.datasets.transforms.AddPFMask
-
-### `AddOPFMask`
-
-::: gridfm_graphkit.datasets.transforms.AddOPFMask
+::: gridfm_graphkit.datasets.transforms.LoadGridParamsFromPath

docs/install/installation.md

@@ -1,14 +1,18 @@
+# Installation
+
 You can install `gridfm-graphkit` directly from PyPI:
 
 ```bash
 pip install gridfm-graphkit
 ```
 
+For GPU support and compatibility with PyTorch Geometric's scatter operations, install PyTorch (and optionally CUDA) first, then install the matching `torch-scatter` wheel. See [PyTorch and torch-scatter](#pytorch-and-torch-scatter-optional) below.
+
 ---
 
 ## Development Setup
 
-To contribute or develop locally, clone the repository and install in editable mode:
+To contribute or develop locally, clone the repository and install in editable mode. Use Python 3.10, 3.11, or 3.12 (3.12 is recommended).
 
 ```bash
 git clone git@github.com:gridfm/gridfm-graphkit.git
@@ -18,6 +22,26 @@ source venv/bin/activate
 pip install -e .
 ```
 
+### PyTorch and torch-scatter (optional)
+
+If you need GPU acceleration or PyTorch Geometric scatter ops (used by the library), install PyTorch and the matching `torch-scatter` wheel:
+
+1. Install PyTorch (see [pytorch.org](https://pytorch.org/) for your platform and CUDA version).
+
+2. Get your Torch + CUDA version string:
+   ```bash
+   TORCH_CUDA_VERSION=$(python -c "import torch; print(torch.__version__ + ('+cpu' if torch.version.cuda is None else ''))")
+   ```
+
+3. Install the correct `torch-scatter` wheel:
+   ```bash
+   pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH_CUDA_VERSION}.html
+   ```
+
+---
+
+## Optional extras
+
 For documentation generation and unit testing, install with the optional `dev` and `test` extras:
 
 ```bash

docs/models/models.md

@@ -1,10 +1,37 @@
 # Models
 
-### `GPSTransformer`
+### `GNS_heterogeneous`
 
-::: gridfm_graphkit.models.gps_transformer.GPSTransformer
+::: gridfm_graphkit.models.gnn_heterogeneous_gns.GNS_heterogeneous
 
+---
 
-### `GNN_TransformerConv`
+## Physics Decoders
 
-::: gridfm_graphkit.models.gnn_transformer.GNN_TransformerConv
+### `PhysicsDecoderOPF`
+
+::: gridfm_graphkit.models.utils.PhysicsDecoderOPF
+
+### `PhysicsDecoderPF`
+
+::: gridfm_graphkit.models.utils.PhysicsDecoderPF
+
+### `PhysicsDecoderSE`
+
+::: gridfm_graphkit.models.utils.PhysicsDecoderSE
+
+---
+
+## Utility Modules
+
+### `ComputeBranchFlow`
+
+::: gridfm_graphkit.models.utils.ComputeBranchFlow
+
+### `ComputeNodeInjection`
+
+::: gridfm_graphkit.models.utils.ComputeNodeInjection
+
+### `ComputeNodeResiduals`
+
+::: gridfm_graphkit.models.utils.ComputeNodeResiduals