JAX-Based Batched AC Power Flow for GPU Acceleration and AI Ecosystem Integration

TL;DR

This paper introduces a JAX-based batched AC power flow solver that significantly accelerates computations and integrates seamlessly with AI tools, enhancing power system analysis and optimization.

Contribution

It presents a novel JAX implementation of power flow algorithms that achieves over 10x speed-up and improves customization and integration with AI ecosystems.

Findings

Achieves over 10x speed-up compared to pandapower and OpenDSS.

Supports Newton-Raphson for transmission and Z-Bus for distribution networks.

Facilitates embedding power flow evaluation within AI workflows.

Abstract

Coordinating growing grid flexibility under uncertainty is becoming increasingly important for efficient and reliable power-system operation. A core computational requirement is the efficient large-scale batched evaluation of AC power flow across candidate operating actions and uncertainty scenarios. Previous work has explored GPU-based batched power-flow evaluation, but has largely relied on hand-written C or CUDA code, creating barriers to customisation, efficient kernel optimisation, and long-term maintenance. JAX is a Python-based framework that enables efficient accelerator execution while keeping implementations in Python. This letter therefore proposes a JAX-based batched AC power-flow solver that uses current JAX functionality to implement Newton--Raphson for transmission networks and Z-Bus power flow for three-phase unbalanced distribution networks, achieving more than 10x speed-ups relative to pandapower and OpenDSS. In addition, JAX integrates seamlessly with the broader JAX-based AI ecosystem, making it straightforward to embed power-flow evaluation within AI methods for future larger-scale and more complex power-system operation.