Grid environments¶

Public API for TransGridEnv and shared transmission infrastructure (DC / AC power flow, OPF helpers). For physics, see Physics → Transmission. For radial distribution envs, see API → Distribution. For the SCUC env used by the TSO benchmark, see API → Unit commitment.

`TransGridEnv`¶

from powerzoojax.envs import TransGridEnv, make_trans_params

Effective modes¶

`physics`	`solver_mode`	Dispatch source	Network	Safety
`0`	`0`	agent	DC PF	MW line limits
`1`	`0`	agent	Newton-Raphson AC PF	apparent-power thermal + voltage bounds
`0`	`1`	DCOPF	DC PF	MW line limits
`1`	`1`	DCOPF	DCOPF + AC PF ex-post check	apparent-power thermal + voltage bounds
`0` or `1`	`2`	ACOPF	ACOPF	apparent-power thermal + voltage bounds

Step contract¶

reward = -reward_scale * gen_cost.
costs = [cost_thermal_overload, cost_voltage_violation, cost_power_balance, cost_resource].
info["cost_sum"] = sum(costs) is the aggregate diagnostic.
DC observation: [line_flow / cap, load / total_cap, unit_p / p_max, sin(t), cos(t), <bundle_obs>].
AC observation: [|S| / cap, vm, load / total_cap, unit_p / p_max, sin(t), cos(t), <bundle_obs>].

`powerzoojax.envs.grid.trans.TransGridEnv` ¶

Bases: Environment

Transmission grid environment — DC / AC Power Flow mode.

The agent dispatches generators every time step. The environment computes power flow (DC or AC), checks safety, and returns an observation of normalised flows, loads, and unit outputs.

`reset(key, params)` ¶

`step(key, state, action, params)` ¶

Run one environment step.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKey`	PRNG key.	required
`state`	`TransGridState`	Current grid state.	required
`action`	`Array`	When `include_unit_dispatch=True`: `(n_units + sum(bundle.action_dim),)` in `[-1, 1]`. When `include_unit_dispatch=False`: `(sum(bundle.action_dim),)` — bundle actions only.	required
`params`	`TransGridParams`	Environment parameters.	required

`observation_space(params)` ¶

`action_space(params)` ¶

`powerzoojax.envs.grid.trans.TransGridState` ¶

Bases: GridState

TransGridEnv state (extends GridState) — pure grid quantities only.

Resource state (battery SOC, etc.) is NOT stored here; it lives in the caller's scan carry. The grid env only sees external_injection_mw passed through step().

`powerzoojax.envs.grid.trans.TransGridParams` ¶

Bases: GridParams

TransGridEnv parameters (extends GridParams).

Attributes:

Name	Type	Description
`case`	`CaseData`	CaseData with PTDF, limits, cost coefficients, etc.
`reward_scale`	`float`	Scaling factor for generation cost in reward.
`acpf_setup`	`ACPFSetup`	Precomputed ACPF data (None for DC-only mode).
`dcopf_setup`	`DCOPFSetup`	Precomputed DCOPF data (None unless solver_mode=1).
`acopf_setup`	`ACOPFSetup`	Precomputed ACOPF data (None unless solver_mode=2).
`resources`	`tuple`	Tuple of ResourceBundle instances for grid attachment. Each bundle's `bus_idx` must be valid internal node indices. Action layout: `[unit_actions \| bundle_0_actions \| bundle_1_actions \| ...]`.

`powerzoojax.envs.grid.trans.make_trans_params(case, load_profiles=None, max_steps=48, delta_t_hours=0.5, steps_per_day=48, cost_thermal_weight=1.0, reward_scale=0.01, physics=0, solver_mode=0, resources=(), include_unit_dispatch=True)` ¶

Create TransGridParams from a CaseData and optional load profiles.

If load_profiles is None, a flat profile at (d_max + d_min) / 2 is used.

Parameters:

Name	Type	Description	Default
`case`	`CaseData`	CaseData instance.	required
`load_profiles`	`Array`	(T, n_loads) load demand array [MW].	`None`
`max_steps`	`int`	Maximum steps per episode.	`48`
`delta_t_hours`	`float`	Time step duration in hours.	`0.5`
`steps_per_day`	`int`	Steps per day.	`48`
`cost_thermal_weight`	`float`	Weight for thermal overload cost.	`1.0`
`reward_scale`	`float`	Multiplier for generation cost in reward.	`0.01`
`physics`	`int`	0=DC, 1=AC.	`0`
`solver_mode`	`int`	0=PF, 1=DCOPF, 2=ACOPF.	`0`
`resources`	`tuple`	Tuple of ResourceBundle instances (e.g. BatteryBundle). Use `make_battery_bundle(case, bus_ids=[...])` to create bundles.	`()`

Shared state containers¶

`powerzoojax.envs.grid.base.GridState` ¶

State for grid environments (immutable pytree).

Attributes:

Name	Type	Description
`time_step`	`Array`	Current simulation step (int32).
`done`	`Array`	Whether the episode has terminated.
`unit_power_mw`	`Array`	Active power of each unit [MW] (n_units,).
`line_flow_mw`	`Array`	Active power flow on each line [MW] (n_lines,).
`node_injection_mw`	`Array`	Net injection at each node [MW] (n_nodes,).
`is_safe`	`Array`	True if all line flows are within limits.
`n_violations`	`Array`	Number of line-limit violations.
`total_cost`	`Array`	Cumulative generation cost this episode [$].
`vm`	`Array`	Voltage magnitudes (n_nodes,) [p.u.] — AC mode only (ones in DC).
`va`	`Array`	Voltage angles (n_nodes,) [rad] — AC mode only (zeros in DC).
`q_gen`	`Array`	Reactive power output (n_units,) [MVAr] — AC mode only (zeros in DC).
`line_flow_q_mw`	`Array`	Reactive line flows (n_lines,) [MVAr] — AC mode only (zeros in DC).

`powerzoojax.envs.grid.base.GridParams` ¶

Parameters for grid environments (immutable pytree).

Attributes:

Name	Type	Description
`load_profiles`	`Array`	Load profiles (T, n_loads) [MW]. Row t gives each load's demand at time step t. Access should always use `state.time_step % T` to prevent out-of-bounds indexing when `max_steps > T`. RL note: the default modulo indexing produces a strictly deterministic cycle. Agents can overfit to this exact period. For generalization benchmarks, consider adding random start-offset sampling or Gaussian noise to the profiles before passing them in.
`max_steps`	`int`	Maximum steps per episode.
`delta_t_hours`	`float`	Time step duration in hours (default 0.5 = 30 min).
`steps_per_day`	`int`	Steps in one day (e.g. 48 for 30-min resolution).
`cost_thermal_weight`	`float`	Per-MW penalty weight for line overload in cost channel.
`physics`	`int`	0=DC, 1=AC (determines which solver to use).
`solver_mode`	`int`	0=PF, 1=OPF (determines action→dispatch mapping).

max_steps, delta_t_hours, steps_per_day, physics, and solver_mode are marked pytree_node=False so they become static under JIT. This prevents ConcretizationTypeError when they are used in control-flow decisions and ensures no wasted tracing.

DC power flow¶

`powerzoojax.envs.grid.power_flow.dc_power_flow(case, unit_power_mw, node_load_mw)` ¶

Compute DC power flow.

The slack bus absorbs the power imbalance. The returned actual_unit_power_mw reflects the true generation including slack adjustment — use it for cost computation and constraint checking so that RL agents cannot exploit "free" slack energy.

Parameters:

Name	Type	Description	Default
`case`	`CaseData`	CaseData with PTDF, nodes_units_map, slack_bus_idx.	required
`unit_power_mw`	`Array`	Generator outputs (n_units,) [MW].	required
`node_load_mw`	`Array`	Nodal load demand (n_nodes,) [MW].	required

Returns:

Name	Type	Description
`line_flow_mw`	`Array`	Line active power flows (n_lines,) [MW].
`node_injection_mw`	`Array`	Net nodal injection (n_nodes,) [MW]. Imbalance is absorbed by adjusting slack-bus load (load-side adjustment); numerically equivalent to adjusting slack-bus generation.
`actual_unit_power_mw`	`Array`	Generator outputs with slack adjustment (n_units,) [MW]. Slack-bus units are treated as an infinite-capacity reference; p_min/p_max are not clamped here. Enforce slack limits in the env or agent if needed.

`powerzoojax.envs.grid.power_flow.dc_power_flow_with_check(case, unit_power_mw, node_load_mw)` ¶

DC power flow followed by safety check.

Parameters:

Name	Type	Description	Default
`case`	`CaseData`	CaseData.	required
`unit_power_mw`	`Array`	Generator outputs (n_units,) [MW].	required
`node_load_mw`	`Array`	Nodal loads (n_nodes,) [MW].	required

Returns:

Name	Type	Description
`line_flow_mw`	`Array`	Line active power flows (n_lines,) [MW].
`node_injection_mw`	`Array`	Net nodal injection (n_nodes,) [MW].
`actual_unit_power_mw`	`Array`	Generator outputs with slack adjustment (n_units,) [MW]. Use this (not the raw `unit_power_mw`) for cost computation and constraint checking so that RL agents cannot exploit "free" slack energy.
`is_safe`	`Array`	True if all lines within limits (bool scalar).
`n_violations`	`Array`	Number of violated lines (int32 scalar).
`cost_thermal`	`Array`	Active-power overload [MW] (DC MW-limit semantics); see :func:`safety_check` for unit and AC-reuse caveats.

`powerzoojax.envs.grid.power_flow.safety_check(line_flow_mw, line_cap, line_floor)` ¶

Check line flow safety and compute thermal overload cost.

Uses active-power (MW) limit semantics for DC power flow. cost_thermal is in MW, not currency; when folding into reward, apply a penalty weight in the env layer (units: $/MW).

Parameters:

Name	Type	Description	Default
`line_flow_mw`	`Array`	Line flows (n_lines,) [MW].	required
`line_cap`	`Array`	Upper capacity limits (n_lines,) [MW].	required
`line_floor`	`Array`	Lower capacity limits (n_lines,) [MW].	required

Returns:

Name	Type	Description
`is_safe`	`Array`	True if all lines within limits (bool scalar).
`n_violations`	`Array`	Number of violated lines (int32 scalar).
`cost_thermal`	`Array`	Total active-power violation [MW] (sum of overrun above cap and underrun below floor); DC active-power limits, not AC MVA thermal limits.

Note

For AC power flow with MVA or current thermal limits, do not reuse this function; implement a separate check to avoid semantic confusion.

`powerzoojax.envs.grid.power_flow.compute_generation_cost(unit_power_mw, cost_a, cost_b, cost_c)` ¶

Compute total generation cost by integrating the marginal cost curve.

Marginal cost: MC(p) = cost_a * p^2 + cost_b * p + cost_c Total cost: TC(p) = (cost_a/3) * p^3 + (cost_b/2) * p^2 + cost_c * p

Parameters:

Name	Type	Description	Default
`unit_power_mw`	`Array`	Generator outputs (n_units,) [MW].	required
`cost_a`	`Array`	Quadratic coefficient in MC(p), i.e. coefficient of p^2 (n_units,). These are marginal-cost (MC) polynomial coefficients, not MATPOWER-style total-cost (TC) coefficients. MATPOWER uses TC = a·p² + b·p + c and MC = 2a·p + b (linear); here MC(p) = cost_a·p² + cost_b·p + cost_c (quadratic) and TC = (a/3)p³ + (b/2)p² + c·p. Convert when importing from MATPOWER.	required
`cost_b`	`Array`	Linear coefficient of MC(p) (n_units,).	required
`cost_c`	`Array`	Constant term of MC(p) [$/MWh], i.e. marginal cost at zero output (n_units,).	required

Returns:

Name	Type	Description
`total_cost`	`Array`	Scalar generation cost [$/h] (assuming a 1-hour dispatch interval; cost_c [$/MWh] × p [MW] = [$/h]).

`powerzoojax.envs.grid.power_flow.proportional_dispatch(total_load_mw, p_min, p_max)` ¶

Dispatch units proportionally to their capacity.

Each unit first produces p_min, then the remaining demand is allocated proportionally to each unit's headroom (p_max - p_min).

If total_load_mw exceeds the total generation capacity (sum(p_max)), all units are clipped to p_max — no error is raised, as this must stay JIT-safe.

Parameters:

Name	Type	Description	Default
`total_load_mw`	`Array`	Total system load [MW] (scalar).	required
`p_min`	`Array`	Minimum output per unit (n_units,) [MW].	required
`p_max`	`Array`	Maximum output per unit (n_units,) [MW].	required

Returns:

Name	Type	Description
`unit_power_mw`	`Array`	Dispatched power (n_units,) [MW], clipped to [p_min, p_max].

AC power flow¶

`powerzoojax.envs.grid.ac_power_flow.prepare_acpf(case)` ¶

Prepare all precomputed data needed for AC power flow.

Extracts bus classification, generator schedules, and builds Y_bus. Called once at env construction time.

Parameters:

Name	Type	Description	Default
`case`	`CaseData`	CaseData with AC fields populated.	required

Returns:

Name	Type	Description
`ACPFSetup`	`ACPFSetup`	frozen struct ready for JIT-compiled NR solver.

`powerzoojax.envs.grid.ac_power_flow.ac_power_flow(setup, max_iter=30, tol=1e-05)` ¶

Newton-Raphson AC power flow solver (pure JAX, JIT-compilable).

Uses lax.while_loop for the NR iteration.

Q-limit enforcement: PV buses whose calculated Q exceeds Qg_min/Qg_max (stored in setup) are switched to PQ mid-iteration (bus_is_pq mask). The Jacobian stays fixed at (2*n_pvpq, 2*n_pvpq) regardless of how many buses switch, satisfying the static-shape requirement of lax.while_loop. PV→PQ switching is one-way within a single solve (no PQ→PV recovery). If Qg_min_bus/Qg_max_bus are ±inf (default when Q data absent), switching never triggers and the solver behaves identically to before.

Note: Under jax.vmap, all batch elements run for max_iter iterations (already-converged elements still execute but do not update). Each iteration contains jnp.linalg.solve(J, mismatch) which is O(n³). For large networks with large batch sizes, consider lax.fori_loop with a fixed iteration count as a performance-tuning alternative.

Parameters:

Name	Type	Description	Default
`setup`	`ACPFSetup`	Precomputed ACPFSetup from prepare_acpf().	required
`max_iter`	`int`	Maximum NR iterations.	`30`
`tol`	`float`	Convergence tolerance on max \|mismatch\| in p.u. Default 1e-5 (float32 precision floor is ~2e-6).	`1e-05`

Returns:

Type	Description
`ACPFResult`	ACPFResult with voltage, power, branch flow, and convergence info.

`powerzoojax.envs.grid.ac_power_flow.ACPFSetup` ¶

Precomputed data for AC power flow (built once at env construction).

`powerzoojax.envs.grid.ac_power_flow.ACPFResult` ¶

Output of AC power flow solver.

OPF modules¶

DCOPF¶

`powerzoojax.envs.grid.dc_opf.prepare_dcopf(case, rho=1.0)` ¶

Prepare DCOPF data from CaseData.

`powerzoojax.envs.grid.dc_opf.dc_opf(setup, node_load_mw, commitment=None, max_iter=100, tol=0.001)` ¶

Solve DC-OPF via ADMM (Alternating Direction Method of Multipliers).

Primal objective: mc_c @ p (linear), matching PowerZoo HiGHS LP.

The ADMM consensus decomposition introduces auxiliary variables z1 (box projection) and z2 (line-flow projection) so that each update step is cheap and JAX-native:

p-update: exact KKT solve (n_units+1 linear system) — balance is enforced exactly at every iteration.
z1-update: element-wise clip to [p_min, p_max].
z2-update: element-wise clip to [line_floor, line_cap].

This handles the "degenerate LP optimal" case (e.g. congested line where the LP optimum requires two generators at non-bound values simultaneously) which merit-order / subgradient methods cannot represent exactly.

Parameters:

Name	Type	Description	Default
`setup`	`DCOPFSetup`	Precomputed DCOPFSetup.	required
`node_load_mw`	`Array`	(n_nodes,) load [MW].	required
`commitment`	`Array`	(n_units,) binary mask. None = all on.	`None`
`max_iter`	`int`	Max ADMM iterations (100 suffices for most cases; use 200 for heavily congested systems).	`100`
`tol`	`float`	Convergence tolerance for primal residuals [MW].	`0.001`

Returns:

Type	Description
`DCOPFResult`	DCOPFResult.

`powerzoojax.envs.grid.dc_opf.DCOPFSetup` ¶

Precomputed data for DCOPF.

Cost convention: mc_a, mc_b, mc_c are marginal cost coefficients. MC(p) = mc_a * p² + mc_b * p + mc_c [$/MWh] TC(p) = (mc_a/3)*p³ + (mc_b/2)*p² + mc_c*p [$]

The dispatch objective uses only mc_c (linear), matching the PowerZoo LP. mc_a and mc_b are retained for total-cost reporting.

A_aug_inv is the precomputed inverse of the (n_units+1) × (n_units+1) ADMM KKT matrix (constant across all ADMM iterations and all env steps). Storing the inverse replaces an O(n³) solve with an O(n²) matmul per ADMM iteration, giving a large speedup for vmap'd training rollouts.

`powerzoojax.envs.grid.dc_opf.DCOPFResult` ¶

Output of DC OPF solver.

ACOPF¶

`powerzoojax.envs.grid.ac_opf.ac_opf(setup, node_load_mw=None, node_load_mvar=None, commitment=None, max_iter=200, tol=0.001, warm_start_unit_power_mw=None, alm_outer_steps=12)` ¶

Solve AC-OPF via reduced-space L-BFGS-B with IFT-differentiable NR.

Parameters:

Name	Type	Description	Default
`setup`	`ACOPFSetup`	ACOPFSetup from `prepare_acopf`.	required
`node_load_mw`	`Array`	(n_bus,) active load [MW]. `None` → use setup.Pd.	`None`
`node_load_mvar`	`Array`	(n_bus,) reactive load [MVAr]. `None` → use setup.Qd.	`None`
`commitment`	`Array`	(n_gen,) binary. `None` → all units on.	`None`
`max_iter`	`int`	Max L-BFGS-B outer iterations (default 200; typically converges in 20–60 iterations for case5).	`200`
`tol`	`float`	L-BFGS-B convergence tolerance.	`0.001`
`warm_start_unit_power_mw`	`Array`	Optional (n_gen,) dispatch guess [MW] (e.g. DC-OPF result). Used to initialise Pg_ctrl.	`None`
`alm_outer_steps`	`int`	Number of augmented-Lagrangian outer iterations (`lax.scan` length). Default 12.	`12`

Returns:

Type	Description
`ACOPFResult`	ACOPFResult with identical field layout to the previous AL-based solver.

`powerzoojax.envs.grid.ac_opf.ACOPFSetup` ¶

Precomputed AC-OPF data including bus-type classification.

Grid environments¶

TransGridEnv¶

Effective modes¶

Step contract¶

powerzoojax.envs.grid.trans.TransGridEnv ¶

reset(key, params) ¶

step(key, state, action, params) ¶

observation_space(params) ¶

action_space(params) ¶

powerzoojax.envs.grid.trans.TransGridState ¶

powerzoojax.envs.grid.trans.TransGridParams ¶

powerzoojax.envs.grid.trans.make_trans_params(case, load_profiles=None, max_steps=48, delta_t_hours=0.5, steps_per_day=48, cost_thermal_weight=1.0, reward_scale=0.01, physics=0, solver_mode=0, resources=(), include_unit_dispatch=True) ¶

Shared state containers¶

powerzoojax.envs.grid.base.GridState ¶

powerzoojax.envs.grid.base.GridParams ¶

DC power flow¶

powerzoojax.envs.grid.power_flow.dc_power_flow(case, unit_power_mw, node_load_mw) ¶

powerzoojax.envs.grid.power_flow.dc_power_flow_with_check(case, unit_power_mw, node_load_mw) ¶

powerzoojax.envs.grid.power_flow.safety_check(line_flow_mw, line_cap, line_floor) ¶

powerzoojax.envs.grid.power_flow.compute_generation_cost(unit_power_mw, cost_a, cost_b, cost_c) ¶

powerzoojax.envs.grid.power_flow.proportional_dispatch(total_load_mw, p_min, p_max) ¶

AC power flow¶

powerzoojax.envs.grid.ac_power_flow.prepare_acpf(case) ¶

powerzoojax.envs.grid.ac_power_flow.ac_power_flow(setup, max_iter=30, tol=1e-05) ¶

powerzoojax.envs.grid.ac_power_flow.ACPFSetup ¶

powerzoojax.envs.grid.ac_power_flow.ACPFResult ¶

OPF modules¶

DCOPF¶

powerzoojax.envs.grid.dc_opf.prepare_dcopf(case, rho=1.0) ¶

powerzoojax.envs.grid.dc_opf.dc_opf(setup, node_load_mw, commitment=None, max_iter=100, tol=0.001) ¶

powerzoojax.envs.grid.dc_opf.DCOPFSetup ¶

powerzoojax.envs.grid.dc_opf.DCOPFResult ¶

ACOPF¶

powerzoojax.envs.grid.ac_opf.ac_opf(setup, node_load_mw=None, node_load_mvar=None, commitment=None, max_iter=200, tol=0.001, warm_start_unit_power_mw=None, alm_outer_steps=12) ¶

powerzoojax.envs.grid.ac_opf.ACOPFSetup ¶

`TransGridEnv`¶

`powerzoojax.envs.grid.trans.TransGridEnv` ¶

`reset(key, params)` ¶

`step(key, state, action, params)` ¶

`observation_space(params)` ¶

`action_space(params)` ¶

`powerzoojax.envs.grid.trans.TransGridState` ¶

`powerzoojax.envs.grid.trans.TransGridParams` ¶

`powerzoojax.envs.grid.trans.make_trans_params(case, load_profiles=None, max_steps=48, delta_t_hours=0.5, steps_per_day=48, cost_thermal_weight=1.0, reward_scale=0.01, physics=0, solver_mode=0, resources=(), include_unit_dispatch=True)` ¶

`powerzoojax.envs.grid.base.GridState` ¶

`powerzoojax.envs.grid.base.GridParams` ¶

`powerzoojax.envs.grid.power_flow.dc_power_flow(case, unit_power_mw, node_load_mw)` ¶

`powerzoojax.envs.grid.power_flow.dc_power_flow_with_check(case, unit_power_mw, node_load_mw)` ¶

`powerzoojax.envs.grid.power_flow.safety_check(line_flow_mw, line_cap, line_floor)` ¶

`powerzoojax.envs.grid.power_flow.compute_generation_cost(unit_power_mw, cost_a, cost_b, cost_c)` ¶

`powerzoojax.envs.grid.power_flow.proportional_dispatch(total_load_mw, p_min, p_max)` ¶

`powerzoojax.envs.grid.ac_power_flow.prepare_acpf(case)` ¶

`powerzoojax.envs.grid.ac_power_flow.ac_power_flow(setup, max_iter=30, tol=1e-05)` ¶

`powerzoojax.envs.grid.ac_power_flow.ACPFSetup` ¶

`powerzoojax.envs.grid.ac_power_flow.ACPFResult` ¶

`powerzoojax.envs.grid.dc_opf.prepare_dcopf(case, rho=1.0)` ¶

`powerzoojax.envs.grid.dc_opf.dc_opf(setup, node_load_mw, commitment=None, max_iter=100, tol=0.001)` ¶

`powerzoojax.envs.grid.dc_opf.DCOPFSetup` ¶

`powerzoojax.envs.grid.dc_opf.DCOPFResult` ¶

`powerzoojax.envs.grid.ac_opf.ac_opf(setup, node_load_mw=None, node_load_mvar=None, commitment=None, max_iter=200, tol=0.001, warm_start_unit_power_mw=None, alm_outer_steps=12)` ¶

`powerzoojax.envs.grid.ac_opf.ACOPFSetup` ¶