Power System Foundation Models

Foundation-style benchmarks and baselines for machine learning research on power systems challenges.

PFMs provides:

• Unified loaders that fetch and standardize popular open datasets into a common machine learning ready format.
• Canonical train/val/test splits spanning state-of-the-art challenges in optimization, control and forecasting problems.
• The first power system foundation model for solving these task.

Overview

1. Quick start
2. Prerequisites
3. Installation
4. Datasets
5. Running Experiments
6. Contributions
7. License

1. Quick start

Clone the repository

git clone https://github.com/ArsamAryandoust/PowerFoundationModels
cd PowerFoundationModels

(Option A) Use the stable ai4climate release (recommended)

pip install -U ai4climate

(Option B) Use the bleeding‑edge version bundled with this repo

pip uninstall -y ai4climate         # ensure the PyPI version is removed
export PYTHONPATH=$PYTHONPATH:$(pwd)/ai4climate

2. Prerequisites

Dependency	Version
Python	≥ 3.9
PyTorch	≥ 2.2
CUDA (opt)	11.8 / 12.X

Create a clean environment (example with conda):

conda create -n psfm python=3.10 pytorch=2.2 cudatoolkit=11.8 -c pytorch -y
conda activate psfm

3. Installation

Install additional project requirements

pip install -r requirements.txt

4. Datasets

All standardized tasks are hosted on Hugging Face Hub and can be downloaded automatically via ai4climate.load. Current coverage:

Task / Dataset	Modality	Docs
OPFData	Graphs (optimal-power-flow)	details
PowerGraph	Transmission-grid graphs	details
SolarCube	Satellite imagery	details
BuildingElectricity	Time-series load profiles	details
WindFarm	SCADA & weather	details

Load a dataset in three lines:

from ai4climate import load

ds = load(
    task_name="OPFData",
    subtask_name="train_small_test_medium",
    root_path="~/AI4Climate"  # local cache directory
)

5. Running Experiments

Experiments are configured via the config.yml file and executed through src/main.py.

python src/main.py

6. Contributions

We welcome pull requests for new datasets, models and evaluation metrics. Please read contributions.md before opening an issue or PR.

7. License

This project is licensed under the MIT License. See the license file for details.