fix(sos): SOS constraints on masked variables (#688)

linopy/io.py

@@ -424,16 +424,23 @@ def sos_to_file(
 
     for name in names:
         var = m.variables[name]
-        sos_type = var.attrs[SOS_TYPE_ATTR]
-        sos_dim = var.attrs[SOS_DIM_ATTR]
+        sos_type = int(var.attrs[SOS_TYPE_ATTR])  # type: ignore[call-overload]
+        sos_dim = str(var.attrs[SOS_DIM_ATTR])
 
         other_dims = [dim for dim in var.labels.dims if dim != sos_dim]
         for var_slice in var.iterate_slices(slice_size, other_dims):
             ds = var_slice.labels.to_dataset()
-            ds["sos_labels"] = ds["labels"].isel({sos_dim: 0})
+            # Per-set id = max member label: unique per set (labels are globally
+            # unique); a fully-masked set reduces to -1 and is dropped below.
+            ds["sos_labels"] = ds["labels"].max(sos_dim)
             ds["weights"] = ds.coords[sos_dim]
             df = to_polars(ds)
 
+            # Drop masked members
+            df = df.filter((pl.col("labels") != -1) & (pl.col("sos_labels") != -1))
+            if df.is_empty():
+                continue
+
             df = df.group_by("sos_labels").agg(
                 pl.concat_str(
                     *print_variable(pl.col("labels")), pl.lit(":"), pl.col("weights")
@@ -593,8 +600,6 @@ def to_file(
     """
     Write out a model to a lp or mps file.
     """
-    m._check_sos_unmasked()
-
     if fn is None:
         fn = Path(m.get_problem_file())
     if isinstance(fn, str):

linopy/model.py

@@ -1300,34 +1300,6 @@ def _resolve_sos_reformulation(
             )
         return not solver_supports(solver_name, SolverFeature.SOS_CONSTRAINTS)
 
-    def _check_sos_unmasked(self) -> None:
-        """
-        Reject the model if any SOS variable has masked entries.
-
-        The SOS plumbing (both direct-API solvers and the LP file writer) treats
-        linopy variable labels as solver column indices / names, which breaks as
-        soon as a label is ``-1`` (linopy's ``FILL_VALUE["labels"]`` for masked
-        slots). The downstream symptoms are solver-specific — ``IndexError`` on
-        gurobipy, ``?404 Invalid column number`` on xpress, parse errors on
-        xpress/cplex LP readers, silent SOS-set corruption on gurobi's LP reader.
-
-        Surface a single clear error until #688 lands the proper fix.
-        """
-        if not self.variables.sos:
-            return
-        affected = [
-            name
-            for name in self.variables.sos
-            if (self.variables[name].labels.values == -1).any()
-        ]
-        if affected:
-            raise NotImplementedError(
-                f"SOS constraints on masked variables are not yet supported "
-                f"(affected: {affected}; "
-                "see https://github.com/PyPSA/linopy/issues/688). "
-                "Pass reformulate_sos=True as a workaround."
-            )
-
     def remove_objective(self) -> None:
         """
         Remove the objective's linear expression from the model.

linopy/solvers.py

@@ -29,7 +29,6 @@
 
 import numpy as np
 import pandas as pd
-import xarray as xr
 from packaging.specifiers import SpecifierSet
 from packaging.version import parse as parse_version
 from scipy.sparse import tril, triu
@@ -106,6 +105,25 @@ def _solution_from_labels(
     return values_to_lookup_array(np.asarray(values, dtype=float), labels, size=size)
 
 
+def _iter_sos_sets(model: Model) -> Iterator[tuple[int, np.ndarray, np.ndarray]]:
+    """Yield ``(sos_type, positions, weights)`` per active SOS set in ``model``."""
+    label_to_pos = model.variables.label_index.label_to_pos
+    for var_name in model.variables.sos:
+        var = model.variables.sos[var_name]
+        sos_type = int(var.attrs[SOS_TYPE_ATTR])  # type: ignore[call-overload]
+        sos_dim = str(var.attrs[SOS_DIM_ATTR])
+
+        labels = var.labels.transpose(sos_dim, ...)
+        weights = labels.coords[sos_dim].values
+        arr = labels.values.reshape(labels.shape[0], -1)
+
+        for i in range(arr.shape[1]):
+            col = arr[:, i]
+            mask = col != -1
+            if mask.any():
+                yield sos_type, label_to_pos[col[mask]], weights[mask]
+
+
 class SolverFeature(Enum):
     """Enumeration of all solver capabilities tracked by linopy."""
 
@@ -518,7 +536,6 @@ def _build(self, **build_kwargs: Any) -> None:
         if self.model is None:
             raise RuntimeError("Solver has no model attached; cannot build.")
         self._validate_model()
-        self.model._check_sos_unmasked()
         if self.io_api == "direct":
             self._build_direct(**build_kwargs)
         else:
@@ -1581,25 +1598,8 @@ def _build_solver_model(
                 names = print_constraints(M.clabels)
                 c.setAttr("ConstrName", names)
 
-        if model.variables.sos:
-            for var_name in model.variables.sos:
-                var = model.variables.sos[var_name]
-                sos_type: int = var.attrs[SOS_TYPE_ATTR]  # type: ignore[assignment]
-                sos_dim: str = var.attrs[SOS_DIM_ATTR]  # type: ignore[assignment]
-
-                def add_sos(s: xr.DataArray, sos_type: int, sos_dim: str) -> None:
-                    s = s.squeeze()
-                    indices = s.values.flatten().tolist()
-                    weights = s.coords[sos_dim].values.tolist()
-                    gm.addSOS(sos_type, x[indices].tolist(), weights)
-
-                others = [dim for dim in var.labels.dims if dim != sos_dim]
-                if not others:
-                    add_sos(var.labels, sos_type, sos_dim)
-                else:
-                    stacked = var.labels.stack(_sos_group=others)
-                    for _, s in stacked.groupby("_sos_group"):
-                        add_sos(s.unstack("_sos_group"), sos_type, sos_dim)
+        for sos_type, positions, weights in _iter_sos_sets(model):
+            gm.addSOS(sos_type, x[positions.tolist()].tolist(), weights.tolist())
 
         gm.update()
         return gm
@@ -2223,25 +2223,8 @@ def _build_solver_model(
             if cnames:
                 problem.addnames(xpress_Namespaces.ROW, cnames, 0, len(cnames) - 1)
 
-        if model.variables.sos:
-            for var_name in model.variables.sos:
-                var = model.variables.sos[var_name]
-                sos_type: int = var.attrs[SOS_TYPE_ATTR]  # type: ignore[assignment]
-                sos_dim: str = var.attrs[SOS_DIM_ATTR]  # type: ignore[assignment]
-
-                def add_sos(s: xr.DataArray, sos_type: int, sos_dim: str) -> None:
-                    s = s.squeeze()
-                    indices = s.values.flatten().tolist()
-                    weights = s.coords[sos_dim].values.tolist()
-                    problem.addSOS(indices, weights, type=sos_type)
-
-                others = [dim for dim in var.labels.dims if dim != sos_dim]
-                if not others:
-                    add_sos(var.labels, sos_type, sos_dim)
-                else:
-                    stacked = var.labels.stack(_sos_group=others)
-                    for _, s in stacked.groupby("_sos_group"):
-                        add_sos(s.unstack("_sos_group"), sos_type, sos_dim)
+        for sos_type, positions, weights in _iter_sos_sets(model):
+            problem.addSOS(positions.tolist(), weights.tolist(), type=sos_type)
 
         return problem
 

test/test_piecewise_constraints.py

@@ -41,7 +41,11 @@
     SEGMENT_DIM,
 )
 from linopy.piecewise import _slopes_to_points
-from linopy.solver_capabilities import SolverFeature, get_available_solvers_with_feature
+from linopy.solver_capabilities import (
+    SolverFeature,
+    get_available_solvers_with_feature,
+    solver_supports,
+)
 
 if TYPE_CHECKING:
     from linopy.piecewise import BreaksLike, _PwlInputs
@@ -52,6 +56,13 @@
 _sos2_solvers = get_available_solvers_with_feature(
     SolverFeature.SOS_CONSTRAINTS, available_solvers
 )
+_sos2_direct_solvers = sorted(
+    s for s in _sos2_solvers if solver_supports(s, SolverFeature.DIRECT_API)
+)
+_SOS_PATHS = [
+    *[pytest.param(s, "direct", id=f"{s}-direct") for s in _sos2_direct_solvers],
+    *[pytest.param(s, "lp", id=f"{s}-lp") for s in sorted(_sos2_solvers)],
+]
 _any_solvers = [
     s for s in ["highs", "gurobi", "glpk", "cplex"] if s in available_solvers
 ]
@@ -2313,23 +2324,37 @@ def test_lp_per_entity_nan_padding(
         Per-entity NaN-padded breakpoints with method='lp': padded
         segments must be masked out so they don't create spurious
         ``y ≤ 0`` constraints (bug-2 regression).
-
-        ``method='sos2'`` would emit a masked SOS lambda variable, which the
-        native SOS path doesn't yet support (#688) — exercised separately in
-        :py:meth:`test_sos2_per_entity_nan_padding_errors`.
         """
         m = nan_padded_pwl_model("lp")
         m.solve()
         # f_b(10) on chord (5,10)→(15,15) is 12.5
         assert abs(float(m.solution.sel({"entity": "b"})["y"]) - 12.5) < 1e-3
 
-    def test_sos2_per_entity_nan_padding_errors(
-        self, nan_padded_pwl_model: Callable[[Method], Model]
+    @pytest.mark.skipif(not _SOS_PATHS, reason="No SOS-capable solver installed")
+    @pytest.mark.parametrize(("solver", "io_api"), _SOS_PATHS)
+    def test_sos2_per_entity_nan_padding(
+        self,
+        nan_padded_pwl_model: Callable[[Method], Model],
+        solver: str,
+        io_api: str,
     ) -> None:
-        """Masked SOS lambdas hit the #688 guard at solve time."""
+        """
+        Per-entity NaN-padded breakpoints with method='sos2': the SOS
+        lambda variable's masked entries must flow through both the
+        direct API (via label→position resolution) and the LP writer
+        (via masked-member filtering) so the solve returns the same
+        answer as ``method='lp'``. Regression for #688.
+
+        Parametrized across every SOS-capable solver × io_api so the
+        bug surfaces no matter which backend handles the SOS section
+        (gurobi-lp masked the bug on master by silently dropping
+        unknown ``x-1`` members; cplex-lp and gurobi-direct surfaced
+        it as a parse / OOB error).
+        """
         m = nan_padded_pwl_model("sos2")
-        with pytest.raises(NotImplementedError, match="masked"):
-            m.solve()
+        m.solve(solver_name=solver, io_api=io_api)
+        # f_b(10) on chord (5,10)→(15,15) is 12.5 — same oracle as lp variant
+        assert abs(float(m.solution.sel({"entity": "b"})["y"]) - 12.5) < 1e-3
 
     def test_lp_rejects_decreasing_x_concave_ge(self) -> None:
         """

test/test_sos_constraints.py

@@ -8,19 +8,6 @@
 import xarray as xr
 
 from linopy import Model, available_solvers
-from linopy.solver_capabilities import (
-    SolverFeature,
-    get_available_solvers_with_feature,
-    solver_supports,
-)
-
-_direct_sos_solvers = [
-    s
-    for s in get_available_solvers_with_feature(
-        SolverFeature.SOS_CONSTRAINTS, available_solvers
-    )
-    if solver_supports(s, SolverFeature.DIRECT_API)
-]
 
 
 def test_add_sos_constraints_registers_variable() -> None:
@@ -209,66 +196,6 @@ def test_qp_sos1_xpress_direct() -> None:
     assert np.isclose(m.objective.value, -25)
 
 
-@pytest.fixture
-def masked_sos_model() -> Model:
-    """Tiny model with a single masked SOS1 variable."""
-    m = Model()
-    coords = pd.Index([0, 1, 2, 3], name="i")
-    mask = pd.Series([True, True, False, True], index=coords)
-    var = m.add_variables(lower=0, upper=1, coords=[coords], mask=mask, name="sos_var")
-    m.add_sos_constraints(var, sos_type=1, sos_dim="i")
-    m.add_objective(-var.sum())
-    return m
-
-
-@pytest.mark.parametrize("solver_name", _direct_sos_solvers)
-def test_direct_api_raises_on_masked_sos(
-    solver_name: str, masked_sos_model: Model
-) -> None:
-    with pytest.raises(NotImplementedError, match="masked"):
-        masked_sos_model.solve(solver_name=solver_name, io_api="direct")
-
-
-def test_lp_writer_raises_on_masked_sos(
-    masked_sos_model: Model, tmp_path: Path
-) -> None:
-    with pytest.raises(NotImplementedError, match="masked"):
-        masked_sos_model.to_file(tmp_path / "sos.lp", io_api="lp")
-
-
-@pytest.mark.parametrize(
-    "solver_name",
-    [
-        pytest.param(
-            "gurobi",
-            marks=pytest.mark.skipif(
-                "gurobi" not in available_solvers, reason="Gurobi not installed"
-            ),
-        ),
-        pytest.param(
-            "highs",
-            marks=pytest.mark.skipif(
-                "highs" not in available_solvers, reason="HiGHS not installed"
-            ),
-        ),
-    ],
-)
-def test_reformulate_sos_true_solves_masked_sos(
-    solver_name: str, masked_sos_model: Model
-) -> None:
-    """The documented workaround for the masked-SOS bug actually solves."""
-    masked_sos_model.solve(solver_name=solver_name, reformulate_sos=True)
-    sol = masked_sos_model.variables["sos_var"].solution.values
-    # SOS1 over 3 unmasked entries, max sum, each in [0, 1]:
-    # one entry == 1, others == 0, masked stays NaN.
-    assert masked_sos_model.objective.value is not None
-    assert np.isclose(masked_sos_model.objective.value, -1.0)
-    assert np.isnan(sol[2])
-    nonzero = np.flatnonzero(~np.isnan(sol) & (sol > 1e-6))
-    assert len(nonzero) == 1
-    assert np.isclose(sol[nonzero[0]], 1.0)
-
-
 @pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobi not installed")
 def test_reformulate_sos_true_reformulates_on_native_solver(tmp_path: Path) -> None:
     """