EnerGNN

A Graph Neural Network library for real-life Energy networks.

---

EnerGNN is a python package based on JAX and Flax, that provides :

• A Hyper Heterogeneous Multi Graph (H2MG) data representation, especially designed for large complex industrial networks (such as an Electrical Power Transmission System);

• A compatible Graph Neural Network (GNN) library, robust to structure variations (outages, construction of new infrastructure, renaming / reordering, etc.);

• An clear interface to help users apply energnn to their own custom use-cases.

It is currently being used in multiple full-scale and real-life use-cases at Réseau de Transport d’Électricité (RTE). If you wish to either contribute or use it for your use cases, feel free to email us at balthazar.donon@rte-france.com.

---

Installation

To install the CPU version.

pip install energnn

Or to install the GPU version.

pip install energnn --extra gpu

---

Basic Usage

This package considers the training of GNN models to solve distributions of optimization problems (which encompasses traditional supervised learning).

from energnn.problem.example import LinearSystemProblemLoader
from energnn.model.ready_to_use import TinyRecurrentEquivariantGNN
from energnn.trainer import Trainer
import optax

problem_loader = LinearSystemProblemLoader(seed=1)
model = TinyRecurrentEquivariantGNN(
    in_structure=problem_loader.context_structure,
    out_structure=problem_loader.decision_structure,
)
trainer = Trainer(model=model, gradient_transformation=optax.adam(1e-3))
trainer.train(train_loader=problem_loader, n_epochs=10)

• The loader LinearSystemProblemLoader allows to iterate over multiple instances of the optimization problem class, which encapsulate the business logic (input, output, objective, etc.).

• The model processes graph data whose structure should match in_structure and returns graphs whose structure is defined as out_structure.

• The trainer iterates over the loader and updates the model weights to improve the model performance.

Once trained, the model can be applied on new problem instances as follows.

test_loader = LinearSystemProblemLoader(seed=3)
for problem_batch in test_loader:
    context_batch, _ = problem_batch.get_context()                  # Extract input
    decision_batch, _ = model.forward_batch(graph=context_batch)    # Infer decisions
    score, _ = problem_batch.get_score(decision=decision_batch)   # Compute score

---

User guides

• Basics
  • Amortized Optimization
  • Implementing your own Use Case
  • Data Representation using H2MG
  • Graph Neural Network Models
  • Trainer
  • Next Steps
• Tutorial
  • Installation
  • Problem Class
  • Problem Batch
  • Problem Loader
  • Graph Neural Network Model
  • Trainer
• Use Case Implementation
  • Overview
  • Step 0 — Define Graph Structures
  • Step 1 — Implement the Problem Interface
  • Step 2 — Handle Batching (ProblemBatch)
  • Step 3 — Data Loading (ProblemLoader)
  • Interface Checklist
  • Summary
  • Next steps

---

API Reference

For detailed description of energnn classes and methods, check out the API reference documentation.

• API reference
  • Graph
  • Model
  • Problem
  • Trainer

---

Supporting Institutions

---

Cite Us

@software{energnn,
  author = {{Committers of EnerGNN}},
  title = {{EnerGNN: A Graph Neural Network library for real-life Energy networks.}},
  url = {https://github.com/energnn},
}