Saved in:
| Main Authors: | , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2605.14103 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866916011841683456 |
|---|---|
| author | Zhou, Yihong Cope, Dylan Foerster, Jakob Morstyn, Thomas |
| author_facet | Zhou, Yihong Cope, Dylan Foerster, Jakob Morstyn, Thomas |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2605_14103 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | JAX-Based Batched AC Power Flow for GPU Acceleration and AI Ecosystem Integration Zhou, Yihong Cope, Dylan Foerster, Jakob Morstyn, Thomas Systems and Control 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. |
| title | JAX-Based Batched AC Power Flow for GPU Acceleration and AI Ecosystem Integration |
| topic | Systems and Control |
| url | https://arxiv.org/abs/2605.14103 |