Expressions container

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

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 <https://github.com/PyPSA/linopy/pull/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 <https://github.com/PyPSA/linopy/pull/566>`__)

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": {

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 += "<empty>\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.