Fixed strange behaviour when creating variables with conflicting coordinates

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
 --------------
 

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
 
 

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

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]

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,)