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/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/expressions.py b/linopy/expressions.py index 21a4160e9..01cff3a2b 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,129 @@ 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 + + dataset_attrs = ["coeffs", "vars", "const"] + + 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/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/model.py b/linopy/model.py index 24594c9cc..cbdd4674b 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, @@ -1381,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/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 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_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() 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)