diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 67399d44..d3ff4dd8 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -31,8 +31,9 @@ Upcoming Version *New feature: constraint softening* -* A constraint can now be softened with ``Constraint.soften(penalty, max_violation=None, name=None)``, which adds a slack variable (a positive/negative pair for equality constraints) to the constraint's ``lhs`` and a penalty term to the objective, returning a ``Slack`` named tuple. It is not supported on frozen constraints or on detached copies from ``.mutable()``, ``.sel()``, or ``.isel()``. ``model.add_constraints(..., penalty=...)`` is a shortcut that softens the constraint right after creation and cannot be combined with ``freeze=True``. -* The slack variable(s) created by ``soften()`` can be retrieved afterwards via the new ``Constraint.slack`` property. +* A constraint can now be softened with ``Constraint.soften(penalty, max_violation=None, name=None)``, which adds a slack variable (a positive/negative pair for equality constraints) to the constraint's ``lhs`` and a penalty term to the objective, returning a ``Slack`` named tuple. It is not supported on detached copies from ``.mutable()``, ``.sel()``, or ``.isel()``. ``model.add_constraints(..., penalty=...)`` is a shortcut that softens the constraint right after creation. +* A frozen ``CSRConstraint`` is softened natively, without densifying: the slack terms are appended to its sparse rows in place, so ``add_constraints(..., penalty=..., freeze=True)`` works as well. The slack is kept by ``to_dense()``/``mutable()``, ``Constraint.freeze()`` and the netcdf round trip. (`#975 `__) +* The slack variable(s) created by ``soften()`` can be retrieved afterwards via the new ``slack`` property of ``Constraint`` and ``CSRConstraint``. *Other* @@ -77,7 +78,7 @@ Upcoming Version * Adding or removing variables after a constraint was added with ``freeze=True`` no longer breaks ``model.matrices``. The frozen constraint stored dense variable positions as its matrix columns, so blocks frozen at different times disagreed on their width and stacking them raised ``ValueError: inconsistent shapes``. Raw variable labels are stored instead and mapped to positions when the matrix is assembled. Frozen constraints in netCDF files written by earlier versions are read as before. (`#926 `__) * A ``CSRConstraint`` built from a CSR-backed expression now keeps its auxiliary coordinates, e.g. the key levels of a multi-key ``groupby(..., observed=True).sum(sparse=True)`` or a scalar coordinate left by a scalar selection, so its ``coords`` equal those of the dense constraint. They are kept by ``to_dense()``/``mutable()`` and written to and read from netcdf, and ``Constraint.freeze()`` keeps them as well. (`#941 `__) * ``linopy.read_netcdf`` no longer fails for a whole model when a frozen constraint has a dimension or MultiIndex level whose name contains ``-``: the coordinate metadata of a ``CSRConstraint`` is now stored under positional variable names with the names as attributes. Files written by earlier versions are read as before. -* ``CSRConstraint.loc``, ``update``, ``from_rule`` and ``soften``, and assigning to ``coeffs``, ``vars`` or ``sign``, now raise an ``AttributeError`` saying the operation is not supported on a frozen constraint and naming ``.mutable()`` as the way out (for ``from_rule``: build with ``Constraint.from_rule`` and ``.freeze()`` the result; for ``soften``: add the constraint with ``freeze=False``), like the read-only ``rhs``, ``lhs`` and ``scaling`` setters, instead of a bare ``AttributeError``. (`#963 `__) +* ``CSRConstraint.loc``, ``update`` and ``from_rule``, and assigning to ``coeffs``, ``vars`` or ``sign``, now raise an ``AttributeError`` saying the operation is not supported on a frozen constraint and naming ``.mutable()`` as the way out (for ``from_rule``: build with ``Constraint.from_rule`` and ``.freeze()`` the result), like the read-only ``rhs``, ``lhs`` and ``scaling`` setters, instead of a bare ``AttributeError``. (`#963 `__) * ``linopy.merge`` with ``join="outer"``, ``"left"`` or ``"right"`` no longer raises an xarray ``MergeError`` when only one operand carries an auxiliary coordinate on a joined dimension. The constant was reindexed with a fill value of ``0`` that was also written into the auxiliary coordinate, so it no longer matched the coefficients' copy. The coordinate now follows its rows onto the joined grid and is NaN where no operand provides it, under both semantics. The same holds for ``.add`` / ``.sub`` / ``.mul`` / ``.div`` with a constant and a reindexing ``join=``. * A frozen constraint caches the label-to-position mapping of its matrix columns by weak reference and only rebuilds it when the constraint or the set of variables changes. While a persistent snapshot holds the arrays, repeated matrix assembly on an unchanged model returns the same objects, so the snapshot diff can again skip the comparison of untouched frozen constraints by object identity; one-off exports such as ``to_file`` retain no extra memory. (`#933 `__) * ``Solver.close()`` no longer leaves dangling native handles behind. The solver model is now dropped before the environment that owns it, instead of after. And the COPT and MindOpt file interfaces no longer hand back a model they already disposed: after a file-based COPT or MindOpt solve, ``model.solver_model`` is ``None`` rather than a handle into freed memory. (`#899 `__) diff --git a/linopy/constraints.py b/linopy/constraints.py index 1a61d8bb..1f0bbdde 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -10,7 +10,15 @@ import warnings import weakref from abc import ABC, abstractmethod -from collections.abc import Callable, Generator, Hashable, ItemsView, Iterator, Sequence +from collections.abc import ( + Callable, + Generator, + Hashable, + ItemsView, + Iterator, + Mapping, + Sequence, +) from dataclasses import dataclass from itertools import product from typing import ( @@ -302,6 +310,128 @@ def active_labels(self) -> np.ndarray: def active_row_mask(self) -> np.ndarray: """Boolean mask over raveled rows selecting active constraint rows.""" + @property + def slack(self) -> Slack | None: + """ + Slack variable(s) added via :meth:`soften`, or ``None`` if the + constraint has never been softened. + """ + names = _slack_names(self.attrs) + if names is None: + return None + positive, negative = names + return Slack( + positive=self.model.variables[positive], + negative=self.model.variables[negative] if negative else None, + ) + + def soften( + self, + penalty: ConstantLike, + *, + max_violation: ConstantLike | None = None, + name: str | None = None, + ) -> Slack: + """ + Soften a constraint, adding a slack variable and a penalty to the objective function. + + Parameters + ---------- + penalty : constant-like + The penalty that will match the slack variable inside the objective function. Must be bigger than 0. + max_violation: constant-like + The max violation possible that caps the slack (upper bound). If None, the slack will be unbounded. + name: string + The name for the slack variable. If None, it will reuse the constraint name and add a '_slack'. + + Returns + ------- + Slack + Named tuple with the `positive` slack variable, and the `negative` one for equality constraints (`None` + for inequality constraints). + + Notes + ----- + On a frozen CSRConstraint the slack terms are appended to the sparse rows in place, without densifying. + A detached copy (from `.mutable()`, `.sel()`, `.isel()` or `.freeze()`) is not registered in + model.constraints, so soften raises ValueError on it. + + Softening an already-softened constraint raises ValueError instead of stacking a second, redundant slack term + onto the same lhs. + + Examples + -------- + >>> from linopy import Model + >>> import pandas as pd + + >>> m = Model() + >>> investments = pd.Index(["A", "B", "C"], name="investments") + >>> expected_return = pd.Series( + ... [0.08, 0.03, 0.1], index=investments, name="expected_return" + ... ) + >>> w = m.add_variables(lower=0, upper=1, coords=[investments], name="weights") + >>> m.add_objective((expected_return * w).sum(), sense="max") + >>> budget_penalty = 2 + + >>> budget_constraint = m.add_constraints(w.sum() == 1, name="budget") + >>> slack = budget_constraint.soften(penalty=budget_penalty) + """ + if not bool(np.all(np.asarray(penalty) > 0)): + raise ValueError("Penalty is not positive.") + + model = self.model + if model.objective.expression.empty: + raise ValueError( + "Objective must be defined via `model.add_objective` before calling `soften` on constraints." + ) + + if model.constraints.data.get(self.name) is not self: + raise ValueError( + f"Constraint {self.name!r} is not the constraint registered in the model, so " + "`soften` would not affect it (it may be a detached copy from `.mutable()`, " + "`.sel()`, `.isel()`, or `.freeze()`). Call `soften` on `model.constraints[name]` " + "directly." + ) + + existing_slack = self.slack + if existing_slack is not None: + raise ValueError( + f"Constraint {self.name!r} was already softened (existing slack " + f"variable {existing_slack.positive.name!r})" + ) + + sign_values = pd.unique(self.sign.values.ravel()) + if len(sign_values) > 1: + raise NotImplementedError( + "Constraint.soften does not support constraints with mixed signs." + ) + sign = sign_values.item() + + name = name or f"{self.name}_slack" + upper = np.inf if max_violation is None else max_violation + + def add_slack(suffix: str) -> variables.Variable: + return model.add_variables( + lower=0, + upper=upper, + coords=self.coords, + mask=self.mask, + name=f"{name}_{suffix}", + ) + + positive = add_slack("pos") + slack = Slack(positive, add_slack("neg") if sign == EQUAL else None) + self._attach_slack(slack, sign) + + violation = positive if slack.negative is None else positive + slack.negative + direction = 1 if model.sense == "min" else -1 + model.objective += direction * (penalty * violation).sum() + return slack + + @abstractmethod + def _attach_slack(self, slack: Slack, sign: str) -> None: + """Add the slack terms to the lhs and record the slack variable names.""" + def __getitem__( self, selector: str | int | slice | list | tuple | dict ) -> Constraint: @@ -599,6 +729,14 @@ def _positional_csr( _UNSUPPORTED = "is not supported on a frozen constraint" +def _slack_names(attrs: Mapping[str, Any]) -> tuple[str, str] | None: + """Names of the positive and negative slack variables recorded by ``soften``.""" + positive = attrs.get("slack_positive") + if positive is None: + return None + return str(positive), str(attrs.get("slack_negative", "")) + + def _frozen_error( attr: str, what: str = "is read-only", @@ -614,6 +752,10 @@ class CSRConstraint(ConstraintBase): """ Frozen constraint backed by a CSR sparse matrix. + The structure is frozen: rows, coeffs, sign and rhs cannot be mutated + (see the raising setters below). The one exception is ``soften``, which + appends slack terms to the CSR matrix in place. + Parameters ---------- csr : scipy.sparse.csr_array @@ -654,6 +796,7 @@ class CSRConstraint(ConstraintBase): "_dual", "_binvar_labels", "_binval", + "_slack", "_positional_cache", ) @@ -671,6 +814,7 @@ def __init__( binvar_labels: np.ndarray | None = None, binval: int | np.ndarray | None = None, scaling: np.ndarray | None = None, + slack: tuple[str, str] | None = None, ) -> None: self._csr = csr self._active_positions = active_positions @@ -688,6 +832,7 @@ def __init__( self._dual = dual self._binvar_labels = binvar_labels self._binval = binval + self._slack = slack self._positional_cache: _PositionalCache | None = None @property @@ -726,6 +871,8 @@ def attrs(self) -> dict[str, Any]: d: dict[str, Any] = {"name": self._name} if self._cindex is not None: d["label_range"] = (self._cindex, self._cindex + self.full_size) + if self._slack is not None: + d["slack_positive"], d["slack_negative"] = self._slack return d @property @@ -779,6 +926,7 @@ def _init_kwargs(self) -> dict[str, Any]: binvar_labels=self._binvar_labels, binval=self._binval, scaling=self._scaling, + slack=self._slack, ) def _replace(self, **changes: Any) -> CSRConstraint: @@ -940,11 +1088,25 @@ def loc(self) -> LocIndexer: def update(self, *args: Any, **kwargs: Any) -> NoReturn: raise _frozen_error("update", _UNSUPPORTED) - def soften(self, *args: Any, **kwargs: Any) -> NoReturn: - raise _frozen_error( - "soften", - _UNSUPPORTED, - "add the constraint with freeze=False to soften it", + def _attach_slack(self, slack: Slack, sign: str) -> None: + slacks = [v for v in slack if v is not None] + sign_coeff = 1.0 if sign == GREATER_EQUAL else -1.0 + coeffs = np.array([sign_coeff, 1.0][: len(slacks)]) + labels = [v.labels.transpose(*self._grid.dims).values.ravel() for v in slacks] + cols = np.stack(labels, axis=1)[self._active_positions].ravel() + n = self.ncons + data = np.tile(coeffs.astype(self._csr.dtype), n) + shape = (n, self._model._xCounter) + indptr = np.arange(n + 1) * len(slacks) + extra = scipy.sparse.csr_array((data, cols, indptr), shape=shape) + widened = scipy.sparse.csr_array( + (self._csr.data, self._csr.indices, self._csr.indptr), shape=shape + ) + self._csr = widened + extra + self._positional_cache = None + self._slack = ( + slack.positive.name, + slack.negative.name if slack.negative is not None else "", ) @classmethod @@ -1174,6 +1336,8 @@ def to_netcdf_ds(self) -> Dataset: } if isinstance(self._sign, str): attrs["sign"] = self._sign + if self._slack is not None: + attrs["slack_positive"], attrs["slack_negative"] = self._slack if self._binvar_labels is not None: attrs["is_indicator"] = True data_vars["_binvar_labels"] = DataArray(self._binvar_labels, dims=["_flat"]) @@ -1229,6 +1393,7 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: if "_binvar_labels" in ds: binvar_labels = ds["_binvar_labels"].values binval = ds["_binval"].values if "_binval" in ds else attrs["binval"] + slack = _slack_names(attrs) return cls( csr, active_positions, @@ -1242,6 +1407,7 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: binvar_labels=binvar_labels, binval=binval, scaling=scaling, + slack=slack, ) def has_labels(self, labels: np.ndarray) -> bool: @@ -1317,7 +1483,7 @@ def sanitize_infinities(self) -> CSRConstraint: return self def freeze(self) -> CSRConstraint: - """Return self (already immutable).""" + """Return self (structure is already frozen).""" return self def to_dense(self) -> Constraint: @@ -1425,7 +1591,10 @@ def from_dense( active_signs = sign_vals[active_mask] unique_signs = np.unique(active_signs) if len(unique_signs) == 0: - sign: str | np.ndarray = "=" + full_unique_signs = np.unique(sign_vals) + sign: str | np.ndarray = ( + str(full_unique_signs.item()) if len(full_unique_signs) == 1 else "=" + ) elif len(unique_signs) == 1: sign = str(unique_signs[0]) else: @@ -1456,6 +1625,7 @@ def from_dense( binvar_labels=binvar_labels, binval=binval, scaling=scaling, + slack=_slack_names(con.attrs), ) @classmethod @@ -1717,21 +1887,6 @@ def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: ) self.update(lhs=value) - @property - def slack(self) -> Slack | None: - """ - Slack variable(s) added via :meth:`soften`, or ``None`` if the - constraint has never been softened. - """ - positive = self.data.attrs.get("slack_positive") - if positive is None: - return None - negative = self.data.attrs.get("slack_negative", "") - return Slack( - positive=self.model.variables[positive], - negative=self.model.variables[negative] if negative else None, - ) - def _assign_lhs( self, expr: expressions.LinearExpression, rhs: DataArray | None = None ) -> None: @@ -2145,135 +2300,17 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai data = lhs.data.assign(sign=sign, rhs=rhs) return cls(data, model=model) - def soften( - self, - penalty: ConstantLike, - *, - max_violation: ConstantLike | None = None, - name: str | None = None, - ) -> Slack: - """ - Soften a constraint, adding a slack variable and a penalty to the objective function. - - Parameters - ---------- - penalty : constant-like - The penalty that will match the slack variable inside the objective function. Must be bigger than 0. - max_violation: constant-like - The max violation possible that caps the slack (upper bound). If None, the slack will be unbounded. - name: string - The name for the slack variable. If None, it well reuse the constraint name and add a '_slack'. - - Returns - ------- - Slack - Named tuple with the `positive` slack variable, and the `negative` one for equality constraints (`None` - for inequality constraints). - - Notes - ----- - Not supported on frozen constraints (e.g. a CSRConstraint from add_constraints(..., freeze=True) or - Model(freeze_constraints=True)). This method is only defined on Constraint and calling it on a CSRConstraint - raises AttributeError. Calling .mutable() first does not help either, since the resulting Constraint is a - detached copy not registered in model.constraints, so soften raises ValueError on it instead. - - Softening an already-softened constraint raises ValueError instead of stacking a second, redundant slack term - onto the same lhs. - - Examples - -------- - >>> from linopy import Model - >>> import pandas as pd - - >>> m = Model() - >>> investments = pd.Index(["A", "B", "C"], name="investments") - >>> expected_return = pd.Series( - ... [0.08, 0.03, 0.1], index=investments, name="expected_return" - ... ) - >>> w = m.add_variables(lower=0, upper=1, coords=[investments], name="weights") - >>> m.add_objective((expected_return * w).sum(), sense="max") - >>> budget_penalty = 2 - - >>> budget_constraint = m.add_constraints(w.sum() == 1, name="budget") - >>> slack = budget_constraint.soften(penalty=budget_penalty) - """ - # Verify valid penalty to continue: - if not bool(np.all(np.asarray(penalty) > 0)): - raise ValueError("Penalty is not positive.") - - # Require the objective function to exist before using soften method (this is to avoid - # `add_objective` overwriting the penalty term added below, since it replaces rather than merges): - model = self.model - if model.objective.expression.empty: - raise ValueError( - "Objective must be defined via `model.add_objective` before calling `soften` on constraints." - ) - - # A detached copy of the constraint (e.g. from `.mutable()`, `.sel()`, `.isel()`) isn't in - # `model.constraints`, so .soften would silently do nothing on the real model. This check is to avoid that: - if model.constraints.data.get(self.name) is not self: - raise ValueError( - f"Constraint {self.name!r} is not the constraint registered in the model, so " - "`soften` would not affect it (it may be a detached copy from `.mutable()`, " - "`.sel()`, or `.isel()`). Call `soften` on `model.constraints[name]` directly." - ) - - if self.slack is not None: - raise ValueError( - f"Constraint {self.name!r} was already softened (existing slack " - f"variable {self.slack.positive.name!r})" - ) - - name = name or f"{self.name}_slack" - upper = np.inf if max_violation is None else max_violation - - sign_values = pd.unique(self.sign.values.ravel()) - if len(sign_values) > 1: - raise NotImplementedError( - "Constraint.soften does not support constraints with mixed signs." - ) - - positive_slack = model.add_variables( - lower=0, - upper=upper, - coords=self.lhs.coords, - mask=self.mask, - name=f"{name}_pos", - ) - negative_slack = None - - # Update left hand side depending on the sign of the constraint: - sign = sign_values.item() - if sign == "<=": - self.update(lhs=self.lhs - positive_slack) - elif sign == ">=": - self.update(lhs=self.lhs + positive_slack) - else: - negative_slack = model.add_variables( - lower=0, - upper=upper, - coords=self.lhs.coords, - mask=self.mask, - name=f"{name}_neg", - ) - self.update(lhs=self.lhs - positive_slack + negative_slack) - - # Update objective function: - constraint_violation = ( - positive_slack + negative_slack - if negative_slack is not None - else positive_slack - ) - direction = 1 if model.sense == "min" else -1 - model.objective += direction * (penalty * constraint_violation).sum() - + def _attach_slack(self, slack: Slack, sign: str) -> None: + positive = slack.positive + lhs = self.lhs + positive if sign == GREATER_EQUAL else self.lhs - positive + if slack.negative is not None: + lhs = lhs + slack.negative + self.update(lhs=lhs) self._data = self._data.assign_attrs( - slack_positive=positive_slack.name, - slack_negative=negative_slack.name if negative_slack is not None else "", + slack_positive=slack.positive.name, + slack_negative=slack.negative.name if slack.negative is not None else "", ) - return Slack(positive=positive_slack, negative=negative_slack) - def to_polars(self) -> pl.DataFrame: """ Convert the constraint to a polars DataFrame. diff --git a/linopy/model.py b/linopy/model.py index 93d80b19..ec6ebd10 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1292,11 +1292,10 @@ def add_constraints( penalty : constant-like, optional If given, soften the constraint right away by calling :meth:`Constraint.soften` with this penalty, adding a slack variable - and a penalty term to the objective. Not allowed together with - ``freeze=True`` (or a model default of ``freeze_constraints=True``), - since softening requires a mutable, registered ``Constraint``. - The resulting Slack is not returned by this shortcut; retrieve - the slack variable(s) from model.variables using the derived + and a penalty term to the objective. A frozen constraint is + softened natively on its sparse rows. The resulting Slack is not + returned by this shortcut; retrieve the slack variable(s) from + model.variables using the derived name f"{name}_slack_pos" (and f"{name}_slack_neg" for equality constraints). @@ -1312,12 +1311,6 @@ def add_constraints( chunked = bool(freeze and self.chunk) freeze = freeze and not self.chunk - if penalty is not None and freeze: - raise ValueError( - "`penalty` cannot be combined with `freeze=True` (or a model default of `freeze_constraints=True`), " - "since `soften` is not supported on frozen constraints." - ) - if isinstance(sign, str): sign = maybe_replace_sign(sign) elif sign is not None: @@ -1346,11 +1339,9 @@ def add_constraints( cindex = self._cCounter self._cCounter += con.full_size con = con.assign_labels(cindex, name, scaling_grid.values.ravel()) - return self.constraints.add(con) + return self._soften_added(self.constraints.add(con), penalty) if isinstance(con, CSRConstraint): - if penalty is not None: - reason = "`penalty` given, softening needs a mutable constraint" - elif chunked: + if chunked: reason = "chunked model, `Model.chunk` adds constraints unfrozen" else: reason = "constraint added unfrozen, `freeze=False`" @@ -1418,10 +1409,15 @@ def add_constraints( enforce_no_multiindex(data, context=f"constraint {name!r}") constraint = Constraint(data, name=name, model=self, skip_broadcast=True) - added = self.constraints.add(constraint, freeze=freeze) + return self._soften_added(self.constraints.add(constraint, freeze), penalty) + + @staticmethod + def _soften_added( + constraint: ConstraintBase, penalty: ConstantLike | None + ) -> ConstraintBase: if penalty is not None: constraint.soften(penalty=penalty) - return added + return constraint def add_indicator_constraints( self, diff --git a/test/test_constraint.py b/test/test_constraint.py index c0f0e41b..13928b2d 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -113,24 +113,6 @@ def test_add_constraints_penalty_softens_constraint( assert_linequal(m.objective.expression, expected_objective) -def test_add_constraints_penalty_with_freeze_true_raises( - m: Model, x: linopy.Variable -) -> None: - with pytest.raises(ValueError, match="`penalty` cannot be combined"): - m.add_constraints(x >= 0, name="frozen_penalized", freeze=True, penalty=10) - - -def test_add_constraints_penalty_with_model_freeze_default_raises() -> None: - """ - `freeze=None` resolves to the model's `freeze_constraints` default, which must - also be checked against `penalty`, not just an explicit `freeze=True`. - """ - m = Model(freeze_constraints=True) - x = m.add_variables(coords=[pd.RangeIndex(10, name="first")], name="x") - with pytest.raises(ValueError, match="`penalty` cannot be combined"): - m.add_constraints(x >= 0, name="frozen_by_default_penalized", penalty=10) - - def test_constraint_name(c: linopy.constraints.CSRConstraint) -> None: assert c.name == "c" diff --git a/test/test_csr.py b/test/test_csr.py index b75f49de..d24e9a8a 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -28,7 +28,12 @@ from linopy.constraints import Constraint, ConstraintBase, CSRConstraint from linopy.csr import CSRLinearExpression, Grid from linopy.semantics import is_v1 -from linopy.testing import assert_conequal, assert_linequal, assert_quadequal +from linopy.testing import ( + assert_conequal, + assert_linequal, + assert_quadequal, + assert_varequal, +) def require_v1() -> None: @@ -1352,11 +1357,6 @@ def add_chunked(e: LinearExpression, c: Case) -> Any: return c.m.add_constraints(e >= 1, freeze=True) -def add_with_penalty(e: LinearExpression, c: Case) -> Any: - c.m.add_objective(1.0 * c.gen_p.sum()) - return c.m.add_constraints(e >= 1, penalty=1.0) - - DENSIFY_OPS: dict[str, tuple[Callable[[LinearExpression, Case], Any], str]] = { "data": (lambda e, c: e.data, "`.data` read"), "merge": (lambda e, c: e + 1.0 * c.flow, "over different dimensions"), @@ -1372,7 +1372,6 @@ def add_with_penalty(e: LinearExpression, c: Case) -> Any: "constraint added unfrozen, `freeze=False`", ), "add-chunked": (add_chunked, "chunked model"), - "add-penalty": (add_with_penalty, "`penalty` given"), "new-dim": (lambda e, c: e * xr.DataArray([1.0, 2.0], coords=[LOC]), "new dim"), "sum-kwargs": (lambda e, c: e.sum(dims="bus"), "`.data` read"), "groupby-fallback": ( @@ -2004,7 +2003,6 @@ def test_frozen_constraint_mutation_names_mutable(op: str) -> None: FROZEN_UNSUPPORTED: dict[str, tuple[Callable[[CSRConstraint], Any], str]] = { - "soften": (lambda con: con.soften(penalty=1.0), r"freeze=False"), "from_rule": ( lambda con: CSRConstraint.from_rule(con.model, lambda m, i: None, [[0]]), r"Constraint\.from_rule .*\.freeze\(\)", @@ -2050,13 +2048,158 @@ def test_expression_rhs_freezes_sparse_and_matches_dense(rhs: str, form: str) -> assert_frozen_equal(con1, con2) -@pytest.mark.parametrize("default", [False, True], ids=["argument", "model-default"]) -def test_penalty_with_freeze_raises(default: bool) -> None: +MASK_KINDS: dict[str, Callable[[Case], xr.DataArray | None]] = { + "full": lambda c: None, + "masked": lambda c: c.load > 4, + "all-masked": lambda c: c.load > np.inf, +} + + +def softened( + sign: str, + mask_kind: str, + route: str, + max_violation: float | None, + freeze: bool, +) -> tuple[Case, ConstraintBase]: + c = base_model() + c.m.add_objective(1.0 * c.gen_p.sum()) + c.m.freeze_constraints = freeze and route == "penalty-default" + mask = MASK_KINDS[mask_kind](c) + args = (c.balance_lhs(sparse=freeze), sign, c.load) + kwargs: dict[str, Any] = dict(name="c", mask=mask) + if route == "soften": + con = c.m.add_constraints(*args, **kwargs, freeze=freeze) + con.soften(penalty=2.0, max_violation=max_violation) + else: + kwargs["freeze"] = None if route == "penalty-default" else freeze + con = c.m.add_constraints(*args, **kwargs, penalty=2.0) + return c, con + + +SOFTEN_ROUTES = [ + ("soften", None), + ("soften", 5.0), + ("penalty", None), + ("penalty-default", None), +] + + +@pytest.mark.parametrize(("route", "max_violation"), SOFTEN_ROUTES) +@pytest.mark.parametrize("mask_kind", list(MASK_KINDS)) +@pytest.mark.parametrize("sign", ["<=", ">=", "=="]) +def test_frozen_soften_matches_dense( + sign: str, mask_kind: str, route: str, max_violation: float | None, tmp_path: Path +) -> None: + require_v1() + dense_case, dense = softened(sign, mask_kind, route, max_violation, freeze=False) + with no_densify(): + case, con = softened(sign, mask_kind, route, max_violation, freeze=True) + case.m.to_netcdf(tmp_path / "m.nc") + assert isinstance(con, CSRConstraint) + assert case.m.constraints["c"] is con + assert_frozen_equal(dense, con) + obj, want_obj = case.m.objective.expression, dense_case.m.objective.expression + assert isinstance(obj, LinearExpression) and isinstance(want_obj, LinearExpression) + assert_cells_equal(obj, want_obj, ()) + assert dense.slack is not None + read = linopy.read_netcdf(tmp_path / "m.nc").constraints["c"] + assert isinstance(read, CSRConstraint) + assert_frozen_equal(dense, read) + for frozen_con in (con, read, con.mutable(), dense.freeze()): + slack = frozen_con.slack + assert slack is not None + assert_varequal(slack.positive, dense.slack.positive) + assert (slack.negative is None) == (dense.slack.negative is None) + if slack.negative is not None and dense.slack.negative is not None: + assert_varequal(slack.negative, dense.slack.negative) + + +def test_frozen_soften_keeps_sparse_objective() -> None: require_v1() c = base_model() - c.m.freeze_constraints = default - lhs = c.balance_lhs(sparse=True) - with pytest.raises(ValueError, match="`penalty` cannot be combined with `freeze"): - c.m.add_constraints( - lhs >= c.load, penalty=1.0, freeze=None if default else True + with no_densify(): + c.m.add_objective((1.0 * c.gen_p).groupby(c.gbus).sum(sparse=True).sum()) + lhs = c.balance_lhs(sparse=True) + c.m.add_constraints(lhs >= c.load, freeze=True, penalty=2.0) + assert c.m.objective.expression.is_sparse + + +def test_frozen_soften_max_sense_with_array_penalty() -> None: + require_v1() + c = base_model() + c.m.add_objective(1.0 * c.gen_p.sum(), sense="max") + with no_densify(): + con = c.m.add_constraints( + c.balance_lhs(sparse=True), ">=", c.load, name="c", freeze=True ) + assert isinstance(con, CSRConstraint) + penalty = xr.full_like(c.load, 2.0) + slack = con.soften(penalty=penalty) + expected_objective = 1.0 * c.gen_p.sum() - (penalty * slack.positive).sum() + obj = c.m.objective.expression + assert isinstance(obj, LinearExpression) + assert_cells_equal(obj, expected_objective, ()) + + +@pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs") +@pytest.mark.parametrize("sign", ["<=", ">=", "=="]) +def test_frozen_soften_solves_like_dense(sign: str) -> None: + require_v1() + values = [] + for freeze in (False, True): + c, _ = softened(sign, "masked", "soften", 20.0, freeze) + for var in (c.gen_p, c.flow): + var.update(lower=0, upper=1) + c.m.solve("highs") + values.append(c.m.objective.value) + assert values[0] == pytest.approx(values[1]) + + +def soften_twice(con: ConstraintBase) -> Any: + con.soften(penalty=1.0) + return con.soften(penalty=1.0) + + +SOFTEN_ERRORS: dict[str, tuple[Callable[[ConstraintBase], Any], type, str]] = { + "twice": (soften_twice, ValueError, "already softened"), + "zero-penalty": (lambda con: con.soften(penalty=0.0), ValueError, "not positive"), + "no-objective": ( + lambda con: con.soften(penalty=1.0), + ValueError, + "Objective must be defined", + ), + "mixed-signs": ( + lambda con: con.soften(penalty=1.0), + NotImplementedError, + "mixed signs", + ), +} + +# "twice", "zero-penalty" and "no-objective" are already covered for dense +# constraints in test/test_constraint.py; only the frozen path needs them here. +FROZEN_ONLY_ERRORS = {"twice", "zero-penalty", "no-objective"} + +SOFTEN_REJECT_CASES = [ + pytest.param(error, freeze, id=f"{error}-{'frozen' if freeze else 'dense'}") + for error in SOFTEN_ERRORS + for freeze in ((True,) if error in FROZEN_ONLY_ERRORS else (False, True)) +] + + +@pytest.mark.parametrize(("error", "freeze"), SOFTEN_REJECT_CASES) +def test_soften_rejects(error: str, freeze: bool) -> None: + require_v1() + c = base_model() + if error != "no-objective": + c.m.add_objective(1.0 * c.gen_p.sum()) + sign: Any = ">=" + if error == "mixed-signs": + sign = xr.DataArray(np.where(c.load > 4, ">=", "<="), coords=c.load.coords) + con = c.m.add_constraints( + c.balance_lhs(sparse=False), sign, c.load, name="c", freeze=freeze + ) + assert isinstance(con, CSRConstraint) == freeze + call, exc, match = SOFTEN_ERRORS[error] + with pytest.raises(exc, match=match): + call(con)