gridfm/gridfm-pf-reconstruction

GridFM — PowerFlow Reconstruction

A heterogeneous graph neural network (GNS_heterogeneous) from gridfm-graphkit for the PowerFlow reconstruction task on power-grid graphs. Buses and generators are nodes; branches and bus↔generator connections are typed edges. Given a power-grid case, the model produces, per node:

• latent embeddings (bus and generator), and
• predictions: bus voltage magnitude/angle (Vm, Va) and generator active power (Pg), denormalized to physical units.

Model details

Architecture	GNS_heterogeneous (heterogeneous GNS), 12 layers, 8 attention heads, hidden size 12
Parameters	~1.3 M
Task	PowerFlow reconstruction
Input dims	bus 15, generator 6, edge 10
Output dims	bus 2 (Vm, Va), generator 1 (Pg)
Normalization	HeteroDataMVANormalizer (fit_on_train)

Training data

Trained on the IEEE case14 grid (case14_ieee) with data generated by gridfm-datakit (PowerModels.jl AC power-flow solves): 2,000 scenarios spanning load, topology, generation-cost and admittance perturbations; 20 epochs; losses LayeredWeightedPhysics (0.1) + MaskedBusMSE (0.9); AdamW, lr 5e-4.

Scope. This is a compact reference checkpoint trained on a single grid (case14) at modest scale. It is intended to demonstrate GridFM PowerFlow reconstruction and vLLM serving; it is not a large-grid production model.

Serving with vLLM

Install gridfm-graphkit with its vLLM extra (registers the GridFMGNS pooling model and the gridfm_pf_reconstruction I/O processor as vLLM entry points):

pip install "gridfm-graphkit[vllm]"

Serve on the /pooling endpoint:

vllm serve gridfm/gridfm-pf-reconstruction \
  --runner pooling \
  --trust-remote-code \
  --skip-tokenizer-init \
  --enforce-eager \
  --io-processor-plugin gridfm_pf_reconstruction \
  --enable-mm-embeds

Then POST a power-grid case (bus / gen / branch record lists) to http://localhost:8000/pooling; the response carries per-node embeddings and denormalized Vm/Va/Pg predictions. See the gridfm-graphkit usage guide for a full client example.

License

Apache-2.0.