fix: ignore masked variable labels when matching constraints

doc/release_notes.rst

@@ -11,6 +11,11 @@ Upcoming Version
 * ``Model.add_expressions`` registers a ``LinearExpression`` or ``QuadraticExpression`` under a name (auto-generated as ``expr0``, ``expr1``, ... if omitted), accessible afterwards via ``Model.expressions`` (an ``Expressions`` container mirroring ``Model.variables``/``Model.constraints``) and removable via ``Model.remove_expressions``.
   Named expressions are persisted by ``Model.to_netcdf``/``linopy.read_netcdf`` and preserved by ``Model.copy``, ``copy.copy``, ``copy.deepcopy``, and pickling.
 
+**Bug fixes**
+
+* ``Model.remove_variables`` no longer removes constraints that never reference the removed variable. A masked variable carries ``-1`` label entries, which matched the ``-1`` that marks an empty term slot in a constraint, so any masked variable looked like it was used by any constraint with padded terms. Models built with ``mask=`` could silently lose constraints and solve to a wrong optimum. (`#883 <https://github.com/PyPSA/linopy/issues/883>`__)
+* ``Model.remove_variables`` now also works on a model with a quadratic objective, where it raised an ``IndexError`` because the term mask was built over the factor dimension as well. Quadratic terms are dropped when any of their factors references the removed variable. (`#883 <https://github.com/PyPSA/linopy/issues/883>`__)
+
 
 Version v0.9.0
 --------------

linopy/common.py

@@ -471,6 +471,40 @@ def replace_by_map(ds: DataArray, mapping: np.ndarray) -> DataArray:
     )
 
 
+def assigned_labels(labels: np.ndarray | DataArray) -> np.ndarray:
+    """
+    Flatten labels and drop the -1 sentinels.
+
+    ``-1`` marks a masked entry in a variable's or constraint's labels, but an
+    empty term slot in the ``vars`` field of a constraint or expression.
+    Matching labels of one object against another must therefore ignore it,
+    otherwise every masked entry compares equal to every empty term slot.
+    """
+    flat = np.asarray(labels).ravel()
+    return flat[flat != -1]
+
+
+def contains_labels(values: np.ndarray, labels: np.ndarray) -> bool:
+    """
+    Whether any entry of ``values`` is one of ``labels``.
+
+    ``labels`` must not contain the ``-1`` sentinel, see :func:`assigned_labels`.
+    Labels are handed out in ascending blocks, so restricting ``values`` to the
+    label range is both a cheap prefilter and, whenever the block is gap-free, the
+    complete answer.
+    """
+    if not labels.size:
+        return False
+    low, high = labels.min(), labels.max()
+    candidates = values[(values >= low) & (values <= high)]
+    if not candidates.size:
+        return False
+    is_gap_free = bool((np.diff(labels) == 1).all())
+    if is_gap_free:
+        return True
+    return bool(np.isin(candidates, labels).any())
+
+
 def to_path(path: str | Path | None) -> Path | None:
     """
     Convert a string to a Path object.

linopy/constraints.py

@@ -38,8 +38,10 @@
     VariableLabelIndex,
     align_lines_by_delimiter,
     assign_multiindex_safe,
+    assigned_labels,
     check_has_nulls,
     check_has_nulls_polars,
+    contains_labels,
     coords_from_dataset,
     coords_to_dataset_vars,
     filter_nulls_polars,
@@ -200,9 +202,23 @@ def data_attrs(self) -> list[str]:
         base = list(Constraints.dataset_attrs)
         return base + ["binary_var", "binary_val"] if self.is_indicator else base
 
-    @abstractmethod
     def has_variable(self, variable: variables.Variable) -> bool:
-        """Check if the constraint references any of the given variable labels."""
+        """
+        Check if the constraint references any of the given variable labels.
+
+        Masked variable entries (label -1) are ignored: they are not part of
+        the model and would otherwise match every empty term slot.
+        """
+        return self.has_labels(assigned_labels(variable.labels))
+
+    @abstractmethod
+    def has_labels(self, labels: np.ndarray) -> bool:
+        """
+        Check if the constraint references any of the given variable labels.
+
+        ``labels`` must not contain the ``-1`` sentinel, see
+        :func:`linopy.common.assigned_labels`.
+        """
 
     @abstractmethod
     def sanitize_zeros(self) -> ConstraintBase:
@@ -940,11 +956,9 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint:
             binval=binval,
         )
 
-    def has_variable(self, variable: variables.Variable) -> bool:
-        vlabels = self._model.variables.label_index.vlabels
-        return bool(
-            np.isin(vlabels[self._csr.indices], variable.labels.values.ravel()).any()
-        )
+    def has_labels(self, labels: np.ndarray) -> bool:
+        label_to_pos = self._model.variables.label_index.label_to_pos
+        return contains_labels(self._csr.indices, label_to_pos[labels])
 
     def to_matrix_with_rhs(
         self, label_index: VariableLabelIndex
@@ -1519,8 +1533,8 @@ def dual(self, value: ConstantLike) -> None:
         value = DataArray(value).broadcast_like(self.labels)
         self._data = assign_multiindex_safe(self.data, dual=value)
 
-    def has_variable(self, variable: variables.Variable) -> bool:
-        return bool(self.data["vars"].isin(variable.labels.values.ravel()).any())
+    def has_labels(self, labels: np.ndarray) -> bool:
+        return contains_labels(self.data["vars"].values.ravel(), labels)
 
     def _matrix_export_data(
         self, label_index: VariableLabelIndex

linopy/model.py

@@ -31,12 +31,14 @@
 from linopy.alignment import as_dataarray, broadcast_to_coords
 from linopy.common import (
     assign_multiindex_safe,
+    assigned_labels,
     best_int,
     maybe_replace_signs,
     replace_by_map,
     to_path,
 )
 from linopy.constants import (
+    FACTOR_DIM,
     GREATER_EQUAL,
     HELPER_DIMS,
     LESS_EQUAL,
@@ -1431,9 +1433,9 @@ def remove_variables(self, name: str) -> None:
 
         self._relaxed_registry.pop(name, None)
 
-        to_remove = [
-            k for k, con in self.constraints.items() if con.has_variable(variable)
-        ]
+        labels = assigned_labels(variable.labels)
+
+        to_remove = [k for k, con in self.constraints.items() if con.has_labels(labels)]
 
         if to_remove:
             warnings.warn(
@@ -1447,9 +1449,11 @@ def remove_variables(self, name: str) -> None:
 
         self.variables.remove(name)
 
-        self.objective = self.objective.sel(
-            {TERM_DIM: ~self.objective.vars.isin(variable.labels)}
-        )
+        referenced = self.objective.vars.isin(labels)
+        if FACTOR_DIM in referenced.dims:
+            referenced = referenced.any(FACTOR_DIM)
+
+        self.objective = self.objective.sel({TERM_DIM: ~referenced})
 
     def remove_constraints(self, name: str | list[str]) -> None:
         """

