Reinsert broadcasted mask (#580)

* reinsert broadcasting of masks

* update release notes

* consolidate broadcast mask into new function, add tests for subsets

* align test logic to broadcasting

* Reinsert broadcasted mask (#581)

* 1. Moved the dimension subset check into broadcast_mask
2. Added a brief docstring to broadcast_mask

* Add tests for superset dims

---------

Co-authored-by: FBumann <117816358+FBumann@users.noreply.github.com>

Author: FabianHofmann

doc/release_notes.rst

@@ -4,6 +4,11 @@ Release Notes
 Upcoming Version
 ----------------
 
+**Fix Regression**
+
+*  Reinsert broadcasting logic of mask object to be fully compatible with performance improvements in version 0.6.2 using `np.where` instead of `xr.where`.
+
+
 Version 0.6.2
 --------------
 

linopy/common.py

@@ -286,6 +286,32 @@ def as_dataarray(
     return arr
 
 
+def broadcast_mask(mask: DataArray, labels: DataArray) -> DataArray:
+    """
+    Broadcast a boolean mask to match the shape of labels.
+
+    Ensures that mask dimensions are a subset of labels dimensions, broadcasts
+    the mask accordingly, and fills any NaN values (from missing coordinates)
+    with False while emitting a FutureWarning.
+    """
+    assert set(mask.dims).issubset(labels.dims), (
+        "Dimensions of mask not a subset of resulting labels dimensions."
+    )
+    mask = mask.broadcast_like(labels)
+    if mask.isnull().any():
+        warn(
+            "Mask contains coordinates not covered by the data dimensions. "
+            "Missing values will be filled with False (masked out). "
+            "In a future version, this will raise an error. "
+            "Use mask.reindex() or `linopy.align()` to explicitly handle missing "
+            "coordinates.",
+            FutureWarning,
+            stacklevel=3,
+        )
+        mask = mask.fillna(False).astype(bool)
+    return mask
+
+
 # TODO: rename to to_pandas_dataframe
 def to_dataframe(
     ds: Dataset,

linopy/model.py

@@ -9,7 +9,6 @@
 import logging
 import os
 import re
-import warnings
 from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
 from tempfile import NamedTemporaryFile, gettempdir
@@ -30,6 +29,7 @@
     as_dataarray,
     assign_multiindex_safe,
     best_int,
+    broadcast_mask,
     maybe_replace_signs,
     replace_by_map,
     set_int_index,
@@ -552,16 +552,7 @@ def add_variables(
 
         if mask is not None:
             mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool)
-            if set(mask.dims) != set(data["labels"].dims):
-                warnings.warn(
-                    f"Mask dimensions {set(mask.dims)} do not match the data "
-                    f"dimensions {set(data['labels'].dims)}. The mask will be "
-                    f"broadcast across the missing dimensions "
-                    f"{set(data['labels'].dims) - set(mask.dims)}. In a future "
-                    "version, this will raise an error.",
-                    FutureWarning,
-                    stacklevel=2,
-                )
+            mask = broadcast_mask(mask, data.labels)
 
         # Auto-mask based on NaN in bounds (use numpy for speed)
         if self.auto_mask:
@@ -582,7 +573,7 @@ def add_variables(
         self._xCounter += data.labels.size
 
         if mask is not None:
-            data.labels.values = data.labels.where(mask, -1).values
+            data.labels.values = np.where(mask.values, data.labels.values, -1)
 
         data = data.assign_attrs(
             label_range=(start, end), name=name, binary=binary, integer=integer
@@ -756,20 +747,7 @@ def add_constraints(
 
         if mask is not None:
             mask = as_dataarray(mask).astype(bool)
-            # TODO: simplify
-            assert set(mask.dims).issubset(data.dims), (
-                "Dimensions of mask not a subset of resulting labels dimensions."
-            )
-            if set(mask.dims) != set(data["labels"].dims):
-                warnings.warn(
-                    f"Mask dimensions {set(mask.dims)} do not match the data "
-                    f"dimensions {set(data['labels'].dims)}. The mask will be "
-                    f"broadcast across the missing dimensions "
-                    f"{set(data['labels'].dims) - set(mask.dims)}. In a future "
-                    "version, this will raise an error.",
-                    FutureWarning,
-                    stacklevel=2,
-                )
+            mask = broadcast_mask(mask, data.labels)
 
         # Auto-mask based on null expressions or NaN RHS (use numpy for speed)
         if self.auto_mask:
@@ -780,11 +758,9 @@ def add_constraints(
             auto_mask_values = ~vars_all_invalid
             if original_rhs_mask is not None:
                 coords, dims, rhs_notnull = original_rhs_mask
-                # Broadcast RHS mask to match data shape if needed
                 if rhs_notnull.shape != auto_mask_values.shape:
                     rhs_da = DataArray(rhs_notnull, coords=coords, dims=dims)
-                    rhs_da, _ = xr.broadcast(rhs_da, data.labels)
-                    rhs_notnull = rhs_da.values
+                    rhs_notnull = rhs_da.broadcast_like(data.labels).values
                 auto_mask_values = auto_mask_values & rhs_notnull
             auto_mask_arr = DataArray(
                 auto_mask_values, coords=data.labels.coords, dims=data.labels.dims
@@ -802,7 +778,7 @@ def add_constraints(
         self._cCounter += data.labels.size
 
         if mask is not None:
-            data.labels.values = data.labels.where(mask, -1).values
+            data.labels.values = np.where(mask.values, data.labels.values, -1)
 
         data = data.assign_attrs(label_range=(start, end), name=name)
 

test/test_constraints.py

@@ -157,54 +157,44 @@ def test_masked_constraints() -> None:
     y = m.add_variables()
 
     mask = pd.Series([True] * 5 + [False] * 5)
-    with pytest.warns(FutureWarning, match="Mask dimensions"):
-        m.add_constraints(1 * x + 10 * y, EQUAL, 0, mask=mask)
+    m.add_constraints(1 * x + 10 * y, EQUAL, 0, mask=mask)
     assert (m.constraints.labels.con0[0:5, :] != -1).all()
     assert (m.constraints.labels.con0[5:10, :] == -1).all()
 
 
 def test_masked_constraints_broadcast() -> None:
-    """Test that a constraint mask with fewer dimensions broadcasts correctly."""
     m: Model = Model()
 
     lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)])
     upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)])
     x = m.add_variables(lower, upper)
     y = m.add_variables()
 
-    # 1D mask applied to 2D constraint — must broadcast over second dim
     mask = pd.Series([True] * 5 + [False] * 5)
-    with pytest.warns(FutureWarning, match="Mask dimensions"):
-        m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc1", mask=mask)
+    m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc1", mask=mask)
     assert (m.constraints.labels.bc1[0:5, :] != -1).all()
     assert (m.constraints.labels.bc1[5:10, :] == -1).all()
 
-    # Mask along second dimension only
     mask2 = xr.DataArray([True] * 5 + [False] * 5, dims=["dim_1"])
-    with pytest.warns(FutureWarning, match="Mask dimensions"):
-        m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc2", mask=mask2)
+    m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc2", mask=mask2)
     assert (m.constraints.labels.bc2[:, 0:5] != -1).all()
     assert (m.constraints.labels.bc2[:, 5:10] == -1).all()
 
-
-def test_constraints_mask_no_warning_when_aligned() -> None:
-    """Test that no FutureWarning is emitted when mask has same dims as data."""
-    m: Model = Model()
-
-    lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)])
-    upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)])
-    x = m.add_variables(lower, upper)
-    y = m.add_variables()
-
-    mask = xr.DataArray(
-        np.array([[True] * 10] * 5 + [[False] * 10] * 5),
-        coords=[range(10), range(10)],
+    mask3 = xr.DataArray(
+        [True, True, False, False, False],
+        dims=["dim_0"],
+        coords={"dim_0": range(5)},
     )
-    import warnings
-
-    with warnings.catch_warnings():
-        warnings.simplefilter("error", FutureWarning)
-        m.add_constraints(1 * x + 10 * y, EQUAL, 0, mask=mask)
+    with pytest.warns(FutureWarning, match="Missing values will be filled"):
+        m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc3", mask=mask3)
+    assert (m.constraints.labels.bc3[0:2, :] != -1).all()
+    assert (m.constraints.labels.bc3[2:5, :] == -1).all()
+    assert (m.constraints.labels.bc3[5:10, :] == -1).all()
+
+    # Mask with extra dimension not in data should raise
+    mask4 = xr.DataArray([True, False], dims=["extra_dim"])
+    with pytest.raises(AssertionError, match="not a subset"):
+        m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc4", mask=mask4)
 
 
 def test_non_aligned_constraints() -> None:

test/test_variable_assignment.py

@@ -227,8 +227,7 @@ def test_variable_assigment_masked() -> None:
     lower = pd.DataFrame(np.zeros((10, 10)))
     upper = pd.Series(np.ones(10))
     mask = pd.Series([True] * 5 + [False] * 5)
-    with pytest.warns(FutureWarning, match="Mask dimensions"):
-        m.add_variables(lower, upper, mask=mask)
+    m.add_variables(lower, upper, mask=mask)
     assert m.variables.labels.var0[-1, -1].item() == -1
 
 

test/test_variables.py

@@ -108,43 +108,36 @@ def test_variables_nvars(m: Model) -> None:
 
 
 def test_variables_mask_broadcast() -> None:
-    """Test that a mask with fewer dimensions broadcasts correctly."""
     m = Model()
 
     lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)])
     upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)])
 
