diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 22c0ffee..7ee6458c 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -9,6 +9,10 @@ Upcoming Version * Improved constraint equality check in `linopy.testing.assert_conequal` to less strict optionally * Minor bugfix for multiplying variables with numpy type constants +**Breaking Changes** +* With this release the behaviour of `m.add_variables` has been changed so that provided coords will always be +respected. Previously data arrays used as lower/upper bounds would override the coords. + Version 0.5.6 -------------- diff --git a/linopy/common.py b/linopy/common.py index 6804ac1e..05f940a8 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -229,6 +229,7 @@ def as_dataarray( arr: Any, coords: CoordsLike | None = None, dims: DimsLike | None = None, + force_broadcast: bool = False, **kwargs: Any, ) -> DataArray: """ @@ -240,8 +241,12 @@ def as_dataarray( The input object. coords (Union[dict, list, None]): The coordinates for the DataArray. If None, default coordinates will be used. + If this are set, constant data will be broadcast to these coordinates. Pandas and xarray type data will not + be broadcast unless force_broadcast is set to True. dims (Union[list, None]): The dimensions for the DataArray. If None, the dimensions will be automatically generated. + force_broadcast (bool): + Ensures that data is broadcast to the given coordinates for any provided arr type. **kwargs: Additional keyword arguments to be passed to the DataArray constructor. @@ -276,6 +281,11 @@ def as_dataarray( ) arr = fill_missing_coords(arr) + + if force_broadcast and coords is not None: + ones = DataArray(data=1.0, coords=coords) + return arr * ones + return arr diff --git a/linopy/model.py b/linopy/model.py index 3982b84d..9af2feba 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -520,8 +520,12 @@ def add_variables( data = Dataset( { - "lower": as_dataarray(lower, coords, **kwargs), - "upper": as_dataarray(upper, coords, **kwargs), + "lower": as_dataarray( + arr=lower, coords=coords, force_broadcast=True, **kwargs + ), + "upper": as_dataarray( + arr=upper, coords=coords, force_broadcast=True, **kwargs + ), "labels": -1, } ) @@ -530,7 +534,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 = as_dataarray( + arr=mask, coords=data.coords, force_broadcast=True, dims=data.dims + ).astype(bool) start = self._xCounter end = start + data.labels.size diff --git a/test/test_common.py b/test/test_common.py index 85059487..41656e90 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -104,6 +104,20 @@ def test_as_dataarray_with_series_override_coords() -> None: assert list(da.coords[target_dim].values) == target_index +def test_as_dataarray_with_series_force_broadcast() -> None: + base = xr.DataArray(np.ones((3, 2)), [("x", ["a", "b", "c"]), ("y", ["M", "N"])]) + x_only = base.sum("y") + + series_index = pd.Index(data=["a", "b", "c"], name="x") + series = pd.Series(index=series_index, data=1.0) + + da_force = as_dataarray(arr=series, coords=base.coords, force_broadcast=True) + assert da_force.coords.to_dataset().equals(base.coords.to_dataset()) + + da_no_force = as_dataarray(arr=series, coords=base.coords, force_broadcast=False) + assert da_no_force.coords.to_dataset().equals(x_only.coords.to_dataset()) + + def test_as_dataarray_with_series_aligned_coords() -> None: """This should not give out a warning even though coords are given.""" target_dim = "dim_0" @@ -120,6 +134,17 @@ def test_as_dataarray_with_series_aligned_coords() -> None: assert list(da.coords[target_dim].values) == target_index +def test_as_dataarray_with_datarray_force_broadcast() -> None: + base = xr.DataArray(np.ones((3, 2)), [("x", ["a", "b", "c"]), ("y", ["M", "N"])]) + x_only = base.sum("y") + + da_force = as_dataarray(arr=x_only, coords=base.coords, force_broadcast=True) + assert da_force.coords.to_dataset().equals(base.coords.to_dataset()) + + da_no_force = as_dataarray(arr=x_only, coords=base.coords, force_broadcast=False) + assert da_no_force.coords.to_dataset().equals(x_only.coords.to_dataset()) + + def test_as_dataarray_dataframe_dims_default() -> None: target_dims = ("dim_0", "dim_1") target_index = [0, 1] diff --git a/test/test_variables.py b/test/test_variables.py index 3984b091..b525ce76 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -9,11 +9,12 @@ import xarray as xr import xarray.core.indexes import xarray.core.utils +from xarray import Coordinates import linopy from linopy import Model from linopy.testing import assert_varequal -from linopy.variables import ScalarVariable +from linopy.variables import ScalarVariable, Variable @pytest.fixture @@ -122,3 +123,27 @@ def test_scalar_variable(m: Model) -> None: x = ScalarVariable(label=0, model=m) assert isinstance(x, ScalarVariable) assert x.__rmul__(x) is NotImplemented # type: ignore + + +def assert_coords_identical(a: Coordinates, b: Coordinates) -> None: + assert a.to_dataset().equals(b.to_dataset()) + + +def test_variable_coordinates(m: Model) -> None: + m = linopy.Model() + base = xr.DataArray(np.ones((3, 2)), [("x", ["a", "b", "c"]), ("y", ["X", "Y"])]) + + foo = m.add_variables(lower=0, upper=base.sum("y"), coords=base.coords, name="foo") + assert_coords_identical(foo.coords, base.coords) + + bar = m.add_variables( + lower=-base.sum("y"), upper=base.sum("y"), coords=base.coords, name="bar" + ) + assert_coords_identical(bar.coords, base.coords) + + my_dim = pd.RangeIndex(2, name="my-dim") + var = m.add_variables( + lower=xr.DataArray(0), upper=xr.DataArray(10), coords=[my_dim] + ) + assert isinstance(var, Variable) + assert var.shape == (2,)