Skip to content

fix: keep coords dimension order for DataArray bounds (#706)#710

Open
FBumann wants to merge 2 commits into
masterfrom
fix/dataarray-bounds-dim-order
Open

fix: keep coords dimension order for DataArray bounds (#706)#710
FBumann wants to merge 2 commits into
masterfrom
fix/dataarray-bounds-dim-order

Conversation

@FBumann
Copy link
Copy Markdown
Collaborator

@FBumann FBumann commented May 21, 2026

Summary

  • add_variables — when a DataArray bound is missing some of the dimensions in coords, it is expanded via DataArray.expand_dims, which prepends the new dimensions and their coordinate variables. Since 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 coords order; two DataArray bounds missing the same dim produced a prepended order). The expanded array is now transposed back to coords order and reconstructed with its coordinate variables in that order, so the variable's dimensions always follow coords. 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 Python set, so the order expand_dims prepended 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.
  • Added regression tests for both.

The fix is a no-op for the common cases: scalar bounds and full-dimensional DataArray bounds never reach the new code path. Only DataArray bounds with missing dimensions pay the (data-size-independent) reordering cost.

Test plan

  • test/test_variable.py — incl. new test_dataarray_broadcast_missing_dim_order
  • test/test_piecewise_constraints.py — incl. new test_broadcast_points_dim_order_follows_exprs
  • test/test_piecewise_feasibility.py, test/test_model.py
  • ruff check / ruff format clean
  • mypy clean

🤖 Generated with Claude Code

FBumann and others added 2 commits May 21, 2026 11:49
…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>
Copy link
Copy Markdown
Contributor

@MaykThewessen MaykThewessen left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Focused fix, full CI green, tests cover the three bound-type permutations that triggered #706 (scalar+da, da+scalar, da+da).

Two notes:

  1. The target reconstruction block in _validate_dataarray_bounds guards against expand_dims leaving non-dim coords in a different order than the transposed dims. Worth a one-line comment explaining why the second DataArray(...) rebuild is needed — purely from reading the diff it looks like dead code after the transpose(*expected) call, since transpose already reorders dims. The case it catches (auxiliary arr.coords keys outside expected) is non-obvious.

  2. The piecewise._broadcast_points fix is the more impactful one for us — silent run-to-run reordering from set iteration 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

add_variables: DataArray bounds with missing dimensions yield wrong dimension order

2 participants