-    # 1D mask applied to 2D variable — must broadcast over second dim
     mask = pd.Series([True] * 5 + [False] * 5)
-    with pytest.warns(FutureWarning, match="Mask dimensions"):
-        x = m.add_variables(lower, upper, name="x", mask=mask)
+    x = m.add_variables(lower, upper, name="x", mask=mask)
     assert (x.labels[0:5, :] != -1).all()
     assert (x.labels[5:10, :] == -1).all()
 
-    # Mask along second dimension only
     mask2 = xr.DataArray([True] * 5 + [False] * 5, dims=["dim_1"])
-    with pytest.warns(FutureWarning, match="Mask dimensions"):
-        y = m.add_variables(lower, upper, name="y", mask=mask2)
+    y = m.add_variables(lower, upper, name="y", mask=mask2)
     assert (y.labels[:, 0:5] != -1).all()
     assert (y.labels[:, 5:10] == -1).all()
 
-
-def test_variables_mask_no_warning_when_aligned() -> None:
-    """Test that no FutureWarning is emitted when mask has same dims as data."""
-    m = Model()
-
-    lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)])
-    upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)])
-
-    mask = xr.DataArray(
-        np.array([[True] * 10] * 5 + [[False] * 10] * 5),
-        coords=[range(10), range(10)],
+    mask3 = xr.DataArray(
+        [True, True, False, False, False],
+        dims=["dim_0"],
+        coords={"dim_0": range(5)},
     )
-    import warnings
-
-    with warnings.catch_warnings():
-        warnings.simplefilter("error", FutureWarning)
-        m.add_variables(lower, upper, name="x", mask=mask)
+    with pytest.warns(FutureWarning, match="Missing values will be filled"):
+        z = m.add_variables(lower, upper, name="z", mask=mask3)
+    assert (z.labels[0:2, :] != -1).all()
+    assert (z.labels[2:5, :] == -1).all()
+    assert (z.labels[5:10, :] == -1).all()
+
+    # Mask with extra dimension not in data should raise
+    mask4 = xr.DataArray([True, False], dims=["extra_dim"])
+    with pytest.raises(AssertionError, match="not a subset"):
+        m.add_variables(lower, upper, name="w", mask=mask4)
 
 
 def test_variables_get_name_by_label(m: Model) -> None: