diff --git a/doc/api.rst b/doc/api.rst index b62351ea..35b7b5d4 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -247,6 +247,61 @@ Post-solve access expressions.Expressions.solution +LazyExpression +============== + +Placeholder for an expression that is built on demand. Returned by +:meth:`Model.add_expressions ` when +`data` is a callable; arithmetic on a ``LazyExpression`` returns another +``LazyExpression`` rather than evaluating immediately. + +.. autosummary:: + :toctree: generated/ + + expressions.LazyExpression + +Evaluation +---------- + +.. autosummary:: + :toctree: generated/ + + expressions.LazyExpression.evaluate + expressions.LazyExpression.promote + expressions.LazyExpression.is_evaluatable + +Arithmetic and constraints +--------------------------- + +Named counterparts of the arithmetic dunders, shared with +``LinearExpression``/``QuadraticExpression`` via +:class:`AbstractExpression `. These +stay lazy where possible; ``to_constraint``/``le``/``ge``/``eq`` (and the +comparison operators) force evaluation and return a ``Constraint``. + +.. autosummary:: + :toctree: generated/ + + expressions.LazyExpression.add + expressions.LazyExpression.sub + expressions.LazyExpression.mul + expressions.LazyExpression.div + expressions.LazyExpression.pow + expressions.LazyExpression.dot + expressions.LazyExpression.le + expressions.LazyExpression.ge + expressions.LazyExpression.eq + expressions.LazyExpression.to_constraint + +Post-solve access +----------------- + +.. autosummary:: + :toctree: generated/ + + expressions.LazyExpression.solution + + LinearExpression ================ @@ -687,3 +742,13 @@ These warning classes can be silenced or filtered via EvolvingAPIWarning PerformanceWarning + NonLinearExpressionWarning + + +Exceptions +========== + +.. autosummary:: + :toctree: generated/ + + NonLinearOperationError diff --git a/doc/release_notes.rst b/doc/release_notes.rst index ada70ec6..24e49002 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -5,6 +5,15 @@ Upcoming Version ---------------- +*Named expressions* + +* ``Model.add_expressions`` also accepts a callable for ``data``, in which case the expression is not built immediately: a ``LazyExpression`` placeholder is registered instead, and the callable (``data(model, **params)``) only runs when the expression is evaluated, via ``.evaluate()``, ``.promote()``, ``.solution``, or a comparison (``<=``, ``>=``, ``==``). ``mask`` accepts a callable too (resolved at the same time as `data`), in addition to a concrete array. Arithmetic between ``LazyExpression`` objects — and between a ``LazyExpression`` and anything else — stays lazy, returning a new, unnamed ``LazyExpression`` that composes the operands rather than evaluating them. + ``Model.to_netcdf`` gained a ``lazy={"evaluate", "skip", "raise"}`` parameter (default ``"evaluate"``) controlling what happens to lazy entries, since an arbitrary evaluator callable cannot itself be serialized to netcdf. +* ``LazyExpression`` now shares its arithmetic/constraint protocol with ``LinearExpression``/``QuadraticExpression`` via a common ``AbstractExpression`` base, and gained the named counterparts (``add``, ``sub``, ``mul``, ``div``, ``pow``, ``dot``, ``le``, ``ge``, ``eq``, ``to_constraint``) that were previously eager-only — a ``LazyExpression`` is now a drop-in substitute for an eager expression wherever those are called, including with a ``join`` argument. ``lazy ** 2`` no longer evaluates the underlying evaluator twice, and ``-lazy`` now matches eager negation exactly (previously it filled masked/NaN coefficients with 0 before negating, like ``lazy * -1`` does). +* Dividing a ``LazyExpression`` by a variable or another expression (e.g. ``cost / output`` for a unit cost), and other operations with no linear/quadratic form (e.g. ``lazy ** 3``), no longer raise immediately: they build a ``LazyExpression`` that is only readable via ``.solution`` once the model has been solved, since every operand is then just a number. ``.evaluate()``, ``.promote()`` and constraint-building still raise — now a dedicated ``linopy.NonLinearOperationError`` (a ``TypeError`` subclass, so existing ``except TypeError`` code keeps working) — pointing at ``.solution`` instead. Where this is already decidable at construction time (as opposed to only inside a leaf callable's body), a ``linopy.NonLinearExpressionWarning`` is raised immediately, and ``LazyExpression.is_evaluatable`` reports it without forcing evaluation. +* A lazy expression's callable may also read post-solve-only data that is not itself an expression — most notably a constraint's ``.dual`` — and return a plain ``DataArray``/constant, or a ``LazyExpression`` built from one. It is then only readable via ``.solution``; ``.promote()`` raises since there is nothing to store as a named expression. + + Version 0.9.1 ------------- @@ -47,7 +56,6 @@ Version v0.9.0 * ``Model.to_netcdf`` now records the writing linopy version in the ``_linopy_version`` dataset attribute. Files written by older versions (without the attribute) continue to read unchanged. (`#780 `__) - *Other* * Default internal integer labels to ``int32``, cutting memory ~25% and speeding up model build 10-35%. Models exceeding the int32 maximum (~2.1 billion labels) widen to ``int64`` automatically with a ``UserWarning``; pass ``Model(dtypes={"labels": np.int64})`` upfront to avoid the mid-build upcast (exposed read-only via ``Model.dtypes``). (`#566 `__) diff --git a/examples/creating-expressions.ipynb b/examples/creating-expressions.ipynb index ce6017ba..8582479f 100644 --- a/examples/creating-expressions.ipynb +++ b/examples/creating-expressions.ipynb @@ -575,6 +575,229 @@ " optimal solution — handy for inspecting derived quantities without\n", " rebuilding the expression by hand." ] + }, + { + "cell_type": "markdown", + "id": "53", + "metadata": {}, + "source": [ + "## Deferred (lazy) expressions\n", + "\n", + "Passing a callable to ``data`` builds the expression on demand instead of\n", + "right away. ``Model.add_expressions`` registers a ``LazyExpression``\n", + "placeholder and only calls the callable once the expression is actually\n", + "needed — via ``.evaluate()``, ``.promote()``, ``.solution``, or a comparison.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "54", + "metadata": {}, + "outputs": [], + "source": [ + "deferred = m.add_expressions(lambda model: model.variables[\"x\"] * 3, name=\"deferred\")\n", + "deferred" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "55", + "metadata": {}, + "outputs": [], + "source": [ + "deferred.evaluate()" + ] + }, + { + "cell_type": "markdown", + "id": "56", + "metadata": {}, + "source": [ + "Extra keyword arguments are forwarded to the callable every time it runs,\n", + "and ``mask`` may itself be a callable, resolved at the same time:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "57", + "metadata": {}, + "outputs": [], + "source": [ + "def scaled(model, factor):\n", + " return model.variables[\"x\"] * factor\n", + "\n", + "\n", + "scaled_expr = m.add_expressions(scaled, name=\"scaled\", factor=5)\n", + "scaled_expr.evaluate()" + ] + }, + { + "cell_type": "markdown", + "id": "58", + "metadata": {}, + "source": [ + "Arithmetic between lazy expressions — and between a lazy expression and\n", + "anything else — stays lazy: it returns a new, unnamed ``LazyExpression``\n", + "that composes the operands, rather than evaluating them immediately.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59", + "metadata": {}, + "outputs": [], + "source": [ + "combined = deferred + scaled_expr\n", + "combined # still a LazyExpression" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "60", + "metadata": {}, + "outputs": [], + "source": [ + "combined.evaluate()" + ] + }, + { + "cell_type": "markdown", + "id": "61", + "metadata": {}, + "source": [ + "Named methods work the same way as on eager expressions, including\n", + "the ``join`` parameter, and stay lazy too:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "62", + "metadata": {}, + "outputs": [], + "source": [ + "combined.add(1, join=\"outer\")" + ] + }, + { + "cell_type": "markdown", + "id": "63", + "metadata": {}, + "source": [ + "``.promote()`` runs the evaluator once and replaces the placeholder\n", + "in-place with the resulting concrete expression:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "64", + "metadata": {}, + "outputs": [], + "source": [ + "deferred.promote()\n", + "m.expressions[\"deferred\"]" + ] + }, + { + "cell_type": "markdown", + "id": "65", + "metadata": {}, + "source": [ + "## Ratios and other post-solve-only expressions\n", + "\n", + "Dividing by a variable or another expression has no linear or quadratic\n", + "form, so it cannot be built into a LinearExpression/QuadraticExpression.\n", + "On a lazy expression it no longer raises outright, though: it defers to a\n", + "LazyExpression that is only readable via .solution, once the model has\n", + "been solved and every operand is just a number.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66", + "metadata": {}, + "outputs": [], + "source": [ + "cost = m.add_expressions(lambda model: 3 * model.variables[\"x\"], name=\"cost\")\n", + "output = m.add_expressions(lambda model: model.variables[\"x\"], name=\"output\")\n", + "unit_cost = cost / output\n", + "unit_cost # still a LazyExpression -- a NonLinearExpressionWarning was also raised" + ] + }, + { + "cell_type": "markdown", + "id": "67", + "metadata": {}, + "source": [ + ".evaluate(), .promote() and constraint-building all raise a\n", + "linopy.NonLinearOperationError (a TypeError subclass) for such an\n", + "expression, pointing at .solution instead:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "68", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " unit_cost.evaluate()\n", + "except Exception as e:\n", + " print(f\"{type(e).__name__}: {e}\")" + ] + }, + { + "cell_type": "markdown", + "id": "69", + "metadata": {}, + "source": [ + "unit_cost.is_evaluatable reports this without forcing an evaluation, and\n", + "once the model is solved, .solution works like any other expression—\n", + "including through m.expressions.solution, since cost and output\n", + "are themselves stored on the model:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "70", + "metadata": {}, + "outputs": [], + "source": [ + "unit_cost.is_evaluatable" + ] + }, + { + "cell_type": "markdown", + "id": "71", + "metadata": {}, + "source": [ + "A lazy expression may also read data that only exists post-solve and is not\n", + "itself an expression — most notably a constraint’s .dual — and return a\n", + "plain DataArray. It is likewise only readable via .solution:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "72", + "metadata": {}, + "outputs": [], + "source": [ + "m.add_constraints(x >= 2, name=\"x_lower_bound\")\n", + "shadow_price = m.add_expressions(\n", + " lambda model: model.constraints[\"x_lower_bound\"].dual, name=\"shadow_price\"\n", + ")\n", + "shadow_price" + ] } ], "metadata": { @@ -597,7 +820,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.3" + "version": "3.11.6" } }, "nbformat": 4, diff --git a/linopy/__init__.py b/linopy/__init__.py index b813f71d..d526cd90 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -19,6 +19,8 @@ GREATER_EQUAL, LESS_EQUAL, EvolvingAPIWarning, + NonLinearExpressionWarning, + NonLinearOperationError, PerformanceWarning, ) from linopy.constraints import ( @@ -27,7 +29,12 @@ Constraints, CSRConstraint, ) -from linopy.expressions import LinearExpression, QuadraticExpression, merge +from linopy.expressions import ( + LazyExpression, + LinearExpression, + QuadraticExpression, + merge, +) from linopy.io import read_netcdf from linopy.model import Model, Variable, Variables from linopy.objective import Objective @@ -56,8 +63,11 @@ "EvolvingAPIWarning", "GREATER_EQUAL", "LESS_EQUAL", + "LazyExpression", "LinearExpression", "Model", + "NonLinearExpressionWarning", + "NonLinearOperationError", "Objective", "OetcHandler", "PiecewiseFormulation", diff --git a/linopy/constants.py b/linopy/constants.py index 7936ef1c..27357ba4 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -22,6 +22,30 @@ class PerformanceWarning(UserWarning): """Warning raised when an operation triggers expensive Dataset reconstruction.""" +class NonLinearOperationError(TypeError): + """ + Raised when an operation would require a non-linear/non-quadratic expression. + + Subclasses :class:`TypeError` so existing ``except TypeError`` handlers (including + Python's own operator dispatch, which relies on ``NotImplemented``/``TypeError``) + keep working unchanged. A :class:`~linopy.expressions.LazyExpression` built from such + an operation can still be read via its ``.solution`` property once the model has been + solved; it just cannot be materialised into a :class:`LinearExpression` or + :class:`QuadraticExpression`. + """ + + +class NonLinearExpressionWarning(UserWarning): + """ + Warned when a :class:`~linopy.expressions.LazyExpression` is built from an operation + that is already known, at construction time, to be non-linear/non-quadratic. + + The resulting expression is still usable through ``.solution`` once the model is + solved; ``.evaluate()``, ``.promote()`` and constraint-building will raise + :class:`NonLinearOperationError`. + """ + + long_EQUAL = "==" short_GREATER_EQUAL = ">" short_LESS_EQUAL = "<" diff --git a/linopy/expressions.py b/linopy/expressions.py index 01cff3a2..298141b7 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -39,7 +39,7 @@ import scipy import xarray as xr import xarray.core.groupby -from numpy import array, nan, ndarray +from numpy import array, nan from pandas.core.frame import DataFrame from pandas.core.series import Series from scipy.sparse import csc_matrix @@ -97,12 +97,16 @@ LESS_EQUAL, STACKED_TERM_DIM, TERM_DIM, + NonLinearExpressionWarning, + NonLinearOperationError, ) from linopy.types import ( CONSTANT_TYPES, ConstantLike, DimsLike, ExpressionLike, + LazySideLike, + MaskLike, SideLike, SignLike, VariableLike, @@ -151,6 +155,34 @@ def _expr_unwrap( return maybe_expr +def _resolve_lazy(value: Any) -> Any: + """ + Evaluate `value` if it is a `LazyExpression`, otherwise return it unchanged. + """ + return value.evaluate() if isinstance(value, LazyExpression) else value + + +def _solution_of(value: Any) -> Any: + """ + Resolve `value` to a post-solve, numeric value. + + A `LazyExpression`, `BaseExpression` or `Variable` is resolved via its `.solution` + property; anything else (a constant, array, or an already-numeric `DataArray` such as a + constraint's `.dual`) is returned unchanged. Used to compose the `.solution` of a derived + `LazyExpression` operand-wise, for operations (e.g. division by a variable) that have no + linear/quadratic form and therefore cannot go through `.evaluate()`. + """ + if isinstance(value, LazyExpression | BaseExpression | variables.Variable): + return value.solution + return value + + +def _as_solution_dataarray(value: Any) -> DataArray: + """Coerce a resolved solution value (a `DataArray`, e.g. a dual, or a plain constant/array) into a named "solution" `DataArray`, matching `BaseExpression.solution`.""" + da = value if isinstance(value, DataArray) else as_dataarray(value) + return da.rename("solution") + + logger = logging.getLogger(__name__) @@ -680,12 +712,214 @@ def sum(self, **kwargs: Any) -> LinearExpression: return LinearExpression(ds, self.model) -class BaseExpression(ABC): - __slots__ = ("_data", "_model") +class AbstractExpression(ABC): + """ + Operator and constraint-building surface shared by eager expressions + (:class:`BaseExpression`) and deferred ones (:class:`LazyExpression`). + + Holds no data and no Dataset machinery: only the numpy/pandas dispatch guards, + the arithmetic protocol and the comparison/constraint surface. `__eq__` returns + a `Constraint` rather than a bool, so instances are deliberately unhashable. + """ + + __slots__ = () __array_ufunc__ = None __array_priority__ = 10000 __pandas_priority__ = 10000 + @abstractmethod + def __add__(self, other: SideLike) -> AbstractExpression: ... + + @abstractmethod + def __radd__(self, other: SideLike) -> AbstractExpression: ... + + @abstractmethod + def __sub__(self, other: SideLike) -> AbstractExpression: ... + + @abstractmethod + def __rsub__(self, other: SideLike) -> AbstractExpression: ... + + @abstractmethod + def __mul__(self, other: SideLike) -> AbstractExpression: ... + + @abstractmethod + def __rmul__(self, other: SideLike) -> AbstractExpression: ... + + @abstractmethod + def __matmul__(self, other: SideLike) -> AbstractExpression: ... + + @abstractmethod + def __pow__(self, other: int) -> AbstractExpression: ... + + @abstractmethod + def __neg__(self) -> AbstractExpression: ... + + @abstractmethod + def __truediv__(self, other: SideLike) -> AbstractExpression: ... + + @abstractmethod + def to_constraint( + self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None + ) -> Constraint: + """ + Turn this expression into a constraint against `rhs` with the given `sign`. + """ + ... + + @abstractmethod + def add( + self, other: SideLike, join: JoinOptions | None = None + ) -> AbstractExpression: + """ + Add an expression to others. + + Parameters + ---------- + other : expression-like + The expression to add. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + ... + + @abstractmethod + def mul( + self, other: SideLike, join: JoinOptions | None = None + ) -> AbstractExpression: + """ + Multiply the expr by a factor. + + Parameters + ---------- + other : expression-like + The factor to multiply by. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + ... + + @abstractmethod + def div( + self, other: VariableLike | ConstantLike, join: JoinOptions | None = None + ) -> AbstractExpression: + """ + Divide the expr by a factor. + + Parameters + ---------- + other : constant-like + The divisor. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + ... + + @abstractmethod + def sub( + self, other: SideLike, join: JoinOptions | None = None + ) -> AbstractExpression: + """ + Subtract others from expression. + + Parameters + ---------- + other : expression-like + The expression to subtract. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + ... + + @abstractmethod + def pow(self, other: int) -> AbstractExpression: + """ + Power of the expression with a coefficient. + """ + ... + + @abstractmethod + def dot(self, other: SideLike) -> AbstractExpression: + """ + Matrix multiplication with other, similar to xarray dot. + """ + ... + + def le(self, rhs: SideLike, join: JoinOptions | None = None) -> Constraint: + """ + Less than or equal constraint. + + Parameters + ---------- + rhs : expression-like + Right-hand side of the constraint. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + return self.to_constraint(LESS_EQUAL, rhs, join=join) + + def ge(self, rhs: SideLike, join: JoinOptions | None = None) -> Constraint: + """ + Greater than or equal constraint. + + Parameters + ---------- + rhs : expression-like + Right-hand side of the constraint. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + return self.to_constraint(GREATER_EQUAL, rhs, join=join) + + def eq(self, rhs: SideLike, join: JoinOptions | None = None) -> Constraint: + """ + Equality constraint. + + Parameters + ---------- + rhs : expression-like + Right-hand side of the constraint. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + return self.to_constraint(EQUAL, rhs, join=join) + + def __le__(self, rhs: SideLike) -> Constraint: + return self.to_constraint(LESS_EQUAL, rhs) + + def __ge__(self, rhs: SideLike) -> Constraint: + return self.to_constraint(GREATER_EQUAL, rhs) + + def __eq__(self, rhs: SideLike) -> Constraint: # type: ignore[override] + return self.to_constraint(EQUAL, rhs) + + def __gt__(self, other: Any) -> NotImplementedType: + raise NotImplementedError( + "Inequalities only ever defined for >= rather than >." + ) + + def __lt__(self, other: Any) -> NotImplementedType: + raise NotImplementedError( + "Inequalities only ever defined for >= rather than >." + ) + + +class BaseExpression(AbstractExpression): + __slots__ = ("_data", "_model") + _fill_value = FILL_VALUE _data: Dataset @@ -829,6 +1063,11 @@ def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None ) print(self) + # Narrower redeclarations of the ABC's abstract dunders: still abstract (no + # body), but pin the return type to what every BaseExpression subclass + # actually guarantees, so methods below (e.g. `add`, `mul`) that call + # `self.__add__`/`self.__mul__` type-check against `Self | QuadraticExpression` + # rather than the ABC's generic `AbstractExpression`. @abstractmethod def __add__(self, other: SideLike) -> Self | QuadraticExpression: ... @@ -1001,13 +1240,17 @@ def _divide_by_constant( return self._apply_constant_op(other, operator.truediv, fill_value=1, join=join) def __div__(self, other: SideLike) -> Self: + # Return NotImplemented (rather than raising) so a lazy divisor gets a chance + # to handle this via its own reflected `__rtruediv__`, deferring to solve time. + if isinstance(other, LazyExpression): + return NotImplemented + if isinstance(other, SUPPORTED_EXPRESSION_TYPES): + raise NonLinearOperationError( + "unsupported operand type(s) for /: " + f"{type(self)} and {type(other)}. " + "Non-linear expressions are not yet supported." + ) try: - if isinstance(other, SUPPORTED_EXPRESSION_TYPES): - raise TypeError( - "unsupported operand type(s) for /: " - f"{type(self)} and {type(other)}" - "Non-linear expressions are not yet supported." - ) return self._divide_by_constant(other) except TypeError: return NotImplemented @@ -1015,25 +1258,6 @@ def __div__(self, other: SideLike) -> Self: def __truediv__(self, other: SideLike) -> Self: return self.__div__(other) - def __le__(self, rhs: SideLike) -> Constraint: - return self.to_constraint(LESS_EQUAL, rhs) - - def __ge__(self, rhs: SideLike) -> Constraint: - return self.to_constraint(GREATER_EQUAL, rhs) - - def __eq__(self, rhs: SideLike) -> Constraint: # type: ignore[override] - return self.to_constraint(EQUAL, rhs) - - def __gt__(self, other: Any) -> NotImplementedType: - raise NotImplementedError( - "Inequalities only ever defined for >= rather than >." - ) - - def __lt__(self, other: Any) -> NotImplementedType: - raise NotImplementedError( - "Inequalities only ever defined for >= rather than >." - ) - def add( self, other: SideLike, @@ -1062,25 +1286,6 @@ def add( other = other.to_quadexpr() return merge([self, other], cls=self.__class__, join=join) - def sub( - self, - other: SideLike, - join: JoinOptions | None = None, - ) -> Self | QuadraticExpression: - """ - Subtract others from expression. - - Parameters - ---------- - other : expression-like - The expression to subtract. - join : str, optional - How to align coordinates. One of "outer", "inner", "left", - "right", "exact", "override". When None (default), uses the - current default behavior. - """ - return self.add(-other, join=join) - def mul( self, other: SideLike, @@ -1126,69 +1331,31 @@ def div( if join is None: return self.__div__(other) if isinstance(other, SUPPORTED_EXPRESSION_TYPES): - raise TypeError( + raise NonLinearOperationError( "unsupported operand type(s) for /: " f"{type(self)} and {type(other)}. " "Non-linear expressions are not yet supported." ) return self._divide_by_constant(other, join=join) - def le( - self, - rhs: SideLike, - join: JoinOptions | None = None, - ) -> Constraint: - """ - Less than or equal constraint. - - Parameters - ---------- - rhs : expression-like - Right-hand side of the constraint. - join : str, optional - How to align coordinates. One of "outer", "inner", "left", - "right", "exact", "override". When None (default), uses the - current default behavior. - """ - return self.to_constraint(LESS_EQUAL, rhs, join=join) - - def ge( - self, - rhs: SideLike, - join: JoinOptions | None = None, - ) -> Constraint: - """ - Greater than or equal constraint. - - Parameters - ---------- - rhs : expression-like - Right-hand side of the constraint. - join : str, optional - How to align coordinates. One of "outer", "inner", "left", - "right", "exact", "override". When None (default), uses the - current default behavior. - """ - return self.to_constraint(GREATER_EQUAL, rhs, join=join) - - def eq( + def sub( self, - rhs: SideLike, + other: SideLike, join: JoinOptions | None = None, - ) -> Constraint: + ) -> Self | QuadraticExpression: """ - Equality constraint. + Subtract others from expression. Parameters ---------- - rhs : expression-like - Right-hand side of the constraint. + other : expression-like + The expression to subtract. join : str, optional How to align coordinates. One of "outer", "inner", "left", "right", "exact", "override". When None (default), uses the current default behavior. """ - return self.to_constraint(EQUAL, rhs, join=join) + return self.add(-other, join=join) def pow(self, other: int) -> QuadraticExpression: """ @@ -1196,7 +1363,7 @@ def pow(self, other: int) -> QuadraticExpression: """ return self.__pow__(other) - def dot(self, other: ndarray) -> Self | QuadraticExpression: + def dot(self, other: SideLike) -> Self | QuadraticExpression: """ Matrix multiplication with other, similar to xarray dot. """ @@ -1980,6 +2147,8 @@ def __add__( Note: If other is a numpy array or pandas object without axes names, dimension names of self will be filled in other """ + if isinstance(other, LazyExpression): + return NotImplemented if isinstance(other, QuadraticExpression): return other.__add__(self) @@ -2039,6 +2208,8 @@ def __mul__( """ Multiply the expr by a factor. """ + if isinstance(other, LazyExpression): + return NotImplemented if isinstance(other, QuadraticExpression): return other.__rmul__(self) @@ -2084,6 +2255,8 @@ def __matmul__( """ Matrix multiplication with other, similar to xarray dot. """ + if isinstance(other, LazyExpression): + return NotImplemented if not isinstance(other, LinearExpression | variables.Variable): other = as_dataarray(other, coords=self.coords, dims=self.coord_dims) @@ -2460,9 +2633,6 @@ class QuadraticExpression(BaseExpression): """ __slots__ = ("_data", "_model") - __array_ufunc__ = None - __array_priority__ = 10000 - __pandas_priority__ = 10000 _fill_value = {"vars": -1, "coeffs": np.nan, "const": np.nan} @@ -2496,8 +2666,14 @@ def __mul__(self, other: SideLike) -> QuadraticExpression: """ Multiply the expr by a factor. """ + # Must run before the SUPPORTED_EXPRESSION_TYPES check below, since + # LazyExpression is now also a member of that tuple: this guard is what + # lets `lazy * quadratic` defer to `LazyExpression.__rmul__` instead of + # hitting the "non-linear expressions" TypeError meant for other cases. + if isinstance(other, LazyExpression): + return NotImplemented if isinstance(other, SUPPORTED_EXPRESSION_TYPES): - raise TypeError( + raise NonLinearOperationError( "unsupported operand type(s) for *: " f"{type(self)} and {type(other)}. " "Higher order non-linear expressions are not yet supported." @@ -2517,6 +2693,8 @@ def __add__(self, other: SideLike) -> QuadraticExpression: Note: If other is a numpy array or pandas object without axes names, dimension names of self will be filled in other """ + if isinstance(other, LazyExpression): + return NotImplemented try: if isinstance(other, CONSTANT_TYPES): return self._add_constant(other) @@ -2558,7 +2736,9 @@ def __rsub__(self, other: SideLike) -> QuadraticExpression: return NotImplemented def __pow__(self, other: SideLike) -> QuadraticExpression: - raise TypeError("Higher order non-linear expressions are not yet supported.") + raise NonLinearOperationError( + "Higher order non-linear expressions are not yet supported." + ) def __matmul__( self, other: ConstantLike | VariableLike | ExpressionLike @@ -2566,8 +2746,12 @@ def __matmul__( """ Matrix multiplication with other, similar to xarray dot. """ + # See the matching comment in __mul__ above: this guard must run first, + # now that LazyExpression is also in SUPPORTED_EXPRESSION_TYPES. + if isinstance(other, LazyExpression): + return NotImplemented if isinstance(other, SUPPORTED_EXPRESSION_TYPES): - raise TypeError( + raise NonLinearOperationError( "Higher order non-linear expressions are not yet supported." ) @@ -2853,13 +3037,388 @@ def merge( return cls(ds, model) +@dataclass(eq=False, repr=False) +class LazyExpression(AbstractExpression): + """ + A placeholder for an expression whose value is computed on demand. + + Unlike :class:`LinearExpression` / :class:`QuadraticExpression`, a `LazyExpression` holds no expression data of its own. + Instead, it stores an `evaluator` callable that, given the model, builds and returns the real expression, and + optionally a `mask` that is resolved and applied at the same time as `evaluator`. + + Arithmetic between `LazyExpression` objects (and between a `LazyExpression` and anything else) stays lazy: it + returns a new, unnamed `LazyExpression` whose evaluator composes the operands. Nothing is built until + `.evaluate()`, `.promote()`, `.solution`, or a comparison (`<=`, `>=`, `==`) is called. + + `LazyExpression` shares its arithmetic and constraint-building protocol with the eager + expression classes via :class:`AbstractExpression`: the named counterparts (`add`, `sub`, + `mul`, `div`, `pow`, `dot`, `le`, `ge`, `eq`) and comparisons all work the same way as on + `LinearExpression`/`QuadraticExpression`, forcing evaluation only where they must. + + Examples + -------- + >>> from linopy import Model + >>> import pandas as pd + >>> m = Model() + >>> time = pd.RangeIndex(10, name="Time") + >>> x = m.add_variables(lower=0, coords=[time], name="x") + >>> lazy = m.add_expressions(lambda m: m.variables["x"] + 1, name="lazy") + >>> lazy.evaluate() # doctest: +SKIP + """ + + # `eq=False` on the dataclass decorator keeps the inherited `__eq__` (which builds a + # Constraint, see AbstractExpression) instead of a generated field-comparison `__eq__`; + # it also leaves `__hash__` alone, so `AbstractExpression.__hash__ = None` applies - + # LazyExpression is deliberately unhashable, just like the eager expression classes. + + model: Model + """Reference to the model the expression belongs to""" + evaluator: Callable[ + ..., + LinearExpression | QuadraticExpression | variables.Variable | DataArray | Any, + ] + """Callable that builds the expression, invoked as ``evaluator(model, **params)``. + May also return a plain ``Variable``, ``DataArray`` or constant -- e.g. a constraint's + ``.dual`` -- for an expression that is only ever read via `.solution`; see `.promote`.""" + name: str | None = None + """Lazy Expression name. `None` for derived expressions produced by arithmetic, which are never + registered in `model.expressions`.""" + mask: MaskLike | Callable[..., MaskLike] | None = None + """Boolean mask applied to the evaluated expression via `.where(mask)`. + A concrete array-like is applied as-is; a callable is invoked as ``mask(model, **params)`` at the + same time `evaluator` runs, so it can depend on data that is only known once the model exists. + A mask that leaves nothing valid produces an all-NaN expression rather than raising.""" + params: dict[str, Any] = field(default_factory=dict) + """Keyword arguments forwarded to `evaluator` (and to `mask`, if callable).""" + input_data: Dataset | None = None + """Pointer to the input data the evaluator reads, kept for introspection only.""" + dims: tuple[Hashable, ...] = () + """Dimensions of the expression, if known in advance. + If not provided, the dimensions are inferred from the evaluated expression.""" + attrs: dict[Any, Any] = field(default_factory=dict) + """Attributes to be assigned to the evaluated expression, if any.""" + source: Any = None + """Optional serialisable description of `evaluator` (e.g. an expression AST produced by a + declarative frontend). Ignored by linopy itself; reserved so that IO can persist a lazy expression + instead of evaluating it, once a frontend that produces such a description exists.""" + mask_source: Any = None + """As `source`, but describing `mask`.""" + _solution_evaluator: Callable[[], Any] | None = None + """Internal. Set by arithmetic (`_combine`, `__neg__`, `__pow__`) on derived expressions: + composes `.solution` operand-wise, as a fallback for operations that have no linear/quadratic + form and so cannot go through `evaluator`/`.evaluate()`. Not part of the public API.""" + _static_nonlinear: bool = False + """Internal. True when this derived expression is already known, at construction time, to + have no linear/quadratic form (e.g. division by an expression). Backs `is_evaluatable`.""" + + def evaluate( + self, + ) -> LinearExpression | QuadraticExpression | variables.Variable | DataArray | Any: + """ + Evaluate the expression using the provided evaluator and mask. + + Note that nothing is cached, so calling this repeatedly will always re-evaluate from scratch. + + Raises + ------ + NonLinearOperationError + If the underlying operation (e.g. division by a variable or another expression) + has no linear/quadratic form. Such an expression can still be read via + `.solution` once the model has been solved. + """ + expr = ( + self.evaluator(self.model, **self.params) + if self.params + else self.evaluator(self.model) + ) + if self.mask is not None: + mask = ( + self.mask(self.model, **self.params) + if callable(self.mask) + else self.mask + ) + mask = as_dataarray(mask, coords=expr.coords, dims=expr.dims).astype(bool) + expr = expr.where(mask) + return expr + + def promote(self) -> LinearExpression | QuadraticExpression: + """ + Materialise this expression in-place, replacing the placeholder. + + If `self.name` no longer refers to this placeholder in `self.model.expressions` (i.e. it has already been promoted), the existing expression is returned unchanged. + + Raises + ------ + ValueError + If this is a derived expression (`self.name is None`), which is not registered in + `self.model.expressions` and therefore cannot be promoted in-place. + """ + if self.name is None: + raise ValueError( + "Cannot promote a derived LazyExpression (name is None); it was produced by " + "arithmetic and is not registered in `model.expressions`. Call `.evaluate()` instead." + ) + current = self.model.expressions.data.get(self.name) + if current is not self and isinstance( + current, LinearExpression | QuadraticExpression + ): + return current + try: + expr = self.evaluate() + except NonLinearOperationError as e: + raise NonLinearOperationError( + f"Cannot promote LazyExpression '{self.name}': {e} " + "It can still be read via `.solution` once the model has been solved." + ) from e + if not isinstance(expr, LinearExpression | QuadraticExpression): + raise NonLinearOperationError( + f"Cannot promote LazyExpression '{self.name}': its evaluator returned " + f"{type(expr)}, not a LinearExpression or QuadraticExpression. " + "It can still be read via `.solution` once the model has been solved." + ) + expr.attrs.update(self.attrs) + expr.attrs["name"] = self.name + self.model.expressions.data[self.name] = expr + return expr + + @property + def is_evaluatable(self) -> bool: + """ + Whether `.evaluate()` can be expected to succeed. + + False once this expression is already known, at construction time, to have no + linear/quadratic form (e.g. built from division by an expression, or from raising + to a power other than 2). A leaf expression whose own `evaluator` callable happens + to build a non-linear result -- or return something other than a `LinearExpression` + / `QuadraticExpression`, e.g. a constraint's `.dual` -- is not detected here; that + only surfaces when `.evaluate()` is actually called. + """ + return not self._static_nonlinear + + @property + def solution(self) -> DataArray: + """ + Get the optimal values of the expression, without promoting it. + + Tries `.evaluate()` first, so a linear/quadratic expression behaves exactly as + before (mask/NaN semantics included). If that fails because the underlying + operation has no linear/quadratic form (e.g. it divides by a variable or another + expression), falls back to composing `.solution` operand-wise instead -- valid once + the model has a solution, since every operand is then just a number. If the + evaluator itself returns something other than an expression (e.g. a constraint's + `.dual`), that value is used as-is. + """ + try: + expr = self.evaluate() + except (NonLinearOperationError, ValueError): + if self._solution_evaluator is None: + raise + return _as_solution_dataarray(self._solution_evaluator()) + return _as_solution_dataarray(_solution_of(expr)) + + @property + def coords(self) -> DatasetCoordinates | dict[Hashable, Any]: + """Coordinates of the expression, if it has already been promoted.""" + current = self.model.expressions.data.get(self.name) if self.name else None + if current is not None and current is not self: + return current.coords + return {} + + @property + def type(self) -> str: + return "LazyExpression" + + def __repr__(self) -> str: + dims = ", ".join(str(d) for d in self.dims) + name = self.name if self.name is not None else "" + return f"LazyExpression '{name}' [{dims}] (not yet evaluated)" + + def __getattr__(self, name: str) -> Any: + # Only reached for attributes not found on the instance/class, i.e. + # everything but the overrides above; forward to a fresh evaluation. + # Names starting with "_" are rejected outright so that pickling/copying + # (which probes dunder/private attributes before __dict__ is populated) + # cannot recurse into `evaluate()` -> `self.evaluator` -> `__getattr__` -> ... + if name.startswith("_"): + raise AttributeError(name) + return getattr(self.evaluate(), name) + + def _combine( + self, + other: Any, + op: Callable[[Any, Any], Any], + swapped: bool = False, + nonlinear_reason: str | None = None, + ) -> LazyExpression: + """ + Build a new, unnamed `LazyExpression` that lazily applies `op` to `self` and `other`. + + `other` is resolved lazily too, if it is itself a `LazyExpression`. + + A parallel, `.solution`-only closure is always attached (see `_solution_of`), used as + a fallback wherever `op` turns out to have no linear/quadratic form. When that is + already known at construction time (`nonlinear_reason` given), a + `NonLinearExpressionWarning` is raised immediately instead of waiting for a failed + `.evaluate()` to discover it. + """ + if nonlinear_reason is not None: + warn( + f"This LazyExpression involves {nonlinear_reason}, which has no " + "linear/quadratic form. It can only be read via `.solution` once the " + "model has been solved; `.evaluate()`, `.promote()` and constraint-building " + "will raise.", + NonLinearExpressionWarning, + stacklevel=3, + ) + + def evaluator(model: Model) -> Any: + left = self.evaluate() + right = _resolve_lazy(other) + return op(right, left) if swapped else op(left, right) + + def solution_evaluator() -> Any: + left = _solution_of(self) + right = _solution_of(other) + return op(right, left) if swapped else op(left, right) + + return LazyExpression( + model=self.model, + evaluator=evaluator, + _solution_evaluator=solution_evaluator, + _static_nonlinear=nonlinear_reason is not None, + ) + + def __add__(self, other: LazySideLike) -> LazyExpression: + return self._combine(other, operator.add) + + def __radd__(self, other: LazySideLike) -> LazyExpression: + return self._combine(other, operator.add) + + def __sub__(self, other: LazySideLike) -> LazyExpression: + return self._combine(other, operator.sub) + + def __rsub__(self, other: LazySideLike) -> LazyExpression: + return self._combine(other, operator.sub, swapped=True) + + def __mul__(self, other: LazySideLike) -> LazyExpression: + return self._combine(other, operator.mul) + + def __rmul__(self, other: LazySideLike) -> LazyExpression: + return self._combine(other, operator.mul) + + def __truediv__(self, other: LazySideLike) -> LazyExpression: + reason = ( + "division by an expression" + if isinstance(other, SUPPORTED_EXPRESSION_TYPES) + else None + ) + return self._combine(other, operator.truediv, nonlinear_reason=reason) + + def __rtruediv__(self, other: LazySideLike) -> LazyExpression: + # Only a constant/array numerator defers here: an eager expression/variable + # numerator returns NotImplemented so the overall operation still raises, matching + # eager-only division (`variable / other_variable`, etc.). + if isinstance(other, SUPPORTED_EXPRESSION_TYPES): + return NotImplemented + return self._combine( + other, + operator.truediv, + swapped=True, + nonlinear_reason="division by an expression", + ) + + def __matmul__(self, other: LazySideLike) -> LazyExpression: + return self._combine(other, operator.matmul) + + def __rmatmul__(self, other: LazySideLike) -> LazyExpression: + return self._combine(other, operator.matmul, swapped=True) + + def __pow__(self, other: int) -> LazyExpression: + # Evaluate the root once and let the eager `__pow__` do the squaring, + # rather than passing `self` as `other` to `_combine` (which would + # evaluate the root twice: once as `left`, once as `right`). + reason = None if other == 2 else f"raising to the power {other}" + return self._combine(other, operator.pow, nonlinear_reason=reason) + + def __neg__(self) -> LazyExpression: + # `operator.neg` matches eager `BaseExpression.__neg__` (negates + # `coeffs`/`const`, preserving NaN), unlike `* -1` which goes through + # `_apply_constant_op` and fills NaN with 0 first. + def evaluator(model: Model) -> Any: + return -self.evaluate() + + def solution_evaluator() -> Any: + return -_solution_of(self) + + return LazyExpression( + model=self.model, + evaluator=evaluator, + _solution_evaluator=solution_evaluator, + ) + + def to_constraint( + self, sign: SignLike, rhs: LazySideLike, join: JoinOptions | None = None + ) -> Constraint: + """ + Turn this expression into a constraint against `rhs` with the given `sign`. + + Forces evaluation of both `self` and (if lazy) `rhs`. + + Raises + ------ + NonLinearOperationError + If `.evaluate()` fails (see `.evaluate`), or its result is not a + `LinearExpression`/`QuadraticExpression` (e.g. it is a constraint's `.dual`). + """ + expr = self.evaluate() + if not isinstance(expr, LinearExpression | QuadraticExpression): + raise NonLinearOperationError( + f"Cannot build a constraint from this LazyExpression: its evaluator " + f"returned {type(expr)}, not a LinearExpression or QuadraticExpression." + ) + return expr.to_constraint(sign, _resolve_lazy(rhs), join=join) + + def add( + self, other: LazySideLike, join: JoinOptions | None = None + ) -> LazyExpression: + return self._combine(other, lambda expr, rhs: expr.add(rhs, join=join)) + + def mul( + self, other: LazySideLike, join: JoinOptions | None = None + ) -> LazyExpression: + return self._combine(other, lambda expr, rhs: expr.mul(rhs, join=join)) + + def div( + self, other: LazySideLike, join: JoinOptions | None = None + ) -> LazyExpression: + reason = ( + "division by an expression" + if isinstance(other, SUPPORTED_EXPRESSION_TYPES) + else None + ) + return self._combine( + other, lambda expr, rhs: expr.div(rhs, join=join), nonlinear_reason=reason + ) + + def sub( + self, other: LazySideLike, join: JoinOptions | None = None + ) -> LazyExpression: + return self.add(-other, join=join) + + def pow(self, other: int) -> LazyExpression: + return self.__pow__(other) + + def dot(self, other: LazySideLike) -> LazyExpression: + return self.__matmul__(other) + + @dataclass(repr=False) class Expressions: """ An expressions container used for storing multiple expression arrays. """ - data: dict[str, LinearExpression | QuadraticExpression] + data: dict[str, LinearExpression | QuadraticExpression | LazyExpression] model: Model dataset_attrs = ["coeffs", "vars", "const"] @@ -2873,19 +3432,23 @@ def _formatted_names(self) -> dict[str, str]: return {format_string_as_variable_name(n): n for n in self} @overload - def __getitem__(self, names: str) -> LinearExpression | QuadraticExpression: ... + def __getitem__( + self, names: str + ) -> LinearExpression | QuadraticExpression | LazyExpression: ... @overload def __getitem__(self, names: list[str]) -> Expressions: ... def __getitem__( self, names: str | list[str] - ) -> LinearExpression | QuadraticExpression | Expressions: + ) -> LinearExpression | QuadraticExpression | LazyExpression | Expressions: if isinstance(names, str): return self.data[names] return Expressions({name: self.data[name] for name in names}, self.model) - def __getattr__(self, name: str) -> LinearExpression | QuadraticExpression: + def __getattr__( + self, name: str + ) -> LinearExpression | QuadraticExpression | LazyExpression: # If name is an attribute of self (including methods and properties), return that if name in self.data: return self.data[name] @@ -2943,7 +3506,9 @@ def __len__(self) -> int: def __iter__(self) -> Iterator[str]: return self.data.__iter__() - def items(self) -> ItemsView[str, LinearExpression | QuadraticExpression]: + def items( + self, + ) -> ItemsView[str, LinearExpression | QuadraticExpression | LazyExpression]: return self.data.items() def _ipython_key_completions_(self) -> list[str]: @@ -2956,15 +3521,22 @@ def _ipython_key_completions_(self) -> list[str]: """ return list(self) - def add(self, expression: LinearExpression | QuadraticExpression) -> None: + def add( + self, expression: LinearExpression | QuadraticExpression | LazyExpression + ) -> None: """ Add an expression to the expressions container. """ + if expression.name is None: + raise ValueError( + "Cannot add a derived LazyExpression (name is None) to `model.expressions`; " + "it was produced by arithmetic between lazy expressions, not by `add_expressions`." + ) self.data[expression.name] = expression def remove(self, name: str) -> None: """ - Remove variable `name` from the variables. + Remove expression `name` from the expressions. """ self.data.pop(name) @@ -2973,7 +3545,10 @@ def solution(self) -> Dataset: """ Get the solution of variables. """ - return save_join(*[v.solution.rename(k) for k, v in self.items()]) + # `list(...)` guards against mutation of `self.data` if a `LazyExpression` + # promotes itself while its `.solution` is being read (it does not, but + # `.solution` deliberately avoids promoting, so this is just a safeguard). + return save_join(*[v.solution.rename(k) for k, v in list(self.items())]) class ScalarLinearExpression: @@ -3136,6 +3711,7 @@ def to_linexpr(self) -> LinearExpression: SUPPORTED_EXPRESSION_TYPES = ( BaseExpression, + LazyExpression, ScalarLinearExpression, variables.Variable, variables.ScalarVariable, diff --git a/linopy/io.py b/linopy/io.py index e6196d47..4a3d5336 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -6,6 +6,7 @@ from __future__ import annotations import copy as _copy +import dataclasses import json import logging import shutil @@ -16,7 +17,7 @@ from io import BufferedWriter from pathlib import Path from tempfile import TemporaryDirectory -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import numpy as np import pandas as pd @@ -26,7 +27,14 @@ from linopy import solvers from linopy.common import sos_weights, to_polars -from linopy.constants import CONCAT_DIM, FACTOR_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR +from linopy.constants import ( + CONCAT_DIM, + FACTOR_DIM, + SOS_DIM_ATTR, + SOS_TYPE_ATTR, + NonLinearOperationError, +) +from linopy.expressions import LazyExpression, LinearExpression, QuadraticExpression from linopy.objective import Objective if TYPE_CHECKING: @@ -920,7 +928,12 @@ def non_bool_dict( return {k: int(v) if isinstance(v, bool) else v for k, v in d.items()} -def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: +def to_netcdf( + m: Model, + *args: Any, + lazy: Literal["evaluate", "skip", "raise"] = "evaluate", + **kwargs: Any, +) -> None: """ Write out the model to a netcdf file. @@ -930,6 +943,23 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: Model to write out. *args Arguments passed to ``xarray.Dataset.to_netcdf``. + lazy : {"evaluate", "skip", "raise"}, default "evaluate" + What to do with :class:`linopy.LazyExpression` entries in ``m.expressions``, + which hold no data of their own and cannot be written as-is: + + - ``"evaluate"``: run each lazy expression's evaluator and write the + result as an ordinary (linear or quadratic) expression. The + placeholder itself, and the fact that it was lazy, are not restored + by :func:`read_netcdf`. Raises if a lazy expression has no + linear/quadratic form (e.g. it divides by a variable or another + expression) or its evaluator returns something other than an + expression (e.g. a constraint's ``.dual``) -- such expressions can + only be read via their ``.solution``, so use ``lazy="skip"`` for + models that hold them. + - ``"skip"``: omit lazy expressions from the file entirely. A warning + names the dropped entries. + - ``"raise"``: raise a :class:`ValueError` naming the lazy entries + instead of writing the file. **kwargs : TYPE Keyword arguments passed to ``xarray.Dataset.to_netcdf``. @@ -938,7 +968,13 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: Variables, constraints, the objective, parameters and named expressions (``Model.expressions``, including their linear/quadratic type) are all persisted and fully restored by - :func:`linopy.io.read_netcdf`. + :func:`linopy.io.read_netcdf`. :class:`LazyExpression` entries are the + exception: they are handled per the `lazy` parameter above, since nothing + in linopy today can serialize an arbitrary evaluator callable. A + lazy expression built from a serialisable description (e.g. an AST + produced by a declarative frontend, attached via `LazyExpression.source`) + could be persisted as such in the future; no such description exists yet, + so every lazy entry currently falls through to the `lazy` policy above. The SOS reformulation lifecycle token lives only on the in-memory Model and is not persisted. If the model has an active SOS @@ -984,13 +1020,55 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: with_prefix(con.to_netcdf_ds(), f"constraints-{name}") for name, con in m.constraints.items() ] - exprs = [ - with_prefix( - expr.data.assign_attrs(name=name, _linopy_expr_type=expr.type), - f"expressions-{name}", - ) - for name, expr in m.expressions.items() + + lazy_names = [ + name for name, expr in m.expressions.items() if isinstance(expr, LazyExpression) ] + if lazy_names: + if lazy == "raise": + raise ValueError( + f"Cannot write lazy expression(s) {lazy_names} to netcdf. " + "Pass lazy='evaluate' to materialise them or lazy='skip' to drop them." + ) + if lazy == "skip": + logger.warning( + f"Dropping lazy expression(s) {lazy_names} from the netcdf file " + "(lazy='skip'); they will not be present after `read_netcdf`." + ) + + exprs = [] + for name, expr_or_lazy in m.expressions.items(): + expr: LinearExpression | QuadraticExpression + if isinstance(expr_or_lazy, LazyExpression): + # Lazy expressions with a serialisable `source` (e.g. an AST produced by a + # declarative frontend) could be persisted here instead of being evaluated. + # Nothing in linopy produces a `source` yet, so every lazy entry falls + # through to the `lazy` policy below. + if lazy == "skip": + continue + try: + evaluated = expr_or_lazy.evaluate() + except NonLinearOperationError as e: + raise NonLinearOperationError( + f"Cannot write lazy expression '{name}' to netcdf with lazy='evaluate': " + f"{e} Pass lazy='skip' to drop it, or read it via `.solution` instead." + ) from e + if not isinstance(evaluated, LinearExpression | QuadraticExpression): + raise TypeError( + f"Cannot write lazy expression '{name}' to netcdf with lazy='evaluate': " + f"its evaluator returned {type(evaluated)}, not a LinearExpression or " + "QuadraticExpression. Pass lazy='skip' to drop it, or read it via " + "`.solution` instead." + ) + expr = evaluated + else: + expr = expr_or_lazy + exprs.append( + with_prefix( + expr.data.assign_attrs(name=name, _linopy_expr_type=expr.type), + f"expressions-{name}", + ) + ) objective = m.objective.data objective = objective.assign_attrs(sense=m.objective.sense) if m.objective.value is not None: @@ -1108,9 +1186,14 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: m._variables = Variables(variables, m) + # Everything written by `to_netcdf` is eager: lazy expressions are either + # evaluated, skipped, or raised on before writing (see the `lazy` parameter + # there). A future `_linopy_lazy_source` attr, persisted from + # `LazyExpression.source`, would be rehydrated into a `LazyExpression` here + # instead of falling into the eager branch below. exprs = [str(k) for k in ds if str(k).startswith("expressions")] expr_names = list({str(k).rsplit("-", 1)[0] for k in exprs}) - expressions: dict[str, LinearExpression | QuadraticExpression] = {} + expressions: dict[str, LinearExpression | QuadraticExpression | LazyExpression] = {} for k in sorted(expr_names): name = remove_prefix(k, "expressions") expr_ds = get_prefix(ds, k) @@ -1211,7 +1294,7 @@ def copy(m: Model, include_solution: bool = False, deep: bool = True) -> Model: A deep or shallow copy of the model. """ from linopy.constraints import Constraint, ConstraintBase, Constraints - from linopy.expressions import Expressions, LinearExpression, QuadraticExpression + from linopy.expressions import Expressions, LinearExpression from linopy.model import Model, Objective from linopy.variables import Variable, Variables @@ -1241,9 +1324,17 @@ def copy(m: Model, include_solution: bool = False, deep: bool = True) -> Model: ) def _copy_expr( - name: str, expr: LinearExpression | QuadraticExpression - ) -> LinearExpression | QuadraticExpression: + name: str, expr: LinearExpression | QuadraticExpression | LazyExpression + ) -> LinearExpression | QuadraticExpression | LazyExpression: # Expressions hold no solve artifacts, so include_solution is irrelevant. + if isinstance(expr, LazyExpression): + # The placeholder itself has no data to copy; just rebind it to the + # new model. `input_data` is the only field that could reasonably + # be deep-copied, since `evaluator`/`mask` are callables. + input_data = expr.input_data + if deep and input_data is not None: + input_data = input_data.copy(deep=True) + return dataclasses.replace(expr, model=new_model, input_data=input_data) new_expr = type(expr)(expr.data.copy(deep=deep), new_model) new_expr.attrs["name"] = name # __init__ resets the name to None return new_expr diff --git a/linopy/model.py b/linopy/model.py index 0729ee71..49e5a39e 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -10,7 +10,7 @@ import os import re import warnings -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Hashable, Mapping, Sequence from pathlib import Path from tempfile import NamedTemporaryFile, gettempdir from types import MappingProxyType @@ -61,6 +61,7 @@ from linopy.dualization import dualize from linopy.expressions import ( Expressions, + LazyExpression, LinearExpression, QuadraticExpression, ScalarLinearExpression, @@ -933,47 +934,120 @@ def add_variables( self.variables.add(variable) return variable + def _next_expression_name(self, name: str | None) -> str: + """Allocate (or validate) the reference name for a new expression.""" + if name is None: + name = f"expr{self._exprnameCounter}" + self._exprnameCounter += 1 + + if name in self.expressions: + raise ValueError(f"Expression '{name}' already assigned to model") + + return name + + @overload + def add_expressions( + self, + data: Callable[ + ..., + LinearExpression + | QuadraticExpression + | Variable + | DataArray + | ConstantLike, + ], + name: str | None = ..., + mask: MaskLike | Callable[..., MaskLike] | None = ..., + dims: tuple[Hashable, ...] = ..., + input_data: Dataset | None = ..., + **params: Any, + ) -> LazyExpression: ... + + @overload def add_expressions( self, data: Variable | LinearExpression | QuadraticExpression | Sequence[tuple[ConstantLike, Variable | str]], + name: str | None = ..., + mask: MaskLike | None = ..., + ) -> LinearExpression | QuadraticExpression: ... + + def add_expressions( + self, + data: Variable + | LinearExpression + | QuadraticExpression + | Sequence[tuple[ConstantLike, Variable | str]] + | Callable[ + ..., + LinearExpression + | QuadraticExpression + | Variable + | DataArray + | ConstantLike, + ], name: str | None = None, - mask: MaskLike | None = None, - ) -> LinearExpression | QuadraticExpression: + mask: MaskLike | Callable[..., MaskLike] | None = None, + dims: tuple[Hashable, ...] = (), + input_data: Dataset | None = None, + **params: Any, + ) -> LinearExpression | QuadraticExpression | LazyExpression: """ Assign a new, possibly multi-dimensional array of expressions to the model. + If `data` is a callable, the expression is not built now: `add_expressions` + registers a :class:`LazyExpression` placeholder that calls `data(self, **params)` + (and, if `mask` is callable, `mask(self, **params)`) only when the expression is + actually evaluated, via `.evaluate()`, `.promote()`, `.solution`, or a comparison. + Arithmetic on the returned `LazyExpression` (e.g. `lazy + 1`) stays lazy too. + Parameters ---------- - data : Variable, LinearExpression, QuadraticExpression, or Sequence of (constant, variable) tuples - The expression(s) to add. - This can be a Variable or LinearExpression, or a sequence of (constant, variable) tuples which will be summed up. - coords : list/xarray.Coordinates, optional - The coords of the expression array. - The default is None. + data : Variable, LinearExpression, QuadraticExpression, Sequence of (constant, variable) tuples, or Callable + The expression(s) to add. This can be a Variable or LinearExpression, a sequence + of (constant, variable) tuples which will be summed up, or a callable + `data(model, **params)` that builds and returns the expression on demand. + A callable may also read post-solve-only data -- e.g. a constraint's `.dual`, + or a ratio that divides by a variable or another expression, which has no + linear/quadratic form -- and return a plain `DataArray`/constant, or a + :class:`LazyExpression` produced by such arithmetic. The result is then only + readable via `.solution`; `.evaluate()`, `.promote()` and constraint-building + raise :class:`~linopy.NonLinearOperationError`. name : str, optional Reference name of the added expressions. The default None results in a name like "expr1", "expr2" etc. - mask : array_like, optional + mask : array_like or Callable, optional Boolean mask with False values for expressions which are skipped. The shape of the mask has to match the shape the added expressions. - Default is None. + If `data` is callable, `mask` may also be a callable `mask(model, **params)`, + resolved at the same time as `data`; a callable `mask` is not accepted + together with a non-callable `data`. Default is None. + dims : tuple of Hashable, optional + Only used when `data` is callable. Dimensions of the eventual expression, + used solely for a cheap `repr` before the expression has been evaluated. + input_data : xr.Dataset, optional + Only used when `data` is callable. Pointer to the input data `data` reads, + kept for introspection only; never copied. + **params : Any + Only used when `data` is callable. Forwarded as keyword arguments to `data` + (and to `mask`, if callable) every time the expression is evaluated. Raises ------ ValueError If neither lower bound and upper bound have coordinates, nor `coords` are directly given. + TypeError + If `mask` is callable but `data` is not. Returns ------- - linopy.LinearExpression | linopy.QuadraticExpression + linopy.LinearExpression | linopy.QuadraticExpression | linopy.LazyExpression Expression which was added to the model. - Examples -------- >>> from linopy import Model @@ -982,13 +1056,39 @@ def add_expressions( >>> time = pd.RangeIndex(10, name="Time") >>> x = m.add_variables(lower=0, coords=[time], name="x") >>> expr = m.add_expressions(x + 1, name="expr") + + A lazily-evaluated expression: + + >>> lazy = m.add_expressions(lambda m: m.variables["x"] + 1, name="lazy") + + A ratio that divides by a variable, only readable via `.solution` once solved: + + >>> unit_cost = m.add_expressions( + ... lambda m: m.variables["x"].sum() / m.variables["x"].sum(), + ... name="unit_cost", + ... ) # doctest: +SKIP """ - if name is None: - name = f"expr{self._exprnameCounter}" - self._exprnameCounter += 1 + if callable(mask) and not callable(data): + raise TypeError( + "A callable mask can only be used with a callable expression; " + "pass a concrete mask or make `data` callable too." + ) - if name in self.expressions: - raise ValueError(f"Expression '{name}' already assigned to model") + if callable(data): + name = self._next_expression_name(name) + lazy = LazyExpression( + model=self, + evaluator=data, + name=name, + mask=mask, + params=params, + input_data=input_data, + dims=dims, + ) + self.expressions.add(lazy) + return lazy + + name = self._next_expression_name(name) expr: LinearExpression | QuadraticExpression if isinstance(data, Variable): diff --git a/linopy/monkey_patch_xarray.py b/linopy/monkey_patch_xarray.py index 1e526c92..fabe41fb 100644 --- a/linopy/monkey_patch_xarray.py +++ b/linopy/monkey_patch_xarray.py @@ -14,6 +14,7 @@ expressions.LinearExpression, expressions.ScalarLinearExpression, expressions.QuadraticExpression, + expressions.LazyExpression, ) diff --git a/linopy/testing.py b/linopy/testing.py index d9c67f7b..b59e68fe 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -4,9 +4,14 @@ import xarray as xr from xarray.testing import assert_equal -from linopy.constants import TERM_DIM +from linopy.constants import TERM_DIM, NonLinearOperationError from linopy.constraints import ConstraintBase, _con_unwrap -from linopy.expressions import LinearExpression, QuadraticExpression, _expr_unwrap +from linopy.expressions import ( + LazyExpression, + LinearExpression, + QuadraticExpression, + _expr_unwrap, +) from linopy.model import Model from linopy.variables import Variable, _var_unwrap @@ -71,16 +76,56 @@ def assert_quadequal( def assert_exprequal( - a: LinearExpression | QuadraticExpression, - b: LinearExpression | QuadraticExpression, + a: LinearExpression | QuadraticExpression | LazyExpression, + b: LinearExpression | QuadraticExpression | LazyExpression, check_name: bool = True, ) -> None: """ - Assert that two expressions are equal, dispatching on linear vs quadratic. + Assert that two expressions are equal, dispatching on linear vs quadratic vs lazy. xarray's assert_equal ignores attrs, so the stored name (which lives in ``attrs["name"]``) is compared explicitly unless ``check_name=False``. + + If either side is a :class:`LazyExpression`, both must be: the placeholder's + `name` and `dims` are compared directly, and the underlying expressions are + compared after calling `.evaluate()` on each (without promoting either). If both + sides raise :class:`NonLinearOperationError` (e.g. both divide by a variable or + another expression), the `name`/`dims` comparison above is treated as sufficient; + if only one side raises, that is a real mismatch and fails. """ + if isinstance(a, LazyExpression) or isinstance(b, LazyExpression): + assert isinstance(a, LazyExpression) and isinstance(b, LazyExpression), ( + f"expression types differ: {type(a)} != {type(b)}" + ) + if check_name: + assert a.name == b.name, ( + f"expression names differ: {a.name!r} != {b.name!r}" + ) + assert a.dims == b.dims, ( + f"lazy expression dims differ: {a.dims!r} != {b.dims!r}" + ) + try: + a_evaluated = a.evaluate() + except NonLinearOperationError as a_error: + try: + b.evaluate() + except NonLinearOperationError: + return + raise AssertionError( + f"only one side raised NonLinearOperationError on `.evaluate()`: {a_error}" + ) from a_error + b_evaluated = b.evaluate() + assert isinstance(a_evaluated, LinearExpression | QuadraticExpression), ( + f"side 'a' evaluated to {type(a_evaluated)}, not a LinearExpression or " + "QuadraticExpression; compare its `.solution` instead" + ) + assert isinstance(b_evaluated, LinearExpression | QuadraticExpression), ( + f"side 'b' evaluated to {type(b_evaluated)}, not a LinearExpression or " + "QuadraticExpression; compare its `.solution` instead" + ) + assert_exprequal(a_evaluated, b_evaluated, check_name=False) + return + assert type(a) is type(b), f"expression types differ: {type(a)} != {type(b)}" if check_name: assert a.name == b.name, f"expression names differ: {a.name!r} != {b.name!r}" diff --git a/linopy/types.py b/linopy/types.py index 6b4cf712..93817a2b 100644 --- a/linopy/types.py +++ b/linopy/types.py @@ -16,6 +16,7 @@ ConstraintBase, ) from linopy.expressions import ( + LazyExpression, LinearExpression, QuadraticExpression, ScalarLinearExpression, @@ -51,3 +52,4 @@ ConstraintLike = Union["ConstraintBase", "AnonymousScalarConstraint"] LinExprLike = Union["Variable", "LinearExpression"] SideLike = Union[ConstantLike, VariableLike, ExpressionLike] # noqa: UP007 +LazySideLike = Union[SideLike, "LazyExpression"] # noqa: UP007 diff --git a/linopy/variables.py b/linopy/variables.py index c2e247bb..eddf5ac9 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -62,6 +62,7 @@ STASHED_LOWER, STASHED_UPPER, TERM_DIM, + NonLinearOperationError, ) from linopy.types import ( ConstantLike, @@ -466,8 +467,19 @@ def __div__( """ Divide variables with a coefficient. """ - if isinstance(other, expressions.LinearExpression | Variable): - raise TypeError( + # Return NotImplemented (rather than raising) so a lazy divisor gets a chance + # to handle this via its own reflected `__rtruediv__`, deferring to solve time. + if isinstance(other, expressions.LazyExpression): + return NotImplemented + if isinstance( + other, + expressions.LinearExpression + | expressions.QuadraticExpression + | expressions.ScalarLinearExpression + | Variable + | ScalarVariable, + ): + raise NonLinearOperationError( "unsupported operand type(s) for /: " f"{type(self)} and {type(other)}. " "Non-linear expressions are not yet supported." @@ -482,6 +494,8 @@ def __truediv__( """ try: return self.__div__(coefficient) + except NonLinearOperationError: + raise except TypeError: return NotImplemented diff --git a/test/test_expressions.py b/test/test_expressions.py index 53c5a57c..75ebe050 100644 --- a/test/test_expressions.py +++ b/test/test_expressions.py @@ -3,14 +3,28 @@ This module aims at testing the correct behavior of the Expressions class. """ +import warnings + +import numpy as np import pandas as pd import pytest import xarray as xr -from linopy import Model -from linopy.expressions import Expressions, LinearExpression, QuadraticExpression +from linopy import Model, Variable +from linopy.constants import ( + LESS_EQUAL, + NonLinearExpressionWarning, + NonLinearOperationError, +) +from linopy.constraints import Constraint +from linopy.expressions import ( + Expressions, + LazyExpression, + LinearExpression, + QuadraticExpression, +) from linopy.solvers import available_solvers -from linopy.testing import assert_linequal +from linopy.testing import assert_conequal, assert_linequal, assert_quadequal @pytest.fixture @@ -81,12 +95,10 @@ def test_add_expressions_from_variable_and_tuples() -> None: expr = m.add_expressions(x, name="from_var") assert isinstance(expr, LinearExpression) assert_linequal(expr, x.to_linexpr()) - assert_linequal(expr, m.expressions["from_var"]) expr = m.add_expressions([(2, x)], name="from_tuples") assert isinstance(expr, LinearExpression) assert_linequal(expr, 2 * x) - assert_linequal(expr, m.expressions["from_tuples"]) def test_add_expressions_quadratic(m: Model) -> None: @@ -141,3 +153,492 @@ def test_expressions_solution() -> None: assert isinstance(sol, xr.Dataset) assert "double_x" in sol assert (sol["double_x"] == 4).all() + + +class TestLazyExpression: + """Tests for the callable-`data` (lazy) path of `Model.add_expressions`.""" + + def test_add_expressions_with_callable_places_placeholder( + self, m: Model, x: Variable, y: Variable + ) -> None: + lazy = m.add_expressions(lambda m: x + y, name="lazy") + assert isinstance(lazy, LazyExpression) + assert m.expressions["lazy"] is lazy + + def test_evaluate_equals_eager_and_stays_lazy( + self, m: Model, x: Variable, y: Variable + ) -> None: + lazy = m.add_expressions(lambda m: x + y, name="lazy") + assert_linequal(lazy.evaluate(), x + y) + assert m.expressions["lazy"] is lazy + + def test_promote_swaps_carries_attrs_and_is_idempotent( + self, m: Model, x: Variable, y: Variable + ) -> None: + lazy = m.add_expressions(lambda m: x + y, name="lazy") + lazy.attrs["references"] = ["x", "y"] + promoted = lazy.promote() + assert isinstance(promoted, LinearExpression) + assert m.expressions["lazy"] is promoted + assert promoted.attrs["references"] == ["x", "y"] + assert promoted.attrs["name"] == "lazy" + assert_linequal(promoted, x + y) + # A second promote (from the stale placeholder) returns the existing entry. + assert lazy.promote() is promoted + + def test_promote_derived_expression_raises( + self, m: Model, x: Variable, y: Variable + ) -> None: + lazy = m.add_expressions(lambda m: x + y, name="lazy") + derived = lazy * 2 + assert derived.name is None + with pytest.raises(ValueError, match="derived"): + derived.promote() + + def test_duplicate_name_raises(self, m: Model, x: Variable) -> None: + m.add_expressions(lambda m: 1 * x, name="lazy") + with pytest.raises(ValueError, match="already assigned to model"): + m.add_expressions(lambda m: 2 * x, name="lazy") + with pytest.raises(ValueError, match="already assigned to model"): + m.add_expressions(1 * x, name="lazy") + + def test_auto_naming(self, m: Model, x: Variable) -> None: + lazy = m.add_expressions(lambda m: 1 * x) + assert lazy.name.startswith("expr") + assert lazy.name in m.expressions + + def test_metadata_without_evaluation(self, m: Model, x: Variable) -> None: + calls = 0 + + def evaluator(model: Model) -> LinearExpression: + nonlocal calls + calls += 1 + return 1 * x + + ds = xr.Dataset({"const": ("dim_0", [1.0, 2.0])}) + lazy = m.add_expressions(evaluator, name="lazy", dims=("dim_0",), input_data=ds) + # The input data is shared by pointer, never copied. + assert lazy.input_data is ds + assert lazy.dims == ("dim_0",) + assert "not yet evaluated" in repr(lazy) + assert calls == 0 + lazy.evaluate() + assert calls == 1 + # Nothing is cached: a second evaluation runs the evaluator again. + lazy.evaluate() + assert calls == 2 + + def test_params_forwarded_to_evaluator(self, m: Model, x: Variable) -> None: + lazy = m.add_expressions( + lambda model, factor: factor * x, name="lazy", factor=3 + ) + assert_linequal(lazy.evaluate(), 3 * x) + + def test_params_forwarded_to_callable_mask(self, m: Model, x: Variable) -> None: + mask_calls = 0 + + def mask(model: Model, threshold: int) -> xr.DataArray: + nonlocal mask_calls + mask_calls += 1 + return x.coords["first"] >= threshold + + lazy = m.add_expressions( + lambda model, threshold: x + 1, name="lazy", mask=mask, threshold=1 + ) + assert mask_calls == 0 + result = lazy.evaluate() + assert mask_calls == 1 + expected = (x + 1).where(x.coords["first"] >= 1) + assert_linequal(result, expected) + + def test_concrete_mask_matches_eager(self, x: Variable) -> None: + m = x.model + mask = x.coords["first"] < 1 + lazy = m.add_expressions(lambda model: x + 1, name="lazy", mask=mask) + eager = m.add_expressions(x + 1, name="eager", mask=mask) + assert_linequal(lazy.evaluate(), eager) + + def test_all_false_mask_yields_empty_expression_no_error(self, x: Variable) -> None: + m = x.model + mask = xr.zeros_like(x.coords["first"], dtype=bool) + lazy = m.add_expressions(lambda model: x + 1, name="lazy", mask=mask) + result = lazy.evaluate() + assert result.const.isnull().all() + + def test_callable_mask_requires_callable_data(self, m: Model, x: Variable) -> None: + with pytest.raises(TypeError, match="callable mask"): + m.add_expressions(x + 1, name="lazy", mask=lambda model: True) + + def test_lazy_algebra_stays_lazy_and_matches_eager( + self, m: Model, x: Variable, y: Variable + ) -> None: + lazy = m.add_expressions(lambda m: x + y, name="lazy") + eager = x + y + + combos = [ + (lazy + lazy, eager + eager), + (lazy - lazy, eager - eager), + (lazy * 2, eager * 2), + (2 * lazy, 2 * eager), + (-lazy, -eager), + (eager + lazy, eager + eager), + (eager - lazy, eager - eager), + (np.array(2) * lazy, eager * 2), + ] + for result, expected in combos: + assert isinstance(result, LazyExpression) + assert_linequal(result.evaluate(), expected) + + def test_lazy_pow_and_matmul(self, m: Model, x: Variable) -> None: + lazy = m.add_expressions(lambda m: 1 * x, name="lazy") + + squared = lazy**2 + assert isinstance(squared, LazyExpression) + assert_quadequal(squared.evaluate(), (1 * x) ** 2) + + arr = xr.DataArray( + np.ones(x.coords["first"].size), coords=x.coords, dims=x.dims + ) + matmul_result = lazy @ arr + assert isinstance(matmul_result, LazyExpression) + assert_linequal(matmul_result.evaluate(), (1 * x) @ arr) + + def test_evaluator_called_once_per_leaf_per_evaluate( + self, m: Model, x: Variable, y: Variable + ) -> None: + calls = {"a": 0, "b": 0} + + def eval_a(model: Model) -> LinearExpression: + calls["a"] += 1 + return 1 * x + + def eval_b(model: Model) -> LinearExpression: + calls["b"] += 1 + return 1 * y + + lazy_a = m.add_expressions(eval_a, name="a") + lazy_b = m.add_expressions(eval_b, name="b") + chain = (lazy_a + lazy_b) * 2 + + assert calls == {"a": 0, "b": 0} + chain.evaluate() + assert calls == {"a": 1, "b": 1} + + def test_solution_raises_before_solve( + self, m: Model, x: Variable, y: Variable + ) -> None: + lazy = m.add_expressions(lambda m: x + y, name="lazy") + with pytest.raises(AttributeError, match="not optimized"): + lazy.solution + + @pytest.mark.skipif(not available_solvers, reason="No solver installed") + def test_solution_matches_eager_after_solve(self) -> None: + m = Model() + time = pd.RangeIndex(3, name="time") + x = m.add_variables(lower=1, coords=[time], name="x") + eager = m.add_expressions(2 * x, name="eager") + lazy = m.add_expressions(lambda m: 2 * m.variables["x"], name="lazy") + m.add_objective(x.sum()) + m.solve(available_solvers[0]) + xr.testing.assert_allclose(lazy.solution, eager.solution) + # Requesting the solution must not promote the placeholder. + assert m.expressions["lazy"] is lazy + + @pytest.mark.skipif(not available_solvers, reason="No solver installed") + def test_expressions_solution_with_lazy_member(self) -> None: + m = Model() + time = pd.RangeIndex(3, name="time") + x = m.add_variables(lower=2, coords=[time], name="x") + m.add_expressions(lambda m: 2 * m.variables["x"], name="lazy") + m.add_objective(x.sum()) + m.solve(available_solvers[0]) + + sol = m.expressions.solution + assert isinstance(sol, xr.Dataset) + assert "lazy" in sol + assert (sol["lazy"] == 4).all() + + def test_named_methods_match_eager( + self, m: Model, x: Variable, y: Variable + ) -> None: + lazy = m.add_expressions(lambda m: x + y, name="lazy") + eager = x + y + + combos = [ + (lazy.add(2), eager.add(2)), + (lazy.sub(2), eager.sub(2)), + (lazy.mul(2), eager.mul(2)), + (lazy.div(2), eager.div(2)), + ] + for result, expected in combos: + assert isinstance(result, LazyExpression) + assert_linequal(result.evaluate(), expected) + + def test_named_methods_pow_and_dot(self, m: Model, x: Variable) -> None: + lazy = m.add_expressions(lambda m: 1 * x, name="lazy") + eager = 1 * x + + squared = lazy.pow(2) + assert isinstance(squared, LazyExpression) + assert_quadequal(squared.evaluate(), eager.pow(2)) + + arr = xr.DataArray( + np.ones(x.coords["first"].size), coords=x.coords, dims=x.dims + ) + dot_result = lazy.dot(arr) + assert isinstance(dot_result, LazyExpression) + assert_linequal(dot_result.evaluate(), eager.dot(arr)) + + def test_named_methods_with_join(self, m: Model, y: Variable) -> None: + lazy = m.add_expressions(lambda m: y + 1, name="lazy") + eager = y + 1 + series = pd.Series([1.0, 2.0, 3.0], index=[1, 2, 4], name="second") + + result = lazy.add(series, join="outer") + assert isinstance(result, LazyExpression) + assert_linequal(result.evaluate(), eager.add(series, join="outer")) + + result = lazy.sub(series, join="outer") + assert_linequal(result.evaluate(), eager.sub(series, join="outer")) + + result = lazy.mul(2, join="override") + assert_linequal(result.evaluate(), eager.mul(2, join="override")) + + result = lazy.div(2, join="outer") + assert_linequal(result.evaluate(), eager.div(2, join="outer")) + + # Joining against another expression is rejected the same way eagerly, + # just deferred to evaluation time. + mul_join_expr = lazy.mul(eager, join="outer") + with pytest.raises(TypeError, match="join parameter"): + mul_join_expr.evaluate() + + def test_div_by_expression_defers(self, m: Model, x: Variable, y: Variable) -> None: + lazy_num = m.add_expressions(lambda m: x + y, name="lazy_num") + lazy_den = m.add_expressions(lambda m: x + 1, name="lazy_den") + eager = x + y + + combos: list[LazyExpression] = [] + with pytest.warns( + NonLinearExpressionWarning, match="division by an expression" + ): + combos.append(lazy_num / lazy_den) + with pytest.warns( + NonLinearExpressionWarning, match="division by an expression" + ): + combos.append(lazy_num / x) + with pytest.warns( + NonLinearExpressionWarning, match="division by an expression" + ): + combos.append(lazy_num / eager) + with pytest.warns( + NonLinearExpressionWarning, match="division by an expression" + ): + combos.append(lazy_num.div(eager)) + + for result in combos: + assert isinstance(result, LazyExpression) + assert result.is_evaluatable is False + with pytest.raises(NonLinearOperationError): + result.evaluate() + # These are unnamed, derived expressions: `.promote()` rejects them for + # that reason first (see test_promote_named_ratio_raises for the + # named/nonlinear case). + with pytest.raises(ValueError, match="derived"): + result.promote() + with pytest.raises(NonLinearOperationError): + result.le(1) + + # A constant numerator over a lazy denominator also defers and warns; its + # `.evaluate()` still fails (the eager classes have no `__rtruediv__` for a bare + # constant numerator, a pre-existing, unrelated limitation), but as a plain + # TypeError rather than NonLinearOperationError. + with pytest.warns( + NonLinearExpressionWarning, match="division by an expression" + ): + const_over_lazy = 2 / lazy_den + assert const_over_lazy.is_evaluatable is False + with pytest.raises(TypeError): + const_over_lazy.evaluate() + + # Ordinary constant division is untouched: no warning, still evaluatable. + with warnings.catch_warnings(): + warnings.simplefilter("error", NonLinearExpressionWarning) + const_div = lazy_num / 2 + assert const_div.is_evaluatable is True + assert_linequal(const_div.evaluate(), eager / 2) + + # An eager numerator divided by a lazy denominator still raises outright: the + # nonlinear-division entry point is the lazy operand, not any eager one. + with pytest.raises(TypeError): + eager / lazy_den + with pytest.raises(TypeError): + x / lazy_den + + def test_promote_named_ratio_raises( + self, m: Model, x: Variable, y: Variable + ) -> None: + # Not statically decidable from the constructor call (the division happens + # inside the callable body), so this only surfaces once `.evaluate()` runs. + ratio = m.add_expressions(lambda m: (x + y) / (x + 1), name="ratio") + assert ratio.is_evaluatable is True + with pytest.raises(NonLinearOperationError, match="ratio"): + ratio.promote() + with pytest.raises(NonLinearOperationError): + ratio.evaluate() + + def test_pow_by_non_square_defers(self, m: Model, x: Variable) -> None: + lazy = m.add_expressions(lambda m: 1 * x, name="lazy") + + with pytest.warns(NonLinearExpressionWarning, match="raising to the power 3"): + cubed = lazy**3 + assert cubed.is_evaluatable is False + # The eager `LinearExpression.__pow__` guard raises `ValueError` (not + # `NonLinearOperationError`) for anything but 2 -- unrelated eager behaviour, + # left untouched. `.solution` (tested via the ratio case elsewhere) still + # falls back correctly since it also catches `ValueError`. + with pytest.raises(ValueError, match="Power must be 2"): + cubed.evaluate() + + # Squaring is unaffected. + with warnings.catch_warnings(): + warnings.simplefilter("error", NonLinearExpressionWarning) + squared = lazy.pow(2) + assert squared.is_evaluatable is True + assert_quadequal(squared.evaluate(), (1 * x) ** 2) + + @pytest.mark.skipif(not available_solvers, reason="No solver available") + def test_ratio_solution_after_solve(self) -> None: + m = Model() + x = m.add_variables( + lower=2, upper=2, coords=[pd.RangeIndex(3, name="time")], name="x" + ) + cost = m.add_expressions(lambda m: 3 * m.variables["x"], name="cost") + output = m.add_expressions(lambda m: m.variables["x"], name="output") + with pytest.warns( + NonLinearExpressionWarning, match="division by an expression" + ): + unit_cost = cost / output + m.add_objective(x.sum()) + m.solve(available_solvers[0]) + + xr.testing.assert_equal(unit_cost.solution, cost.solution / output.solution) + assert (unit_cost.solution == 3).all() + + sol = m.expressions.solution + assert "cost" in sol and "output" in sol + + @pytest.mark.skipif(not available_solvers, reason="No solver available") + def test_derived_linear_solution_still_goes_through_evaluate(self) -> None: + m = Model() + time = pd.RangeIndex(3, name="time") + x = m.add_variables(lower=1, coords=[time], name="x") + mask = x.coords["time"] < 2 + lazy = m.add_expressions(lambda m: m.variables["x"] + 1, name="lazy", mask=mask) + derived = ( + lazy + 1 + ) # a derived node, exercising `_combine`'s solution fallback path + assert derived.is_evaluatable is True + m.add_objective(x.sum()) + m.solve(available_solvers[0]) + assert lazy.solution.isnull().any() + # `derived.solution` must come from `derived.evaluate().solution` (which fills the + # masked NaN with 0 before adding, per `_add_constant`), not from silently falling + # back to solution-composition (which would let the NaN propagate through instead). + xr.testing.assert_allclose(derived.solution, derived.evaluate().solution) + assert not derived.solution.isnull().any() + + @pytest.mark.skipif(not available_solvers, reason="No solver available") + def test_dual_reading_lazy_expression(self) -> None: + m = Model() + x = m.add_variables(lower=0, coords=[pd.RangeIndex(3, name="time")], name="x") + m.add_constraints(x >= 2, name="c") + m.add_objective(x.sum()) + + dual_only = m.add_expressions( + lambda m: m.constraints["c"].dual, name="dual_only" + ) + with pytest.raises(AttributeError, match="not optimized"): + dual_only.solution + + weighted = m.add_expressions( + lambda m: m.constraints["c"].dual * m.variables["x"], name="weighted" + ) + + m.solve(available_solvers[0]) + + xr.testing.assert_equal( + dual_only.solution, m.constraints["c"].dual.rename("solution") + ) + assert_linequal(weighted.evaluate(), m.constraints["c"].dual * x) + + with pytest.raises(NonLinearOperationError): + dual_only.promote() + + sol = m.expressions.solution + assert "dual_only" in sol and "weighted" in sol + + def test_constraints_from_lazy(self, m: Model, x: Variable, y: Variable) -> None: + lazy = m.add_expressions(lambda m: x + y, name="lazy") + eager = x + y + + assert_conequal( + lazy.to_constraint(LESS_EQUAL, 5), eager.to_constraint(LESS_EQUAL, 5) + ) + assert_conequal(lazy.le(5, join="outer"), eager.le(5, join="outer")) + assert_conequal(lazy.ge(5), eager.ge(5)) + assert_conequal(lazy.eq(5), eager.eq(5)) + + for con, expected in [(lazy <= 5, eager <= 5), (lazy == 5, eager == 5)]: + assert isinstance(con, Constraint) + assert_conequal(con, expected) + + with pytest.raises(NotImplementedError): + lazy < 5 + with pytest.raises(NotImplementedError): + lazy > 5 + + def test_pow_evaluates_leaf_once(self, m: Model, x: Variable) -> None: + calls = 0 + + def evaluator(model: Model) -> LinearExpression: + nonlocal calls + calls += 1 + return 1 * x + + lazy = m.add_expressions(evaluator, name="lazy") + squared = lazy**2 + assert calls == 0 + squared.evaluate() + assert calls == 1 + + def test_neg_matches_eager_with_mask(self, x: Variable) -> None: + m = x.model + mask = x.coords["first"] < 1 + lazy = m.add_expressions(lambda model: x + 1, name="lazy", mask=mask) + eager = m.add_expressions(x + 1, name="eager", mask=mask) + assert_linequal((-lazy).evaluate(), -eager) + + def test_eager_operands_defer_to_lazy( + self, m: Model, x: Variable, y: Variable + ) -> None: + lazy = m.add_expressions(lambda m: x + y, name="lazy") + eager = x + y + + for result in ( + eager + lazy, + eager - lazy, + eager * lazy, + (x * y) + lazy, + pd.Series([1.0, 2.0, 3.0], index=[1, 2, 3], name="second") * lazy, + xr.DataArray(y.coords["second"].values, coords=y.coords) + lazy, + ): + assert isinstance(result, LazyExpression) + + with pytest.raises(TypeError): + eager / lazy + + def test_lazy_is_unhashable(self, m: Model, x: Variable) -> None: + lazy = m.add_expressions(lambda m: 1 * x, name="lazy") + with pytest.raises(TypeError): + hash(lazy) diff --git a/test/test_io.py b/test/test_io.py index 1842dd10..6bec5bbd 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -18,7 +18,7 @@ import xarray as xr from linopy import LESS_EQUAL, Model, available_solvers, read_netcdf -from linopy.constants import FACTOR_DIM +from linopy.constants import FACTOR_DIM, NonLinearOperationError from linopy.expressions import LinearExpression, QuadraticExpression from linopy.io import signed_number from linopy.testing import assert_exprequal, assert_model_equal @@ -392,6 +392,65 @@ def test_model_to_netcdf_preserves_exprname_counter( assert new_expr.name == "expr2" +@pytest.fixture +def model_with_lazy_expressions() -> Model: + m = Model() + x = m.add_variables(4, pd.Series([8, 10]), name="x") + m.add_expressions(lambda m: m.variables["x"] + 1, name="lazy_lin") + # The division happens inside the callable body, not via LazyExpression + # arithmetic, so it is not statically decidable and only fails at `.evaluate()`. + m.add_expressions(lambda m: m.variables["x"] / m.variables["x"], name="ratio") + m.add_objective(x.sum()) + return m + + +def test_model_to_netcdf_lazy_evaluate( + model_with_lazy_expressions: Model, tmp_path: Path +) -> None: + m = model_with_lazy_expressions + m.remove_expressions("ratio") + fn = tmp_path / "test.nc" + m.to_netcdf(fn, lazy="evaluate") + p = read_netcdf(fn) + + assert "lazy_lin" in p.expressions + assert_exprequal( + p.expressions["lazy_lin"], + m.expressions["lazy_lin"].evaluate(), + check_name=False, + ) + + +def test_model_to_netcdf_lazy_evaluate_raises_for_nonlinear_ratio( + model_with_lazy_expressions: Model, tmp_path: Path +) -> None: + m = model_with_lazy_expressions + fn = tmp_path / "test.nc" + with pytest.raises(NonLinearOperationError, match="ratio"): + m.to_netcdf(fn, lazy="evaluate") + + +def test_model_to_netcdf_lazy_skip( + model_with_lazy_expressions: Model, tmp_path: Path +) -> None: + m = model_with_lazy_expressions + fn = tmp_path / "test.nc" + m.to_netcdf(fn, lazy="skip") + p = read_netcdf(fn) + + assert "lazy_lin" not in p.expressions + assert "ratio" not in p.expressions + + +def test_model_to_netcdf_lazy_raise( + model_with_lazy_expressions: Model, tmp_path: Path +) -> None: + m = model_with_lazy_expressions + fn = tmp_path / "test.nc" + with pytest.raises(ValueError, match="lazy_lin"): + m.to_netcdf(fn, lazy="raise") + + def test_pickle_model_with_expressions( model_with_expressions: Model, tmp_path: Path ) -> None: