fix: keep coords dimension order for DataArray bounds (#706)#710
fix: keep coords dimension order for DataArray bounds (#706)#710FBumann wants to merge 2 commits into
Conversation
…unds When add_variables expands a DataArray bound that is missing some of the dimensions in `coords`, `DataArray.expand_dims` prepends the new dimensions and their coordinate variables. Because a Dataset derives its dimension order from coordinate insertion order, the resulting variable's dimension order depended on the type of the bounds: a scalar bound kept the `coords` order, while two DataArray bounds missing the same dimension produced a different (prepended) order. Transpose the expanded array back to the `coords` order and reconstruct it with the coordinate variables in that order, so the variable's dimensions always follow `coords` regardless of the bounds' type. Closes #706. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
_broadcast_points collected the dimensions missing from the breakpoint array into a Python set, so the order in which expand_dims prepended them depended on hash randomization and varied between processes. Build the expansion map by iterating the expressions and their dimensions in order, giving a stable, reproducible dimension order. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
MaykThewessen
left a comment
There was a problem hiding this comment.
LGTM. Focused fix, full CI green, tests cover the three bound-type permutations that triggered #706 (scalar+da, da+scalar, da+da).
Two notes:
-
The
targetreconstruction block in_validate_dataarray_boundsguards againstexpand_dimsleaving non-dim coords in a different order than the transposed dims. Worth a one-line comment explaining why the secondDataArray(...)rebuild is needed — purely from reading the diff it looks like dead code after thetranspose(*expected)call, sincetransposealready reorders dims. The case it catches (auxiliaryarr.coordskeys outsideexpected) is non-obvious. -
The
piecewise._broadcast_pointsfix is the more impactful one for us — silent run-to-run reordering fromsetiteration is exactly the kind of thing that bites multi-process scenario sweeps. Glad it has a regression test.
Useful for downstream PyPSA-Eur consumers that pass DataArray bounds via add_variables in custom constraint modules. Tested mentally against our state_of_charge_set and PST tap-extendable code paths — both bypass this path (scalar bounds), so no breakage expected.
Summary
add_variables— when aDataArraybound is missing some of the dimensions incoords, it is expanded viaDataArray.expand_dims, which prepends the new dimensions and their coordinate variables. Since aDatasetderives its dimension order from coordinate insertion order, the resulting variable's dimension order depended on the type of the bounds (a scalar bound keptcoordsorder; twoDataArraybounds missing the same dim produced a prepended order). The expanded array is now transposed back tocoordsorder and reconstructed with its coordinate variables in that order, so the variable's dimensions always followcoords. Closes add_variables: DataArray bounds with missing dimensions yield wrong dimension order #706._broadcast_points(piecewise) — the dimensions missing from the breakpoint array were collected into a Pythonset, so the orderexpand_dimsprepended them depended on hash randomization and varied between processes. The expansion map is now built by iterating expressions and their dimensions in order, giving a stable, reproducible dimension order.The fix is a no-op for the common cases: scalar bounds and full-dimensional
DataArraybounds never reach the new code path. OnlyDataArraybounds with missing dimensions pay the (data-size-independent) reordering cost.Test plan
test/test_variable.py— incl. newtest_dataarray_broadcast_missing_dim_ordertest/test_piecewise_constraints.py— incl. newtest_broadcast_points_dim_order_follows_exprstest/test_piecewise_feasibility.py,test/test_model.pyruff check/ruff formatcleanmypyclean🤖 Generated with Claude Code