From d3cadb1c9cf6698927fb92c5fee51445b770c670 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Tue, 20 Jan 2026 20:26:53 +0100 Subject: [PATCH 01/18] Fix as_dataarray to apply coords parameter for DataArray input Previously, when a DataArray was passed to as_dataarray(), the coords parameter was silently ignored. This was inconsistent with other input types (numpy, pandas) where coords are applied. Now, when coords is provided as a dict and the input is a DataArray, the function will reindex the array to match the provided coordinates. This ensures consistent behavior across all input types. Co-Authored-By: Claude Opus 4.5 --- linopy/common.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 09f67355..190371e4 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -9,7 +9,7 @@ import operator import os -from collections.abc import Callable, Generator, Hashable, Iterable, Sequence +from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence from functools import partial, reduce, wraps from pathlib import Path from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload @@ -245,7 +245,13 @@ def as_dataarray( arr = DataArray(float(arr), coords=coords, dims=dims, **kwargs) elif isinstance(arr, int | float | str | bool | list): arr = DataArray(arr, coords=coords, dims=dims, **kwargs) - + elif isinstance(arr, DataArray): + # Apply coords via reindex if provided as dict (for consistency with other input types) + if coords is not None and isinstance(coords, Mapping): + # Only reindex dimensions that exist in both arr and coords + reindex_coords = {k: v for k, v in coords.items() if k in arr.dims} + if reindex_coords: + arr = arr.reindex(reindex_coords) elif not isinstance(arr, DataArray): supported_types = [ np.number, From 4d32c3414121d03088657d017ce09fabc0093c39 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Tue, 20 Jan 2026 20:46:11 +0100 Subject: [PATCH 02/18] 1. Reindexes existing dims to match coord order 2. Expands to new dims from coords (broadcast) Summary: - as_dataarray now consistently applies coords for all input types - DataArrays with fewer dims are expanded to match the full coords specification - This fixes the inconsistency when creating variables with DataArray bounds --- linopy/common.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 190371e4..d560930c 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -246,12 +246,16 @@ def as_dataarray( elif isinstance(arr, int | float | str | bool | list): arr = DataArray(arr, coords=coords, dims=dims, **kwargs) elif isinstance(arr, DataArray): - # Apply coords via reindex if provided as dict (for consistency with other input types) + # Apply coords via reindex/expand if provided as dict (for consistency with other input types) if coords is not None and isinstance(coords, Mapping): - # Only reindex dimensions that exist in both arr and coords + # Reindex dimensions that exist in both arr and coords reindex_coords = {k: v for k, v in coords.items() if k in arr.dims} if reindex_coords: arr = arr.reindex(reindex_coords) + # Expand to new dimensions from coords + expand_coords = {k: v for k, v in coords.items() if k not in arr.dims} + if expand_coords: + arr = arr.expand_dims(expand_coords) elif not isinstance(arr, DataArray): supported_types = [ np.number, From ce7c7e2b7dff5186dccbde87696b7086e42819d5 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 07:43:07 +0100 Subject: [PATCH 03/18] Raise on coord mismatch instead of silent reindex in as_dataarray Replace reindex with a strict equality check for DataArray inputs. Silent reindexing is dangerous as it introduces NaNs for missing indices and drops unmatched ones, masking user bugs. Now raises ValueError if coords don't match, while still allowing expand_dims for broadcasting to new dimensions. --- linopy/common.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index d560930c..66d4cd0c 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -246,13 +246,17 @@ def as_dataarray( elif isinstance(arr, int | float | str | bool | list): arr = DataArray(arr, coords=coords, dims=dims, **kwargs) elif isinstance(arr, DataArray): - # Apply coords via reindex/expand if provided as dict (for consistency with other input types) if coords is not None and isinstance(coords, Mapping): - # Reindex dimensions that exist in both arr and coords - reindex_coords = {k: v for k, v in coords.items() if k in arr.dims} - if reindex_coords: - arr = arr.reindex(reindex_coords) - # Expand to new dimensions from coords + for k, v in coords.items(): + if k in arr.dims: + expected = pd.Index(v) + actual = pd.Index(arr.coords[k].values) + if not actual.equals(expected): + raise ValueError( + f"Coordinates for dimension '{k}' do not match: " + f"expected {expected.tolist()}, got {actual.tolist()}" + ) + # Expand to new dimensions from coords (broadcast) expand_coords = {k: v for k, v in coords.items() if k not in arr.dims} if expand_coords: arr = arr.expand_dims(expand_coords) From fa8605d757ada868e8f2f97c13f0dac24007003f Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 07:49:19 +0100 Subject: [PATCH 04/18] Add allow_extra_dims flag to as_dataarray Strict by default: raises ValueError if a DataArray has dimensions not present in coords. Call sites that need broadcasting (multiply, dot, add) opt in with allow_extra_dims=True. Structural call sites like add_variables bounds/mask remain strict. --- linopy/common.py | 11 +++++++++++ linopy/expressions.py | 12 +++++++++--- linopy/variables.py | 4 +++- 3 files changed, 23 insertions(+), 4 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 66d4cd0c..6bb483f4 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -214,6 +214,7 @@ def as_dataarray( arr: Any, coords: CoordsLike | None = None, dims: DimsLike | None = None, + allow_extra_dims: bool = False, **kwargs: Any, ) -> DataArray: """ @@ -227,6 +228,10 @@ def as_dataarray( The coordinates for the DataArray. If None, default coordinates will be used. dims (Union[list, None]): The dimensions for the DataArray. If None, the dimensions will be automatically generated. + allow_extra_dims: + If False (default), raise ValueError when a DataArray input has + dimensions not present in coords. Set to True to allow extra + dimensions for broadcasting. **kwargs: Additional keyword arguments to be passed to the DataArray constructor. @@ -247,6 +252,12 @@ def as_dataarray( arr = DataArray(arr, coords=coords, dims=dims, **kwargs) elif isinstance(arr, DataArray): if coords is not None and isinstance(coords, Mapping): + if not allow_extra_dims: + extra_dims = set(arr.dims) - set(coords) + if extra_dims: + raise ValueError( + f"DataArray has extra dimensions not in coords: {extra_dims}" + ) for k, v in coords.items(): if k in arr.dims: expected = pd.Index(v) diff --git a/linopy/expressions.py b/linopy/expressions.py index d2ae9022..7307fd96 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -597,7 +597,9 @@ def _add_constant( # so that missing data does not silently propagate through arithmetic. if np.isscalar(other) and join is None: return self.assign(const=self.const.fillna(0) + other) - da = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + da = as_dataarray( + other, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True + ) self_const, da, needs_data_reindex = self._align_constant( da, fill_value=0, join=join ) @@ -626,7 +628,9 @@ def _apply_constant_op( - factor (other) is filled with fill_value (0 for mul, 1 for div) - coeffs and const are filled with 0 (additive identity) """ - factor = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + factor = as_dataarray( + other, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True + ) self_const, factor, needs_data_reindex = self._align_constant( factor, fill_value=fill_value, join=join ) @@ -1705,7 +1709,9 @@ def __matmul__( Matrix multiplication with other, similar to xarray dot. """ if not isinstance(other, LinearExpression | variables.Variable): - other = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + other = as_dataarray( + other, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True + ) common_dims = list(set(self.coord_dims).intersection(other.dims)) return (self * other).sum(dim=common_dims) diff --git a/linopy/variables.py b/linopy/variables.py index f99fb938..fd6eef2e 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -315,7 +315,9 @@ def to_linexpr( linopy.LinearExpression Linear expression with the variables and coefficients. """ - coefficient = as_dataarray(coefficient, coords=self.coords, dims=self.dims) + coefficient = as_dataarray( + coefficient, coords=self.coords, dims=self.dims, allow_extra_dims=True + ) coefficient = coefficient.reindex_like(self.labels, fill_value=0) coefficient = coefficient.fillna(0) ds = Dataset({"coeffs": coefficient, "vars": self.labels}).expand_dims( From 6be142faa61ec472acd46200368820f8217a5a02 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 07:59:52 +0100 Subject: [PATCH 05/18] Normalize sequence coords to dict for DataArray validation When coords is a sequence (e.g. from add_variables), convert it to a dict using dims or Index names so the same validation applies. This closes the gap where sequence coords were silently ignored for DataArray inputs. Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/linopy/common.py b/linopy/common.py index 6bb483f4..1526e299 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -251,6 +251,14 @@ def as_dataarray( elif isinstance(arr, int | float | str | bool | list): arr = DataArray(arr, coords=coords, dims=dims, **kwargs) elif isinstance(arr, DataArray): + if coords is not None and not isinstance(coords, Mapping): + # Normalize sequence coords to a dict + seq = list(coords) + if dims is not None: + dim_names = [dims] if isinstance(dims, str) else list(dims) + coords = dict(zip(dim_names, seq)) + else: + coords = {c.name: c for c in seq if hasattr(c, "name") and c.name} if coords is not None and isinstance(coords, Mapping): if not allow_extra_dims: extra_dims = set(arr.dims) - set(coords) From b6c1758b015d28fe01f56fefa51f1286fea9fca8 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 08:02:35 +0100 Subject: [PATCH 06/18] Extract _coords_to_mapping helper for coords normalization Unifies the sequence-to-dict coords conversion used in pandas_to_dataarray, numpy_to_dataarray, and the DataArray branch of as_dataarray into a single helper. Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 48 +++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 1526e299..b735dcc6 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -128,6 +128,22 @@ def get_from_iterable(lst: DimsLike | None, index: int) -> Any | None: return lst[index] if 0 <= index < len(lst) else None +def _coords_to_mapping(coords: CoordsLike, dims: DimsLike | None = None) -> Mapping: + """ + Normalize coords to a Mapping. + + If coords is already a Mapping, return as-is. If it's a sequence, + convert to a dict using dims (if provided) or Index names. + """ + if isinstance(coords, Mapping): + return coords + seq = list(coords) + if dims is not None: + dim_names = [dims] if isinstance(dims, str) else list(dims) + return dict(zip(dim_names, seq)) + return {c.name: c for c in seq if hasattr(c, "name") and c.name} + + def pandas_to_dataarray( arr: pd.DataFrame | pd.Series, coords: CoordsLike | None = None, @@ -161,6 +177,26 @@ def pandas_to_dataarray( axis.name or get_from_iterable(dims, i) or f"dim_{i}" for i, axis in enumerate(arr.axes) ] + if coords is not None: + pandas_coords = dict(zip(dims, arr.axes)) + coords = _coords_to_mapping(coords, dims) + shared_dims = set(pandas_coords.keys()) & set(coords.keys()) + non_aligned = [] + for dim in shared_dims: + coord = coords[dim] + if not isinstance(coord, pd.Index): + coord = pd.Index(coord) + if not pandas_coords[dim].equals(coord): + non_aligned.append(dim) + if any(non_aligned): + warn( + f"coords for dimension(s) {non_aligned} is not aligned with the pandas object. " + "Previously, the indexes of the pandas were ignored and overwritten in " + "these cases. Now, the pandas object's coordinates are taken considered" + " for alignment." + ) + + return DataArray(arr, coords=None, dims=dims, **kwargs) @@ -203,7 +239,7 @@ def numpy_to_dataarray( if dims is not None and len(dims) and coords is not None: if isinstance(coords, list): - coords = dict(zip(dims, coords[: arr.ndim])) + coords = _coords_to_mapping(coords, dims) elif is_dict_like(coords): coords = {k: v for k, v in coords.items() if k in dims} @@ -251,14 +287,8 @@ def as_dataarray( elif isinstance(arr, int | float | str | bool | list): arr = DataArray(arr, coords=coords, dims=dims, **kwargs) elif isinstance(arr, DataArray): - if coords is not None and not isinstance(coords, Mapping): - # Normalize sequence coords to a dict - seq = list(coords) - if dims is not None: - dim_names = [dims] if isinstance(dims, str) else list(dims) - coords = dict(zip(dim_names, seq)) - else: - coords = {c.name: c for c in seq if hasattr(c, "name") and c.name} + if coords is not None: + coords = _coords_to_mapping(coords, dims) if coords is not None and isinstance(coords, Mapping): if not allow_extra_dims: extra_dims = set(arr.dims) - set(coords) From fa16ba2f1d08234e2777e5339e25729483a16b72 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 08:10:51 +0100 Subject: [PATCH 07/18] Update docstrings for as_dataarray and add_variables Document the coord validation and broadcasting behavior from the user perspective. --- linopy/common.py | 10 ++++++++++ linopy/model.py | 14 +++++++++----- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index b735dcc6..6ed495a5 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -275,6 +275,16 @@ def as_dataarray( ------- DataArray: The converted DataArray. + + Raises + ------ + ValueError + If arr is a DataArray and coords is provided: raised when + coordinates for shared dimensions do not match, or when the + DataArray has dimensions not covered by coords (unless + allow_extra_dims is True). The DataArray's dimensions may be + a subset of coords — missing dimensions are added via + expand_dims. """ if isinstance(arr, pd.Series | pd.DataFrame): arr = pandas_to_dataarray(arr, coords=coords, dims=dims, **kwargs) diff --git a/linopy/model.py b/linopy/model.py index f1d7e5ef..a58c20c6 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -489,10 +489,12 @@ def add_variables( Upper bound of the variable(s). Ignored if `binary` is True. The default is inf. coords : list/xarray.Coordinates, optional - The coords of the variable array. - These are directly passed to the DataArray creation of - `lower` and `upper`. For every single combination of - coordinates a optimization variable is added to the model. + The coords of the variable array. For every single + combination of coordinates an optimization variable is + added to the model. Data for `lower`, `upper` and `mask` + is fitted to these coords: shared dimensions must have + matching coordinates, and missing dimensions are broadcast. + A ValueError is raised if the data is not compatible. The default is None. name : str, optional Reference name of the added variables. The default None results in @@ -514,7 +516,9 @@ def add_variables( ------ ValueError If neither lower bound and upper bound have coordinates, nor - `coords` are directly given. + `coords` are directly given. Also raised if `lower` or + `upper` are DataArrays whose coordinates do not match the + provided `coords`. Returns ------- From e46a19edf2de1b276b4b21f36e5a47858df17725 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 08:27:09 +0100 Subject: [PATCH 08/18] Extract dataarray path into method --- linopy/common.py | 65 ++++++++++++++++++++++++++++++++---------------- 1 file changed, 43 insertions(+), 22 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 6ed495a5..96fd93af 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -144,6 +144,48 @@ def _coords_to_mapping(coords: CoordsLike, dims: DimsLike | None = None) -> Mapp return {c.name: c for c in seq if hasattr(c, "name") and c.name} +def _validate_dataarray_coords( + arr: DataArray, + coords: CoordsLike | None = None, + dims: DimsLike | None = None, + allow_extra_dims: bool = False, +) -> DataArray: + """ + Validate a DataArray's coordinates against expected coords. + + For shared dimensions, coordinates must match exactly. Extra + dimensions in the DataArray raise unless allow_extra_dims is True. + Missing dimensions are broadcast via expand_dims. + """ + if coords is None: + return arr + + expected = _coords_to_mapping(coords, dims) + if not expected: + return arr + + if not allow_extra_dims: + extra = set(arr.dims) - set(expected) + if extra: + raise ValueError(f"DataArray has extra dimensions not in coords: {extra}") + + for k, v in expected.items(): + if k in arr.dims: + expected_idx = pd.Index(v) + actual_idx = pd.Index(arr.coords[k].values) + if not actual_idx.equals(expected_idx): + raise ValueError( + f"Coordinates for dimension '{k}' do not match: " + f"expected {expected_idx.tolist()}, got {actual_idx.tolist()}" + ) + + expand = {k: v for k, v in expected.items() if k not in arr.dims} + if expand: + arr = arr.expand_dims(expand) + + return arr + + def pandas_to_dataarray( arr: pd.DataFrame | pd.Series, coords: CoordsLike | None = None, @@ -297,28 +339,7 @@ def as_dataarray( elif isinstance(arr, int | float | str | bool | list): arr = DataArray(arr, coords=coords, dims=dims, **kwargs) elif isinstance(arr, DataArray): - if coords is not None: - coords = _coords_to_mapping(coords, dims) - if coords is not None and isinstance(coords, Mapping): - if not allow_extra_dims: - extra_dims = set(arr.dims) - set(coords) - if extra_dims: - raise ValueError( - f"DataArray has extra dimensions not in coords: {extra_dims}" - ) - for k, v in coords.items(): - if k in arr.dims: - expected = pd.Index(v) - actual = pd.Index(arr.coords[k].values) - if not actual.equals(expected): - raise ValueError( - f"Coordinates for dimension '{k}' do not match: " - f"expected {expected.tolist()}, got {actual.tolist()}" - ) - # Expand to new dimensions from coords (broadcast) - expand_coords = {k: v for k, v in coords.items() if k not in arr.dims} - if expand_coords: - arr = arr.expand_dims(expand_coords) + arr = _validate_dataarray_coords(arr, coords, dims, allow_extra_dims) elif not isinstance(arr, DataArray): supported_types = [ np.number, From 2905fd444c1f18457d4d0a1e02335f99a6fac2ef Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 08:29:49 +0100 Subject: [PATCH 09/18] Improve performance --- linopy/common.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 96fd93af..b73addcd 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -169,17 +169,19 @@ def _validate_dataarray_coords( if extra: raise ValueError(f"DataArray has extra dimensions not in coords: {extra}") + expand = {} for k, v in expected.items(): - if k in arr.dims: - expected_idx = pd.Index(v) - actual_idx = pd.Index(arr.coords[k].values) - if not actual_idx.equals(expected_idx): - raise ValueError( - f"Coordinates for dimension '{k}' do not match: " - f"expected {expected_idx.tolist()}, got {actual_idx.tolist()}" - ) + if k not in arr.dims: + expand[k] = v + continue + actual = arr.coords[k] + v_idx = v if isinstance(v, pd.Index) else pd.Index(v) + if not actual.to_index().equals(v_idx): + raise ValueError( + f"Coordinates for dimension '{k}' do not match: " + f"expected {v_idx.tolist()}, got {actual.values.tolist()}" + ) - expand = {k: v for k, v in expected.items() if k not in arr.dims} if expand: arr = arr.expand_dims(expand) @@ -339,7 +341,8 @@ def as_dataarray( elif isinstance(arr, int | float | str | bool | list): arr = DataArray(arr, coords=coords, dims=dims, **kwargs) elif isinstance(arr, DataArray): - arr = _validate_dataarray_coords(arr, coords, dims, allow_extra_dims) + if coords is not None: + arr = _validate_dataarray_coords(arr, coords, dims, allow_extra_dims) elif not isinstance(arr, DataArray): supported_types = [ np.number, From fea8f9b1c49dde14a038742e82a34a919f52594d Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 09:23:02 +0100 Subject: [PATCH 10/18] ensure mask broadcasting still works as expected, even with misaligned coords --- linopy/model.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/linopy/model.py b/linopy/model.py index a58c20c6..11b954fe 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -574,7 +574,10 @@ def add_variables( self._check_valid_dim_names(data) if mask is not None: - mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool) + # TODO: Simplify when removing Future Warning from broadcast_mask + if not isinstance(mask, DataArray): + mask = as_dataarray(mask, coords=data.coords, dims=data.dims) + mask = mask.astype(bool) mask = broadcast_mask(mask, data.labels) # Auto-mask based on NaN in bounds (use numpy for speed) @@ -799,7 +802,10 @@ def add_constraints( (data,) = xr.broadcast(data, exclude=[TERM_DIM]) if mask is not None: - mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool) + # TODO: Simplify when removing Future Warning from broadcast_mask + if not isinstance(mask, DataArray): + mask = as_dataarray(mask, coords=data.coords, dims=data.dims) + mask = mask.astype(bool) mask = broadcast_mask(mask, data.labels) # Auto-mask based on null expressions or NaN RHS (use numpy for speed) From 65bcfb2a702eed32bda8e4d1ac1dfe23edb958bb Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 09:28:59 +0100 Subject: [PATCH 11/18] fix types --- linopy/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/linopy/common.py b/linopy/common.py index b73addcd..76ee08a5 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -139,7 +139,7 @@ def _coords_to_mapping(coords: CoordsLike, dims: DimsLike | None = None) -> Mapp return coords seq = list(coords) if dims is not None: - dim_names = [dims] if isinstance(dims, str) else list(dims) + dim_names: list[str] = [dims] if isinstance(dims, str) else list(dims) # type: ignore[arg-type] return dict(zip(dim_names, seq)) return {c.name: c for c in seq if hasattr(c, "name") and c.name} From be936638cb13e11826f0fefb271b20d23b9a4101 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 19 Feb 2026 09:57:11 +0100 Subject: [PATCH 12/18] Update Changelog and add some tests to ensure behaviour --- doc/release_notes.rst | 2 +- test/test_common.py | 80 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 1 deletion(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 0697e8a2..d6173086 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -10,6 +10,7 @@ Upcoming Version - Comparison operators (``==``, ``<=``, ``>=``) fill missing RHS coords with NaN (no constraint created) - Fixes crash on ``subset + var`` / ``subset + expr`` reverse addition - Fixes superset DataArrays expanding result coords beyond the variable's coordinate space +* When passing DataArray bounds to ``add_variables`` with explicit ``coords``, the ``coords`` parameter now defines the variable's coordinates. DataArray bounds are validated against these coords (raises ``ValueError`` on mismatch) and broadcast to missing dimensions. Previously, the ``coords`` parameter was silently ignored for DataArray inputs. * Add ``add_piecewise_constraints()`` with SOS2, incremental, LP, and disjunctive formulations (``linopy.piecewise(x, x_pts, y_pts) == y``). * Add ``linopy.piecewise()`` to create piecewise linear function descriptors (`PiecewiseExpression`) from separate x/y breakpoint arrays. * Add ``linopy.breakpoints()`` factory for convenient breakpoint construction from lists, Series, DataFrames, DataArrays, or dicts. Supports slopes mode. @@ -25,7 +26,6 @@ Version 0.6.5 * Expose the knitro context to allow for more flexible use of the knitro python API. - Version 0.6.4 -------------- diff --git a/test/test_common.py b/test/test_common.py index f1190024..fa4869be 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -409,6 +409,86 @@ def test_as_dataarray_with_dataarray_default_dims_coords() -> None: assert list(da_out.coords["dim2"].values) == list(da_in.coords["dim2"].values) +def test_as_dataarray_with_dataarray_coord_mismatch() -> None: + da_in = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) + with pytest.raises(ValueError, match="do not match"): + as_dataarray(da_in, coords={"x": [10, 20, 40]}) + + +def test_as_dataarray_with_dataarray_extra_dims() -> None: + da_in = DataArray([[1, 2], [3, 4]], dims=["x", "y"]) + with pytest.raises(ValueError, match="extra dimensions"): + as_dataarray(da_in, coords={"x": [0, 1]}) + + +def test_as_dataarray_with_dataarray_extra_dims_allowed() -> None: + da_in = DataArray( + [[1, 2], [3, 4]], + dims=["x", "y"], + coords={"x": [0, 1], "y": [0, 1]}, + ) + da_out = as_dataarray(da_in, coords={"x": [0, 1]}, allow_extra_dims=True) + assert da_out.dims == da_in.dims + assert da_out.shape == da_in.shape + + +def test_as_dataarray_with_dataarray_broadcast() -> None: + da_in = DataArray([1, 2], dims=["x"], coords={"x": ["a", "b"]}) + da_out = as_dataarray( + da_in, coords={"x": ["a", "b"], "y": [1, 2, 3]}, dims=["x", "y"] + ) + assert set(da_out.dims) == {"x", "y"} + assert da_out.sizes["y"] == 3 + + +def test_as_dataarray_with_dataarray_no_coords() -> None: + da_in = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) + da_out = as_dataarray(da_in) + assert_equal(da_out, da_in) + + +def test_as_dataarray_with_dataarray_sequence_coords() -> None: + da_in = DataArray([1, 2], dims=["x"], coords={"x": [0, 1]}) + idx = pd.RangeIndex(2, name="x") + da_out = as_dataarray(da_in, coords=[idx], dims=["x"]) + assert list(da_out.coords["x"].values) == [0, 1] + + +def test_as_dataarray_with_dataarray_sequence_coords_mismatch() -> None: + da_in = DataArray([1, 2], dims=["x"], coords={"x": [0, 1]}) + idx = pd.RangeIndex(3, name="x") + with pytest.raises(ValueError, match="do not match"): + as_dataarray(da_in, coords=[idx], dims=["x"]) + + +def test_add_variables_with_dataarray_bounds_and_coords() -> None: + model = Model() + time = pd.RangeIndex(5, name="time") + lower = DataArray([0, 0, 0, 0, 0], dims=["time"], coords={"time": range(5)}) + var = model.add_variables(lower=lower, coords=[time], name="x") + assert var.shape == (5,) + assert list(var.data.coords["time"].values) == list(range(5)) + + +def test_add_variables_with_dataarray_bounds_coord_mismatch() -> None: + model = Model() + time = pd.RangeIndex(5, name="time") + lower = DataArray([0, 0, 0], dims=["time"], coords={"time": [0, 1, 2]}) + with pytest.raises(ValueError, match="do not match"): + model.add_variables(lower=lower, coords=[time], name="x") + + +def test_add_variables_with_dataarray_bounds_broadcast() -> None: + model = Model() + time = pd.RangeIndex(3, name="time") + space = pd.Index(["a", "b"], name="space") + lower = DataArray([0, 0, 0], dims=["time"], coords={"time": range(3)}) + var = model.add_variables(lower=lower, coords=[time, space], name="x") + assert set(var.data.dims) == {"time", "space"} + assert var.data.sizes["time"] == 3 + assert var.data.sizes["space"] == 2 + + def test_as_dataarray_with_unsupported_type() -> None: with pytest.raises(TypeError): as_dataarray(lambda x: 1, dims=["dim1"], coords=[["a"]]) From 9508117820f07b8e7a4990ac76dbf75314d1bb17 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Wed, 11 Mar 2026 13:50:05 +0100 Subject: [PATCH 13/18] Fix DataArray validation for upstream compatibility - Skip coord validation for DataArray inputs in arithmetic contexts (allow_extra_dims=True) to preserve xarray's native alignment - Add allow_extra_dims=True to comparison operator and quadratic dot as_dataarray calls for consistent broadcasting - Handle MultiIndex levels in expand_dims guard Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 19 +++++++++++-------- linopy/expressions.py | 8 ++++++-- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 76ee08a5..778255d6 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -172,15 +172,19 @@ def _validate_dataarray_coords( expand = {} for k, v in expected.items(): if k not in arr.dims: + # Skip coords that already exist (e.g. MultiIndex levels) + if k in arr.coords: + continue expand[k] = v continue - actual = arr.coords[k] - v_idx = v if isinstance(v, pd.Index) else pd.Index(v) - if not actual.to_index().equals(v_idx): - raise ValueError( - f"Coordinates for dimension '{k}' do not match: " - f"expected {v_idx.tolist()}, got {actual.values.tolist()}" - ) + if not allow_extra_dims: + actual = arr.coords[k] + v_idx = v if isinstance(v, pd.Index) else pd.Index(v) + if not actual.to_index().equals(v_idx): + raise ValueError( + f"Coordinates for dimension '{k}' do not match: " + f"expected {v_idx.tolist()}, got {actual.values.tolist()}" + ) if expand: arr = arr.expand_dims(expand) @@ -240,7 +244,6 @@ def pandas_to_dataarray( " for alignment." ) - return DataArray(arr, coords=None, dims=dims, **kwargs) diff --git a/linopy/expressions.py b/linopy/expressions.py index 7307fd96..7865e6fe 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1146,7 +1146,9 @@ def to_constraint( ) if isinstance(rhs, SUPPORTED_CONSTANT_TYPES): - rhs = as_dataarray(rhs, coords=self.coords, dims=self.coord_dims) + rhs = as_dataarray( + rhs, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True + ) extra_dims = set(rhs.dims) - set(self.coord_dims) if extra_dims: @@ -2197,7 +2199,9 @@ def __matmul__( "Higher order non-linear expressions are not yet supported." ) - other = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + other = as_dataarray( + other, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True + ) common_dims = list(set(self.coord_dims).intersection(other.dims)) return (self * other).sum(dim=common_dims) From 83ad30b8076fc7bb0184060c59f0320d4f19c6b2 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Wed, 11 Mar 2026 14:08:23 +0100 Subject: [PATCH 14/18] Simplify: validate DataArray coords only in add_variables MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move validation out of as_dataarray into model.add_variables directly. This removes the allow_extra_dims flag and all changes to expressions.py and variables.py — arithmetic call sites are unaffected. Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 54 +++++++++----------------------- linopy/expressions.py | 20 +++--------- linopy/model.py | 17 +++++----- linopy/variables.py | 4 +-- test/test_common.py | 73 ++++++++++++++++++++++++------------------- 5 files changed, 70 insertions(+), 98 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 778255d6..62c1e6ef 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -146,28 +146,23 @@ def _coords_to_mapping(coords: CoordsLike, dims: DimsLike | None = None) -> Mapp def _validate_dataarray_coords( arr: DataArray, - coords: CoordsLike | None = None, + coords: CoordsLike, dims: DimsLike | None = None, - allow_extra_dims: bool = False, ) -> DataArray: """ - Validate a DataArray's coordinates against expected coords. + Validate and broadcast a DataArray against expected coords. - For shared dimensions, coordinates must match exactly. Extra - dimensions in the DataArray raise unless allow_extra_dims is True. - Missing dimensions are broadcast via expand_dims. + - Shared dimensions must have matching coordinates (raises ValueError). + - Extra dimensions in the DataArray raise ValueError. + - Missing dimensions are broadcast via expand_dims. """ - if coords is None: - return arr - expected = _coords_to_mapping(coords, dims) if not expected: return arr - if not allow_extra_dims: - extra = set(arr.dims) - set(expected) - if extra: - raise ValueError(f"DataArray has extra dimensions not in coords: {extra}") + extra = set(arr.dims) - set(expected) + if extra: + raise ValueError(f"DataArray has extra dimensions not in coords: {extra}") expand = {} for k, v in expected.items(): @@ -177,14 +172,13 @@ def _validate_dataarray_coords( continue expand[k] = v continue - if not allow_extra_dims: - actual = arr.coords[k] - v_idx = v if isinstance(v, pd.Index) else pd.Index(v) - if not actual.to_index().equals(v_idx): - raise ValueError( - f"Coordinates for dimension '{k}' do not match: " - f"expected {v_idx.tolist()}, got {actual.values.tolist()}" - ) + actual = arr.coords[k] + v_idx = v if isinstance(v, pd.Index) else pd.Index(v) + if not actual.to_index().equals(v_idx): + raise ValueError( + f"Coordinates for dimension '{k}' do not match: " + f"expected {v_idx.tolist()}, got {actual.values.tolist()}" + ) if expand: arr = arr.expand_dims(expand) @@ -297,7 +291,6 @@ def as_dataarray( arr: Any, coords: CoordsLike | None = None, dims: DimsLike | None = None, - allow_extra_dims: bool = False, **kwargs: Any, ) -> DataArray: """ @@ -311,10 +304,6 @@ def as_dataarray( The coordinates for the DataArray. If None, default coordinates will be used. dims (Union[list, None]): The dimensions for the DataArray. If None, the dimensions will be automatically generated. - allow_extra_dims: - If False (default), raise ValueError when a DataArray input has - dimensions not present in coords. Set to True to allow extra - dimensions for broadcasting. **kwargs: Additional keyword arguments to be passed to the DataArray constructor. @@ -322,16 +311,6 @@ def as_dataarray( ------- DataArray: The converted DataArray. - - Raises - ------ - ValueError - If arr is a DataArray and coords is provided: raised when - coordinates for shared dimensions do not match, or when the - DataArray has dimensions not covered by coords (unless - allow_extra_dims is True). The DataArray's dimensions may be - a subset of coords — missing dimensions are added via - expand_dims. """ if isinstance(arr, pd.Series | pd.DataFrame): arr = pandas_to_dataarray(arr, coords=coords, dims=dims, **kwargs) @@ -343,9 +322,6 @@ def as_dataarray( arr = DataArray(float(arr), coords=coords, dims=dims, **kwargs) elif isinstance(arr, int | float | str | bool | list): arr = DataArray(arr, coords=coords, dims=dims, **kwargs) - elif isinstance(arr, DataArray): - if coords is not None: - arr = _validate_dataarray_coords(arr, coords, dims, allow_extra_dims) elif not isinstance(arr, DataArray): supported_types = [ np.number, diff --git a/linopy/expressions.py b/linopy/expressions.py index 7865e6fe..d2ae9022 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -597,9 +597,7 @@ def _add_constant( # so that missing data does not silently propagate through arithmetic. if np.isscalar(other) and join is None: return self.assign(const=self.const.fillna(0) + other) - da = as_dataarray( - other, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True - ) + da = as_dataarray(other, coords=self.coords, dims=self.coord_dims) self_const, da, needs_data_reindex = self._align_constant( da, fill_value=0, join=join ) @@ -628,9 +626,7 @@ def _apply_constant_op( - factor (other) is filled with fill_value (0 for mul, 1 for div) - coeffs and const are filled with 0 (additive identity) """ - factor = as_dataarray( - other, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True - ) + factor = as_dataarray(other, coords=self.coords, dims=self.coord_dims) self_const, factor, needs_data_reindex = self._align_constant( factor, fill_value=fill_value, join=join ) @@ -1146,9 +1142,7 @@ def to_constraint( ) if isinstance(rhs, SUPPORTED_CONSTANT_TYPES): - rhs = as_dataarray( - rhs, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True - ) + rhs = as_dataarray(rhs, coords=self.coords, dims=self.coord_dims) extra_dims = set(rhs.dims) - set(self.coord_dims) if extra_dims: @@ -1711,9 +1705,7 @@ def __matmul__( Matrix multiplication with other, similar to xarray dot. """ if not isinstance(other, LinearExpression | variables.Variable): - other = as_dataarray( - other, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True - ) + other = as_dataarray(other, coords=self.coords, dims=self.coord_dims) common_dims = list(set(self.coord_dims).intersection(other.dims)) return (self * other).sum(dim=common_dims) @@ -2199,9 +2191,7 @@ def __matmul__( "Higher order non-linear expressions are not yet supported." ) - other = as_dataarray( - other, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True - ) + other = as_dataarray(other, coords=self.coords, dims=self.coord_dims) common_dims = list(set(self.coord_dims).intersection(other.dims)) return (self * other).sum(dim=common_dims) diff --git a/linopy/model.py b/linopy/model.py index 11b954fe..83c676f4 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -26,6 +26,7 @@ from linopy import solvers from linopy.common import ( + _validate_dataarray_coords, as_dataarray, assign_multiindex_safe, best_int, @@ -562,6 +563,12 @@ def add_variables( else: lower, upper = 0, 1 + if coords is not None: + if isinstance(lower, DataArray): + lower = _validate_dataarray_coords(lower, coords) + if isinstance(upper, DataArray): + upper = _validate_dataarray_coords(upper, coords) + data = Dataset( { "lower": as_dataarray(lower, coords, **kwargs), @@ -574,10 +581,7 @@ def add_variables( self._check_valid_dim_names(data) if mask is not None: - # TODO: Simplify when removing Future Warning from broadcast_mask - if not isinstance(mask, DataArray): - mask = as_dataarray(mask, coords=data.coords, dims=data.dims) - mask = mask.astype(bool) + mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool) mask = broadcast_mask(mask, data.labels) # Auto-mask based on NaN in bounds (use numpy for speed) @@ -802,10 +806,7 @@ def add_constraints( (data,) = xr.broadcast(data, exclude=[TERM_DIM]) if mask is not None: - # TODO: Simplify when removing Future Warning from broadcast_mask - if not isinstance(mask, DataArray): - mask = as_dataarray(mask, coords=data.coords, dims=data.dims) - mask = mask.astype(bool) + mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool) mask = broadcast_mask(mask, data.labels) # Auto-mask based on null expressions or NaN RHS (use numpy for speed) diff --git a/linopy/variables.py b/linopy/variables.py index fd6eef2e..f99fb938 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -315,9 +315,7 @@ def to_linexpr( linopy.LinearExpression Linear expression with the variables and coefficients. """ - coefficient = as_dataarray( - coefficient, coords=self.coords, dims=self.dims, allow_extra_dims=True - ) + coefficient = as_dataarray(coefficient, coords=self.coords, dims=self.dims) coefficient = coefficient.reindex_like(self.labels, fill_value=0) coefficient = coefficient.fillna(0) ds = Dataset({"coeffs": coefficient, "vars": self.labels}).expand_dims( diff --git a/test/test_common.py b/test/test_common.py index fa4869be..13aa3e77 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -409,56 +409,63 @@ def test_as_dataarray_with_dataarray_default_dims_coords() -> None: assert list(da_out.coords["dim2"].values) == list(da_in.coords["dim2"].values) -def test_as_dataarray_with_dataarray_coord_mismatch() -> None: +def test_as_dataarray_with_dataarray_no_coords() -> None: da_in = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) + da_out = as_dataarray(da_in) + assert_equal(da_out, da_in) + + +def test_validate_dataarray_coords_match() -> None: + from linopy.common import _validate_dataarray_coords + + da = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) + out = _validate_dataarray_coords(da, coords={"x": [10, 20, 30]}) + assert_equal(out, da) + + +def test_validate_dataarray_coords_mismatch() -> None: + from linopy.common import _validate_dataarray_coords + + da = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) with pytest.raises(ValueError, match="do not match"): - as_dataarray(da_in, coords={"x": [10, 20, 40]}) + _validate_dataarray_coords(da, coords={"x": [10, 20, 40]}) -def test_as_dataarray_with_dataarray_extra_dims() -> None: - da_in = DataArray([[1, 2], [3, 4]], dims=["x", "y"]) - with pytest.raises(ValueError, match="extra dimensions"): - as_dataarray(da_in, coords={"x": [0, 1]}) +def test_validate_dataarray_coords_extra_dims() -> None: + from linopy.common import _validate_dataarray_coords + da = DataArray([[1, 2], [3, 4]], dims=["x", "y"]) + with pytest.raises(ValueError, match="extra dimensions"): + _validate_dataarray_coords(da, coords={"x": [0, 1]}) -def test_as_dataarray_with_dataarray_extra_dims_allowed() -> None: - da_in = DataArray( - [[1, 2], [3, 4]], - dims=["x", "y"], - coords={"x": [0, 1], "y": [0, 1]}, - ) - da_out = as_dataarray(da_in, coords={"x": [0, 1]}, allow_extra_dims=True) - assert da_out.dims == da_in.dims - assert da_out.shape == da_in.shape +def test_validate_dataarray_coords_broadcast() -> None: + from linopy.common import _validate_dataarray_coords -def test_as_dataarray_with_dataarray_broadcast() -> None: - da_in = DataArray([1, 2], dims=["x"], coords={"x": ["a", "b"]}) - da_out = as_dataarray( - da_in, coords={"x": ["a", "b"], "y": [1, 2, 3]}, dims=["x", "y"] + da = DataArray([1, 2], dims=["x"], coords={"x": ["a", "b"]}) + out = _validate_dataarray_coords( + da, coords={"x": ["a", "b"], "y": [1, 2, 3]}, dims=["x", "y"] ) - assert set(da_out.dims) == {"x", "y"} - assert da_out.sizes["y"] == 3 - + assert set(out.dims) == {"x", "y"} + assert out.sizes["y"] == 3 -def test_as_dataarray_with_dataarray_no_coords() -> None: - da_in = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) - da_out = as_dataarray(da_in) - assert_equal(da_out, da_in) +def test_validate_dataarray_coords_sequence() -> None: + from linopy.common import _validate_dataarray_coords -def test_as_dataarray_with_dataarray_sequence_coords() -> None: - da_in = DataArray([1, 2], dims=["x"], coords={"x": [0, 1]}) + da = DataArray([1, 2], dims=["x"], coords={"x": [0, 1]}) idx = pd.RangeIndex(2, name="x") - da_out = as_dataarray(da_in, coords=[idx], dims=["x"]) - assert list(da_out.coords["x"].values) == [0, 1] + out = _validate_dataarray_coords(da, coords=[idx], dims=["x"]) + assert list(out.coords["x"].values) == [0, 1] + +def test_validate_dataarray_coords_sequence_mismatch() -> None: + from linopy.common import _validate_dataarray_coords -def test_as_dataarray_with_dataarray_sequence_coords_mismatch() -> None: - da_in = DataArray([1, 2], dims=["x"], coords={"x": [0, 1]}) + da = DataArray([1, 2], dims=["x"], coords={"x": [0, 1]}) idx = pd.RangeIndex(3, name="x") with pytest.raises(ValueError, match="do not match"): - as_dataarray(da_in, coords=[idx], dims=["x"]) + _validate_dataarray_coords(da, coords=[idx], dims=["x"]) def test_add_variables_with_dataarray_bounds_and_coords() -> None: From 07868392626ab30c9619cf17b39ad66927ae0d49 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 12 Mar 2026 10:33:44 +0100 Subject: [PATCH 15/18] Fix typo in warning message and add TODO for mask validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix "coordinates are taken considered" → "are considered" in pandas_to_dataarray warning - Add TODO noting mask DataArray validation is intentionally skipped to preserve broadcast_mask's fill-with-False behavior - Clarify coords docstring for add_variables Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 2 +- linopy/model.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/linopy/common.py b/linopy/common.py index 62c1e6ef..b2e526f3 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -234,7 +234,7 @@ def pandas_to_dataarray( warn( f"coords for dimension(s) {non_aligned} is not aligned with the pandas object. " "Previously, the indexes of the pandas were ignored and overwritten in " - "these cases. Now, the pandas object's coordinates are taken considered" + "these cases. Now, the pandas object's coordinates are considered" " for alignment." ) diff --git a/linopy/model.py b/linopy/model.py index c7a8fdaf..d32925e2 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -528,6 +528,8 @@ def add_variables( is fitted to these coords: shared dimensions must have matching coordinates, and missing dimensions are broadcast. A ValueError is raised if the data is not compatible. + When ``coords`` is ``None``, coordinates are inferred from + the bounds (``lower``/``upper``) if they are DataArrays. The default is None. name : str, optional Reference name of the added variables. The default None results in @@ -613,6 +615,9 @@ def add_variables( lower = _validate_dataarray_coords(lower, coords) if isinstance(upper, DataArray): upper = _validate_dataarray_coords(upper, coords) + # TODO: validate mask DataArray coords too, but this would + # change behavior — broadcast_mask currently fills mismatched + # coords with False silently. Needs a deprecation path. data = Dataset( { From b99dc11111da9ea8c7b8beac21ab9e3dc2b8e746 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 12 Mar 2026 19:13:28 +0100 Subject: [PATCH 16/18] Redesign as_dataarray: clean public API + internal _coerce_to_dataarray helper Split as_dataarray into two functions with distinct responsibilities: - as_dataarray (public): strict coord validation when coords provided, rejects extra dims, expands missing dims, raises on coord mismatch - _coerce_to_dataarray (internal): pure type conversion using coords only as construction hints, no validation or expansion This removes the allow_extra_dims flag and makes the API predictable: callers that need lenient conversion (expression arithmetic, masks) use _coerce_to_dataarray, while add_variables bounds use as_dataarray. Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 212 ++++++++++++++++++++++++++++++------------ linopy/expressions.py | 13 ++- linopy/model.py | 19 ++-- linopy/variables.py | 6 +- test/test_common.py | 68 +++++++++++--- 5 files changed, 226 insertions(+), 92 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index b2e526f3..d6e624ba 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -148,29 +148,39 @@ def _validate_dataarray_coords( arr: DataArray, coords: CoordsLike, dims: DimsLike | None = None, -) -> DataArray: +) -> None: """ - Validate and broadcast a DataArray against expected coords. + Validate a DataArray against expected coords (strict). - Shared dimensions must have matching coordinates (raises ValueError). - - Extra dimensions in the DataArray raise ValueError. - - Missing dimensions are broadcast via expand_dims. + - Extra dimensions in the DataArray always raise ValueError. + + This function only validates — it does not modify ``arr``. + Expansion of missing dims should be done before calling this. """ expected = _coords_to_mapping(coords, dims) if not expected: - return arr + return - extra = set(arr.dims) - set(expected) + # Filter to dimension coordinates only — skip non-dimension coords + # like MultiIndex levels (level1, level2) that share a dimension + # with their parent index. + if hasattr(coords, "dims"): + # xarray Coordinates object — only keep entries that are dimensions + dim_coords = {k: v for k, v in expected.items() if k in coords.dims} + else: + dim_coords = dict(expected) + + extra = set(arr.dims) - set(dim_coords) if extra: raise ValueError(f"DataArray has extra dimensions not in coords: {extra}") - expand = {} - for k, v in expected.items(): + for k, v in dim_coords.items(): if k not in arr.dims: - # Skip coords that already exist (e.g. MultiIndex levels) - if k in arr.coords: - continue - expand[k] = v + continue + # Skip validation for multiindex dimensions — the level coords + # cannot be compared directly via pd.Index.equals + if isinstance(arr.indexes.get(k), pd.MultiIndex): continue actual = arr.coords[k] v_idx = v if isinstance(v, pd.Index) else pd.Index(v) @@ -180,11 +190,6 @@ def _validate_dataarray_coords( f"expected {v_idx.tolist()}, got {actual.values.tolist()}" ) - if expand: - arr = arr.expand_dims(expand) - - return arr - def pandas_to_dataarray( arr: pd.DataFrame | pd.Series, @@ -219,25 +224,6 @@ def pandas_to_dataarray( axis.name or get_from_iterable(dims, i) or f"dim_{i}" for i, axis in enumerate(arr.axes) ] - if coords is not None: - pandas_coords = dict(zip(dims, arr.axes)) - coords = _coords_to_mapping(coords, dims) - shared_dims = set(pandas_coords.keys()) & set(coords.keys()) - non_aligned = [] - for dim in shared_dims: - coord = coords[dim] - if not isinstance(coord, pd.Index): - coord = pd.Index(coord) - if not pandas_coords[dim].equals(coord): - non_aligned.append(dim) - if any(non_aligned): - warn( - f"coords for dimension(s) {non_aligned} is not aligned with the pandas object. " - "Previously, the indexes of the pandas were ignored and overwritten in " - "these cases. Now, the pandas object's coordinates are considered" - " for alignment." - ) - return DataArray(arr, coords=None, dims=dims, **kwargs) @@ -287,30 +273,17 @@ def numpy_to_dataarray( return DataArray(arr, coords=coords, dims=dims, **kwargs) -def as_dataarray( +def _type_dispatch( arr: Any, coords: CoordsLike | None = None, dims: DimsLike | None = None, **kwargs: Any, ) -> DataArray: """ - Convert an object to a DataArray. + Convert arr to a DataArray via type dispatch. - Parameters - ---------- - arr: - The input object. - coords (Union[dict, list, None]): - The coordinates for the DataArray. If None, default coordinates will be used. - dims (Union[list, None]): - The dimensions for the DataArray. If None, the dimensions will be automatically generated. - **kwargs: - Additional keyword arguments to be passed to the DataArray constructor. - - Returns - ------- - DataArray: - The converted DataArray. + This is the shared conversion logic used by both ``as_dataarray`` + and ``_coerce_to_dataarray``. It does NOT validate or expand dims. """ if isinstance(arr, pd.Series | pd.DataFrame): arr = pandas_to_dataarray(arr, coords=coords, dims=dims, **kwargs) @@ -318,10 +291,21 @@ def as_dataarray( arr = numpy_to_dataarray(arr, coords=coords, dims=dims, **kwargs) elif isinstance(arr, pl.Series): arr = numpy_to_dataarray(arr.to_numpy(), coords=coords, dims=dims, **kwargs) - elif isinstance(arr, np.number): - arr = DataArray(float(arr), coords=coords, dims=dims, **kwargs) - elif isinstance(arr, int | float | str | bool | list): - arr = DataArray(arr, coords=coords, dims=dims, **kwargs) + elif isinstance(arr, np.number | int | float | str | bool | list): + if isinstance(arr, np.number): + arr = float(arr) + # For scalars with coords but no dims, infer dims from coords + # to avoid xarray's CoordinateValidationError + if coords is not None and dims is None and np.ndim(arr) == 0: + inferred = _coords_to_mapping(coords) + if inferred: + arr = DataArray( + arr, coords=coords, dims=list(inferred.keys()), **kwargs + ) + else: + arr = DataArray(arr, coords=coords, dims=dims, **kwargs) + else: + arr = DataArray(arr, coords=coords, dims=dims, **kwargs) elif not isinstance(arr, DataArray): supported_types = [ np.number, @@ -339,10 +323,122 @@ def as_dataarray( f"Unsupported type of arr: {type(arr)}. Supported types are: {supported_types_str}" ) - arr = fill_missing_coords(arr) + return fill_missing_coords(arr) + + +def _expand_missing_dims( + arr: DataArray, + coords: CoordsLike, + dims: DimsLike | None = None, +) -> DataArray: + """Broadcast missing dims via expand_dims.""" + expected = _coords_to_mapping(coords, dims) + if not expected: + return arr + + # Filter to dimension coordinates only — skip non-dimension coords + # like MultiIndex levels (level1, level2) that share a dimension + # with their parent index. + if hasattr(coords, "dims"): + dim_coords = {k: v for k, v in expected.items() if k in coords.dims} + else: + dim_coords = dict(expected) + + expand = {k: v for k, v in dim_coords.items() if k not in arr.dims} + if expand: + arr = arr.expand_dims(expand) return arr +def _coerce_to_dataarray( + arr: Any, + coords: CoordsLike, + dims: DimsLike | None = None, + **kwargs: Any, +) -> DataArray: + """ + Internal: convert any type to DataArray using coords for construction. + + Only does type conversion — no coord validation, no expand_dims. + Callers handle alignment and broadcasting themselves. + + Parameters + ---------- + arr: + The input object. + coords: + Coordinates used as construction hints for types without + their own coords (numpy, scalar, list). + dims: + Dimension names. + **kwargs: + Additional keyword arguments passed to the DataArray constructor. + + Returns + ------- + DataArray + """ + return _type_dispatch(arr, coords=coords, dims=dims, **kwargs) + + +def as_dataarray( + arr: Any, + coords: CoordsLike | None = None, + dims: DimsLike | None = None, + **kwargs: Any, +) -> DataArray: + """ + Convert an object to a DataArray with strict coord validation. + + Behavior depends on whether ``coords`` is provided: + + - **No coords**: just type conversion (scalar->DA, numpy->DA, + pandas->DA, DataArray->passthrough). + - **With coords** (strict): shared dims must have matching + coordinates (ValueError if not), extra dims in input are + rejected (ValueError), missing dims are broadcast via + ``expand_dims``. + + Parameters + ---------- + arr: + The input object. + coords: + The coordinates for the DataArray. If None, only type + conversion is performed. + dims: + The dimensions for the DataArray. If None, the dimensions + will be automatically generated. + **kwargs: + Additional keyword arguments passed to the DataArray constructor. + + Returns + ------- + DataArray: + The converted DataArray. + """ + if coords is not None: + # Inputs that already have their own coordinates (DataArray, + # pandas) or scalars need validation against coords. Raw arrays + # (numpy, list) get coords applied during construction, so + # validation and expand_dims are not needed. + needs_validation = isinstance(arr, DataArray | pd.Series | pd.DataFrame) or ( + not isinstance(arr, np.ndarray | pl.Series | list) and np.ndim(arr) == 0 + ) + + if needs_validation: + arr = _type_dispatch(arr, coords=coords, dims=dims, **kwargs) + arr = _expand_missing_dims(arr, coords, dims) + _validate_dataarray_coords(arr, coords, dims) + else: + arr = _type_dispatch(arr, coords=coords, dims=dims, **kwargs) + + return arr + + # No coords — just type conversion + return _type_dispatch(arr, coords=None, dims=dims, **kwargs) + + def broadcast_mask(mask: DataArray, labels: DataArray) -> DataArray: """ Broadcast a boolean mask to match the shape of labels. diff --git a/linopy/expressions.py b/linopy/expressions.py index d2ae9022..542ac685 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -46,6 +46,7 @@ from linopy.common import ( EmptyDeprecationWrapper, LocIndexer, + _coerce_to_dataarray, as_dataarray, assign_multiindex_safe, check_common_keys_values, @@ -597,7 +598,7 @@ def _add_constant( # so that missing data does not silently propagate through arithmetic. if np.isscalar(other) and join is None: return self.assign(const=self.const.fillna(0) + other) - da = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + da = _coerce_to_dataarray(other, coords=self.coords, dims=self.coord_dims) self_const, da, needs_data_reindex = self._align_constant( da, fill_value=0, join=join ) @@ -626,7 +627,7 @@ def _apply_constant_op( - factor (other) is filled with fill_value (0 for mul, 1 for div) - coeffs and const are filled with 0 (additive identity) """ - factor = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + factor = _coerce_to_dataarray(other, coords=self.coords, dims=self.coord_dims) self_const, factor, needs_data_reindex = self._align_constant( factor, fill_value=fill_value, join=join ) @@ -1142,7 +1143,7 @@ def to_constraint( ) if isinstance(rhs, SUPPORTED_CONSTANT_TYPES): - rhs = as_dataarray(rhs, coords=self.coords, dims=self.coord_dims) + rhs = _coerce_to_dataarray(rhs, coords=self.coords, dims=self.coord_dims) extra_dims = set(rhs.dims) - set(self.coord_dims) if extra_dims: @@ -1705,7 +1706,9 @@ def __matmul__( Matrix multiplication with other, similar to xarray dot. """ if not isinstance(other, LinearExpression | variables.Variable): - other = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + other = _coerce_to_dataarray( + other, coords=self.coords, dims=self.coord_dims + ) common_dims = list(set(self.coord_dims).intersection(other.dims)) return (self * other).sum(dim=common_dims) @@ -2191,7 +2194,7 @@ def __matmul__( "Higher order non-linear expressions are not yet supported." ) - other = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + other = _coerce_to_dataarray(other, coords=self.coords, dims=self.coord_dims) common_dims = list(set(self.coord_dims).intersection(other.dims)) return (self * other).sum(dim=common_dims) diff --git a/linopy/model.py b/linopy/model.py index d32925e2..b260d587 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -26,7 +26,7 @@ from linopy import solvers from linopy.common import ( - _validate_dataarray_coords, + _coerce_to_dataarray, as_dataarray, assign_multiindex_safe, best_int, @@ -610,15 +610,6 @@ def add_variables( "Semi-continuous variables require a positive scalar lower bound." ) - if coords is not None: - if isinstance(lower, DataArray): - lower = _validate_dataarray_coords(lower, coords) - if isinstance(upper, DataArray): - upper = _validate_dataarray_coords(upper, coords) - # TODO: validate mask DataArray coords too, but this would - # change behavior — broadcast_mask currently fills mismatched - # coords with False silently. Needs a deprecation path. - data = Dataset( { "lower": as_dataarray(lower, coords, **kwargs), @@ -631,7 +622,9 @@ def add_variables( self._check_valid_dim_names(data) if mask is not None: - mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool) + mask = _coerce_to_dataarray( + mask, coords=data.coords, dims=data.dims + ).astype(bool) mask = broadcast_mask(mask, data.labels) # Auto-mask based on NaN in bounds (use numpy for speed) @@ -860,7 +853,9 @@ def add_constraints( (data,) = xr.broadcast(data, exclude=[TERM_DIM]) if mask is not None: - mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool) + mask = _coerce_to_dataarray( + mask, coords=data.coords, dims=data.dims + ).astype(bool) mask = broadcast_mask(mask, data.labels) # Auto-mask based on null expressions or NaN RHS (use numpy for speed) diff --git a/linopy/variables.py b/linopy/variables.py index 4332a037..9d45f0da 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -33,7 +33,7 @@ from linopy.common import ( LabelPositionIndex, LocIndexer, - as_dataarray, + _coerce_to_dataarray, assign_multiindex_safe, check_has_nulls, check_has_nulls_polars, @@ -321,7 +321,9 @@ def to_linexpr( linopy.LinearExpression Linear expression with the variables and coefficients. """ - coefficient = as_dataarray(coefficient, coords=self.coords, dims=self.dims) + coefficient = _coerce_to_dataarray( + coefficient, coords=self.coords, dims=self.dims + ) coefficient = coefficient.reindex_like(self.labels, fill_value=0) coefficient = coefficient.fillna(0) ds = Dataset({"coeffs": coefficient, "vars": self.labels}).expand_dims( diff --git a/test/test_common.py b/test/test_common.py index 13aa3e77..75a0c201 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -419,8 +419,8 @@ def test_validate_dataarray_coords_match() -> None: from linopy.common import _validate_dataarray_coords da = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) - out = _validate_dataarray_coords(da, coords={"x": [10, 20, 30]}) - assert_equal(out, da) + # Should not raise + _validate_dataarray_coords(da, coords={"x": [10, 20, 30]}) def test_validate_dataarray_coords_mismatch() -> None: @@ -439,24 +439,13 @@ def test_validate_dataarray_coords_extra_dims() -> None: _validate_dataarray_coords(da, coords={"x": [0, 1]}) -def test_validate_dataarray_coords_broadcast() -> None: - from linopy.common import _validate_dataarray_coords - - da = DataArray([1, 2], dims=["x"], coords={"x": ["a", "b"]}) - out = _validate_dataarray_coords( - da, coords={"x": ["a", "b"], "y": [1, 2, 3]}, dims=["x", "y"] - ) - assert set(out.dims) == {"x", "y"} - assert out.sizes["y"] == 3 - - def test_validate_dataarray_coords_sequence() -> None: from linopy.common import _validate_dataarray_coords da = DataArray([1, 2], dims=["x"], coords={"x": [0, 1]}) idx = pd.RangeIndex(2, name="x") - out = _validate_dataarray_coords(da, coords=[idx], dims=["x"]) - assert list(out.coords["x"].values) == [0, 1] + # Should not raise + _validate_dataarray_coords(da, coords=[idx], dims=["x"]) def test_validate_dataarray_coords_sequence_mismatch() -> None: @@ -468,6 +457,55 @@ def test_validate_dataarray_coords_sequence_mismatch() -> None: _validate_dataarray_coords(da, coords=[idx], dims=["x"]) +def test_coerce_to_dataarray_scalar() -> None: + from linopy.common import _coerce_to_dataarray + + da = _coerce_to_dataarray(1, coords={"x": [10, 20]}, dims=["x"]) + assert isinstance(da, DataArray) + assert da.dims == ("x",) + assert list(da.coords["x"].values) == [10, 20] + + +def test_coerce_to_dataarray_numpy() -> None: + from linopy.common import _coerce_to_dataarray + + arr = np.array([1, 2, 3]) + da = _coerce_to_dataarray(arr, coords={"x": [10, 20, 30]}, dims=["x"]) + assert isinstance(da, DataArray) + assert da.dims == ("x",) + assert list(da.coords["x"].values) == [10, 20, 30] + + +def test_coerce_to_dataarray_no_expand() -> None: + from linopy.common import _coerce_to_dataarray + + da_in = DataArray([1, 2], dims=["x"], coords={"x": ["a", "b"]}) + # Missing dim "y" should NOT be expanded — coerce only does type conversion + da = _coerce_to_dataarray( + da_in, coords={"x": ["a", "b"], "y": [1, 2, 3]}, dims=["x", "y"] + ) + assert da.dims == ("x",) + assert "y" not in da.dims + + +def test_coerce_to_dataarray_allows_extra_dims() -> None: + from linopy.common import _coerce_to_dataarray + + da_in = DataArray([[1, 2], [3, 4]], dims=["x", "y"]) + # Should NOT raise even though "y" is not in coords + da = _coerce_to_dataarray(da_in, coords={"x": [0, 1]}, dims=["x"]) + assert set(da.dims) == {"x", "y"} + + +def test_coerce_to_dataarray_no_coord_validation() -> None: + from linopy.common import _coerce_to_dataarray + + da_in = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) + # Should NOT raise even though coords don't match + da = _coerce_to_dataarray(da_in, coords={"x": [10, 20, 40]}, dims=["x"]) + assert list(da.coords["x"].values) == [10, 20, 30] + + def test_add_variables_with_dataarray_bounds_and_coords() -> None: model = Model() time = pd.RangeIndex(5, name="time") From bd955ef2a9cadd308da8478be6faef840b70da4c Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:03:07 +0100 Subject: [PATCH 17/18] Rename _coerce_to_dataarray to ensure_dataarray, revert as_dataarray to optional coords MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename internal helper to ensure_dataarray (no underscore, consistent with other common.py conventions like broadcast_mask) - Revert as_dataarray coords back to optional — when coords is None, delegates to _type_dispatch for pure type conversion - Deduplicate _type_dispatch call in as_dataarray (always called, validation conditional) - Update all call sites in expressions.py, model.py, variables.py - Update test names to match new function name Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 82 +++++++++++++++++++++---------------------- linopy/expressions.py | 23 ++++++------ linopy/model.py | 20 +++++------ linopy/variables.py | 6 ++-- test/test_common.py | 32 ++++++++--------- 5 files changed, 79 insertions(+), 84 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index d6e624ba..07f2aaf0 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -283,7 +283,7 @@ def _type_dispatch( Convert arr to a DataArray via type dispatch. This is the shared conversion logic used by both ``as_dataarray`` - and ``_coerce_to_dataarray``. It does NOT validate or expand dims. + and ``ensure_dataarray``. It does NOT validate or expand dims. """ if isinstance(arr, pd.Series | pd.DataFrame): arr = pandas_to_dataarray(arr, coords=coords, dims=dims, **kwargs) @@ -350,9 +350,9 @@ def _expand_missing_dims( return arr -def _coerce_to_dataarray( +def ensure_dataarray( arr: Any, - coords: CoordsLike, + coords: CoordsLike | None = None, dims: DimsLike | None = None, **kwargs: Any, ) -> DataArray: @@ -368,7 +368,7 @@ def _coerce_to_dataarray( The input object. coords: Coordinates used as construction hints for types without - their own coords (numpy, scalar, list). + their own coords (numpy, scalar, list). Optional. dims: Dimension names. **kwargs: @@ -388,55 +388,55 @@ def as_dataarray( **kwargs: Any, ) -> DataArray: """ - Convert an object to a DataArray with strict coord validation. + Convert an object to a DataArray with optional strict coord validation. + + When ``coords`` is provided, performs type conversion plus + validation: - Behavior depends on whether ``coords`` is provided: + - For inputs with their own coordinates (DataArray, pandas, scalars): shared dims must match exactly + (``ValueError`` if not), extra dims are rejected (``ValueError``), missing dims are broadcast via ``expand_dims``. + - For raw arrays (numpy, list, polars): ``coords`` is applied + during construction (no validation needed). - - **No coords**: just type conversion (scalar->DA, numpy->DA, - pandas->DA, DataArray->passthrough). - - **With coords** (strict): shared dims must have matching - coordinates (ValueError if not), extra dims in input are - rejected (ValueError), missing dims are broadcast via - ``expand_dims``. + When ``coords`` is ``None``, performs pure type conversion only + (accepts any input type, including numpy/list/polars). Parameters ---------- - arr: - The input object. - coords: - The coordinates for the DataArray. If None, only type - conversion is performed. - dims: - The dimensions for the DataArray. If None, the dimensions - will be automatically generated. - **kwargs: - Additional keyword arguments passed to the DataArray constructor. + arr : + The input object. + coords : CoordsLike, optional + Expected coordinates. When provided, used for construction + (numpy, scalar, list) and for validation (DataArray, pandas). + When ``None``, only type conversion is performed. + dims : DimsLike, optional + Dimension names. + **kwargs : + Additional keyword arguments passed to the DataArray + constructor. Returns ------- - DataArray: - The converted DataArray. + DataArray """ - if coords is not None: - # Inputs that already have their own coordinates (DataArray, - # pandas) or scalars need validation against coords. Raw arrays - # (numpy, list) get coords applied during construction, so - # validation and expand_dims are not needed. - needs_validation = isinstance(arr, DataArray | pd.Series | pd.DataFrame) or ( - not isinstance(arr, np.ndarray | pl.Series | list) and np.ndim(arr) == 0 - ) + if coords is None: + return _type_dispatch(arr, coords=None, dims=dims, **kwargs) - if needs_validation: - arr = _type_dispatch(arr, coords=coords, dims=dims, **kwargs) - arr = _expand_missing_dims(arr, coords, dims) - _validate_dataarray_coords(arr, coords, dims) - else: - arr = _type_dispatch(arr, coords=coords, dims=dims, **kwargs) + # Inputs that already have their own coordinates (DataArray, + # pandas) or scalars need validation against coords. Raw arrays + # (numpy, list) get coords applied during construction, so + # validation and expand_dims are not needed. + needs_validation = isinstance(arr, DataArray | pd.Series | pd.DataFrame) or ( + not isinstance(arr, np.ndarray | pl.Series | list) and np.ndim(arr) == 0 + ) - return arr + arr = _type_dispatch(arr, coords=coords, dims=dims, **kwargs) - # No coords — just type conversion - return _type_dispatch(arr, coords=None, dims=dims, **kwargs) + if needs_validation: + arr = _expand_missing_dims(arr, coords, dims) + _validate_dataarray_coords(arr, coords, dims) + + return arr def broadcast_mask(mask: DataArray, labels: DataArray) -> DataArray: diff --git a/linopy/expressions.py b/linopy/expressions.py index 542ac685..c8e5d61f 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -46,12 +46,11 @@ from linopy.common import ( EmptyDeprecationWrapper, LocIndexer, - _coerce_to_dataarray, - as_dataarray, assign_multiindex_safe, check_common_keys_values, check_has_nulls, check_has_nulls_polars, + ensure_dataarray, fill_missing_coords, filter_nulls_polars, forward_as_properties, @@ -356,7 +355,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: da = xr.DataArray([], dims=[TERM_DIM]) data = Dataset({"coeffs": da, "vars": da, "const": 0.0}) elif isinstance(data, SUPPORTED_CONSTANT_TYPES): - const = as_dataarray(data) + const = ensure_dataarray(data) da = xr.DataArray([], dims=[TERM_DIM]) data = Dataset({"coeffs": da, "vars": da, "const": const}) elif not isinstance(data, Dataset): @@ -598,7 +597,7 @@ def _add_constant( # so that missing data does not silently propagate through arithmetic. if np.isscalar(other) and join is None: return self.assign(const=self.const.fillna(0) + other) - da = _coerce_to_dataarray(other, coords=self.coords, dims=self.coord_dims) + da = ensure_dataarray(other, coords=self.coords, dims=self.coord_dims) self_const, da, needs_data_reindex = self._align_constant( da, fill_value=0, join=join ) @@ -627,7 +626,7 @@ def _apply_constant_op( - factor (other) is filled with fill_value (0 for mul, 1 for div) - coeffs and const are filled with 0 (additive identity) """ - factor = _coerce_to_dataarray(other, coords=self.coords, dims=self.coord_dims) + factor = ensure_dataarray(other, coords=self.coords, dims=self.coord_dims) self_const, factor, needs_data_reindex = self._align_constant( factor, fill_value=fill_value, join=join ) @@ -1143,7 +1142,7 @@ def to_constraint( ) if isinstance(rhs, SUPPORTED_CONSTANT_TYPES): - rhs = _coerce_to_dataarray(rhs, coords=self.coords, dims=self.coord_dims) + rhs = ensure_dataarray(rhs, coords=self.coords, dims=self.coord_dims) extra_dims = set(rhs.dims) - set(self.coord_dims) if extra_dims: @@ -1706,9 +1705,7 @@ def __matmul__( Matrix multiplication with other, similar to xarray dot. """ if not isinstance(other, LinearExpression | variables.Variable): - other = _coerce_to_dataarray( - other, coords=self.coords, dims=self.coord_dims - ) + other = ensure_dataarray(other, coords=self.coords, dims=self.coord_dims) common_dims = list(set(self.coord_dims).intersection(other.dims)) return (self * other).sum(dim=common_dims) @@ -2042,7 +2039,7 @@ def process_one( DeprecationWarning, ) # assume that the element is a constant - const = as_dataarray(t[0]) + const = ensure_dataarray(t[0]) if model is None: raise ValueError("Model must be provided when using constants.") return LinearExpression(const, model) @@ -2070,7 +2067,7 @@ def from_constant(cls, model: Model, constant: ConstantLike) -> LinearExpression linopy.LinearExpression A linear expression representing the constant value. """ - const_da = as_dataarray(constant) + const_da = ensure_dataarray(constant) return LinearExpression(const_da, model) @@ -2194,7 +2191,7 @@ def __matmul__( "Higher order non-linear expressions are not yet supported." ) - other = _coerce_to_dataarray(other, coords=self.coords, dims=self.coord_dims) + other = ensure_dataarray(other, coords=self.coords, dims=self.coord_dims) common_dims = list(set(self.coord_dims).intersection(other.dims)) return (self * other).sum(dim=common_dims) @@ -2329,7 +2326,7 @@ def as_expression( return obj.to_linexpr() else: try: - obj = as_dataarray(obj, **kwargs) + obj = ensure_dataarray(obj, **kwargs) except ValueError as e: raise ValueError("Cannot convert to LinearExpression") from e return LinearExpression(obj, model) diff --git a/linopy/model.py b/linopy/model.py index b260d587..0d9c86ef 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -26,11 +26,11 @@ from linopy import solvers from linopy.common import ( - _coerce_to_dataarray, as_dataarray, assign_multiindex_safe, best_int, broadcast_mask, + ensure_dataarray, maybe_replace_signs, replace_by_map, set_int_index, @@ -622,9 +622,9 @@ def add_variables( self._check_valid_dim_names(data) if mask is not None: - mask = _coerce_to_dataarray( - mask, coords=data.coords, dims=data.dims - ).astype(bool) + mask = ensure_dataarray(mask, coords=data.coords, dims=data.dims).astype( + bool + ) mask = broadcast_mask(mask, data.labels) # Auto-mask based on NaN in bounds (use numpy for speed) @@ -786,14 +786,14 @@ def add_constraints( name = f"con{self._connameCounter}" self._connameCounter += 1 if sign is not None: - sign = maybe_replace_signs(as_dataarray(sign)) + sign = maybe_replace_signs(ensure_dataarray(sign)) # Capture original RHS for auto-masking before constraint creation # (NaN values in RHS are lost during constraint creation) # Use numpy for speed instead of xarray's notnull() original_rhs_mask = None if self.auto_mask and rhs is not None: - rhs_da = as_dataarray(rhs) + rhs_da = ensure_dataarray(rhs) original_rhs_mask = (rhs_da.coords, rhs_da.dims, ~np.isnan(rhs_da.values)) if isinstance(lhs, LinearExpression): @@ -846,16 +846,16 @@ def add_constraints( mask = ( rhs_mask if mask is None - else (as_dataarray(mask).astype(bool) & rhs_mask) + else (ensure_dataarray(mask).astype(bool) & rhs_mask) ) data["labels"] = -1 (data,) = xr.broadcast(data, exclude=[TERM_DIM]) if mask is not None: - mask = _coerce_to_dataarray( - mask, coords=data.coords, dims=data.dims - ).astype(bool) + mask = ensure_dataarray(mask, coords=data.coords, dims=data.dims).astype( + bool + ) mask = broadcast_mask(mask, data.labels) # Auto-mask based on null expressions or NaN RHS (use numpy for speed) diff --git a/linopy/variables.py b/linopy/variables.py index 9d45f0da..8f3be95b 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -33,10 +33,10 @@ from linopy.common import ( LabelPositionIndex, LocIndexer, - _coerce_to_dataarray, assign_multiindex_safe, check_has_nulls, check_has_nulls_polars, + ensure_dataarray, filter_nulls_polars, format_string_as_variable_name, generate_indices_for_printout, @@ -321,9 +321,7 @@ def to_linexpr( linopy.LinearExpression Linear expression with the variables and coefficients. """ - coefficient = _coerce_to_dataarray( - coefficient, coords=self.coords, dims=self.dims - ) + coefficient = ensure_dataarray(coefficient, coords=self.coords, dims=self.dims) coefficient = coefficient.reindex_like(self.labels, fill_value=0) coefficient = coefficient.fillna(0) ds = Dataset({"coeffs": coefficient, "vars": self.labels}).expand_dims( diff --git a/test/test_common.py b/test/test_common.py index 75a0c201..2eea1de3 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -457,52 +457,52 @@ def test_validate_dataarray_coords_sequence_mismatch() -> None: _validate_dataarray_coords(da, coords=[idx], dims=["x"]) -def test_coerce_to_dataarray_scalar() -> None: - from linopy.common import _coerce_to_dataarray +def test_ensure_dataarray_scalar() -> None: + from linopy.common import ensure_dataarray - da = _coerce_to_dataarray(1, coords={"x": [10, 20]}, dims=["x"]) + da = ensure_dataarray(1, coords={"x": [10, 20]}, dims=["x"]) assert isinstance(da, DataArray) assert da.dims == ("x",) assert list(da.coords["x"].values) == [10, 20] -def test_coerce_to_dataarray_numpy() -> None: - from linopy.common import _coerce_to_dataarray +def test_ensure_dataarray_numpy() -> None: + from linopy.common import ensure_dataarray arr = np.array([1, 2, 3]) - da = _coerce_to_dataarray(arr, coords={"x": [10, 20, 30]}, dims=["x"]) + da = ensure_dataarray(arr, coords={"x": [10, 20, 30]}, dims=["x"]) assert isinstance(da, DataArray) assert da.dims == ("x",) assert list(da.coords["x"].values) == [10, 20, 30] -def test_coerce_to_dataarray_no_expand() -> None: - from linopy.common import _coerce_to_dataarray +def test_ensure_dataarray_no_expand() -> None: + from linopy.common import ensure_dataarray da_in = DataArray([1, 2], dims=["x"], coords={"x": ["a", "b"]}) - # Missing dim "y" should NOT be expanded — coerce only does type conversion - da = _coerce_to_dataarray( + # Missing dim "y" should NOT be expanded — ensure only does type conversion + da = ensure_dataarray( da_in, coords={"x": ["a", "b"], "y": [1, 2, 3]}, dims=["x", "y"] ) assert da.dims == ("x",) assert "y" not in da.dims -def test_coerce_to_dataarray_allows_extra_dims() -> None: - from linopy.common import _coerce_to_dataarray +def test_ensure_dataarray_allows_extra_dims() -> None: + from linopy.common import ensure_dataarray da_in = DataArray([[1, 2], [3, 4]], dims=["x", "y"]) # Should NOT raise even though "y" is not in coords - da = _coerce_to_dataarray(da_in, coords={"x": [0, 1]}, dims=["x"]) + da = ensure_dataarray(da_in, coords={"x": [0, 1]}, dims=["x"]) assert set(da.dims) == {"x", "y"} -def test_coerce_to_dataarray_no_coord_validation() -> None: - from linopy.common import _coerce_to_dataarray +def test_ensure_dataarray_no_coord_validation() -> None: + from linopy.common import ensure_dataarray da_in = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) # Should NOT raise even though coords don't match - da = _coerce_to_dataarray(da_in, coords={"x": [10, 20, 40]}, dims=["x"]) + da = ensure_dataarray(da_in, coords={"x": [10, 20, 40]}, dims=["x"]) assert list(da.coords["x"].values) == [10, 20, 30] From 31733c71cd9e1999c891c2e44c488132b1c2c9c4 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:11:33 +0100 Subject: [PATCH 18/18] Add validation to _coords_to_mapping for edge cases - Raise ValueError when dims length doesn't match coords sequence length - Raise ValueError on duplicate .name in coords sequence - Warn on unnamed items in coords sequence (when dims is None) Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 22 +++++++++++++++++++++- test/test_common.py | 24 ++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/linopy/common.py b/linopy/common.py index 07f2aaf0..513c935f 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -140,8 +140,28 @@ def _coords_to_mapping(coords: CoordsLike, dims: DimsLike | None = None) -> Mapp seq = list(coords) if dims is not None: dim_names: list[str] = [dims] if isinstance(dims, str) else list(dims) # type: ignore[arg-type] + if len(dim_names) != len(seq): + raise ValueError( + f"Length of dims ({len(dim_names)}) does not match " + f"length of coords sequence ({len(seq)})." + ) return dict(zip(dim_names, seq)) - return {c.name: c for c in seq if hasattr(c, "name") and c.name} + result: dict = {} + for c in seq: + if not hasattr(c, "name") or not c.name: + warn( + f"Coordinate {c!r} has no .name attribute and will be ignored. " + "Pass dims explicitly or use named indexes (e.g. pd.RangeIndex).", + UserWarning, + stacklevel=3, + ) + continue + if c.name in result: + raise ValueError( + f"Duplicate coordinate name '{c.name}' in coords sequence." + ) + result[c.name] = c + return result def _validate_dataarray_coords( diff --git a/test/test_common.py b/test/test_common.py index 2eea1de3..4feee38d 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -457,6 +457,30 @@ def test_validate_dataarray_coords_sequence_mismatch() -> None: _validate_dataarray_coords(da, coords=[idx], dims=["x"]) +def test_coords_to_mapping_dims_length_mismatch() -> None: + from linopy.common import _coords_to_mapping + + with pytest.raises(ValueError, match="Length of dims"): + _coords_to_mapping([pd.RangeIndex(3), pd.RangeIndex(5)], dims=["x"]) + + +def test_coords_to_mapping_unnamed_index_warns() -> None: + from linopy.common import _coords_to_mapping + + with pytest.warns(UserWarning, match="no .name attribute"): + result = _coords_to_mapping([np.array([1, 2, 3])]) + assert result == {} + + +def test_coords_to_mapping_duplicate_name() -> None: + from linopy.common import _coords_to_mapping + + idx1 = pd.RangeIndex(3, name="x") + idx2 = pd.RangeIndex(5, name="x") + with pytest.raises(ValueError, match="Duplicate coordinate name"): + _coords_to_mapping([idx1, idx2]) + + def test_ensure_dataarray_scalar() -> None: from linopy.common import ensure_dataarray