GridFM App Distribution

GridFM Distribution appkitgridfm-app-distribution is a Python library to train, finetune, and predict a foundation model of an electrical power grid. Specialised training tasks can be developed for specific distribution applications.

Installation

Create and activate a virtual environment (make sure you use the right python version = 3.10, 3.11, 3.12 or 3.13.)

python -m venv venv
source venv/bin/activate

Install gridfm-app-distribution from PyPI

pip install gridfm-app-distribution

torch-scatter is a required dependency. It cannot be bundled in pyproject.toml because the correct wheel depends on your PyTorch and CUDA versions, so it must be installed separately.

Get PyTorch + CUDA version for torch-scatter

TORCH_CUDA_VERSION=$(python -c "import torch; print(torch.__version__ + ('+cpu' if torch.version.cuda is None else ''))")

Install the correct torch-scatter wheel

pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH_CUDA_VERSION}.html
pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH_CUDA_VERSION}.html

For documentation generation and unit testing, install with the optional dev and test extras:

pip install "gridfm-app-distribution[dev,test]"

Generating training data (contingency_data_gen)

The repository ships a second application, contingency_data_gen, that builds distribution networks and runs N-k contingency analyses with pandapower, writing the results as gridfm-datakit-format partitioned parquet under <output_dir>/raw/ (bus_data, gen_data, branch_data). That is the layout gridfm_app_dist reads from --data_path, so generation chains directly into training:

contingency_data_gen --config data_gen_1.yaml          # writes data/scenario_33radial/raw/, ...
gridfm_app_dist train --config configs/scenario_33radial.yaml --data_path data

Jobs are described in a YAML config (one entry per scenario). Networks can come from an Excel workbook (flag: 1|2|3, network_file_path), a YAML network definition built with the bundled netlib builder (flag: 9, builtin_network: netbuilder/urban_LV.yaml), a MATPOWER file (flag: 10) or a pandapower.networks factory (flag: null, builtin_network: case33bw). Relative paths resolve against the config's folder and then the package's configs/ folder; output_dir resolves against the working directory.

contingency_data_gen --list-configs                      # configs shipped with the package
contingency_data_gen --config my_jobs.yaml --dry-run     # validate and print the plan
contingency_data_gen --config data_gen_1.yaml --only case33radial

Network builder examples (examples/datakit/) show how to assemble MV/LV networks from the standard-type catalogue in contingency_data_gen.netlib.

CLI commands

Interface to train, fine-tune, evaluate, and run inference on GridFM models of distribution networks using YAML configs and MLflow tracking.

gridfm_app_dist <command> [OPTIONS]

Available commands:

• train - Train a new model from scratch
• finetune - Fine-tune an existing pre-trained model
• evaluate - Evaluate model performance on a dataset
• predict - Run inference and save predictions

Training Models

gridfm_app_dist --config path/to/config.yaml

Arguments

Examples

Standard Training:

gridfm_app_dist train --config examples/config/rural_MVLV_test.yaml --data_path examples/data

Fine-Tuning Models

gridfm_app_dist finetune --config path/to/config.yaml --model_path path/to/model.pt

Arguments

Evaluating Models

gridfm_app_dist evaluate --config path/to/eval.yaml --model_path path/to/model.pt

Arguments

Example with saved normalizer stats

When evaluating a model on a dataset, you can pass the normalizer statistics from the original training run to ensure the same normalization parameters are used:

gridfm_app_dist evaluate \
  --config examples/config/HGNS_PF_datakit_case118.yaml \
  --model_path mlruns/<experiment_id>/<run_id>/artifacts/model/best_model_state_dict.pt \
  --normalizer_stats mlruns/<experiment_id>/<run_id>/artifacts/stats/normalizer_stats.pt \
  --data_path data

Note: The --normalizer_stats flag only affects normalizers with fit_strategy = "fit_on_train" (e.g. HeteroDataMVANormalizer). Per-sample normalizers (HeteroDataPerSampleMVANormalizer) always recompute their statistics from the current dataset regardless of this flag.

Running Predictions

gridfm_app_dist predict --config path/to/config.yaml --model_path path/to/model.pt

Arguments

Benchmarking Dataloader Throughput

gridfm_app_dist benchmark --config path/to/config.yaml

Arguments

Use built-in help for full command details:

gridfm_app_dist --help
gridfm_app_dist <command> --help

Running Tests

Unit and Integration Tests

Install the test dependencies first (if not already done):

pip install -e .[dev,test]

Run the full unit test suite:

pytest ./tests

Run the base set integration tests:

pytest ./integrationtests/test_base_set.py

Running Base Set Tests on an LSF Cluster (GPU)

To submit the base set integration tests as an interactive LSF job with GPU access, use bsub. Adjust the paths to match your environment:

bsub -gpu "num=1" \
     -n 16 \
     -R "rusage[mem=32GB] span[hosts=1]" \
     -Is \
     -J gridfm_base_set_tests \
     /bin/bash -c "
       cd /path/to/gridfm-app-distribution && \
       export PATH=/path/to/cuda/bin:\$PATH \
               CUDA_HOME=/path/to/cuda \
               LD_LIBRARY_PATH=/path/to/cuda/lib64:\$LD_LIBRARY_PATH && \
       source /path/to/venv/bin/activate && \
       pytest ./integrationtests/test_base_set.py
     "

Key bsub options used above:

Concrete example (adapt paths to your cluster setup):

bsub -gpu "num=1" -n 16 -R "rusage[mem=32GB] span[hosts=1]" -Is -J hpo_trial_190 /bin/bash -c "cd /dccstor/terratorch/users/rkie/gitco/gridfm-app-distribution && export PATH=/opt/share/cuda-12.8.1/bin:\$PATH CUDA_HOME=/opt/share/cuda-12.8.1 LD_LIBRARY_PATH=/opt/share/cuda-12.8.1/lib64:\$LD_LIBRARY_PATH && source /u/rkie/venvs/venv_gridfm-graphkit/bin/activate && pytest ./integrationtests/test_base_set.py"