From a32cf4bc3fcba8af9f60d490b6a9e4c9c6c2027c Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:18:32 +0100 Subject: [PATCH 01/13] Fix add_variables silently ignoring coords for DataArray bounds When DataArray bounds were passed to add_variables with explicit coords, the coords parameter was silently ignored because as_dataarray skips conversion for DataArray inputs. Now validates DataArray bounds against coords: raises ValueError on mismatched or extra dimensions, and broadcasts missing dimensions via expand_dims. Co-Authored-By: Claude Opus 4.6 --- linopy/model.py | 68 +++++++++++++++++++++++++++++++++++++++++++++ test/test_common.py | 55 ++++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+) diff --git a/linopy/model.py b/linopy/model.py index 54334411..bb6b1ac3 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -94,6 +94,70 @@ logger = logging.getLogger(__name__) +def _coords_to_dict( + coords: Sequence[Sequence | pd.Index | DataArray] | Mapping, +) -> dict[str, Any]: + """Normalize coords to a dict mapping dim names to coordinate values.""" + if isinstance(coords, Mapping): + return dict(coords) + # xarray Coordinates objects + if hasattr(coords, "dims"): + return {k: v for k, v in coords.items() if k in coords.dims} + # Sequence of indexes + result: dict[str, Any] = {} + for c in coords: + if isinstance(c, pd.Index) and c.name: + result[c.name] = c + return result + + +def _validate_dataarray_bounds(arr: Any, coords: Any) -> Any: + """ + Validate and expand DataArray bounds against explicit coords. + + If ``arr`` is not a DataArray, return it unchanged (``as_dataarray`` + will handle conversion). For DataArray inputs: + + - Raises ``ValueError`` if the array has dimensions not in coords. + - Raises ``ValueError`` if shared dimension coordinates don't match. + - Expands missing dimensions via ``expand_dims``. + """ + if not isinstance(arr, DataArray): + return arr + + expected = _coords_to_dict(coords) + if not expected: + return arr + + extra = set(arr.dims) - set(expected) + if extra: + raise ValueError(f"DataArray has extra dimensions not in coords: {extra}") + + for dim, coord_values in expected.items(): + if dim not in arr.dims: + continue + if isinstance(arr.indexes.get(dim), pd.MultiIndex): + continue + expected_idx = ( + coord_values + if isinstance(coord_values, pd.Index) + else pd.Index(coord_values) + ) + actual_idx = arr.coords[dim].to_index() + if not actual_idx.equals(expected_idx): + raise ValueError( + f"Coordinates for dimension '{dim}' do not match: " + f"expected {expected_idx.tolist()}, got {actual_idx.tolist()}" + ) + + # Expand missing dimensions + expand = {k: v for k, v in expected.items() if k not in arr.dims} + if expand: + arr = arr.expand_dims(expand) + + return arr + + class Model: """ Linear optimization model. @@ -603,6 +667,10 @@ def add_variables( "Semi-continuous variables require a positive scalar lower bound." ) + if coords is not None: + lower = _validate_dataarray_bounds(lower, coords) + upper = _validate_dataarray_bounds(upper, coords) + data = Dataset( { "lower": as_dataarray(lower, coords, **kwargs), diff --git a/test/test_common.py b/test/test_common.py index f1190024..eea6a01e 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -409,6 +409,61 @@ 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_add_variables_with_dataarray_bounds_and_coords() -> None: + from linopy import Model + + 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: + from linopy import Model + + 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_extra_dims() -> None: + from linopy import Model + + model = Model() + lower = DataArray([[1, 2], [3, 4]], dims=["x", "y"]) + with pytest.raises(ValueError, match="extra dimensions"): + model.add_variables(lower=lower, coords={"x": [0, 1]}, name="x") + + +def test_add_variables_with_dataarray_bounds_broadcast() -> None: + from linopy import Model + + 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_add_variables_with_non_dataarray_bounds_unchanged() -> None: + """Non-DataArray bounds (scalars, numpy) should work as before.""" + from linopy import Model + + model = Model() + time = pd.RangeIndex(3, name="time") + var = model.add_variables( + lower=0, upper=np.array([1, 2, 3]), coords=[time], name="x" + ) + assert var.shape == (3,) + + def test_as_dataarray_with_unsupported_type() -> None: with pytest.raises(TypeError): as_dataarray(lambda x: 1, dims=["dim1"], coords=[["a"]]) From cb438bdbcf819181fe4fa05d992f0785cda3aa6f Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:24:28 +0100 Subject: [PATCH 02/13] Another test --- test/test_common.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/test/test_common.py b/test/test_common.py index eea6a01e..0addd796 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -464,6 +464,21 @@ def test_add_variables_with_non_dataarray_bounds_unchanged() -> None: assert var.shape == (3,) +def test_add_variables_with_mixed_dataarray_bounds() -> None: + """Non-DataArray bounds (scalars, numpy) should work as before.""" + from linopy import Model + + model = Model() + time = pd.RangeIndex(3, name="time") + space = pd.Index(["a", "b"], name="space") + lower = DataArray([1, 1, 1], dims=["time"], coords={"time": range(3)}) + upper = xr.DataArray(0) + var = model.add_variables(lower=lower, upper=upper, 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 1655fa0eb152eb8b876c8eceeff5bbfa3b8ad7da Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:31:09 +0100 Subject: [PATCH 03/13] Add additional test coverage for DataArray bounds validation Test MultiIndex coords (validation skip), xarray Coordinates object, dims-only DataArrays, and upper bound mismatch detection. Co-Authored-By: Claude Opus 4.6 --- test/test_common.py | 47 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/test/test_common.py b/test/test_common.py index 0addd796..555976d7 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -479,6 +479,53 @@ def test_add_variables_with_mixed_dataarray_bounds() -> None: assert var.data.sizes["space"] == 2 +def test_add_variables_with_dataarray_upper_mismatch() -> None: + """Validation should catch mismatched upper bounds too, not just lower.""" + from linopy import Model + + model = Model() + time = pd.RangeIndex(5, name="time") + upper = DataArray([1, 2, 3], dims=["time"], coords={"time": [0, 1, 2]}) + with pytest.raises(ValueError, match="do not match"): + model.add_variables(upper=upper, coords=[time], name="x") + + +def test_add_variables_with_multiindex_coords() -> None: + """MultiIndex with scalar bounds should still work (existing pattern).""" + from linopy import Model + + model = Model() + idx = pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=("level1", "level2")) + idx.name = "multi" + var = model.add_variables(lower=0, upper=1, coords=[idx], name="x") + assert var.shape == (4,) + + +def test_add_variables_with_xarray_coordinates() -> None: + """Coords passed as xarray Coordinates object should work.""" + from linopy import Model + + model = Model() + time = pd.RangeIndex(3, name="time") + existing_var = model.add_variables(lower=0, coords=[time], name="base") + xr_coords = existing_var.data.coords + lower = DataArray([1, 1, 1], dims=["time"], coords={"time": range(3)}) + var = model.add_variables(lower=lower, coords=xr_coords, name="x2") + assert var.shape == (3,) + + +def test_add_variables_with_dims_only_dataarray() -> None: + """DataArray with dims but no explicit coord values should still work.""" + from linopy import Model + + model = Model() + lower = DataArray([0, 0, 0], dims=["x"]) + var = model.add_variables( + lower=lower, coords=[pd.RangeIndex(3, name="x")], name="x" + ) + assert var.shape == (3,) + + def test_as_dataarray_with_unsupported_type() -> None: with pytest.raises(TypeError): as_dataarray(lambda x: 1, dims=["dim1"], coords=[["a"]]) From ecca2f8463d662c36944d1f357cb01fb63858e78 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:33:38 +0100 Subject: [PATCH 04/13] Add TODO noting as_dataarray fails for scalars with dict coords Co-Authored-By: Claude Opus 4.6 --- linopy/model.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/linopy/model.py b/linopy/model.py index bb6b1ac3..f95e6ac4 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -671,6 +671,10 @@ def add_variables( lower = _validate_dataarray_bounds(lower, coords) upper = _validate_dataarray_bounds(upper, coords) + # TODO: as_dataarray fails for scalars with dict coords (e.g. + # as_dataarray(inf, coords={"x": [0,1,2]})) because dims is not + # inferred from the dict keys. Sequence coords work fine. + # See https://github.com/PyPSA/linopy/pull/551 data = Dataset( { "lower": as_dataarray(lower, coords, **kwargs), From e0683577c940e035cc66ff22087f137b0cfab9b4 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:35:53 +0100 Subject: [PATCH 05/13] Fix as_dataarray for scalars with dict coords Infer dims from dict keys when dims is None and the input is a scalar. Previously this raised xarray's CoordinateValidationError because xarray can't broadcast a 0-dim value to coords without explicit dims. Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 4 ++++ linopy/model.py | 4 ---- test/test_common.py | 10 ++++++++++ 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 09f67355..7fb76288 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -242,8 +242,12 @@ def as_dataarray( elif isinstance(arr, pl.Series): arr = numpy_to_dataarray(arr.to_numpy(), coords=coords, dims=dims, **kwargs) elif isinstance(arr, np.number): + if dims is None and is_dict_like(coords) and np.ndim(arr) == 0: + dims = list(coords.keys()) arr = DataArray(float(arr), coords=coords, dims=dims, **kwargs) elif isinstance(arr, int | float | str | bool | list): + if dims is None and is_dict_like(coords) and np.ndim(arr) == 0: + dims = list(coords.keys()) arr = DataArray(arr, coords=coords, dims=dims, **kwargs) elif not isinstance(arr, DataArray): diff --git a/linopy/model.py b/linopy/model.py index f95e6ac4..bb6b1ac3 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -671,10 +671,6 @@ def add_variables( lower = _validate_dataarray_bounds(lower, coords) upper = _validate_dataarray_bounds(upper, coords) - # TODO: as_dataarray fails for scalars with dict coords (e.g. - # as_dataarray(inf, coords={"x": [0,1,2]})) because dims is not - # inferred from the dict keys. Sequence coords work fine. - # See https://github.com/PyPSA/linopy/pull/551 data = Dataset( { "lower": as_dataarray(lower, coords, **kwargs), diff --git a/test/test_common.py b/test/test_common.py index 555976d7..06d345e0 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -514,6 +514,16 @@ def test_add_variables_with_xarray_coordinates() -> None: assert var.shape == (3,) +def test_add_variables_with_scalar_and_dict_coords() -> None: + """Scalar bounds with dict coords should infer dims from dict keys.""" + from linopy import Model + + model = Model() + var = model.add_variables(lower=0, upper=10, coords={"x": [0, 1, 2]}, name="x") + assert var.shape == (3,) + assert list(var.data.coords["x"].values) == [0, 1, 2] + + def test_add_variables_with_dims_only_dataarray() -> None: """DataArray with dims but no explicit coord values should still work.""" from linopy import Model From 51b10b6ae31c0b0a99950d9a6fc9b1b4afb81170 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:40:31 +0100 Subject: [PATCH 06/13] Replace individual tests with parameterized test suite Consolidate add_variables tests into TestAddVariablesBoundsWithCoords class with parameterized tests covering all bound types (scalar, np.number, numpy, pandas, list, DataArray, DataArray-no-coords) x both coord formats (sequence, dict). Also fixes as_dataarray for scalars with dict coords by inferring dims from dict keys. Co-Authored-By: Claude Opus 4.6 --- test/test_common.py | 201 +++++++++++++++++++------------------------- 1 file changed, 85 insertions(+), 116 deletions(-) diff --git a/test/test_common.py b/test/test_common.py index 06d345e0..86307bd7 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -409,131 +409,100 @@ 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_add_variables_with_dataarray_bounds_and_coords() -> None: - from linopy import Model - - 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: - from linopy import Model - - 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") +class TestAddVariablesBoundsWithCoords: + """Test that add_variables correctly handles all bound types with coords.""" + SEQ_COORDS = [pd.RangeIndex(3, name="x")] + DICT_COORDS = {"x": [0, 1, 2]} -def test_add_variables_with_dataarray_bounds_extra_dims() -> None: - from linopy import Model + @pytest.fixture() + def model(self) -> "Model": + from linopy import Model - model = Model() - lower = DataArray([[1, 2], [3, 4]], dims=["x", "y"]) - with pytest.raises(ValueError, match="extra dimensions"): - model.add_variables(lower=lower, coords={"x": [0, 1]}, name="x") - - -def test_add_variables_with_dataarray_bounds_broadcast() -> None: - from linopy import Model - - 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 + return Model() + # -- All bound types should work with both coord formats --------------- -def test_add_variables_with_non_dataarray_bounds_unchanged() -> None: - """Non-DataArray bounds (scalars, numpy) should work as before.""" - from linopy import Model - - model = Model() - time = pd.RangeIndex(3, name="time") - var = model.add_variables( - lower=0, upper=np.array([1, 2, 3]), coords=[time], name="x" + @pytest.mark.parametrize( + "lower", + [ + pytest.param(0, id="scalar"), + pytest.param(np.float64(0), id="np.number"), + pytest.param(np.array([0, 0, 0]), id="numpy"), + pytest.param( + pd.Series([0, 0, 0], index=pd.RangeIndex(3, name="x")), id="pandas" + ), + pytest.param([0, 0, 0], id="list"), + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + id="dataarray", + ), + pytest.param(DataArray([0, 0, 0], dims=["x"]), id="dataarray-no-coords"), + ], ) - assert var.shape == (3,) - - -def test_add_variables_with_mixed_dataarray_bounds() -> None: - """Non-DataArray bounds (scalars, numpy) should work as before.""" - from linopy import Model - - model = Model() - time = pd.RangeIndex(3, name="time") - space = pd.Index(["a", "b"], name="space") - lower = DataArray([1, 1, 1], dims=["time"], coords={"time": range(3)}) - upper = xr.DataArray(0) - var = model.add_variables(lower=lower, upper=upper, 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_add_variables_with_dataarray_upper_mismatch() -> None: - """Validation should catch mismatched upper bounds too, not just lower.""" - from linopy import Model - - model = Model() - time = pd.RangeIndex(5, name="time") - upper = DataArray([1, 2, 3], dims=["time"], coords={"time": [0, 1, 2]}) - with pytest.raises(ValueError, match="do not match"): - model.add_variables(upper=upper, coords=[time], name="x") - - -def test_add_variables_with_multiindex_coords() -> None: - """MultiIndex with scalar bounds should still work (existing pattern).""" - from linopy import Model - - model = Model() - idx = pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=("level1", "level2")) - idx.name = "multi" - var = model.add_variables(lower=0, upper=1, coords=[idx], name="x") - assert var.shape == (4,) - - -def test_add_variables_with_xarray_coordinates() -> None: - """Coords passed as xarray Coordinates object should work.""" - from linopy import Model - - model = Model() - time = pd.RangeIndex(3, name="time") - existing_var = model.add_variables(lower=0, coords=[time], name="base") - xr_coords = existing_var.data.coords - lower = DataArray([1, 1, 1], dims=["time"], coords={"time": range(3)}) - var = model.add_variables(lower=lower, coords=xr_coords, name="x2") - assert var.shape == (3,) - - -def test_add_variables_with_scalar_and_dict_coords() -> None: - """Scalar bounds with dict coords should infer dims from dict keys.""" - from linopy import Model - - model = Model() - var = model.add_variables(lower=0, upper=10, coords={"x": [0, 1, 2]}, name="x") - assert var.shape == (3,) - assert list(var.data.coords["x"].values) == [0, 1, 2] - + @pytest.mark.parametrize( + "coords", + [ + pytest.param([pd.RangeIndex(3, name="x")], id="seq-coords"), + pytest.param({"x": [0, 1, 2]}, id="dict-coords"), + ], + ) + def test_bound_types_with_coords(self, model: "Model", lower, coords) -> None: + var = model.add_variables(lower=lower, coords=coords, name="x") + assert var.shape == (3,) + assert var.dims == ("x",) + assert list(var.data.coords["x"].values) == [0, 1, 2] -def test_add_variables_with_dims_only_dataarray() -> None: - """DataArray with dims but no explicit coord values should still work.""" - from linopy import Model + # -- DataArray validation: mismatch and extra dims --------------------- - model = Model() - lower = DataArray([0, 0, 0], dims=["x"]) - var = model.add_variables( - lower=lower, coords=[pd.RangeIndex(3, name="x")], name="x" + @pytest.mark.parametrize( + "coords", + [ + pytest.param([pd.RangeIndex(5, name="x")], id="seq-coords"), + pytest.param({"x": [0, 1, 2, 3, 4]}, id="dict-coords"), + ], ) - assert var.shape == (3,) + def test_dataarray_coord_mismatch(self, model: "Model", coords) -> None: + lower = DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}) + with pytest.raises(ValueError, match="do not match"): + model.add_variables(lower=lower, coords=coords, name="x") + + def test_dataarray_coord_mismatch_upper(self, model: "Model") -> None: + upper = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) + with pytest.raises(ValueError, match="do not match"): + model.add_variables(upper=upper, coords=self.SEQ_COORDS, name="x") + + def test_dataarray_extra_dims(self, model: "Model") -> None: + lower = DataArray([[1, 2], [3, 4]], dims=["x", "y"]) + with pytest.raises(ValueError, match="extra dimensions"): + model.add_variables(lower=lower, coords=self.DICT_COORDS, name="x") + + # -- Broadcasting missing dims ----------------------------------------- + + def test_dataarray_broadcast_missing_dim(self, model: "Model") -> None: + 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, "space": 2} + + # -- Special coord formats --------------------------------------------- + + def test_multiindex_coords(self, model: "Model") -> None: + idx = pd.MultiIndex.from_product( + [[1, 2], ["a", "b"]], names=("level1", "level2") + ) + idx.name = "multi" + var = model.add_variables(lower=0, upper=1, coords=[idx], name="x") + assert var.shape == (4,) + + def test_xarray_coordinates_object(self, model: "Model") -> None: + time = pd.RangeIndex(3, name="time") + base = model.add_variables(lower=0, coords=[time], name="base") + lower = DataArray([1, 1, 1], dims=["time"], coords={"time": range(3)}) + var = model.add_variables(lower=lower, coords=base.data.coords, name="x2") + assert var.shape == (3,) def test_as_dataarray_with_unsupported_type() -> None: From ff52cb1cb67c09f59730e24e76aa9c6902f9da4c Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:41:34 +0100 Subject: [PATCH 07/13] Assert broadcast test checks actual values, not NaN Co-Authored-By: Claude Opus 4.6 --- test/test_common.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/test/test_common.py b/test/test_common.py index 86307bd7..ef0387e3 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -482,10 +482,14 @@ def test_dataarray_extra_dims(self, model: "Model") -> None: def test_dataarray_broadcast_missing_dim(self, model: "Model") -> None: time = pd.RangeIndex(3, name="time") space = pd.Index(["a", "b"], name="space") - lower = DataArray([0, 0, 0], dims=["time"], coords={"time": range(3)}) + lower = DataArray([1, 2, 3], 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, "space": 2} + # Verify broadcast filled with actual values, not NaN + assert not var.data.lower.isnull().any() + assert (var.data.lower.sel(space="a") == [1, 2, 3]).all() + assert (var.data.lower.sel(space="b") == [1, 2, 3]).all() # -- Special coord formats --------------------------------------------- From 2d22ad16b18c9c33854b30b3f07efcbae51ea6a9 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:43:36 +0100 Subject: [PATCH 08/13] Add mixed bound type combination and edge case tests Test DataArray+numpy, DataArray+scalar, DataArray+DataArray combos for lower/upper. Also test both bounds covering different dim subsets with broadcast, and that only the mismatched bound raises ValueError. Co-Authored-By: Claude Opus 4.6 --- test/test_common.py | 69 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/test/test_common.py b/test/test_common.py index ef0387e3..d4bfb7ae 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -508,6 +508,75 @@ def test_xarray_coordinates_object(self, model: "Model") -> None: var = model.add_variables(lower=lower, coords=base.data.coords, name="x2") assert var.shape == (3,) + # -- Mixed bound type combinations ------------------------------------ + + @pytest.mark.parametrize( + "lower, upper", + [ + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + np.array([1, 1, 1]), + id="da-lower+numpy-upper", + ), + pytest.param( + np.array([0, 0, 0]), + DataArray([1, 1, 1], dims=["x"], coords={"x": [0, 1, 2]}), + id="numpy-lower+da-upper", + ), + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + DataArray([1, 1, 1], dims=["x"], coords={"x": [0, 1, 2]}), + id="da-lower+da-upper", + ), + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + 10, + id="da-lower+scalar-upper", + ), + pytest.param( + 0, + DataArray([1, 1, 1], dims=["x"], coords={"x": [0, 1, 2]}), + id="scalar-lower+da-upper", + ), + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + xr.DataArray(10), + id="da-lower+scalar-da-upper", + ), + ], + ) + def test_mixed_bound_types(self, model: "Model", lower, upper) -> None: + var = model.add_variables( + lower=lower, upper=upper, coords=self.SEQ_COORDS, name="x" + ) + assert var.shape == (3,) + assert var.dims == ("x",) + assert not var.data.lower.isnull().any() + assert not var.data.upper.isnull().any() + + def test_both_dataarray_different_dim_subsets(self, model: "Model") -> None: + """Lower and upper cover different subsets of dims, both broadcast.""" + time = pd.RangeIndex(3, name="time") + space = pd.Index(["a", "b"], name="space") + lower = DataArray([0, 0, 0], dims=["time"], coords={"time": range(3)}) + upper = DataArray([10, 20], dims=["space"], coords={"space": ["a", "b"]}) + var = model.add_variables( + lower=lower, upper=upper, coords=[time, space], name="x" + ) + assert var.data.sizes == {"time": 3, "space": 2} + assert not var.data.lower.isnull().any() + assert not var.data.upper.isnull().any() + assert (var.data.upper.sel(time=0) == [10, 20]).all() + + def test_one_dataarray_mismatches_other_ok(self, model: "Model") -> None: + """Only the mismatched bound should raise, regardless of the other.""" + lower = DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}) + upper = DataArray([1, 1], dims=["x"], coords={"x": [10, 20]}) + with pytest.raises(ValueError, match="do not match"): + model.add_variables( + lower=lower, upper=upper, coords=self.SEQ_COORDS, name="x" + ) + def test_as_dataarray_with_unsupported_type() -> None: with pytest.raises(TypeError): From 3c29405eb9ed30c862ecba9173cf69806d0e40b7 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:44:53 +0100 Subject: [PATCH 09/13] Add 0-dim bound types and fix numpy_to_dataarray with dict coords Add numpy-0d and dataarray-0d to the parameterized bound type tests. Fix numpy_to_dataarray to infer dims from dict keys for 0-dim arrays, matching the scalar fix in as_dataarray. Co-Authored-By: Claude Opus 4.6 --- linopy/common.py | 2 ++ test/test_common.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/linopy/common.py b/linopy/common.py index 7fb76288..b9e36024 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -191,6 +191,8 @@ def numpy_to_dataarray( """ # fallback case for zero dim arrays if arr.ndim == 0: + if dims is None and is_dict_like(coords): + dims = list(coords.keys()) return DataArray(arr.item(), coords=coords, dims=dims, **kwargs) if isinstance(dims, Iterable | Sequence): diff --git a/test/test_common.py b/test/test_common.py index d4bfb7ae..93e1b2d6 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -428,7 +428,8 @@ def model(self) -> "Model": [ pytest.param(0, id="scalar"), pytest.param(np.float64(0), id="np.number"), - pytest.param(np.array([0, 0, 0]), id="numpy"), + pytest.param(np.array(0), id="numpy-0d"), + pytest.param(np.array([0, 0, 0]), id="numpy-1d"), pytest.param( pd.Series([0, 0, 0], index=pd.RangeIndex(3, name="x")), id="pandas" ), @@ -438,6 +439,7 @@ def model(self) -> "Model": id="dataarray", ), pytest.param(DataArray([0, 0, 0], dims=["x"]), id="dataarray-no-coords"), + pytest.param(xr.DataArray(0), id="dataarray-0d"), ], ) @pytest.mark.parametrize( From b76c4c3479a9e69a8c87f8263da70088b7860a22 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:46:49 +0100 Subject: [PATCH 10/13] Add tests for inferred coords, multi-dim, string/datetime coords Cover three gaps: coords inferred from bounds (no coords arg) for DataArray and pandas, multi-dimensional coord specifications with both scalar and DataArray bounds, and real-world coordinate types (string regions, datetime index) including mismatch detection. Co-Authored-By: Claude Opus 4.6 --- test/test_common.py | 95 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/test/test_common.py b/test/test_common.py index 93e1b2d6..03b4f528 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -579,6 +579,101 @@ def test_one_dataarray_mismatches_other_ok(self, model: "Model") -> None: lower=lower, upper=upper, coords=self.SEQ_COORDS, name="x" ) + # -- Coords inferred from bounds (no coords arg) ---------------------- + + @pytest.mark.parametrize( + "lower", + [ + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [10, 20, 30]}), + id="dataarray", + ), + pytest.param( + pd.Series([0, 0, 0], index=pd.Index([10, 20, 30], name="x")), + id="pandas", + ), + ], + ) + def test_coords_inferred_from_bounds(self, model: "Model", lower) -> None: + """When coords is None, dims/coords are inferred from the bounds.""" + var = model.add_variables(lower=lower, name="x") + assert var.dims == ("x",) + assert list(var.data.coords["x"].values) == [10, 20, 30] + + def test_coords_inferred_multidim(self, model: "Model") -> None: + lower = DataArray( + np.zeros((3, 2)), + dims=["time", "space"], + coords={"time": [0, 1, 2], "space": ["a", "b"]}, + ) + var = model.add_variables(lower=lower, name="x") + assert set(var.dims) == {"time", "space"} + assert var.data.sizes == {"time": 3, "space": 2} + + # -- Multi-dimensional coords ----------------------------------------- + + @pytest.mark.parametrize( + "coords", + [ + pytest.param( + [pd.RangeIndex(3, name="time"), pd.Index(["a", "b"], name="space")], + id="seq-coords", + ), + pytest.param( + {"time": [0, 1, 2], "space": ["a", "b"]}, + id="dict-coords", + ), + ], + ) + def test_multidim_coords_with_scalar(self, model: "Model", coords) -> None: + var = model.add_variables(lower=0, upper=1, coords=coords, name="x") + assert set(var.dims) == {"time", "space"} + assert var.data.sizes == {"time": 3, "space": 2} + + def test_multidim_dataarray_with_coords(self, model: "Model") -> None: + lower = DataArray( + np.zeros((3, 2)), + dims=["time", "space"], + coords={"time": [0, 1, 2], "space": ["a", "b"]}, + ) + coords = [pd.RangeIndex(3, name="time"), pd.Index(["a", "b"], name="space")] + var = model.add_variables(lower=lower, coords=coords, name="x") + assert set(var.dims) == {"time", "space"} + assert var.data.sizes == {"time": 3, "space": 2} + assert not var.data.lower.isnull().any() + + # -- String and datetime coordinates ----------------------------------- + + def test_string_coordinates(self, model: "Model") -> None: + coords = {"region": ["north", "south", "east"]} + lower = DataArray( + [0, 0, 0], + dims=["region"], + coords={"region": ["north", "south", "east"]}, + ) + var = model.add_variables(lower=lower, coords=coords, name="x") + assert var.dims == ("region",) + assert list(var.data.coords["region"].values) == ["north", "south", "east"] + + def test_datetime_coordinates(self, model: "Model") -> None: + dates = pd.date_range("2025-01-01", periods=3) + coords = [dates.rename("time")] + lower = DataArray([0, 0, 0], dims=["time"], coords={"time": dates}) + var = model.add_variables(lower=lower, coords=coords, name="x") + assert var.dims == ("time",) + assert var.shape == (3,) + + def test_string_coords_mismatch(self, model: "Model") -> None: + lower = DataArray( + [0, 0], dims=["region"], coords={"region": ["north", "south"]} + ) + with pytest.raises(ValueError, match="do not match"): + model.add_variables( + lower=lower, + coords={"region": ["north", "south", "east"]}, + name="x", + ) + def test_as_dataarray_with_unsupported_type() -> None: with pytest.raises(TypeError): From c7a0b8458dc53470c565699527e879c64f696d66 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:49:09 +0100 Subject: [PATCH 11/13] Add test for bounds with different dimension order Verify lower(time, space) and upper(space, time) align correctly via xarray broadcast. Co-Authored-By: Claude Opus 4.6 --- test/test_common.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/test/test_common.py b/test/test_common.py index 03b4f528..3207ac09 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -642,6 +642,27 @@ def test_multidim_dataarray_with_coords(self, model: "Model") -> None: assert var.data.sizes == {"time": 3, "space": 2} assert not var.data.lower.isnull().any() + def test_bounds_with_different_dim_order(self, model: "Model") -> None: + """Lower (time, space) and upper (space, time) should align correctly.""" + time = pd.RangeIndex(3, name="time") + space = pd.Index(["a", "b"], name="space") + lower = DataArray( + np.zeros((3, 2)), + dims=["time", "space"], + coords={"time": range(3), "space": ["a", "b"]}, + ) + upper = DataArray( + np.ones((2, 3)), + dims=["space", "time"], + coords={"space": ["a", "b"], "time": range(3)}, + ) + var = model.add_variables( + lower=lower, upper=upper, coords=[time, space], name="x" + ) + assert var.data.sizes == {"time": 3, "space": 2} + assert (var.data.lower.values == 0).all() + assert (var.data.upper.values == 1).all() + # -- String and datetime coordinates ----------------------------------- def test_string_coordinates(self, model: "Model") -> None: From 308a0ed1fe2a9b92700a0238716f9e47d6cbdd22 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:51:03 +0100 Subject: [PATCH 12/13] Reindex DataArray bounds with reordered coordinates When a DataArray bound has the same coordinate values as coords but in a different order, reindex to match instead of raising ValueError. Still raises when the values actually differ (not just reordered). Co-Authored-By: Claude Opus 4.6 --- linopy/model.py | 14 ++++++++++---- test/test_common.py | 16 ++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/linopy/model.py b/linopy/model.py index bb6b1ac3..462f9cb1 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -145,10 +145,16 @@ def _validate_dataarray_bounds(arr: Any, coords: Any) -> Any: ) actual_idx = arr.coords[dim].to_index() if not actual_idx.equals(expected_idx): - raise ValueError( - f"Coordinates for dimension '{dim}' do not match: " - f"expected {expected_idx.tolist()}, got {actual_idx.tolist()}" - ) + # Same values, different order → reindex to match expected order + if len(actual_idx) == len(expected_idx) and set(actual_idx) == set( + expected_idx + ): + arr = arr.reindex({dim: expected_idx}) + else: + raise ValueError( + f"Coordinates for dimension '{dim}' do not match: " + f"expected {expected_idx.tolist()}, got {actual_idx.tolist()}" + ) # Expand missing dimensions expand = {k: v for k, v in expected.items() if k not in arr.dims} diff --git a/test/test_common.py b/test/test_common.py index 3207ac09..e333bd50 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -663,6 +663,22 @@ def test_bounds_with_different_dim_order(self, model: "Model") -> None: assert (var.data.lower.values == 0).all() assert (var.data.upper.values == 1).all() + # -- Reordered coordinates --------------------------------------------- + + def test_reordered_coords_reindexed(self, model: "Model") -> None: + """Same coord values in different order should reindex, not raise.""" + lower = DataArray([10, 20, 30], dims=["x"], coords={"x": ["c", "a", "b"]}) + var = model.add_variables(lower=lower, coords={"x": ["a", "b", "c"]}, name="x") + assert list(var.data.coords["x"].values) == ["a", "b", "c"] + # Values must follow the reindexed order, not the original + assert list(var.data.lower.values) == [20, 30, 10] + + def test_reordered_coords_different_values_raises(self, model: "Model") -> None: + """Overlapping but not identical coord sets must still raise.""" + lower = DataArray([10, 20], dims=["x"], coords={"x": ["a", "b"]}) + with pytest.raises(ValueError, match="do not match"): + model.add_variables(lower=lower, coords={"x": ["a", "c"]}, name="x") + # -- String and datetime coordinates ----------------------------------- def test_string_coordinates(self, model: "Model") -> None: From f53d1968e53d70962eb208b01c3f938f693f0083 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 13 Mar 2026 10:11:06 +0100 Subject: [PATCH 13/13] Fix mypy errors: remove dead code branch, add type annotations Remove unreachable hasattr(coords, "dims") branch in _coords_to_dict (xarray Coordinates are Mappings, caught by isinstance check above). Add Any type annotations to parameterized test arguments. Co-Authored-By: Claude Opus 4.6 --- linopy/model.py | 3 --- test/test_common.py | 14 +++++++++----- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/linopy/model.py b/linopy/model.py index 462f9cb1..d8c86bd4 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -100,9 +100,6 @@ def _coords_to_dict( """Normalize coords to a dict mapping dim names to coordinate values.""" if isinstance(coords, Mapping): return dict(coords) - # xarray Coordinates objects - if hasattr(coords, "dims"): - return {k: v for k, v in coords.items() if k in coords.dims} # Sequence of indexes result: dict[str, Any] = {} for c in coords: diff --git a/test/test_common.py b/test/test_common.py index e333bd50..c85cfcc5 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -5,6 +5,8 @@ @author: fabian """ +from typing import Any + import numpy as np import pandas as pd import polars as pl @@ -449,7 +451,9 @@ def model(self) -> "Model": pytest.param({"x": [0, 1, 2]}, id="dict-coords"), ], ) - def test_bound_types_with_coords(self, model: "Model", lower, coords) -> None: + def test_bound_types_with_coords( + self, model: "Model", lower: Any, coords: Any + ) -> None: var = model.add_variables(lower=lower, coords=coords, name="x") assert var.shape == (3,) assert var.dims == ("x",) @@ -464,7 +468,7 @@ def test_bound_types_with_coords(self, model: "Model", lower, coords) -> None: pytest.param({"x": [0, 1, 2, 3, 4]}, id="dict-coords"), ], ) - def test_dataarray_coord_mismatch(self, model: "Model", coords) -> None: + def test_dataarray_coord_mismatch(self, model: "Model", coords: Any) -> None: lower = DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}) with pytest.raises(ValueError, match="do not match"): model.add_variables(lower=lower, coords=coords, name="x") @@ -547,7 +551,7 @@ def test_xarray_coordinates_object(self, model: "Model") -> None: ), ], ) - def test_mixed_bound_types(self, model: "Model", lower, upper) -> None: + def test_mixed_bound_types(self, model: "Model", lower: Any, upper: Any) -> None: var = model.add_variables( lower=lower, upper=upper, coords=self.SEQ_COORDS, name="x" ) @@ -594,7 +598,7 @@ def test_one_dataarray_mismatches_other_ok(self, model: "Model") -> None: ), ], ) - def test_coords_inferred_from_bounds(self, model: "Model", lower) -> None: + def test_coords_inferred_from_bounds(self, model: "Model", lower: Any) -> None: """When coords is None, dims/coords are inferred from the bounds.""" var = model.add_variables(lower=lower, name="x") assert var.dims == ("x",) @@ -625,7 +629,7 @@ def test_coords_inferred_multidim(self, model: "Model") -> None: ), ], ) - def test_multidim_coords_with_scalar(self, model: "Model", coords) -> None: + def test_multidim_coords_with_scalar(self, model: "Model", coords: Any) -> None: var = model.add_variables(lower=0, upper=1, coords=coords, name="x") assert set(var.dims) == {"time", "space"} assert var.data.sizes == {"time": 3, "space": 2}