fix: apply coords-as-truth rule to mask in add_variables/add_constraints

doc/release_notes.rst

@@ -53,6 +53,7 @@ Most users should keep calling ``model.solve(...)``. If you want more control, y
 **Bug Fixes**
 
 * ``Model.add_variables``: 0.7.0 made ``coords`` (dims, order, and values) the source of truth for ``DataArray`` bounds; this release closes the two remaining gaps. Pandas ``Series`` / ``DataFrame`` bounds missing a dimension are broadcast to ``coords`` instead of being silently dropped (`#709 <https://github.com/PyPSA/linopy/issues/709>`__), and the variable's dimension order always follows ``coords`` regardless of bound type (`#706 <https://github.com/PyPSA/linopy/issues/706>`__).
+* ``add_variables`` / ``add_constraints``: the same rule now applies to ``mask`` — pandas ``Series`` / ``DataFrame`` masks missing a dimension are broadcast to the variable/constraint shape. As previously announced via ``FutureWarning``, masks whose coordinates are a sparse subset of the data's coordinates now raise ``ValueError`` rather than silently filling missing entries with ``False``; masks with dimensions not in the data raise ``ValueError`` instead of ``AssertionError``.
 * ``add_piecewise_formulation`` now produces a reproducible dimension order in the broadcast breakpoint array. The previous set-based expansion gave a hash-randomized order that varied between processes.
 * SOS constraints on masked variables no longer cause solver-specific failures (Gurobi ``IndexError``, Xpress ``?404 Invalid column number``, LP parse errors, silent set corruption). ``Model.solve()`` and ``Model.to_file()`` now raise a clear ``NotImplementedError`` referring users to `#688 <https://github.com/PyPSA/linopy/issues/688>`__; pass ``reformulate_sos=True`` as a workaround.
 * ``Model.solve(..., reformulate_sos=True)`` now actually reformulates SOS constraints even when the solver supports them natively. Previously it was silently ignored with a warning.

linopy/model.py

@@ -20,7 +20,7 @@
 import pandas as pd
 import xarray as xr
 from deprecation import deprecated
-from numpy import inf, ndarray
+from numpy import inf
 from pandas.core.frame import DataFrame
 from pandas.core.series import Series
 from xarray import DataArray, Dataset
@@ -32,7 +32,6 @@
     as_dataarray_in_coords,
     assign_multiindex_safe,
     best_int,
-    broadcast_mask,
     maybe_replace_signs,
     replace_by_map,
     to_path,
@@ -593,7 +592,7 @@ def add_variables(
         upper: Any = inf,
         coords: Sequence[Sequence | pd.Index] | Mapping | None = None,
         name: str | None = None,
-        mask: DataArray | ndarray | Series | None = None,
+        mask: MaskLike | None = None,
         binary: bool = False,
         integer: bool = False,
         semi_continuous: bool = False,
@@ -713,8 +712,7 @@ 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 = broadcast_mask(mask, data.labels)
+            mask = as_dataarray_in_coords(mask, data.coords).astype(bool)
 
         # Auto-mask based on NaN in bounds (use numpy for speed)
         if self.auto_mask:
@@ -978,8 +976,7 @@ def add_constraints(
         (data,) = xr.broadcast(data, exclude=[TERM_DIM])
 
         if mask is not None:
-            mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool)
-            mask = broadcast_mask(mask, data.labels)
+            mask = as_dataarray_in_coords(mask, data.coords).astype(bool)
 
         # Auto-mask based on null expressions or NaN RHS (use numpy for speed)
         if self.auto_mask:

test/test_constraints.py

@@ -258,20 +258,27 @@ def test_masked_constraints_broadcast() -> None:
     assert (m.constraints.labels.bc2[:, 0:5] != -1).all()
     assert (m.constraints.labels.bc2[:, 5:10] == -1).all()
 
+    # Pandas Series with named index missing a dim is broadcast to data.coords.
+    mask_pd = pd.Series(
+        [True, False, True] + [False] * 7, index=pd.RangeIndex(10, name="dim_0")
+    )
+    m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc_pd", mask=mask_pd)
+    assert (m.constraints.labels.bc_pd[[0, 2], :] != -1).all()
+    assert (m.constraints.labels.bc_pd[[1, 3, 4, 5, 6, 7, 8, 9], :] == -1).all()
+
+    # Mask with sparse coords (subset of data's coords) now raises instead of
+    # emitting a FutureWarning — the rule from the bounds path applies here too.
     mask3 = xr.DataArray(
         [True, True, False, False, False],
         dims=["dim_0"],
         coords={"dim_0": range(5)},
     )
-    with pytest.warns(FutureWarning, match="Missing values will be filled"):
+    with pytest.raises(ValueError, match="Coordinates for dimension 'dim_0'"):
         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"):
+    with pytest.raises(ValueError, match="extra dimensions"):
         m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc4", mask=mask4)
 
 

test/test_variables.py

@@ -123,20 +123,27 @@ def test_variables_mask_broadcast() -> None:
     assert (y.labels[:, 0:5] != -1).all()
     assert (y.labels[:, 5:10] == -1).all()
 
+    # Pandas Series with named index missing a dim is broadcast to data.coords.
+    mask_pd = pd.Series(
+        [True, False, True] + [False] * 7, index=pd.RangeIndex(10, name="dim_0")
+    )
+    v = m.add_variables(lower, upper, name="v", mask=mask_pd)
+    assert (v.labels[[0, 2], :] != -1).all()
+    assert (v.labels[[1, 3, 4, 5, 6, 7, 8, 9], :] == -1).all()
+
+    # Mask with sparse coords (subset of data's coords) now raises instead of
+    # emitting a FutureWarning — the rule from the bounds path applies here too.
     mask3 = xr.DataArray(
         [True, True, False, False, False],
         dims=["dim_0"],
         coords={"dim_0": range(5)},
     )
-    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()
+    with pytest.raises(ValueError, match="Coordinates for dimension 'dim_0'"):
+        m.add_variables(lower, upper, name="z", mask=mask3)
 
     # 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"):
+    with pytest.raises(ValueError, match="extra dimensions"):
         m.add_variables(lower, upper, name="w", mask=mask4)