diff --git a/linopy/common.py b/linopy/common.py index 09f67355..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): @@ -242,8 +244,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 54334411..d8c86bd4 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -94,6 +94,73 @@ 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) + # 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): + # 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} + if expand: + arr = arr.expand_dims(expand) + + return arr + + class Model: """ Linear optimization model. @@ -603,6 +670,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..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 @@ -409,6 +411,311 @@ def test_as_dataarray_with_dataarray_default_dims_coords() -> None: assert list(da_out.coords["dim2"].values) == list(da_in.coords["dim2"].values) +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]} + + @pytest.fixture() + def model(self) -> "Model": + from linopy import Model + + return Model() + + # -- All bound types should work with both coord formats --------------- + + @pytest.mark.parametrize( + "lower", + [ + pytest.param(0, id="scalar"), + pytest.param(np.float64(0), id="np.number"), + 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" + ), + 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"), + pytest.param(xr.DataArray(0), id="dataarray-0d"), + ], + ) + @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: Any, coords: Any + ) -> 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] + + # -- DataArray validation: mismatch and extra dims --------------------- + + @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"), + ], + ) + 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") + + 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([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 --------------------------------------------- + + 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,) + + # -- 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: Any, upper: Any) -> 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" + ) + + # -- 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: 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",) + 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: 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} + + 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() + + 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() + + # -- 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: + 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): as_dataarray(lambda x: 1, dims=["dim1"], coords=[["a"]])