diff --git a/doc/release_notes.rst b/doc/release_notes.rst index d9bc95aa..4df695c7 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -4,6 +4,7 @@ Release Notes .. Upcoming Version * Fix LP file writing for negative zero (-0.0) values that produced invalid syntax like "+-0.0" rejected by Gurobi +* Add SOS1 and SOS2 reformulations for solvers not supporting them. Version 0.6.0 -------------- diff --git a/doc/sos-constraints.rst b/doc/sos-constraints.rst index 37dd72d2..a2731400 100644 --- a/doc/sos-constraints.rst +++ b/doc/sos-constraints.rst @@ -75,7 +75,7 @@ Method Signature .. code-block:: python - Model.add_sos_constraints(variable, sos_type, sos_dim) + Model.add_sos_constraints(variable, sos_type, sos_dim, big_m=None) **Parameters:** @@ -85,6 +85,8 @@ Method Signature Type of SOS constraint (1 or 2) - ``sos_dim`` : str Name of the dimension along which the SOS constraint applies +- ``big_m`` : float | None + Custom Big-M value for reformulation (see :ref:`sos-reformulation`) **Requirements:** @@ -254,6 +256,83 @@ SOS constraints are supported by most modern mixed-integer programming solvers t - MOSEK - MindOpt +For unsupported solvers, use automatic reformulation (see below). + +.. _sos-reformulation: + +SOS Reformulation +----------------- + +For solvers without native SOS support, linopy can reformulate SOS constraints +as binary + linear constraints using the Big-M method. + +.. code-block:: python + + # Automatic reformulation during solve + m.solve(solver_name="highs", reformulate_sos=True) + + # Or reformulate manually + m.reformulate_sos_constraints() + m.solve(solver_name="highs") + +**Requirements:** + +- Variables must have **non-negative lower bounds** (lower >= 0) +- Big-M values are derived from variable upper bounds +- For infinite upper bounds, specify custom values via the ``big_m`` parameter + +.. code-block:: python + + # Finite bounds (default) + x = m.add_variables(lower=0, upper=100, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + # Infinite upper bounds: specify Big-M + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=10) + +The reformulation uses the tighter of ``big_m`` and variable upper bound. + +Mathematical Formulation +~~~~~~~~~~~~~~~~~~~~~~~~ + +**SOS1 Reformulation:** + +Original constraint: :math:`\text{SOS1}(\{x_1, x_2, \ldots, x_n\})` means at most one +:math:`x_i` can be non-zero. + +Given :math:`x = (x_1, \ldots, x_n) \in \mathbb{R}^n_+`, introduce binary +:math:`y = (y_1, \ldots, y_n) \in \{0,1\}^n`: + +.. math:: + + x_i &\leq M_i \cdot y_i \quad \forall i \in \{1, \ldots, n\} \\ + x_i &\geq 0 \quad \forall i \in \{1, \ldots, n\} \\ + \sum_{i=1}^{n} y_i &\leq 1 \\ + y_i &\in \{0, 1\} \quad \forall i \in \{1, \ldots, n\} + +where :math:`M_i \geq \sup\{x_i\}` (upper bound on :math:`x_i`). + +**SOS2 Reformulation:** + +Original constraint: :math:`\text{SOS2}(\{x_1, x_2, \ldots, x_n\})` means at most two +:math:`x_i` can be non-zero, and if two are non-zero, they must have consecutive indices. + +Given :math:`x = (x_1, \ldots, x_n) \in \mathbb{R}^n_+`, introduce binary +:math:`y = (y_1, \ldots, y_{n-1}) \in \{0,1\}^{n-1}`: + +.. math:: + + x_1 &\leq M_1 \cdot y_1 \\ + x_i &\leq M_i \cdot (y_{i-1} + y_i) \quad \forall i \in \{2, \ldots, n-1\} \\ + x_n &\leq M_n \cdot y_{n-1} \\ + x_i &\geq 0 \quad \forall i \in \{1, \ldots, n\} \\ + \sum_{i=1}^{n-1} y_i &\leq 1 \\ + y_i &\in \{0, 1\} \quad \forall i \in \{1, \ldots, n-1\} + +where :math:`M_i \geq \sup\{x_i\}`. Interpretation: :math:`y_i = 1` activates interval +:math:`[i, i+1]`, allowing :math:`x_i` and :math:`x_{i+1}` to be non-zero. + Common Patterns --------------- diff --git a/linopy/constants.py b/linopy/constants.py index 021a9a10..2e1ef47a 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -49,6 +49,11 @@ CV_DIM, ] +# SOS constraint attribute keys +SOS_TYPE_ATTR = "sos_type" +SOS_DIM_ATTR = "sos_dim" +SOS_BIG_M_ATTR = "big_m_upper" + class ModelStatus(Enum): """ diff --git a/linopy/io.py b/linopy/io.py index 56fe033d..7dc4177e 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -25,7 +25,7 @@ from linopy import solvers from linopy.common import to_polars -from linopy.constants import CONCAT_DIM +from linopy.constants import CONCAT_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR from linopy.objective import Objective if TYPE_CHECKING: @@ -374,8 +374,8 @@ def sos_to_file( for name in names: var = m.variables[name] - sos_type = var.attrs["sos_type"] - sos_dim = var.attrs["sos_dim"] + sos_type = var.attrs[SOS_TYPE_ATTR] + sos_dim = 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): @@ -773,8 +773,8 @@ def to_gurobipy( if m.variables.sos: for var_name in m.variables.sos: var = m.variables.sos[var_name] - sos_type: int = var.attrs["sos_type"] # type: ignore[assignment] - sos_dim: str = var.attrs["sos_dim"] # type: ignore[assignment] + 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() diff --git a/linopy/model.py b/linopy/model.py index 657b2d45..8828d819 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -38,6 +38,9 @@ GREATER_EQUAL, HELPER_DIMS, LESS_EQUAL, + SOS_BIG_M_ATTR, + SOS_DIM_ATTR, + SOS_TYPE_ATTR, TERM_DIM, ModelStatus, TerminationCondition, @@ -65,6 +68,7 @@ IO_APIS, available_solvers, ) +from linopy.sos_reformulation import reformulate_all_sos from linopy.types import ( ConstantLike, ConstraintLike, @@ -556,6 +560,7 @@ def add_sos_constraints( variable: Variable, sos_type: Literal[1, 2], sos_dim: str, + big_m: float | None = None, ) -> None: """ Add an sos1 or sos2 constraint for one dimension of a variable @@ -569,15 +574,26 @@ def add_sos_constraints( Type of SOS sos_dim : str Which dimension of variable to add SOS constraint to + big_m : float | None, optional + Big-M value for SOS reformulation. Only used when reformulating + SOS constraints for solvers that don't support them natively. + + - None (default): Use variable upper bounds as Big-M + - float: Custom Big-M value + + The reformulation uses the tighter of big_m and variable upper bound: + M = min(big_m, var.upper). + + Tighter Big-M values improve LP relaxation quality and solve time. """ if sos_type not in (1, 2): raise ValueError(f"sos_type must be 1 or 2, got {sos_type}") if sos_dim not in variable.dims: raise ValueError(f"sos_dim must name a variable dimension, got {sos_dim}") - if "sos_type" in variable.attrs or "sos_dim" in variable.attrs: - existing_sos_type = variable.attrs.get("sos_type") - existing_sos_dim = variable.attrs.get("sos_dim") + if SOS_TYPE_ATTR in variable.attrs or SOS_DIM_ATTR in variable.attrs: + existing_sos_type = variable.attrs.get(SOS_TYPE_ATTR) + existing_sos_dim = variable.attrs.get(SOS_DIM_ATTR) raise ValueError( f"variable already has an sos{existing_sos_type} constraint on {existing_sos_dim}" ) @@ -589,7 +605,12 @@ def add_sos_constraints( f"but got {variable.coords[sos_dim].dtype}" ) - variable.attrs.update(sos_type=sos_type, sos_dim=sos_dim) + # Process and store big_m value + attrs_update: dict[str, Any] = {SOS_TYPE_ATTR: sos_type, SOS_DIM_ATTR: sos_dim} + if big_m is not None: + attrs_update[SOS_BIG_M_ATTR] = float(big_m) + + variable.attrs.update(attrs_update) def add_constraints( self, @@ -830,18 +851,65 @@ def remove_sos_constraints(self, variable: Variable) -> None: ------- None. """ - if "sos_type" not in variable.attrs or "sos_dim" not in variable.attrs: + if SOS_TYPE_ATTR not in variable.attrs or SOS_DIM_ATTR not in variable.attrs: raise ValueError(f"Variable '{variable.name}' has no SOS constraints") - sos_type = variable.attrs["sos_type"] - sos_dim = variable.attrs["sos_dim"] + sos_type = variable.attrs[SOS_TYPE_ATTR] + sos_dim = variable.attrs[SOS_DIM_ATTR] - del variable.attrs["sos_type"], variable.attrs["sos_dim"] + del variable.attrs[SOS_TYPE_ATTR], variable.attrs[SOS_DIM_ATTR] + + # Also remove big_m attribute if present + variable.attrs.pop(SOS_BIG_M_ATTR, None) logger.debug( f"Removed sos{sos_type} constraint on {sos_dim} from {variable.name}" ) + def reformulate_sos_constraints(self, prefix: str = "_sos_reform_") -> list[str]: + """ + Reformulate SOS constraints as binary + linear constraints. + + This converts SOS1 and SOS2 constraints into equivalent binary variable + formulations using the Big-M method. This allows solving models with SOS + constraints using solvers that don't support them natively (e.g., HiGHS, GLPK). + + Big-M values are determined as follows: + 1. If custom big_m was specified in add_sos_constraints(), use that + 2. Otherwise, use the variable bounds (tightest valid Big-M) + + Parameters + ---------- + prefix : str, optional + Prefix for naming auxiliary variables and constraints. + Default is "_sos_reform_". + + Returns + ------- + list[str] + List of variable names that were reformulated. + + Raises + ------ + ValueError + If any SOS variable has infinite bounds and no custom big_m was specified. + + Examples + -------- + >>> m = Model() + >>> x = m.add_variables( + ... lower=0, upper=1, coords=[pd.Index([0, 1, 2], name="i")], name="x" + ... ) + >>> m.add_sos_constraints(x, sos_type=1, sos_dim="i") + >>> m.reformulate_sos_constraints() # Now solvable with HiGHS + ['x'] + + With custom big_m for tighter relaxation: + + >>> m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=0.5) + """ + return reformulate_all_sos(self, prefix=prefix) + def remove_objective(self) -> None: """ Remove the objective's linear expression from the model. @@ -1126,6 +1194,7 @@ def solve( remote: RemoteHandler | OetcHandler = None, # type: ignore progress: bool | None = None, mock_solve: bool = False, + reformulate_sos: bool = False, **solver_options: Any, ) -> tuple[str, str]: """ @@ -1195,6 +1264,11 @@ def solve( than 10000 variables and constraints. mock_solve : bool, optional Whether to run a mock solve. This will skip the actual solving. Variables will be set to have dummy values + reformulate_sos : bool, optional + Whether to automatically reformulate SOS constraints as binary + linear + constraints for solvers that don't support them natively. + This uses the Big-M method and requires all SOS variables to have finite bounds. + Default is False. **solver_options : kwargs Options passed to the solver. @@ -1293,10 +1367,17 @@ def solve( ) # SOS constraints are not supported by all solvers - if self.variables.sos and not solver_supports( - solver_name, SolverFeature.SOS_CONSTRAINTS - ): - raise ValueError(f"Solver {solver_name} does not support SOS constraints.") + if self.variables.sos: + if reformulate_sos and not solver_supports( + solver_name, SolverFeature.SOS_CONSTRAINTS + ): + logger.info(f"Reformulating SOS constraints for solver {solver_name}") + self.reformulate_sos_constraints() + elif not solver_supports(solver_name, SolverFeature.SOS_CONSTRAINTS): + raise ValueError( + f"Solver {solver_name} does not support SOS constraints. " + "Use reformulate_sos=True or a solver that supports SOS (gurobi, cplex)." + ) try: solver_class = getattr(solvers, f"{solvers.SolverName(solver_name).name}") diff --git a/linopy/sos_reformulation.py b/linopy/sos_reformulation.py new file mode 100644 index 00000000..beb83ab5 --- /dev/null +++ b/linopy/sos_reformulation.py @@ -0,0 +1,206 @@ +""" +SOS constraint reformulation using Big-M method. + +Converts SOS1/SOS2 constraints to binary + linear constraints for solvers +that don't support them natively. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import numpy as np +import pandas as pd +from xarray import DataArray + +from linopy.constants import SOS_BIG_M_ATTR, SOS_DIM_ATTR, SOS_TYPE_ATTR + +if TYPE_CHECKING: + from linopy.model import Model + from linopy.variables import Variable + +logger = logging.getLogger(__name__) + + +def compute_big_m_values(var: Variable) -> DataArray: + """ + Compute Big-M values from variable bounds and custom big_m attribute. + + Uses the tighter of variable upper bound and custom big_m to ensure + the best possible LP relaxation. + + Parameters + ---------- + var : Variable + Variable with bounds (and optionally big_m_upper attr). + + Returns + ------- + DataArray + M_upper for reformulation constraints: x <= M_upper * y + + Raises + ------ + ValueError + If variable has negative lower bounds or infinite upper bounds. + """ + # SOS reformulation requires non-negative variables + if (var.lower < 0).any(): + raise ValueError( + f"Variable '{var.name}' has negative lower bounds. " + "SOS reformulation requires non-negative variables (lower >= 0)." + ) + + big_m_upper = var.attrs.get(SOS_BIG_M_ATTR) + M_upper = var.upper + + if big_m_upper is not None: + M_upper = M_upper.clip(max=big_m_upper) # type: ignore[arg-type] + + # Validate finiteness + if np.isinf(M_upper).any(): + raise ValueError( + f"Variable '{var.name}' has infinite upper bounds. " + "Set finite bounds or specify big_m in add_sos_constraints()." + ) + + return M_upper + + +def reformulate_sos1(model: Model, var: Variable, prefix: str) -> None: + """ + Reformulate SOS1 constraint as binary + linear constraints. + + For each x[i] with upper bound M[i]: + - Add binary indicator y[i] + - x[i] <= M[i] * y[i] + - sum(y) <= 1 + + Parameters + ---------- + model : Model + Model to add reformulation constraints to. + var : Variable + Variable with SOS1 constraint (must have non-negative lower bounds). + prefix : str + Prefix for naming auxiliary variables and constraints. + """ + sos_dim = str(var.attrs[SOS_DIM_ATTR]) + name = var.name + M = compute_big_m_values(var) + + coords = [var.coords[d] for d in var.dims] + y = model.add_variables(coords=coords, name=f"{prefix}{name}_y", binary=True) + + model.add_constraints(var <= M * y, name=f"{prefix}{name}_upper") + model.add_constraints(y.sum(dim=sos_dim) <= 1, name=f"{prefix}{name}_card") + + +def reformulate_sos2(model: Model, var: Variable, prefix: str) -> None: + """ + Reformulate SOS2 constraint as binary + linear constraints. + + For ordered x[0..n-1] with upper bounds M[i]: + - Add n-1 binary segment indicators z[i] + - x[0] <= M[0] * z[0] + - x[i] <= M[i] * (z[i-1] + z[i]) for middle elements + - x[n-1] <= M[n-1] * z[n-2] + - sum(z) <= 1 + + Parameters + ---------- + model : Model + Model to add reformulation constraints to. + var : Variable + Variable with SOS2 constraint (must have non-negative lower bounds). + prefix : str + Prefix for naming auxiliary variables and constraints. + """ + sos_dim = str(var.attrs[SOS_DIM_ATTR]) + name = var.name + n = var.sizes[sos_dim] + + if n <= 1: + return + + M = compute_big_m_values(var) + + # Create n-1 segment indicators + z_coords = [ + pd.Index(var.coords[sos_dim].values[:-1], name=sos_dim) + if d == sos_dim + else var.coords[d] + for d in var.dims + ] + z = model.add_variables(coords=z_coords, name=f"{prefix}{name}_z", binary=True) + + # Convert to expressions to avoid SOS attr validation on isel + x, z_expr = 1 * var, 1 * z + + # First: x[0] <= M[0] * z[0] + model.add_constraints( + x.isel({sos_dim: 0}) <= M.isel({sos_dim: 0}) * z_expr.isel({sos_dim: 0}), + name=f"{prefix}{name}_upper_first", + ) + # Middle: x[i] <= M[i] * (z[i-1] + z[i]) + for i in range(1, n - 1): + model.add_constraints( + x.isel({sos_dim: i}) + <= M.isel({sos_dim: i}) + * (z_expr.isel({sos_dim: i - 1}) + z_expr.isel({sos_dim: i})), + name=f"{prefix}{name}_upper_mid_{i}", + ) + # Last: x[n-1] <= M[n-1] * z[n-2] + model.add_constraints( + x.isel({sos_dim: n - 1}) + <= M.isel({sos_dim: n - 1}) * z_expr.isel({sos_dim: n - 2}), + name=f"{prefix}{name}_upper_last", + ) + + model.add_constraints(z.sum(dim=sos_dim) <= 1, name=f"{prefix}{name}_card") + + +def reformulate_all_sos(model: Model, prefix: str = "_sos_reform_") -> list[str]: + """ + Reformulate all SOS constraints in the model. + + Parameters + ---------- + model : Model + Model containing SOS constraints to reformulate. + prefix : str, optional + Prefix for auxiliary variables and constraints. Default: "_sos_reform_" + + Returns + ------- + list[str] + Names of variables that were reformulated. + """ + reformulated = [] + + for var_name in list(model.variables.sos): + var = model.variables[var_name] + sos_type = var.attrs.get(SOS_TYPE_ATTR) + sos_dim = var.attrs.get(SOS_DIM_ATTR) + + if sos_type is None or sos_dim is None: + continue + if var.sizes[sos_dim] <= 1: + continue + + # Check if fixed to zero + M = compute_big_m_values(var) + if (M == 0).all(): + continue + + if sos_type == 1: + reformulate_sos1(model, var, prefix) + elif sos_type == 2: + reformulate_sos2(model, var, prefix) + + model.remove_sos_constraints(var) + reformulated.append(var_name) + + logger.info(f"Reformulated {len(reformulated)} SOS constraint(s)") + return reformulated diff --git a/linopy/variables.py b/linopy/variables.py index e2570b5d..75349c4b 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -52,7 +52,7 @@ to_polars, ) from linopy.config import options -from linopy.constants import HELPER_DIMS, TERM_DIM +from linopy.constants import HELPER_DIMS, SOS_DIM_ATTR, SOS_TYPE_ATTR, TERM_DIM from linopy.solver_capabilities import SolverFeature, solver_supports from linopy.types import ( ConstantLike, @@ -195,10 +195,10 @@ def __init__( if "label_range" not in data.attrs: data.assign_attrs(label_range=(data.labels.min(), data.labels.max())) - if "sos_type" in data.attrs or "sos_dim" in data.attrs: - if (sos_type := data.attrs.get("sos_type")) not in (1, 2): + if SOS_TYPE_ATTR in data.attrs or SOS_DIM_ATTR in data.attrs: + if (sos_type := data.attrs.get(SOS_TYPE_ATTR)) not in (1, 2): raise ValueError(f"sos_type must be 1 or 2, got {sos_type}") - if (sos_dim := data.attrs.get("sos_dim")) not in data.dims: + if (sos_dim := data.attrs.get(SOS_DIM_ATTR)) not in data.dims: raise ValueError( f"sos_dim must name a variable dimension, got {sos_dim}" ) @@ -328,8 +328,8 @@ def __repr__(self) -> str: dim_names = self.coord_names dim_sizes = list(self.sizes.values()) masked_entries = (~self.mask).sum().values - sos_type = self.attrs.get("sos_type") - sos_dim = self.attrs.get("sos_dim") + sos_type = self.attrs.get(SOS_TYPE_ATTR) + sos_dim = self.attrs.get(SOS_DIM_ATTR) lines = [] if dims: @@ -1244,8 +1244,8 @@ def __repr__(self) -> str: if ds.coords else "" ) - if (sos_type := ds.attrs.get("sos_type")) in (1, 2) and ( - sos_dim := ds.attrs.get("sos_dim") + if (sos_type := ds.attrs.get(SOS_TYPE_ATTR)) in (1, 2) and ( + sos_dim := ds.attrs.get(SOS_DIM_ATTR) ): coords += f" - sos{sos_type} on {sos_dim}" r += f" * {name}{coords}\n" @@ -1401,8 +1401,8 @@ def sos(self) -> Variables: { name: self.data[name] for name in self - if self[name].attrs.get("sos_dim") - and self[name].attrs.get("sos_type") in (1, 2) + if self[name].attrs.get(SOS_DIM_ATTR) + and self[name].attrs.get(SOS_TYPE_ATTR) in (1, 2) }, self.model, ) diff --git a/test/test_sos_reformulation.py b/test/test_sos_reformulation.py new file mode 100644 index 00000000..c6445b03 --- /dev/null +++ b/test/test_sos_reformulation.py @@ -0,0 +1,621 @@ +"""Tests for SOS constraint reformulation.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from linopy import Model, available_solvers +from linopy.sos_reformulation import ( + compute_big_m_values, + reformulate_all_sos, + reformulate_sos1, + reformulate_sos2, +) + + +class TestValidateBounds: + """Tests for bound validation in compute_big_m_values.""" + + def test_finite_bounds_pass(self) -> None: + """Finite non-negative bounds should pass validation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + compute_big_m_values(x) # Should not raise + + def test_infinite_upper_bounds_raise(self) -> None: + """Infinite upper bounds should raise ValueError.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + with pytest.raises(ValueError, match="infinite upper bounds"): + compute_big_m_values(x) + + def test_negative_lower_bounds_raise(self) -> None: + """Negative lower bounds should raise ValueError.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=-1, upper=1, coords=[idx], name="x") + with pytest.raises(ValueError, match="negative lower bounds"): + compute_big_m_values(x) + + def test_mixed_negative_lower_bounds_raise(self) -> None: + """Mixed finite/negative lower bounds should raise ValueError.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables( + lower=np.array([0, -1, 0]), + upper=np.array([1, 1, 1]), + coords=[idx], + name="x", + ) + with pytest.raises(ValueError, match="negative lower bounds"): + compute_big_m_values(x) + + +class TestComputeBigM: + """Tests for compute_big_m_values.""" + + def test_positive_bounds(self) -> None: + """Test Big-M computation with positive bounds.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=10, coords=[idx], name="x") + M = compute_big_m_values(x) + assert np.allclose(M.values, [10, 10, 10]) + + def test_varying_bounds(self) -> None: + """Test Big-M computation with varying upper bounds.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables( + lower=np.array([0, 0, 0]), + upper=np.array([1, 2, 3]), + coords=[idx], + name="x", + ) + M = compute_big_m_values(x) + assert np.allclose(M.values, [1, 2, 3]) + + def test_custom_big_m_scalar(self) -> None: + """Test Big-M uses tighter of custom value and bounds.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=100, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=10) + M = compute_big_m_values(x) + # M = min(10, 100) = 10 (custom is tighter) + assert np.allclose(M.values, [10, 10, 10]) + + def test_custom_big_m_allows_infinite_bounds(self) -> None: + """Test that custom big_m allows variables with infinite bounds.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=10) + # Should not raise - custom big_m makes result finite + M = compute_big_m_values(x) + assert np.allclose(M.values, [10, 10, 10]) + + +class TestSOS1Reformulation: + """Tests for SOS1 reformulation.""" + + def test_basic_sos1(self) -> None: + """Test basic SOS1 reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + # Reformulate + reformulate_sos1(m, x, "_test_") + m.remove_sos_constraints(x) + + # Check auxiliary variables and constraints were added + assert "_test_x_y" in m.variables + assert "_test_x_upper" in m.constraints + assert "_test_x_card" in m.constraints + + # Binary variable should have same dimensions + y = m.variables["_test_x_y"] + assert y.dims == x.dims + assert y.sizes == x.sizes + + def test_sos1_multidimensional(self) -> None: + """Test SOS1 reformulation with multi-dimensional variables.""" + m = Model() + idx_i = pd.Index([0, 1, 2], name="i") + idx_j = pd.Index([0, 1], name="j") + x = m.add_variables(lower=0, upper=1, coords=[idx_i, idx_j], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulate_sos1(m, x, "_test_") + m.remove_sos_constraints(x) + + # Binary variable should have same dimensions + y = m.variables["_test_x_y"] + assert set(y.dims) == {"i", "j"} + + # Cardinality constraint should have reduced dimensions (summed over i) + card_con = m.constraints["_test_x_card"] + assert "j" in card_con.dims + + +class TestSOS2Reformulation: + """Tests for SOS2 reformulation.""" + + def test_basic_sos2(self) -> None: + """Test basic SOS2 reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + reformulate_sos2(m, x, "_test_") + m.remove_sos_constraints(x) + + # Check auxiliary variables and constraints were added + assert "_test_x_z" in m.variables + assert "_test_x_upper_first" in m.constraints + assert "_test_x_upper_last" in m.constraints + assert "_test_x_card" in m.constraints + + # Segment indicators should have n-1 elements + z = m.variables["_test_x_z"] + assert z.sizes["i"] == 2 # n-1 = 3-1 = 2 + + def test_sos2_trivial_single_element(self) -> None: + """Test SOS2 with single element (trivially satisfied, skipped).""" + m = Model() + idx = pd.Index([0], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + # Should not add anything for trivial case + reformulate_sos2(m, x, "_test_") + + assert "_test_x_z" not in m.variables + + def test_sos2_two_elements(self) -> None: + """Test SOS2 with two elements (one segment indicator).""" + m = Model() + idx = pd.Index([0, 1], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + reformulate_sos2(m, x, "_test_") + m.remove_sos_constraints(x) + + # Should have 1 segment indicator + z = m.variables["_test_x_z"] + assert z.sizes["i"] == 1 + + def test_sos2_with_middle_constraints(self) -> None: + """Test SOS2 with 4+ elements to ensure middle constraints are created.""" + m = Model() + idx = pd.Index([0, 1, 2, 3], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + reformulate_sos2(m, x, "_test_") + m.remove_sos_constraints(x) + + # Should have first, middle, and last constraints + assert "_test_x_upper_first" in m.constraints + assert "_test_x_upper_mid_1" in m.constraints # middle element at index 1 + assert "_test_x_upper_mid_2" in m.constraints # middle element at index 2 + assert "_test_x_upper_last" in m.constraints + + def test_sos2_multidimensional(self) -> None: + """Test SOS2 reformulation with multi-dimensional variables.""" + m = Model() + idx_i = pd.Index([0, 1, 2], name="i") + idx_j = pd.Index([0, 1], name="j") + x = m.add_variables(lower=0, upper=1, coords=[idx_i, idx_j], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + reformulate_sos2(m, x, "_test_") + m.remove_sos_constraints(x) + + # Segment indicator should have (n-1) elements in i dimension, same j dimension + z = m.variables["_test_x_z"] + assert set(z.dims) == {"i", "j"} + assert z.sizes["i"] == 2 # n-1 = 3-1 = 2 + assert z.sizes["j"] == 2 + + # Cardinality constraint should have j dimension preserved + card_con = m.constraints["_test_x_card"] + assert "j" in card_con.dims + + +class TestReformulateAllSOS: + """Tests for reformulate_all_sos.""" + + def test_reformulate_single_sos1(self) -> None: + """Test reformulating a single SOS1 variable.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulated = reformulate_all_sos(m) + + assert reformulated == ["x"] + assert len(list(m.variables.sos)) == 0 # No more SOS variables + + def test_reformulate_multiple_sos(self) -> None: + """Test reformulating multiple SOS variables.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + y = m.add_variables(lower=0, upper=2, coords=[idx], name="y") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_sos_constraints(y, sos_type=2, sos_dim="i") + + reformulated = reformulate_all_sos(m) + + assert set(reformulated) == {"x", "y"} + assert len(list(m.variables.sos)) == 0 + + def test_reformulate_skips_single_element(self) -> None: + """Test that single element SOS is skipped.""" + m = Model() + idx = pd.Index([0], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulated = reformulate_all_sos(m) + + # Single element SOS is skipped (trivially satisfied) + assert reformulated == [] + + def test_reformulate_skips_zero_bounds(self) -> None: + """Test that variables with zero bounds are skipped.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=0, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulated = reformulate_all_sos(m) + + # Zero bounds means variable is fixed to 0, no reformulation needed + assert reformulated == [] + + def test_reformulate_raises_on_infinite_bounds(self) -> None: + """Test that infinite bounds raise ValueError.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + with pytest.raises(ValueError, match="infinite"): + reformulate_all_sos(m) + + def test_reformulate_raises_on_negative_lower_bounds(self) -> None: + """Test that negative lower bounds raise ValueError.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=-1, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + with pytest.raises(ValueError, match="negative lower bounds"): + reformulate_all_sos(m) + + +class TestModelReformulateSOS: + """Tests for Model.reformulate_sos_constraints method.""" + + def test_reformulate_inplace(self) -> None: + """Test inplace reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + result = m.reformulate_sos_constraints() + + assert result == ["x"] + assert len(list(m.variables.sos)) == 0 + assert "_sos_reform_x_y" in m.variables + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestSolveWithReformulation: + """Tests for solving with SOS reformulation.""" + + def test_sos1_maximize_with_highs(self) -> None: + """Test SOS1 maximize problem with HiGHS using reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # Should maximize by choosing x[2] = 1 + assert np.isclose(x.solution.values[2], 1, atol=1e-5) + assert np.isclose(x.solution.values[0], 0, atol=1e-5) + assert np.isclose(x.solution.values[1], 0, atol=1e-5) + assert m.objective.value is not None + assert np.isclose(m.objective.value, 3, atol=1e-5) + + def test_sos1_minimize_with_highs(self) -> None: + """Test SOS1 minimize problem with HiGHS using reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([3, 2, 1]), sense="min") + + m.solve(solver_name="highs", reformulate_sos=True) + + # Should minimize to 0 by setting all x = 0 + assert m.objective.value is not None + assert np.isclose(m.objective.value, 0, atol=1e-5) + + def test_sos2_maximize_with_highs(self) -> None: + """Test SOS2 maximize problem with HiGHS using reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # SOS2 allows two adjacent non-zeros, so x[1] and x[2] can both be 1 + # Maximum is 2 + 3 = 5 + assert m.objective.value is not None + assert np.isclose(m.objective.value, 5, atol=1e-5) + # Check that at most two adjacent variables are non-zero + nonzero_count = (np.abs(x.solution.values) > 1e-5).sum() + assert nonzero_count <= 2 + + def test_sos2_different_coefficients(self) -> None: + """Test SOS2 with different coefficients.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + m.add_objective(x * np.array([2, 1, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # Best is x[1]=1 and x[2]=1 giving 1+3=4 + # or x[0]=1 and x[1]=1 giving 2+1=3 + assert m.objective.value is not None + assert np.isclose(m.objective.value, 4, atol=1e-5) + + def test_reformulate_sos_false_raises_error(self) -> None: + """Test that HiGHS without reformulate_sos raises error.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + with pytest.raises(ValueError, match="does not support SOS"): + m.solve(solver_name="highs", reformulate_sos=False) + + def test_multidimensional_sos1_with_highs(self) -> None: + """Test multi-dimensional SOS1 with HiGHS.""" + m = Model() + idx_i = pd.Index([0, 1, 2], name="i") + idx_j = pd.Index([0, 1], name="j") + x = m.add_variables(lower=0, upper=1, coords=[idx_i, idx_j], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # For each j, at most one x[i, j] can be non-zero + # Maximum is achieved by one non-zero per j column: 2 total + assert m.objective.value is not None + assert np.isclose(m.objective.value, 2, atol=1e-5) + + # Check SOS1 is satisfied for each j + for j in idx_j: + nonzero_count = (np.abs(x.solution.sel(j=j).values) > 1e-5).sum() + assert nonzero_count <= 1 + + def test_multidimensional_sos2_with_highs(self) -> None: + """Test multi-dimensional SOS2 with HiGHS.""" + m = Model() + idx_i = pd.Index([0, 1, 2], name="i") + idx_j = pd.Index([0, 1], name="j") + x = m.add_variables(lower=0, upper=1, coords=[idx_i, idx_j], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # For each j, at most two adjacent x[i, j] can be non-zero + # Maximum is achieved by two adjacent non-zeros per j column: 4 total + assert m.objective.value is not None + assert np.isclose(m.objective.value, 4, atol=1e-5) + + # Check SOS2 is satisfied for each j + for j in idx_j: + sol_j = x.solution.sel(j=j).values + nonzero_indices = np.where(np.abs(sol_j) > 1e-5)[0] + # At most 2 non-zeros + assert len(nonzero_indices) <= 2 + # If 2 non-zeros, they must be adjacent + if len(nonzero_indices) == 2: + assert abs(nonzero_indices[1] - nonzero_indices[0]) == 1 + + +@pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobi not installed") +class TestEquivalenceWithGurobi: + """Tests comparing reformulated solutions with native Gurobi SOS.""" + + def test_sos1_equivalence(self) -> None: + """Test that reformulated SOS1 gives same result as native Gurobi.""" + gurobipy = pytest.importorskip("gurobipy") + + # Native Gurobi solution + m1 = Model() + idx = pd.Index([0, 1, 2], name="i") + x1 = m1.add_variables(lower=0, upper=1, coords=[idx], name="x") + m1.add_sos_constraints(x1, sos_type=1, sos_dim="i") + m1.add_objective(x1 * np.array([1, 2, 3]), sense="max") + + try: + m1.solve(solver_name="gurobi") + except gurobipy.GurobiError as exc: + pytest.skip(f"Gurobi environment unavailable: {exc}") + + # Reformulated solution with HiGHS + m2 = Model() + x2 = m2.add_variables(lower=0, upper=1, coords=[idx], name="x") + m2.add_sos_constraints(x2, sos_type=1, sos_dim="i") + m2.add_objective(x2 * np.array([1, 2, 3]), sense="max") + + if "highs" in available_solvers: + m2.solve(solver_name="highs", reformulate_sos=True) + assert m1.objective.value is not None + assert m2.objective.value is not None + assert np.isclose(m1.objective.value, m2.objective.value, atol=1e-5) + + def test_sos2_equivalence(self) -> None: + """Test that reformulated SOS2 gives same result as native Gurobi.""" + gurobipy = pytest.importorskip("gurobipy") + + # Native Gurobi solution + m1 = Model() + idx = pd.Index([0, 1, 2], name="i") + x1 = m1.add_variables(lower=0, upper=1, coords=[idx], name="x") + m1.add_sos_constraints(x1, sos_type=2, sos_dim="i") + m1.add_objective(x1 * np.array([1, 2, 3]), sense="max") + + try: + m1.solve(solver_name="gurobi") + except gurobipy.GurobiError as exc: + pytest.skip(f"Gurobi environment unavailable: {exc}") + + # Reformulated solution with HiGHS + m2 = Model() + x2 = m2.add_variables(lower=0, upper=1, coords=[idx], name="x") + m2.add_sos_constraints(x2, sos_type=2, sos_dim="i") + m2.add_objective(x2 * np.array([1, 2, 3]), sense="max") + + if "highs" in available_solvers: + m2.solve(solver_name="highs", reformulate_sos=True) + assert m1.objective.value is not None + assert m2.objective.value is not None + assert np.isclose(m1.objective.value, m2.objective.value, atol=1e-5) + + +class TestEdgeCases: + """Tests for edge cases.""" + + def test_preserves_non_sos_variables(self) -> None: + """Test that non-SOS variables are preserved.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_variables(lower=0, upper=2, coords=[idx], name="y") # No SOS + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulate_all_sos(m) + + # y should be unchanged + assert "y" in m.variables + assert "sos_type" not in m.variables["y"].attrs + + def test_custom_prefix(self) -> None: + """Test custom prefix for reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulate_all_sos(m, prefix="_custom_") + + assert "_custom_x_y" in m.variables + assert "_custom_x_upper" in m.constraints + assert "_custom_x_card" in m.constraints + + def test_constraints_with_sos_variables(self) -> None: + """Test that existing constraints with SOS variables work after reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + y = m.add_variables(lower=0, upper=10, name="y") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + # Add constraint involving SOS variable + m.add_constraints(x.sum() <= y, name="linking") + + # Reformulate + reformulate_all_sos(m) + + # Original constraint should still exist + assert "linking" in m.constraints + + def test_float_coordinates(self) -> None: + """Test SOS with float coordinates (common for piecewise linear).""" + m = Model() + breakpoints = pd.Index([0.0, 0.5, 1.0], name="bp") + lambdas = m.add_variables(lower=0, upper=1, coords=[breakpoints], name="lambda") + m.add_sos_constraints(lambdas, sos_type=2, sos_dim="bp") + + reformulate_all_sos(m) + + # Should work with float coordinates + assert "_sos_reform_lambda_z" in m.variables + z = m.variables["_sos_reform_lambda_z"] + # Segment indicators have n-1 = 2 elements + assert z.sizes["bp"] == 2 + + def test_custom_big_m_removed_on_remove_sos(self) -> None: + """Test that custom big_m attribute is removed with SOS constraint.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=100, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=10) + + assert "big_m_upper" in x.attrs + + m.remove_sos_constraints(x) + + assert "big_m_upper" not in x.attrs + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestCustomBigM: + """Tests for custom Big-M functionality.""" + + def test_solve_with_custom_big_m(self) -> None: + """Test solving with custom big_m value.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + # Large bounds but tight effective constraint + x = m.add_variables(lower=0, upper=1000, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=1) + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # With big_m=1, maximum should be 3 (x[2]=1) + assert m.objective.value is not None + assert np.isclose(m.objective.value, 3, atol=1e-5) + + def test_solve_with_infinite_bounds_and_custom_big_m(self) -> None: + """Test solving with infinite bounds but custom big_m.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=5) + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + # Should work because big_m is specified + m.solve(solver_name="highs", reformulate_sos=True) + + assert m.objective.value is not None + # Max is 15 (x[2]=5) + assert np.isclose(m.objective.value, 15, atol=1e-5)