refactor: unify as_dataarray; split broadcasting from coords validation#726
Open
FBumann wants to merge 1 commit into
Open
refactor: unify as_dataarray; split broadcasting from coords validation#726FBumann wants to merge 1 commit into
FBumann wants to merge 1 commit into
Conversation
Closes #723. Folds the body of `as_dataarray_in_coords` into `as_dataarray` and extracts the contract checks into `assert_compatible_with_coords`, so linopy now has one broadcasting primitive and one validation companion. `as_dataarray(arr, coords)` aligns the result against `coords` for every input type: labels positional inputs (numpy / unnamed pandas / scalar) by position, reindexes same-values-different-order, expands missing dims, and transposes to coords order. Extra dims and disagreeing value sets on shared dims pass through unchanged, so xarray broadcasting in expression arithmetic keeps working. `assert_compatible_with_coords(arr, coords)` enforces the strict contract (`arr.dims ⊆ coords.dims`, plus exact coord-value equality on shared dims). `add_variables` and `add_constraints` now call it after `as_dataarray` for `lower` / `upper` / `mask`, replacing the deleted `as_dataarray_in_coords` helper. `_coords_to_dict` filters MultiIndex level coords out of `xarray.Coordinates` inputs so the new strict-by-default path treats `station` (and not its derived `letter` / `num` levels) as the dim. Test suite: 3698 passed (no regressions). Two existing tests were updated to reflect the new "coords is source of truth" semantics: `test_as_dataarray_with_ndarray_coords_dict_set_dims_not_aligned` (extra coord entries now broadcast in) and `test_dataarray_extra_dims` (now triggers the subset check rather than the value-mismatch check). Microbenchmark in dev-scripts/benchmark_as_dataarray.py shows flat timings vs the base branch on both add_variables-heavy and arithmetic- heavy workloads. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #723. Stacked on #725 (which is stacked on #722).
What changes
as_dataarray_in_coordsis folded intoas_dataarrayand split along the seam between "broadcastarragainstcoords" and "enforce thecoordscontract":as_dataarray(arr, coords)— the broadcasting primitive. For every input type, the result is aligned withcoords: positional inputs (numpy / unnamed pandas / scalar) are labeled by position, shared-dim coords are reindexed when values are equal in a different order, dims present incoordsbut not inarrare expanded, and the result is transposed tocoordsorder. Extra dims and disagreeing value sets on shared dims pass through, so xarray broadcasting in expression arithmetic keeps working.assert_compatible_with_coords(arr, coords)— the validation companion. Raises ifarrintroduces dims not incoords(was:as_dataarray_in_coords's extra-dim raise) or if a shared dim has disagreeing coord values (was: its "do not match" raise).add_variablesandadd_constraintsnow callas_dataarrayfollowed byassert_compatible_with_coordsforlower/upper/mask. The previousas_dataarray_in_coordshelper is deleted._coords_to_dictnow filters MultiIndex level coords out ofxarray.Coordinatesinputs, so the new strict-by-default path treats e.g.station(and not its derivedletter/numlevels) as the dim. This was already a latent issue once strict semantics governedas_dataarray's coords arg.Audit summary (the 12 call sites listed in #723)
model.py:705/706lower/upperinadd_variablesas_dataarray+assert_compatible_with_coordsmodel.py:715maskinadd_variablesmodel.py:979maskinadd_constraintsexpressions.py:584/613/1105/1668/2154arithmeticvariables.py:330to_linexpr(coefficient)expressions.py:341/2002/2030/2289,model.py:912/919/972,variables.py:1369Breaking changes (relative to PR #725's strict semantics)
as_dataarrayno longer raises on shared-dim value-set mismatch. Disagreeing values are passed through for downstream xarray alignment to handle. Callers that want the old behavior should callassert_compatible_with_coordsafter the conversion (add_variables/add_constraintsalready do).as_dataarrayno longer raises on extra dims. Same migration: useassert_compatible_with_coords.Two existing tests updated to reflect the new "coords is source of truth, extras broadcast in" semantics:
test_as_dataarray_with_ndarray_coords_dict_set_dims_not_aligned: extra coord entries now expand into the result.test_dataarray_extra_dims: rewritten so the subset check fires (rather than the value-mismatch check from the old order).Test plan
test/test_common.py— added five tests pinning the new split (extra-dim preservation, disjoint shared-dim values,assert_compatible_with_coordsextra-dim raise, value-mismatch raise, subset-dims allowed).pre-commit run --all-filesclean (ruff/format/blackdoc/codespell).dev-scripts/benchmark_as_dataarray.py, untracked): flat timings vs the base branch on both add_variables-heavy (≈22 ms / 50–57 ms mean) and arithmetic-heavy (≈80–82 ms) workloads.🤖 Generated with Claude Code