From eafbf77db5289144ebd49a4ac315861a5980b6f9 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:09:58 +0100 Subject: [PATCH 1/7] Add expressions container and adder method --- linopy/expressions.py | 151 +++++++++++++++++++++++++++++++++++++++++- linopy/model.py | 94 ++++++++++++++++++++++++++ 2 files changed, 243 insertions(+), 2 deletions(-) diff --git a/linopy/expressions.py b/linopy/expressions.py index 21a4160e9..7a7d5a8b3 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -11,10 +11,26 @@ import logging import operator from abc import ABC, abstractmethod -from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence +from collections.abc import ( + Callable, + Hashable, + ItemsView, + Iterable, + Iterator, + Mapping, + Sequence, +) from dataclasses import dataclass, field from itertools import product, zip_longest -from typing import TYPE_CHECKING, Any, Self, TypeAlias, TypeVar, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + Self, + TypeAlias, + TypeVar, + cast, + overload, +) from warnings import warn import numpy as np @@ -55,6 +71,7 @@ filter_nulls_polars, format_coord, format_single_expression, + format_string_as_variable_name, forward_as_properties, generate_indices_for_printout, get_dims_with_index_levels, @@ -64,6 +81,7 @@ is_constant, iterate_slices, maybe_group_terms_polars, + save_join, to_dataframe, to_polars, ) @@ -735,6 +753,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: # TODO: add a warning here, routines should be safe against this data = data.drop_vars(drop_dims) + data = data.assign_attrs(name=None) self._model = model self._data = cast(Dataset, data) @@ -1235,6 +1254,13 @@ def loc(self) -> LocIndexer: def type(self) -> str: return "LinearExpression" + @property + def name(self) -> str: + """ + Return the name of the variable. + """ + return str(self.attrs["name"]) + @property def data(self) -> Dataset: return self._data @@ -2827,6 +2853,127 @@ def merge( return cls(ds, model) +@dataclass(repr=False) +class Expressions: + """ + An expressions container used for storing multiple expression arrays. + """ + + data: dict[str, LinearExpression | QuadraticExpression] + model: Model + + def _formatted_names(self) -> dict[str, str]: + """ + Get a dictionary of formatted names to the proper variable names. + This map enables a attribute like accession of variable names which + are not valid python variable names. + """ + return {format_string_as_variable_name(n): n for n in self} + + @overload + def __getitem__(self, names: str) -> LinearExpression | QuadraticExpression: ... + + @overload + def __getitem__(self, names: list[str]) -> Expressions: ... + + def __getitem__( + self, names: str | list[str] + ) -> LinearExpression | QuadraticExpression | 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: + # If name is an attribute of self (including methods and properties), return that + if name in self.data: + return self.data[name] + else: + if name in (formatted_names := self._formatted_names()): + return self.data[formatted_names[name]] + raise AttributeError( + f"Expressions has no attribute `{name}` or the attribute is not accessible / raises an error." + ) + + def __getstate__(self) -> dict: + return self.__dict__ + + def __setstate__(self, d: dict) -> None: + self.__dict__.update(d) + + def __dir__(self) -> list[str]: + base_attributes = list(super().__dir__()) + formatted_names = [ + n for n in self._formatted_names() if n not in base_attributes + ] + return base_attributes + formatted_names + + def _format_items(self, exclude: set[str] | None = None) -> str: + """Format expression items, optionally excluding names in a group.""" + r = "" + count = 0 + for name, ds in self.items(): + if exclude and name in exclude: + continue + count += 1 + coords = ( + " (" + ", ".join(str(coord) for coord in ds.coords) + ")" + if ds.coords + else "" + ) + r += f" * {name}{coords}\n" + if count == 0: + r += "\n" + return r + + def __repr__(self) -> str: + """ + Return a string representation of the expressions container. + """ + r = "linopy.model.Expressions" + line = "-" * len(r) + r += f"\n{line}\n" + r += self._format_items() + return r + + def __len__(self) -> int: + return self.data.__len__() + + def __iter__(self) -> Iterator[str]: + return self.data.__iter__() + + def items(self) -> ItemsView[str, LinearExpression | QuadraticExpression]: + return self.data.items() + + def _ipython_key_completions_(self) -> list[str]: + """ + Provide method for the key-autocompletions in IPython. + + See + http://ipython.readthedocs.io/en/stable/config/integrating.html#tab-completion + For the details. + """ + return list(self) + + def add(self, expression: LinearExpression | QuadraticExpression) -> None: + """ + Add an expression to the expressions container. + """ + self.data[expression.name] = expression + + def remove(self, name: str) -> None: + """ + Remove variable `name` from the variables. + """ + self.data.pop(name) + + @property + def solution(self) -> Dataset: + """ + Get the solution of variables. + """ + return save_join(*[v.solution.rename(k) for k, v in self.items()]) + + class ScalarLinearExpression: """ A scalar linear expression container. diff --git a/linopy/model.py b/linopy/model.py index 24594c9cc..136ac0adb 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -57,6 +57,7 @@ ) from linopy.dualization import dualize from linopy.expressions import ( + Expressions, LinearExpression, QuadraticExpression, ScalarLinearExpression, @@ -133,6 +134,7 @@ class Model: _solver: solvers.Solver | None _variables: Variables + _expressions: Expressions _constraints: Constraints _objective: Objective _parameters: Dataset @@ -144,6 +146,7 @@ class Model: _cCounter: int _dtypes: dict[DtypeKey, type[np.signedinteger]] _varnameCounter: int + _exprnameCounter: int _connameCounter: int _pwlCounter: int _blocks: DataArray | None @@ -155,6 +158,7 @@ class Model: __slots__ = ( # containers "_variables", + "_expressions", "_constraints", "_objective", "_parameters", @@ -168,6 +172,7 @@ class Model: "_cCounter", "_dtypes", "_varnameCounter", + "_exprnameCounter", "_connameCounter", "_pwlCounter", "_blocks", @@ -258,6 +263,7 @@ def __init__( dtypes ) self._variables: Variables = Variables({}, model=self) + self._expressions: Expressions = Expressions({}, model=self) self._constraints: Constraints = Constraints({}, model=self) self._objective: Objective = Objective(LinearExpression(None, self), self) self._parameters: Dataset = Dataset() @@ -267,6 +273,7 @@ def __init__( self._xCounter: int = 0 self._cCounter: int = 0 self._varnameCounter: int = 0 + self._exprnameCounter: int = 0 self._connameCounter: int = 0 self._pwlCounter: int = 0 self._blocks: DataArray | None = None @@ -326,6 +333,13 @@ def variables(self) -> Variables: """ return self._variables + @property + def expressions(self) -> Expressions: + """ + Expressions assigned to the model. + """ + return self._expressions + @property def constraints(self) -> Constraints: """ @@ -572,6 +586,7 @@ def scalar_attrs(self) -> list[str]: "_xCounter", "_cCounter", "_varnameCounter", + "_exprnameCounter", "_connameCounter", "_pwlCounter", "force_dim_names", @@ -590,11 +605,13 @@ def __repr__(self) -> str: var_names, con_names = _get_piecewise_groups(self) var_string = self.variables._format_items(exclude=var_names) con_string = self.constraints._format_items(exclude=con_names) + expr_string = self.expressions._format_items() model_string = f"Linopy {self.type} model" return ( f"{model_string}\n{'=' * len(model_string)}\n\n" f"Variables:\n----------\n{var_string}\n" + f"Expressions:\n------------\n{expr_string}\n" f"Constraints:\n------------\n{con_string}" f"{pwl_repr_summary(self)}" f"\nStatus:\n-------\n{self.status}" @@ -913,6 +930,83 @@ def add_variables( self.variables.add(variable) return variable + def add_expressions( + self, + data: Variable + | LinearExpression + | QuadraticExpression + | Sequence[tuple[ConstantLike, Variable | str]], + name: str | None = None, + mask: MaskLike | None = None, + ) -> LinearExpression | QuadraticExpression: + """ + Assign a new, possibly multi-dimensional array of expressions to the + model. + + 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. + name : str, optional + Reference name of the added expressions. The default None results in + a name like "expr1", "expr2" etc. + mask : array_like, 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. + + Raises + ------ + ValueError + If neither lower bound and upper bound have coordinates, nor + `coords` are directly given. + + Returns + ------- + linopy.LinearExpression | linopy.QuadraticExpression + Expression which was added to the model. + + + 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") + >>> expr = m.add_expressions(x + 1, name="expr") + """ + 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") + + expr: LinearExpression | QuadraticExpression + if isinstance(data, Variable): + expr = data.to_linexpr() + elif isinstance(data, Sequence): + expr = self.linexpr(*data) + else: + expr = data + self.check_force_dim_names(expr.data) + self._check_valid_dim_names(expr.data) + + if mask is not None: + mask = as_dataarray(mask, coords=expr.coords, dims=expr.dims).astype(bool) + expr = expr.where(mask) + if self.chunk: + expr = expr.chunk(self.chunk) + + expr.attrs["name"] = name + self.expressions.add(expr) + return expr + def add_sos_constraints( self, variable: Variable, From 800c18e08a83aea28d9e7d840fdd48e6d30cef91 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:12:18 +0100 Subject: [PATCH 2/7] Add expressions tests & docs; add remove_expression method Co-authored-by: Claude --- examples/creating-expressions.ipynb | 90 ++++++++++++++++++ linopy/model.py | 21 +++++ test/test_expressions.py | 141 ++++++++++++++++++++++++++++ 3 files changed, 252 insertions(+) create mode 100644 test/test_expressions.py diff --git a/examples/creating-expressions.ipynb b/examples/creating-expressions.ipynb index cb41a2c66..ce6017ba4 100644 --- a/examples/creating-expressions.ipynb +++ b/examples/creating-expressions.ipynb @@ -485,6 +485,96 @@ "source": [ "x.rolling(time=3).sum()" ] + }, + { + "cell_type": "markdown", + "id": "45", + "metadata": {}, + "source": [ + "## Storing expressions on the model\n", + "\n", + "The expressions we have built so far are ordinary Python objects: they live in\n", + "a notebook variable but are not attached to the model in any way. Sometimes\n", + "you want to reuse the same expression in several constraints, in the\n", + "objective, or inspect it after solving — for that, `m.add_expressions`\n", + "registers an expression under a name on the model, similar to how\n", + "`m.add_variables` registers a variable.\n", + "\n", + "If you don't pass a `name`, one is generated automatically (`expr0`, `expr1`,\n", + "...). Both `LinearExpression` and `QuadraticExpression` can be stored." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "46", + "metadata": {}, + "outputs": [], + "source": [ + "total = m.add_expressions(x + y, name=\"total\")\n", + "total" + ] + }, + { + "cell_type": "markdown", + "id": "47", + "metadata": {}, + "source": [ + "The stored expressions are reachable through `m.expressions`, which behaves\n", + "like a dict of expressions:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "48", + "metadata": {}, + "outputs": [], + "source": [ + "m.expressions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "49", + "metadata": {}, + "outputs": [], + "source": [ + "# equivalent to m.expressions.total\n", + "m.expressions[\"total\"]" + ] + }, + { + "cell_type": "markdown", + "id": "50", + "metadata": {}, + "source": [ + "Stored expressions also show up in the model's overview, next to the\n", + "variables and constraints:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "51", + "metadata": {}, + "outputs": [], + "source": [ + "m" + ] + }, + { + "cell_type": "markdown", + "id": "52", + "metadata": {}, + "source": [ + ".. tip::\n", + " After solving the model, ``m.expressions.solution`` returns an\n", + " `xarray.Dataset` with one entry per stored expression, evaluated at the\n", + " optimal solution — handy for inspecting derived quantities without\n", + " rebuilding the expression by hand." + ] } ], "metadata": { diff --git a/linopy/model.py b/linopy/model.py index 136ac0adb..cbdd4674b 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1475,6 +1475,27 @@ def remove_constraints(self, name: str | list[str]) -> None: logger.debug(f"Removed constraint: {name}") self.constraints.remove(name) + def remove_expressions(self, name: str | list[str]) -> None: + """ + Remove all expressions stored under reference name 'name' from the + model. + + Parameters + ---------- + name : str or list of str + Reference name(s) of the expressions to remove. If a single name is + provided, only that expression will be removed. If a list of names + is provided, all expressions with those names will be removed. + + Returns + ------- + None. + """ + names = [name] if isinstance(name, str) else name + for n in names: + logger.debug(f"Removed expression: {n}") + self.expressions.remove(n) + def remove_sos_constraints(self, variable: Variable) -> None: """ Remove all sos constraints from a given variable. diff --git a/test/test_expressions.py b/test/test_expressions.py new file mode 100644 index 000000000..cb49af8e1 --- /dev/null +++ b/test/test_expressions.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +""" +This module aims at testing the correct behavior of the Expressions class. +""" + +import pandas as pd +import pytest +import xarray as xr + +from linopy import Model +from linopy.expressions import Expressions, LinearExpression, QuadraticExpression +from linopy.solvers import available_solvers +from linopy.testing import assert_linequal + + +@pytest.fixture +def m() -> Model: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(10, name="first")], name="x") + y = m.add_variables(coords=[pd.Index([1, 2, 3], name="second")], name="y") + m.add_expressions(x + 1, name="expr_x") + m.add_expressions(x * y, name="expr_xy") + return m + + +def test_expressions_repr(m: Model) -> None: + m.expressions.__repr__() + repr(Model()) + + +def test_expressions_getitem(m: Model) -> None: + assert isinstance(m.expressions["expr_x"], LinearExpression) + + subset = m.expressions[["expr_x"]] + assert isinstance(subset, Expressions) + assert len(subset) == 1 + + +def test_expressions_getattr(m: Model) -> None: + assert_linequal(m.expressions.expr_x, m.expressions["expr_x"]) + + with pytest.raises(AttributeError): + m.expressions.does_not_exist + + +def test_expressions_getattr_formatted() -> None: + m = Model() + x = m.add_variables(name="x") + m.add_expressions(x + 1, name="e-0") + assert_linequal(m.expressions.e_0, m.expressions["e-0"]) + + +def test_expressions_dict_protocol(m: Model) -> None: + assert len(m.expressions) == 2 + assert set(iter(m.expressions)) == {"expr_x", "expr_xy"} + assert set(dict(m.expressions.items())) == {"expr_x", "expr_xy"} + assert "expr_x" in m.expressions + assert m.expressions._ipython_key_completions_() == list(m.expressions) + assert "expr_x" in dir(m.expressions) + + +def test_expressions_name_counter() -> None: + m = Model() + x = m.add_variables(name="x") + m.add_expressions(x + 1) + m.add_expressions(x + 1) + assert "expr0" in m.expressions + assert "expr1" in m.expressions + + +def test_expressions_duplicate_name_raises(m: Model) -> None: + x = m.variables["x"] + with pytest.raises(ValueError, match="already assigned"): + m.add_expressions(x + 1, name="expr_x") + + +def test_add_expressions_from_variable_and_tuples() -> None: + m = Model() + x = m.add_variables(name="x") + + expr = m.add_expressions(x, name="from_var") + assert isinstance(expr, LinearExpression) + assert_linequal(expr, x.to_linexpr()) + + expr = m.add_expressions([(2, x)], name="from_tuples") + assert isinstance(expr, LinearExpression) + assert_linequal(expr, 2 * x) + + +def test_add_expressions_quadratic(m: Model) -> None: + assert isinstance(m.expressions["expr_xy"], QuadraticExpression) + + +def test_add_expressions_mask() -> None: + m = Model() + idx = pd.RangeIndex(10, name="first") + x = m.add_variables(coords=[idx], name="x") + mask = xr.DataArray([True] * 5 + [False] * 5, coords=[idx]) + + expr = m.add_expressions(x + 1, name="masked", mask=mask) + assert_linequal(expr, (x + 1).where(mask)) + + +def test_expressions_remove(m: Model) -> None: + m.expressions.remove("expr_x") + assert "expr_x" not in m.expressions + + with pytest.raises(KeyError): + m.expressions.remove("expr_x") + + +def test_remove_expressions(m: Model) -> None: + m.remove_expressions("expr_x") + assert "expr_x" not in m.expressions + assert "expr_xy" in m.expressions + + +def test_remove_expressions_with_list(m: Model) -> None: + m.remove_expressions(["expr_x", "expr_xy"]) + assert len(m.expressions) == 0 + + +def test_model_repr_contains_expressions(m: Model) -> None: + r = repr(m) + assert "Expressions:" in r + assert "* expr_x" in r + + +@pytest.mark.skipif(not available_solvers, reason="No solver available") +def test_expressions_solution() -> None: + m = Model() + x = m.add_variables(lower=0, coords=[pd.RangeIndex(3, name="first")], name="x") + m.add_constraints(x >= 2) + m.add_expressions(2 * x, name="double_x") + m.add_objective(x.sum()) + m.solve(available_solvers[0]) + + sol = m.expressions.solution + assert isinstance(sol, xr.Dataset) + assert "double_x" in sol + assert (sol["double_x"] == 4).all() From e9aa8aa3f2c96d927f9bc67aa1eb430257dd1283 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:13:46 +0100 Subject: [PATCH 3/7] Fix IO issue caused by nameless objective --- linopy/objective.py | 1 + 1 file changed, 1 insertion(+) diff --git a/linopy/objective.py b/linopy/objective.py index a51b22076..67d141c8a 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -192,6 +192,7 @@ def expression( if (expr.const != 0.0) and not np.isnan(expr.const): raise ValueError("Constant values in objective function not supported.") + expr.attrs["name"] = "objective" self._expression = expr @property From c7e79f815cb11e52042f7936ee22c066fb4e4c1f Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:19:08 +0100 Subject: [PATCH 4/7] Include expressions in IO. Co-authored-by: Claude --- doc/api.rst | 34 +++++++ doc/release_notes.rst | 5 + linopy/expressions.py | 2 + linopy/io.py | 53 ++++++++++- linopy/testing.py | 25 +++++ test/test_io.py | 208 +++++++++++++++++++++++++++++++++++++++++- test/test_model.py | 38 ++++++++ test/test_testing.py | 48 +++++++++- 8 files changed, 407 insertions(+), 6 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index 973915893..b62351ea4 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -42,6 +42,7 @@ Building a model model.Model.add_variables model.Model.add_constraints model.Model.add_objective + model.Model.add_expressions model.Model.add_sos_constraints model.Model.add_piecewise_formulation @@ -53,6 +54,7 @@ Inspecting a model model.Model.variables model.Model.constraints + model.Model.expressions model.Model.objective model.Model.sense model.Model.type @@ -67,6 +69,7 @@ Modifying a model model.Model.remove_variables model.Model.remove_constraints + model.Model.remove_expressions model.Model.remove_objective model.Model.remove_sos_constraints model.Model.copy @@ -215,6 +218,35 @@ Inventory variables.Variables.sos +Expressions +=========== + +Container for the collection of named expressions on a model. Accessed via +``model.expressions``. + +.. autosummary:: + :toctree: generated/ + + expressions.Expressions + +Modification +------------ + +.. autosummary:: + :toctree: generated/ + + expressions.Expressions.add + expressions.Expressions.remove + +Post-solve access +----------------- + +.. autosummary:: + :toctree: generated/ + + expressions.Expressions.solution + + LinearExpression ================ @@ -251,6 +283,7 @@ Structure .. autosummary:: :toctree: generated/ + expressions.LinearExpression.name expressions.LinearExpression.vars expressions.LinearExpression.coeffs expressions.LinearExpression.const @@ -292,6 +325,7 @@ Structure .. autosummary:: :toctree: generated/ + expressions.QuadraticExpression.name expressions.QuadraticExpression.vars expressions.QuadraticExpression.coeffs expressions.QuadraticExpression.const diff --git a/doc/release_notes.rst b/doc/release_notes.rst index cc14abf92..c98061cb6 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -17,6 +17,11 @@ Upcoming Version * ``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 `__) +*Named expressions* + +* ``Model.add_expressions`` registers a ``LinearExpression`` or ``QuadraticExpression`` under a name (auto-generated as ``expr0``, ``expr1``, ... if omitted), accessible afterwards via ``Model.expressions`` (an ``Expressions`` container mirroring ``Model.variables``/``Model.constraints``) and removable via ``Model.remove_expressions``. + Named expressions are persisted by ``Model.to_netcdf``/``linopy.read_netcdf`` and preserved by ``Model.copy``, ``copy.copy``, ``copy.deepcopy``, and pickling. + *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/linopy/expressions.py b/linopy/expressions.py index 7a7d5a8b3..01cff3a2b 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -2862,6 +2862,8 @@ class Expressions: data: dict[str, LinearExpression | QuadraticExpression] model: Model + dataset_attrs = ["coeffs", "vars", "const"] + def _formatted_names(self) -> dict[str, str]: """ Get a dictionary of formatted names to the proper variable names. diff --git a/linopy/io.py b/linopy/io.py index 462fa5b8f..aca62e06a 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -26,7 +26,7 @@ from linopy import solvers from linopy.common import to_polars -from linopy.constants import CONCAT_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR +from linopy.constants import CONCAT_DIM, FACTOR_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR from linopy.objective import Objective if TYPE_CHECKING: @@ -934,6 +934,11 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: Notes ----- + 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`. + The SOS reformulation lifecycle token lives only on the in-memory Model and is not persisted. If the model has an active SOS reformulation at serialization time, the netcdf contains the @@ -978,6 +983,13 @@ 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() + ] objective = m.objective.data objective = objective.assign_attrs(sense=m.objective.sense) if m.objective.value is not None: @@ -986,7 +998,7 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: params = [with_prefix(m.parameters, "parameters")] scalars = {k: getattr(m, k) for k in m.scalar_attrs} - ds = xr.merge(vars + cons + obj + params, combine_attrs="drop_conflicts") + ds = xr.merge(vars + cons + exprs + obj + params, combine_attrs="drop_conflicts") ds = ds.assign_attrs(scalars) ds.attrs[NETCDF_VERSION_ATTR] = version("linopy") if m._relaxed_registry: @@ -1039,7 +1051,7 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: Constraints, CSRConstraint, ) - from linopy.expressions import LinearExpression + from linopy.expressions import Expressions, LinearExpression, QuadraticExpression from linopy.model import Model from linopy.variables import Variable, Variables @@ -1095,6 +1107,26 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: m._variables = Variables(variables, m) + 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] = {} + for k in sorted(expr_names): + name = remove_prefix(k, "expressions") + expr_ds = get_prefix(ds, k) + expr_type = expr_ds.attrs.pop("_linopy_expr_type", None) + expr_ds.attrs.pop("name", None) # re-attached below, after construction + expr: LinearExpression | QuadraticExpression + if expr_type == "QuadraticExpression" or ( + expr_type is None and FACTOR_DIM in expr_ds.dims + ): + expr = QuadraticExpression(expr_ds, m) + else: + expr = LinearExpression(expr_ds, m) + expr.attrs["name"] = name + expressions[name] = expr + + m._expressions = Expressions(expressions, m) + cons = [str(k) for k in ds if str(k).startswith("constraints")] con_names = list({str(k).rsplit("-", 1)[0] for k in cons}) constraints: dict[str, ConstraintBase] = {} @@ -1178,7 +1210,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 LinearExpression + from linopy.expressions import Expressions, LinearExpression, QuadraticExpression from linopy.model import Model, Objective from linopy.variables import Variable, Variables @@ -1207,6 +1239,19 @@ def copy(m: Model, include_solution: bool = False, deep: bool = True) -> Model: new_model, ) + def _copy_expr( + name: str, expr: LinearExpression | QuadraticExpression + ) -> LinearExpression | QuadraticExpression: + # Expressions hold no solve artifacts, so include_solution is irrelevant. + 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 + + new_model._expressions = Expressions( + {name: _copy_expr(name, expr) for name, expr in m.expressions.items()}, + new_model, + ) + def _copy_con_data(con: ConstraintBase) -> xr.Dataset: d = con.mutable().data if include_solution: diff --git a/linopy/testing.py b/linopy/testing.py index 5e88f2a9b..d9c67f7be 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -70,6 +70,26 @@ def assert_quadequal( return assert_equal(_expr_unwrap(a), _expr_unwrap(b)) +def assert_exprequal( + a: LinearExpression | QuadraticExpression, + b: LinearExpression | QuadraticExpression, + check_name: bool = True, +) -> None: + """ + Assert that two expressions are equal, dispatching on linear vs quadratic. + + xarray's assert_equal ignores attrs, so the stored name (which lives in + ``attrs["name"]``) is compared explicitly unless ``check_name=False``. + """ + 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}" + if isinstance(a, QuadraticExpression): + assert_quadequal(a, b) + else: + assert_linequal(a, b) + + def assert_conequal(a: ConstraintBase, b: ConstraintBase, strict: bool = True) -> None: """ Assert that two constraints are equal. @@ -105,6 +125,11 @@ def assert_model_equal(a: Model, b: Model) -> None: for c in a.constraints: assert_conequal(a.constraints[c], b.constraints[c]) + assert set(a.expressions) == set(b.expressions) + + for e in a.expressions: + assert_exprequal(a.expressions[e], b.expressions[e]) + assert_linequal(a.objective.expression, b.objective.expression) assert a.objective.sense == b.objective.sense assert a.objective.value == b.objective.value diff --git a/test/test_io.py b/test/test_io.py index 27cba396b..1842dd10b 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -18,8 +18,10 @@ import xarray as xr from linopy import LESS_EQUAL, Model, available_solvers, read_netcdf +from linopy.constants import FACTOR_DIM +from linopy.expressions import LinearExpression, QuadraticExpression from linopy.io import signed_number -from linopy.testing import assert_model_equal +from linopy.testing import assert_exprequal, assert_model_equal HAS_NETCDF4 = importlib.util.find_spec("netCDF4") is not None @@ -74,6 +76,37 @@ def model_with_multiindex() -> Model: return m +@pytest.fixture +def model_with_expressions() -> Model: + m = Model() + + x = m.add_variables(4, pd.Series([8, 10]), name="x") + y = m.add_variables(0, pd.DataFrame([[1, 2], [3, 4]]), name="y") + + m.add_expressions(x + 1, name="lin") + m.add_expressions(x * y, name="quad") + m.add_expressions(2 * x + 3 * y, name="mixed-dims-expr") + + m.add_constraints(x + y, LESS_EQUAL, 10) + m.add_objective(m.expressions["mixed-dims-expr"]) + + return m + + +@pytest.fixture +def model_with_masked_expression() -> Model: + m = Model() + + idx = pd.RangeIndex(6, name="i") + x = m.add_variables(coords=[idx], name="x") + mask = xr.DataArray([True, True, True, False, False, False], coords=[idx]) + m.add_expressions(x + 1, name="masked", mask=mask) + + m.add_objective(x.sum()) + + return m + + def test_model_to_netcdf(model: Model, tmp_path: Path) -> None: m = model fn = tmp_path / "test.nc" @@ -202,6 +235,179 @@ def test_model_to_netcdf_with_multiindex_scipy_engine( assert_model_equal(m, read_netcdf(fn)) +def test_model_to_netcdf_with_expressions( + model_with_expressions: Model, tmp_path: Path +) -> None: + m = model_with_expressions + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert set(p.expressions) == {"lin", "quad", "mixed-dims-expr"} + assert_model_equal(m, p) + + +def test_model_to_netcdf_linear_expression( + model_with_expressions: Model, tmp_path: Path +) -> None: + m = model_with_expressions + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert isinstance(p.expressions["lin"], LinearExpression) + assert not isinstance(p.expressions["lin"], QuadraticExpression) + assert_exprequal(m.expressions["lin"], p.expressions["lin"]) + + +def test_model_to_netcdf_quadratic_expression( + model_with_expressions: Model, tmp_path: Path +) -> None: + m = model_with_expressions + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert isinstance(p.expressions["quad"], QuadraticExpression) + assert p.expressions["quad"].data.sizes[FACTOR_DIM] == 2 + assert_exprequal(m.expressions["quad"], p.expressions["quad"]) + + +def test_model_to_netcdf_expression_dash_name( + model_with_expressions: Model, tmp_path: Path +) -> None: + m = model_with_expressions + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert "mixed-dims-expr" in p.expressions + assert p.expressions["mixed-dims-expr"].name == "mixed-dims-expr" + + +def test_model_to_netcdf_masked_expression( + model_with_masked_expression: Model, tmp_path: Path +) -> None: + m = model_with_masked_expression + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert_model_equal(m, p) + + masked = p.expressions["masked"] + np.testing.assert_array_equal( + masked.vars.values, m.expressions["masked"].vars.values + ) + assert np.isnan(masked.coeffs.values[3:]).all() + + +def test_model_to_netcdf_expression_with_multiindex( + model_with_multiindex: Model, tmp_path: Path +) -> None: + m = model_with_multiindex + x = m.variables["x-var"] + y = m.variables["y-var"] + m.add_expressions(x + y, name="mi-expr") + + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert_model_equal(m, p) + index = p.expressions["mi-expr"].indexes["dim_0"] + assert isinstance(index, pd.MultiIndex) + assert list(index.names) == ["first", "second"] + + +def test_model_to_netcdf_expression_with_multiindex_scipy_engine( + model_with_multiindex: Model, tmp_path: Path +) -> None: + m = model_with_multiindex + x = m.variables["x-var"] + y = m.variables["y-var"] + m.add_expressions(x + y, name="mi-expr") + + fn = tmp_path / "test.nc" + m.to_netcdf(fn, engine="scipy") + + raw_attrs = xr.load_dataset(fn).attrs + expr_multiindex_attrs = { + k: v + for k, v in raw_attrs.items() + if k.startswith("expressions-mi-expr") and k.endswith("_multiindex") + } + assert expr_multiindex_attrs + for k, v in expr_multiindex_attrs.items(): + assert isinstance(v, str), f"{k!r}: {v!r}" + + assert_model_equal(m, read_netcdf(fn)) + + +def test_model_to_netcdf_expression_labels_stay_valid( + model_with_expressions: Model, tmp_path: Path +) -> None: + m = model_with_expressions + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + valid_labels = set(p.variables.flat.labels) + labels = p.expressions["lin"].vars.values.ravel() + assert all(label == -1 or label in valid_labels for label in labels) + + # "lin" is `x + 1`, so its single term per element should equal x's own labels. + x_labels = p.variables["x"].labels + lin_labels = p.expressions["lin"].vars.isel(_term=0) + xr.testing.assert_equal(lin_labels.rename(None), x_labels.rename(None)) + + +def test_model_to_netcdf_empty_expressions(model: Model, tmp_path: Path) -> None: + m = model + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert len(p.expressions) == 0 + assert_model_equal(m, p) + + raw = xr.load_dataset(fn) + assert not any(str(k).startswith("expressions") for k in raw) + + +def test_model_to_netcdf_preserves_exprname_counter( + model: Model, tmp_path: Path +) -> None: + m = model + x = m.variables["x"] + m.add_expressions(x + 1) + m.add_expressions(x + 2) + + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert p._exprnameCounter == m._exprnameCounter == 2 + new_expr = p.add_expressions(p.variables["x"] + 3) + assert new_expr.name == "expr2" + + +def test_pickle_model_with_expressions( + model_with_expressions: Model, tmp_path: Path +) -> None: + m = model_with_expressions + fn = tmp_path / "test.pkl" + + with open(fn, "wb") as f: + pickle.dump(m, f) + + with open(fn, "rb") as f: + p = pickle.load(f) + + assert_model_equal(m, p) + assert p.expressions["lin"].model is p + + @pytest.mark.skipif(not HAS_NETCDF4, reason="legacy format requires netCDF4 backend") def test_read_netcdf_with_multiindex_legacy_list_attr( model_with_multiindex: Model, tmp_path: Path diff --git a/test/test_model.py b/test/test_model.py index 6b9e31576..a246f3bfb 100644 --- a/test/test_model.py +++ b/test/test_model.py @@ -319,6 +319,44 @@ def test_model_deepcopy_protocol(copy_test_model: Model) -> None: assert m.objective.sense == original_sense +@pytest.fixture(scope="module") +def copy_test_model_with_expressions() -> Model: + """Representative model with named linear and quadratic expressions.""" + m: Model = Model() + + lower: xr.DataArray = xr.DataArray( + np.zeros((10, 10)), coords=[range(10), range(10)] + ) + upper: xr.DataArray = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)]) + x = m.add_variables(lower, upper, name="x") + y = m.add_variables(name="y") + + m.add_expressions(x + 1, name="lin") + m.add_expressions(x * y, name="quad") + + m.add_constraints(1 * x + 10 * y, EQUAL, 0) + m.add_objective((10 * x + 5 * y).sum()) + + return m + + +def test_copy_model_with_expressions( + copy_test_model_with_expressions: Model, +) -> None: + """Model.copy(), copy.copy() and copy.deepcopy() all preserve expressions.""" + m = copy_test_model_with_expressions.copy(deep=True) + + for c in (m.copy(), pycopy.copy(m), pycopy.deepcopy(m)): + assert_model_equal(m, c) + assert c.expressions["lin"].model is c + assert c.expressions["quad"].model is c + + deep = pycopy.deepcopy(m) + original_coeff = m.expressions["lin"].coeffs.values.flat[0].item() + deep.expressions["lin"].coeffs.values.flat[0] = original_coeff + 42 + assert m.expressions["lin"].coeffs.values.flat[0] == original_coeff + + @pytest.mark.skipif(not available_solvers, reason="No solver installed") class TestModelCopySolved: def test_model_deepcopy_protocol_excludes_solution( diff --git a/test/test_testing.py b/test/test_testing.py index d0fabc86e..274f31ae6 100644 --- a/test/test_testing.py +++ b/test/test_testing.py @@ -2,7 +2,7 @@ import pytest from linopy import Model -from linopy.testing import assert_linequal +from linopy.testing import assert_exprequal, assert_linequal, assert_model_equal @pytest.fixture @@ -34,3 +34,49 @@ def test_assert_linequal_still_detects_real_differences(model: Model) -> None: assert_linequal(1 * a, 1 * c) # different dimension sets with pytest.raises(AssertionError): assert_linequal(1 * a, 2 * a) # different coefficients + + +def test_assert_exprequal_detects_type_mismatch(model: Model) -> None: + """A linear and a quadratic expression must never compare equal.""" + a = model.add_variables(coords=[pd.Index([0, 1], name="i")], name="a") + b = model.add_variables(coords=[pd.Index([0, 1], name="i")], name="b") + + with pytest.raises(AssertionError, match="expression types differ"): + assert_exprequal(a + 1, a * b) + + +def test_assert_exprequal_detects_name_mismatch(model: Model) -> None: + """Expressions with identical values but different stored names differ.""" + a = model.add_variables(coords=[pd.Index([0, 1], name="i")], name="a") + + lhs = model.add_expressions(a + 1, name="first") + rhs = model.add_expressions(a + 1, name="second") + + with pytest.raises(AssertionError, match="expression names differ"): + assert_exprequal(lhs, rhs) + + # names deliberately ignored + assert_exprequal(lhs, rhs, check_name=False) + + +def test_assert_model_equal_detects_expression_difference() -> None: + """assert_model_equal must fail when expressions differ between models.""" + m1 = Model() + a1 = m1.add_variables(coords=[pd.Index([0, 1], name="i")], name="a") + m1.add_expressions(a1 + 1, name="expr") + m1.add_objective(a1.sum()) + + m2 = Model() + a2 = m2.add_variables(coords=[pd.Index([0, 1], name="i")], name="a") + m2.add_expressions(a2 + 2, name="expr") # different coefficients + m2.add_objective(a2.sum()) + + with pytest.raises(AssertionError): + assert_model_equal(m1, m2) + + m3 = Model() + a3 = m3.add_variables(coords=[pd.Index([0, 1], name="i")], name="a") + m3.add_objective(a3.sum()) # no "expr" at all + + with pytest.raises(AssertionError): + assert_model_equal(m1, m3) From 6f0a220025e5e05706c7c0da889129704542829b Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:37:57 +0100 Subject: [PATCH 5/7] Add lazy expression object Co-authored-by: Claude --- doc/api.rst | 31 ++++ doc/release_notes.rst | 2 + examples/creating-expressions.ipynb | 111 +++++++++++- linopy/__init__.py | 8 +- linopy/expressions.py | 264 +++++++++++++++++++++++++++- linopy/io.py | 87 +++++++-- linopy/model.py | 109 ++++++++++-- linopy/testing.py | 31 +++- test/test_expressions.py | 216 ++++++++++++++++++++++- 9 files changed, 812 insertions(+), 47 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index b62351ea4..c54ecc0e6 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -247,6 +247,37 @@ 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 + +Post-solve access +----------------- + +.. autosummary:: + :toctree: generated/ + + expressions.LazyExpression.solution + + LinearExpression ================ diff --git a/doc/release_notes.rst b/doc/release_notes.rst index c98061cb6..d7bc9a1ed 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -21,6 +21,8 @@ Upcoming Version * ``Model.add_expressions`` registers a ``LinearExpression`` or ``QuadraticExpression`` under a name (auto-generated as ``expr0``, ``expr1``, ... if omitted), accessible afterwards via ``Model.expressions`` (an ``Expressions`` container mirroring ``Model.variables``/``Model.constraints``) and removable via ``Model.remove_expressions``. Named expressions are persisted by ``Model.to_netcdf``/``linopy.read_netcdf`` and preserved by ``Model.copy``, ``copy.copy``, ``copy.deepcopy``, and pickling. +* ``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. *Other* diff --git a/examples/creating-expressions.ipynb b/examples/creating-expressions.ipynb index ce6017ba4..e1acba3ed 100644 --- a/examples/creating-expressions.ipynb +++ b/examples/creating-expressions.ipynb @@ -575,6 +575,115 @@ " 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": [ + "``.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": "62", + "metadata": {}, + "outputs": [], + "source": [ + "deferred.promote()\n", + "m.expressions[\"deferred\"]" + ] } ], "metadata": { @@ -597,7 +706,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 b813f71d5..d7664d6f4 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -27,7 +27,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,6 +61,7 @@ "EvolvingAPIWarning", "GREATER_EQUAL", "LESS_EQUAL", + "LazyExpression", "LinearExpression", "Model", "Objective", diff --git a/linopy/expressions.py b/linopy/expressions.py index 01cff3a2b..1bdc3a1f4 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -103,6 +103,7 @@ ConstantLike, DimsLike, ExpressionLike, + MaskLike, SideLike, SignLike, VariableLike, @@ -1980,6 +1981,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 +2042,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 +2089,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) @@ -2496,6 +2503,8 @@ def __mul__(self, other: SideLike) -> QuadraticExpression: """ Multiply the expr by a factor. """ + if isinstance(other, LazyExpression): + return NotImplemented if isinstance(other, SUPPORTED_EXPRESSION_TYPES): raise TypeError( "unsupported operand type(s) for *: " @@ -2517,6 +2526,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) @@ -2566,6 +2577,8 @@ def __matmul__( """ Matrix multiplication with other, similar to xarray dot. """ + if isinstance(other, LazyExpression): + return NotImplemented if isinstance(other, SUPPORTED_EXPRESSION_TYPES): raise TypeError( "Higher order non-linear expressions are not yet supported." @@ -2853,13 +2866,232 @@ def merge( return cls(ds, model) +@dataclass +class LazyExpression: + """ + 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. + + 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 + """ + + # Guard attributes that make numpy/pandas defer arithmetic to LazyExpression's + # own dunders instead of broadcasting element-wise into an object array; + # mirrors `BaseExpression.__array_ufunc__` / `__array_priority__` above. Plain + # (unannotated) class attributes, so `@dataclass` does not treat them as fields. + __array_ufunc__ = None + __array_priority__ = 10000 + __pandas_priority__ = 10000 + + model: Model + """Reference to the model the expression belongs to""" + evaluator: Callable[..., LinearExpression | QuadraticExpression] + """Callable that builds the expression, invoked as ``evaluator(model, **params)``.""" + 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`.""" + + def evaluate(self) -> LinearExpression | QuadraticExpression: + """ + Evaluate the expression using the provided evaluator and mask. + + Note that nothing is cached, so calling this repeatedly will always re-evaluate from scratch. + """ + 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 + expr = self.evaluate() + expr.attrs.update(self.attrs) + expr.attrs["name"] = self.name + self.model.expressions.data[self.name] = expr + return expr + + @property + def solution(self) -> DataArray: + """ + Get the optimal values of the expression, without promoting it. + """ + return self.evaluate().solution + + @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 + ) -> LazyExpression: + """ + Build a new, unnamed `LazyExpression` that lazily applies `op` to `self` and `other`. + + `other` is evaluated lazily too, if it is itself a `LazyExpression`. + """ + + def evaluator(model: Model) -> LinearExpression | QuadraticExpression: + left = self.evaluate() + right = other.evaluate() if isinstance(other, LazyExpression) else other + return op(right, left) if swapped else op(left, right) + + return LazyExpression(model=self.model, evaluator=evaluator) + + def __add__(self, other: SideLike | LazyExpression) -> LazyExpression: + return self._combine(other, operator.add) + + def __radd__(self, other: SideLike | LazyExpression) -> LazyExpression: + return self._combine(other, operator.add) + + def __sub__(self, other: SideLike | LazyExpression) -> LazyExpression: + return self._combine(other, operator.sub) + + def __rsub__(self, other: SideLike | LazyExpression) -> LazyExpression: + return self._combine(other, operator.sub, swapped=True) + + def __mul__(self, other: SideLike | LazyExpression) -> LazyExpression: + return self._combine(other, operator.mul) + + def __rmul__(self, other: SideLike | LazyExpression) -> LazyExpression: + return self._combine(other, operator.mul) + + def __truediv__(self, other: SideLike | LazyExpression) -> LazyExpression: + if isinstance(other, (LazyExpression, *SUPPORTED_EXPRESSION_TYPES)): + raise TypeError( + f"unsupported operand type(s) for /: {type(self)} and {type(other)}. " + "Expressions cannot be used as a divisor." + ) + return self._combine(other, operator.truediv) + + def __matmul__(self, other: SideLike | LazyExpression) -> LazyExpression: + return self._combine(other, operator.matmul) + + def __rmatmul__(self, other: SideLike | LazyExpression) -> LazyExpression: + return self._combine(other, operator.matmul, swapped=True) + + def __pow__(self, other: int) -> LazyExpression: + if other != 2: + raise ValueError("Power must be 2.") + return self._combine(self, operator.mul) + + def __neg__(self) -> LazyExpression: + return self._combine(-1, operator.mul) + + def __le__(self, other: SideLike) -> Constraint: + return self.evaluate() <= other + + def __ge__(self, other: SideLike) -> Constraint: + return self.evaluate() >= other + + def __eq__(self, other: SideLike) -> Constraint: # type: ignore[override] + return self.evaluate() == other + + def __lt__(self, other: Any) -> NotImplementedType: + raise NotImplementedError( + "Inequalities only ever defined for >= rather than >." + ) + + def __gt__(self, other: Any) -> NotImplementedType: + raise NotImplementedError( + "Inequalities only ever defined for >= rather than >." + ) + + @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 +3105,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 +3179,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 +3194,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 +3218,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: diff --git a/linopy/io.py b/linopy/io.py index aca62e06a..505fc2c5f 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 @@ -27,6 +28,7 @@ from linopy import solvers from linopy.common import to_polars from linopy.constants import CONCAT_DIM, FACTOR_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR +from linopy.expressions import LazyExpression from linopy.objective import Objective if TYPE_CHECKING: @@ -919,7 +921,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. @@ -929,6 +936,18 @@ 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`. + - ``"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``. @@ -937,7 +956,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 @@ -983,13 +1008,38 @@ 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 in m.expressions.items(): + if isinstance(expr, 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 + expr = expr.evaluate() + 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: @@ -1107,9 +1157,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) @@ -1240,9 +1295,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 cbdd4674b..dfe8fc208 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 @@ -58,6 +58,7 @@ from linopy.dualization import dualize from linopy.expressions import ( Expressions, + LazyExpression, LinearExpression, QuadraticExpression, ScalarLinearExpression, @@ -930,47 +931,100 @@ 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], + 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], 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. 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 @@ -979,13 +1033,32 @@ 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") """ - 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/testing.py b/linopy/testing.py index d9c67f7be..ae7847c0c 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -6,7 +6,12 @@ from linopy.constants import TERM_DIM 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,34 @@ 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 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}" + ) + assert_exprequal(a.evaluate(), b.evaluate(), 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/test/test_expressions.py b/test/test_expressions.py index cb49af8e1..52da61adf 100644 --- a/test/test_expressions.py +++ b/test/test_expressions.py @@ -3,14 +3,20 @@ This module aims at testing the correct behavior of the Expressions class. """ +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.expressions import ( + Expressions, + LazyExpression, + LinearExpression, + QuadraticExpression, +) from linopy.solvers import available_solvers -from linopy.testing import assert_linequal +from linopy.testing import assert_linequal, assert_quadequal @pytest.fixture @@ -139,3 +145,207 @@ 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() From a6ec54397171652c3a0892578e89baee1e60ea6d Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Thu, 13 Aug 2026 08:53:33 +0100 Subject: [PATCH 6/7] Allow non-linear expressions and constraint duals in solution expressions Co-authored-by: Claude --- doc/api.rst | 34 ++ doc/release_notes.rst | 3 + examples/creating-expressions.ipynb | 116 ++++- linopy/__init__.py | 4 + linopy/constants.py | 24 ++ linopy/expressions.py | 632 +++++++++++++++++++++------- linopy/io.py | 42 +- linopy/model.py | 31 +- linopy/monkey_patch_xarray.py | 1 + linopy/testing.py | 28 +- linopy/types.py | 2 + linopy/variables.py | 18 +- test/test_expressions.py | 295 ++++++++++++- test/test_io.py | 61 ++- 14 files changed, 1122 insertions(+), 169 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index c54ecc0e6..35b7b5d46 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -268,6 +268,30 @@ Evaluation 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 ----------------- @@ -718,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 d7bc9a1ed..51fd640d5 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -23,6 +23,9 @@ Upcoming Version Named expressions are persisted by ``Model.to_netcdf``/``linopy.read_netcdf`` and preserved by ``Model.copy``, ``copy.copy``, ``copy.deepcopy``, and pickling. * ``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. *Other* diff --git a/examples/creating-expressions.ipynb b/examples/creating-expressions.ipynb index e1acba3ed..8582479f1 100644 --- a/examples/creating-expressions.ipynb +++ b/examples/creating-expressions.ipynb @@ -669,6 +669,25 @@ "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" @@ -677,13 +696,108 @@ { "cell_type": "code", "execution_count": null, - "id": "62", + "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": { diff --git a/linopy/__init__.py b/linopy/__init__.py index d7664d6f4..d526cd909 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 ( @@ -64,6 +66,8 @@ "LazyExpression", "LinearExpression", "Model", + "NonLinearExpressionWarning", + "NonLinearOperationError", "Objective", "OetcHandler", "PiecewiseFormulation", diff --git a/linopy/constants.py b/linopy/constants.py index 7936ef1c2..27357ba48 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 1bdc3a1f4..298141b72 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,15 @@ LESS_EQUAL, STACKED_TERM_DIM, TERM_DIM, + NonLinearExpressionWarning, + NonLinearOperationError, ) from linopy.types import ( CONSTANT_TYPES, ConstantLike, DimsLike, ExpressionLike, + LazySideLike, MaskLike, SideLike, SignLike, @@ -152,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__) @@ -681,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 @@ -830,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: ... @@ -1002,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 @@ -1016,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, @@ -1063,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, @@ -1127,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: """ @@ -1197,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. """ @@ -2467,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} @@ -2503,10 +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." @@ -2569,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 @@ -2577,10 +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." ) @@ -2866,8 +3037,8 @@ def merge( return cls(ds, model) -@dataclass -class LazyExpression: +@dataclass(eq=False, repr=False) +class LazyExpression(AbstractExpression): """ A placeholder for an expression whose value is computed on demand. @@ -2879,6 +3050,11 @@ class LazyExpression: 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 @@ -2890,18 +3066,20 @@ class LazyExpression: >>> lazy.evaluate() # doctest: +SKIP """ - # Guard attributes that make numpy/pandas defer arithmetic to LazyExpression's - # own dunders instead of broadcasting element-wise into an object array; - # mirrors `BaseExpression.__array_ufunc__` / `__array_priority__` above. Plain - # (unannotated) class attributes, so `@dataclass` does not treat them as fields. - __array_ufunc__ = None - __array_priority__ = 10000 - __pandas_priority__ = 10000 + # `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] - """Callable that builds the expression, invoked as ``evaluator(model, **params)``.""" + 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`.""" @@ -2925,12 +3103,28 @@ class LazyExpression: instead of evaluating it, once a frontend that produces such a description exists.""" mask_source: Any = None """As `source`, but describing `mask`.""" - - def evaluate(self) -> LinearExpression | QuadraticExpression: + _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) @@ -2969,18 +3163,58 @@ def promote(self) -> LinearExpression | QuadraticExpression: current, LinearExpression | QuadraticExpression ): return current - expr = self.evaluate() + 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. """ - return self.evaluate().solution + 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]: @@ -3010,80 +3244,173 @@ def __getattr__(self, name: str) -> Any: return getattr(self.evaluate(), name) def _combine( - self, other: Any, op: Callable[[Any, Any], Any], swapped: bool = False + 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 evaluated lazily too, if it is itself a `LazyExpression`. + `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) -> LinearExpression | QuadraticExpression: + def evaluator(model: Model) -> Any: left = self.evaluate() - right = other.evaluate() if isinstance(other, LazyExpression) else other + right = _resolve_lazy(other) return op(right, left) if swapped else op(left, right) - return LazyExpression(model=self.model, evaluator=evaluator) + 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: SideLike | LazyExpression) -> LazyExpression: + def __add__(self, other: LazySideLike) -> LazyExpression: return self._combine(other, operator.add) - def __radd__(self, other: SideLike | LazyExpression) -> LazyExpression: + def __radd__(self, other: LazySideLike) -> LazyExpression: return self._combine(other, operator.add) - def __sub__(self, other: SideLike | LazyExpression) -> LazyExpression: + def __sub__(self, other: LazySideLike) -> LazyExpression: return self._combine(other, operator.sub) - def __rsub__(self, other: SideLike | LazyExpression) -> LazyExpression: + def __rsub__(self, other: LazySideLike) -> LazyExpression: return self._combine(other, operator.sub, swapped=True) - def __mul__(self, other: SideLike | LazyExpression) -> LazyExpression: + def __mul__(self, other: LazySideLike) -> LazyExpression: return self._combine(other, operator.mul) - def __rmul__(self, other: SideLike | LazyExpression) -> LazyExpression: + def __rmul__(self, other: LazySideLike) -> LazyExpression: return self._combine(other, operator.mul) - def __truediv__(self, other: SideLike | LazyExpression) -> LazyExpression: - if isinstance(other, (LazyExpression, *SUPPORTED_EXPRESSION_TYPES)): - raise TypeError( - f"unsupported operand type(s) for /: {type(self)} and {type(other)}. " - "Expressions cannot be used as a divisor." - ) - return self._combine(other, operator.truediv) + 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 __matmul__(self, other: SideLike | LazyExpression) -> LazyExpression: + 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: SideLike | LazyExpression) -> LazyExpression: + def __rmatmul__(self, other: LazySideLike) -> LazyExpression: return self._combine(other, operator.matmul, swapped=True) def __pow__(self, other: int) -> LazyExpression: - if other != 2: - raise ValueError("Power must be 2.") - return self._combine(self, operator.mul) + # 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: - return self._combine(-1, operator.mul) + # `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 __le__(self, other: SideLike) -> Constraint: - return self.evaluate() <= other + 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`. - def __ge__(self, other: SideLike) -> Constraint: - return self.evaluate() >= other + Forces evaluation of both `self` and (if lazy) `rhs`. - def __eq__(self, other: SideLike) -> Constraint: # type: ignore[override] - return self.evaluate() == other + 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 __lt__(self, other: Any) -> NotImplementedType: - raise NotImplementedError( - "Inequalities only ever defined for >= rather than >." - ) + def add( + self, other: LazySideLike, join: JoinOptions | None = None + ) -> LazyExpression: + return self._combine(other, lambda expr, rhs: expr.add(rhs, join=join)) - def __gt__(self, other: Any) -> NotImplementedType: - raise NotImplementedError( - "Inequalities only ever defined for >= rather than >." + 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: @@ -3384,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 505fc2c5f..8eaa00dc6 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -27,8 +27,14 @@ from linopy import solvers from linopy.common import to_polars -from linopy.constants import CONCAT_DIM, FACTOR_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR -from linopy.expressions import LazyExpression +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: @@ -943,7 +949,12 @@ def to_netcdf( - ``"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`. + 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 @@ -1025,15 +1036,32 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ) exprs = [] - for name, expr in m.expressions.items(): - if isinstance(expr, LazyExpression): + 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 - expr = expr.evaluate() + 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), @@ -1265,7 +1293,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 diff --git a/linopy/model.py b/linopy/model.py index dfe8fc208..cf9558f8d 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -945,7 +945,14 @@ def _next_expression_name(self, name: str | None) -> str: @overload def add_expressions( self, - data: Callable[..., LinearExpression | QuadraticExpression], + data: Callable[ + ..., + LinearExpression + | QuadraticExpression + | Variable + | DataArray + | ConstantLike, + ], name: str | None = ..., mask: MaskLike | Callable[..., MaskLike] | None = ..., dims: tuple[Hashable, ...] = ..., @@ -970,7 +977,14 @@ def add_expressions( | LinearExpression | QuadraticExpression | Sequence[tuple[ConstantLike, Variable | str]] - | Callable[..., LinearExpression | QuadraticExpression], + | Callable[ + ..., + LinearExpression + | QuadraticExpression + | Variable + | DataArray + | ConstantLike, + ], name: str | None = None, mask: MaskLike | Callable[..., MaskLike] | None = None, dims: tuple[Hashable, ...] = (), @@ -993,6 +1007,12 @@ def add_expressions( 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. @@ -1037,6 +1057,13 @@ def add_expressions( 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 callable(mask) and not callable(data): raise TypeError( diff --git a/linopy/monkey_patch_xarray.py b/linopy/monkey_patch_xarray.py index 1e526c927..fabe41fb0 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 ae7847c0c..b59e68fec 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -4,7 +4,7 @@ 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 ( LazyExpression, @@ -88,7 +88,10 @@ def assert_exprequal( 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). + 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), ( @@ -101,7 +104,26 @@ def assert_exprequal( assert a.dims == b.dims, ( f"lazy expression dims differ: {a.dims!r} != {b.dims!r}" ) - assert_exprequal(a.evaluate(), b.evaluate(), check_name=False) + 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)}" diff --git a/linopy/types.py b/linopy/types.py index 6b4cf712d..93817a2b9 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 c2e247bb0..eddf5ac9f 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 52da61adf..75ebe0507 100644 --- a/test/test_expressions.py +++ b/test/test_expressions.py @@ -3,12 +3,20 @@ 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, Variable +from linopy.constants import ( + LESS_EQUAL, + NonLinearExpressionWarning, + NonLinearOperationError, +) +from linopy.constraints import Constraint from linopy.expressions import ( Expressions, LazyExpression, @@ -16,7 +24,7 @@ QuadraticExpression, ) from linopy.solvers import available_solvers -from linopy.testing import assert_linequal, assert_quadequal +from linopy.testing import assert_conequal, assert_linequal, assert_quadequal @pytest.fixture @@ -349,3 +357,288 @@ def test_expressions_solution_with_lazy_member(self) -> None: 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 1842dd10b..6bec5bbdf 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: From 2fc3767b1a6fa71609349c260ae7ba2e5b4efa18 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:08:55 +0100 Subject: [PATCH 7/7] Fix merge error --- linopy/io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/linopy/io.py b/linopy/io.py index 3e0de4e7c..4a3d5336c 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -26,7 +26,7 @@ from tqdm import tqdm from linopy import solvers -from linopy.common import to_polars +from linopy.common import sos_weights, to_polars from linopy.constants import ( CONCAT_DIM, FACTOR_DIM,