JAX-Based Batched AC Power Flow for GPU Acceleration and AI Ecosystem Integration
For power system operators, it provides a fast, customizable, and maintainable tool for large-scale batched power flow evaluation, crucial for uncertainty and flexibility management.
This work proposes a JAX-based batched AC power-flow solver that achieves over 10x speed-up compared to pandapower and OpenDSS for transmission and distribution networks, while enabling seamless integration with AI frameworks.
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.