LUMINA-2M

LUMINA-2M

Model Description

LUMINA-2M is a heterogeneous graph neural network (HGT) designed for fast, high-fidelity prediction of AC Optimal Power Flow (ACOPF) solutions. The model encodes power system components—buses, generators, loads, and transmission lines—as distinct node and edge types, capturing structural and physical heterogeneity in electric grids.

Trained on ACOPF datasets spanning multiple network topologies, LUMINA-2M generalizes across diverse grid configurations and operating conditions. It provides efficient approximations of OPF variables such as voltages, power injections, and line flows, serving as a scalable surrogate for traditional optimization solvers.

Model Details

• Version: v0.1.5-alpha
• Model Architecture: HGT
• Total Parameters: 3,755,642 (3.8M)
• Trainable Parameters: 3,755,642 (3.8M)
• Training Case: case14_ieee, case30_ieee, case57_ieee, case118_ieee, case500_goc, case2000_goc, case4661_sdet, case6470_rte, case10000_goc, case13659_pegase
• Training Data Size: 270,000 samples per case
• Training Date: 2026-09-23
• Training Epochs: 86
• Loss Type: Augmented Lagrangian
• Final Validation Loss: 0.072916

Note: The validation loss is only comparable across releases that share the same Loss Type. For example, an Augmented Lagrangian loss includes constraint-violation penalties and Lagrangian terms, so its magnitude is not directly comparable to a plain MSE loss.

Input Channels

• Bus: 7 features
• Generator: 11 features
• Load: 2 features
• Shunt: 2 features
• Grid: 4 features

Model Files

This repository contains the model in multiple formats:

• config.json - Model configuration and metadata
• model.pt - Complete PyTorch checkpoint with metadata and config
• model.safetensors - Model weights in SafeTensors format (recommended)
• requirements.txt - Required Python libraries
• pareto_front.png - Accuracy vs. constraint-violation plot across all released checkpoints (see Release History)

Installation

1. Clone or pull the latest lumina-inference repository:

   git clone [email protected]:argonne-gridfm/lumina-inference.git
   

2. Install the package in editable mode:

   cd lumina-inference
   pip install -e .
   

Basic Usage

Model Setup

import json

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from lumina_inference import Modeler
from lumina_inference.dataset import OPFDataset
from lumina_inference.loader import DataLoader

# Download model artifacts from Hugging Face
config_path = hf_hub_download(repo_id="argonne/LUMINA-2M", filename="config.json")

# Load config
with open(config_path, "r") as f:
    config_data = json.load(f)

# Set up device and modeler
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
modeler = Modeler(device)

Load Weights Option 1: Using SafeTensors (Recommended)

# Load weights from SafeTensors
safetensors_path = hf_hub_download(
    repo_id="argonne/LUMINA-2M", filename="model.safetensors"
)
state_dict = load_file(safetensors_path)
modeler.load_model(config_data, state_dict)

Load Weights Option 2: Using PyTorch Checkpoint

# Load weights from PyTorch Checkpoint on CPU
model_path = hf_hub_download(repo_id="argonne/LUMINA-2M", filename="model.pt")
checkpoint = torch.load(model_path, map_location="cpu")
state_dict = checkpoint.get("model_state_dict")
modeler.load_model(config_data, state_dict)

Load Data and Run Model (Batch)

# Simple argument defaults
case_name = config_data.get("case_name", "pglib_opf_case14_ieee")
batch_size = 1
max_batches = 50

# Loading OPF dataset
dataset = OPFDataset(root="./opf_data", case_name=case_name)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)

# Run model across data batches
preds = modeler.run_predictions(loader, max_batches=max_batches)

# Inspect first prediction
if preds:
    predictions_cpu, batch_cpu = preds[0]
    print(f"Prediction keys: {list(predictions_cpu.keys())}")
    for key, tensor in predictions_cpu.items():
        if isinstance(tensor, torch.Tensor):
            print(f"  {key}: shape={tensor.shape}, dtype={tensor.dtype}")

Single-Sample Prediction (No DataLoader)

from lumina_inference import load_from_json_file, load_from_dict

# From a JSON file
data = load_from_json_file("path/to/opf_sample.json")
result = modeler.predict_single(data)

# From a Python dictionary
data = load_from_dict(opf_dict)
result = modeler.predict_single(data)

Data Source

Google DeepMind. (2024). OPFData: Large-scale datasets for AC optimal power flow. Google Cloud Storage. Retrieved November 12, 2025, from storage.googleapis.com, under the MIT License. https://doi.org/10.48550/arXiv.2406.07234.

Release History

Citation

If you use LUMINA-2M in your research, please cite the LUMINA paper:

@inproceedings{li2026lumina,
  title={LUMINA: Foundation Models for Topology Transferable ACOPF},
  author={Li, Yijiang and Memon, Zeeshan and Jin, Hongwei and Fenu, Stefano and Song, Keunju and Sharma, Sunash B and Gasana, Parfait and Kim, Hongseok and Zhao, Liang and Kim, Kibaek},
  booktitle={ICLR 2026 Workshop on Foundation Models for Science: Real-World Impact and Science-First Design},
  doi={10.48550/arXiv.2603.04300},
  url={https://doi.org/10.48550/arXiv.2603.04300}
}

Acknowledgements

This material is based upon work supported by Laboratory Directed Research and Development (LDRD) funding from Argonne National Laboratory, provided by the Director, Office of Science, of the U.S. Department of Energy under contract DE-AC02-06CH11357.

An award of computer time was provided by the ASCR Leadership Computing Challenge (ALCC) program. This research used resources of the Argonne Leadership Computing Facility, which is a U.S. Department of Energy Office of Science User Facility operated under contract DE-AC02-06CH11357.

This research used resources of the National Energy Research Scientific Computing Center (NERSC), a Department of Energy User Facility using NERSC award ALCC-ERCAP0038201.

License

Released under the Apache License, Version 2.0. See the bundled LICENSE and NOTICE files for the full terms and attribution.

Copyright 2026 UChicago Argonne, LLC.