From 619b07c2b136069a29f730b52138008fbf0bc036 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/12] 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 7dd97b65..8554a9c1 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 @@ -264,7 +264,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 44a94b5c5335151f7004cafe9914645f5d585c45 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/12] 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 8554a9c1..ab8ea6f2 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -265,12 +265,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 727eb2c78b2d88aa7a840dbe35e05f4893580a6f 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/12] 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 ab8ea6f2..9294155c 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -265,13 +265,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 31642cdfcdfe2867cc9ab15e497714af5359e55f 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/12] 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 | 8 ++++++-- linopy/variables.py | 4 +++- 3 files changed, 20 insertions(+), 3 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 9294155c..8904c6d3 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -233,6 +233,7 @@ def as_dataarray( arr: Any, coords: CoordsLike | None = None, dims: DimsLike | None = None, + allow_extra_dims: bool = False, **kwargs: Any, ) -> DataArray: """ @@ -246,6 +247,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. @@ -266,6 +271,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 10e243de..8af037b9 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -529,7 +529,9 @@ def _multiply_by_linear_expression( def _multiply_by_constant( self: GenericExpression, other: ConstantLike ) -> GenericExpression: - multiplier = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + multiplier = as_dataarray( + other, coords=self.coords, dims=self.coord_dims, allow_extra_dims=True + ) coeffs = self.coeffs * multiplier assert all(coeffs.sizes[d] == s for d, s in self.coeffs.sizes.items()) const = self.const * multiplier @@ -1406,7 +1408,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 e2570b5d..64fcf340 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -313,7 +313,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 + ) ds = Dataset({"coeffs": coefficient, "vars": self.labels}).expand_dims( TERM_DIM, -1 ) From 4351e93131fd0a58201a5b2c7fd406704d15b669 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/12] 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 8904c6d3..9ed95f8c 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -270,6 +270,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 047f288802a8de32ee44b8dfe78bd93c6235867a 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/12] 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 | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 9ed95f8c..7948d1a5 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, @@ -163,8 +179,7 @@ def pandas_to_dataarray( ] if coords is not None: pandas_coords = dict(zip(dims, arr.axes)) - if isinstance(coords, Sequence): - coords = dict(zip(dims, coords)) + coords = _coords_to_mapping(coords, dims) shared_dims = set(pandas_coords.keys()) & set(coords.keys()) non_aligned = [] for dim in shared_dims: @@ -224,7 +239,7 @@ def numpy_to_dataarray( dims = [get_from_iterable(dims, i) or f"dim_{i}" for i in range(ndim)] if isinstance(coords, list) and dims is not None and len(dims): - coords = dict(zip(dims, coords)) + coords = _coords_to_mapping(coords, dims) return DataArray(arr, coords=coords, dims=dims, **kwargs) @@ -270,14 +285,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 73e85d2976f97822cc0a24b749b08c0bffb2807f 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/12] 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 7948d1a5..1c99bf2a 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -273,6 +273,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 657b2d45..90d57393 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -449,10 +449,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 @@ -474,7 +476,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 5f2e94df2a5040f0eb79090f51c69d39bb755e4e 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/12] 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 1c99bf2a..ea29531b 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, @@ -295,28 +337,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 42a036cb7405ec297a59eb0a7e8ed70fdc3f2b12 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/12] Improve performance --- linopy/common.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index ea29531b..2f81fd52 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) @@ -337,7 +339,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 1bab4966be73ad9a99f6b401b74d3ece292e746e 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/12] 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 3081323d..52770dfe 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -555,7 +555,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) @@ -750,7 +753,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 53e449842896c3726794b5f05d5dce48690b7060 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/12] fix types --- linopy/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/linopy/common.py b/linopy/common.py index 8ad4e517..d6a480aa 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 75b0610ce551ed504b261741b6b5143c9a28107a 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/12] Update Changelog and add some tests to ensure behaviour --- doc/release_notes.rst | 2 ++ test/test_common.py | 80 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 7731443b..7a6abbb8 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -4,6 +4,8 @@ Release Notes Upcoming Version ---------------- +* 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. + Version 0.6.4 -------------- diff --git a/test/test_common.py b/test/test_common.py index c3500155..d73bda2d 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -432,6 +432,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"]])