test/test_model.py

@@ -11,6 +11,7 @@
 from tempfile import gettempdir
 
 import numpy as np
+import pandas as pd
 import pytest
 import xarray as xr
 
@@ -160,6 +161,42 @@ def test_remove_variable() -> None:
     assert not m.objective.vars.isin(x.labels).any()
 
 
+@pytest.mark.parametrize("freeze", [False, True])
+@pytest.mark.parametrize("quadratic", [False, True])
+def test_remove_masked_variable_keeps_unrelated_constraints(
+    freeze: bool, quadratic: bool
+) -> None:
+    # https://github.com/PyPSA/linopy/issues/883
+    m: Model = Model(freeze_constraints=freeze)
+
+    i = pd.Index(range(3), name="i")
+    mask = [True, False, True]
+    a = m.add_variables(coords=[i], name="a", mask=mask)
+    b = m.add_variables(coords=[i], name="b", mask=mask)
+    c = m.add_variables(coords=[i], name="c")
+
+    # `b` is masked, so the constraint carries empty term slots (-1), but it
+    # never references `a`
+    without_a = m.add_constraints(b.sum() + c, EQUAL, 0, name="without_a")
+    assert not without_a.has_variable(a)
+
+    with_a = m.add_constraints(a.sum() + c, EQUAL, 0, name="with_a")
+    assert with_a.has_variable(a)
+
+    if quadratic:
+        m.add_objective((a * c).sum() + (c * c).sum())
+    else:
+        m.add_objective((1 * a).sum() + (1 * c).sum())
+
+    with pytest.warns(UserWarning, match="with_a"):
+        m.remove_variables("a")
+
+    assert "without_a" in m.constraints
+    assert "with_a" not in m.constraints
+    assert not m.objective.vars.isin(a.labels[a.labels != -1]).any()
+    assert m.objective.vars.isin(c.labels).any()
+
+
 def test_remove_constraint() -> None:
     m: Model = Model()