diff --git a/doc/api.rst b/doc/api.rst index ed43b625..0f024637 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -129,6 +129,7 @@ data. Requires the ``spec`` dependency group. model.Model.from_spec model.Model.spec spec.ModelSpec + spec.Layer spec.NamedExpressions spec.NamedExpression spec.Declaration diff --git a/doc/release_notes.rst b/doc/release_notes.rst index f77fe7fb..5764e169 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -25,6 +25,12 @@ Upcoming Version * ``Model.from_spec`` / ``model.add_spec`` build a model from a `math-spec `__ YAML program attached to data, and ``model.spec`` (a ``linopy.spec.ModelSpec``) reads it back. Requires the ``spec`` dependency group (``uv sync --group spec`` / ``uv pip install --group spec``, Python >= 3.12) and v1 semantics. Data is attached onto the spec's dimensions and parameters with ``linopy.spec.attach``, raising a ``linopy.spec.SpecDataError`` on mismatched or missing data; ``linopy.spec.Attached`` carries the attached result. The parameters the spec retains live in ``model.spec.parameters``, its own dataset; ``model.parameters`` stays the caller's and a build never writes to it. ``retain`` decides what a netcdf file holds, not what a session can read: a parameter it dropped is resolved from the sources the model was built with, and only a model read back from a file can run out of data. The spec API emits an :class:`linopy.EvolvingAPIWarning` once per session while it stabilises. See :doc:`building-models-from-specs` for a worked example. +* A spec can extend a model that already holds variables. ``model.add_spec`` builds into this model whether it is empty or not, and a linopy ``Variable`` passed under a declared variable's name in ``sources`` *binds* it: the spec reads the model's variable instead of building one, provided the declaration matches it (same dimension names, default bounds, no ``where``, same domain). Dimensions match by name, so an extending spec uses the model's own axis names; a dimension no source keys takes its labels from an earlier layer or from a bound variable, and every claimant to one dimension name must label it the same way, a bound variable at most spanning a subset of the master. A spec introducing a variable, constraint or named expression the model already holds is refused, as is a spec declaring an objective on a model that already has one, or a special-ordered set on a variable that already carries one. + +* ``model.spec`` holds named *layers*, one per ``add_spec`` (``name=`` names the layer, else the file's stem, else ``"spec"``), each a ``linopy.spec.Layer`` with its own ``program``, ``text``, ``parameters``, ``coords``, ``lookups`` and ``names`` (spec name to model name for every bound variable). ``model.spec[name]`` reads one layer; ``model.spec.expressions``, ``declaration`` and ``evaluate`` dispatch to the layer that declares the name. With a single layer ``model.spec.program``, ``text``, ``parameters``, ``coords`` and ``lookups`` still read through the accessor; with several they raise a ``ValueError`` naming the layers, and ``model.spec[name].parameters`` is the way in. ``model.spec.whole`` says whether the layers describe the whole model or extend a hand-built one, ``model.spec.objective_owner`` which layer's objective the model holds, and typesetting several layers joins their renderings and refuses ``standalone=True``. + +* The netcdf layout follows the layers: each is written under a ``spec-`` prefix with its text and bindings in per-layer attributes, and the layer order, ``whole`` and ``objective_owner`` are attributes of the file. A file written by an earlier linopy under the bare ``spec`` prefix still reads, as one layer named ``"spec"``. + * ``model.spec.expressions`` (a ``linopy.spec.NamedExpressions`` mapping) returns a ``linopy.spec.NamedExpression`` for each declared name, with three views: ``.node`` (the lowered formula), ``.expression`` (the unsolved linopy expression — a ``LinearExpression``, bare ``Variable``, array or scalar) and ``.solution`` (the expression folded over the solved model). ``model.spec.evaluate(name, sources)`` returns the same object with its parameters attached afresh. * ``model.spec.typeset(fmt)`` typesets the spec in any format math-spec knows, with ``.to_latex`` / ``.to_markdown`` / ``.to_typst`` spelling the three it knows today, and ``model.spec.declaration(name)`` returns a ``linopy.spec.Declaration`` whose same three methods typeset one named expression, constraint or variable as a single line (math only, no document); a ``NamedExpression`` carries those methods too. A ``ModelSpec``, a ``Declaration`` and a ``NamedExpression`` all render as Markdown in a notebook. diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb index 35a8da65..9d7717e8 100644 --- a/examples/building-models-from-specs.ipynb +++ b/examples/building-models-from-specs.ipynb @@ -203,8 +203,10 @@ "`dict`, or a `math_spec.Spec`. (A pre-lowered `Program` is refused — it has no\n", "YAML form to keep on the model.)\n", "\n", - "`add_spec` builds into an *empty* model; `from_spec` is sugar that makes the\n", - "model for you and forwards any `Model(...)` keyword arguments." + "`from_spec` is sugar over `add_spec`, which builds into *this* model and adds\n", + "the spec as a named **layer** (`name=`, else the file's stem, else `\"spec\"`).\n", + "Section 11 shows the other use of `add_spec`: extending a model you built by\n", + "hand." ] }, { @@ -1021,26 +1023,194 @@ "cell_type": "markdown", "id": "52", "metadata": {}, + "source": [ + "## 11. Extending a hand-built model\n", + "\n", + "A spec does not have to own the whole model. Take a dispatch model built by\n", + "hand — the same `p`, balance and cost as `DISPATCH`, but written as plain\n", + "linopy calls, the way a large model such as PyPSA builds its core." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "53", + "metadata": {}, + "outputs": [], + "source": [ + "base = Model()\n", + "p = base.add_variables(\n", + " lower=0,\n", + " upper=dispatch_data[\"p_max\"].to_xarray(),\n", + " coords=[snapshot, generator],\n", + " name=\"p\",\n", + ")\n", + "base.add_constraints(\n", + " p.sum(\"generator\") == dispatch_data[\"load\"].to_xarray(), name=\"power_balance\"\n", + ")\n", + "base.add_objective((p * dispatch_data[\"cost\"].to_xarray()).sum())\n", + "base" + ] + }, + { + "cell_type": "markdown", + "id": "54", + "metadata": {}, + "source": [ + "An emissions cap can now be added as a spec **layer**. The spec is complete on\n", + "its own — it declares every dimension, parameter and variable it uses — and\n", + "`add_spec` builds it into the existing model. The one new rule: a linopy\n", + "`Variable` passed under a declared variable's name in `sources` **binds** that\n", + "declaration to the existing variable instead of building a new one. The\n", + "declaration must agree with the model variable (same dimension names, default\n", + "bounds, no `where:`, same domain), and the layer's dimension names are the\n", + "model's own axis names." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "55", + "metadata": {}, + "outputs": [], + "source": [ + "CO2 = \"\"\"\n", + "description: Emission cap on a dispatch fleet.\n", + "\n", + "dimensions:\n", + " snapshot: { dtype: int }\n", + " generator: {}\n", + "\n", + "parameters:\n", + " emission_factor: { dims: [generator], description: t CO2 per unit of output }\n", + " co2_cap: { dims: [], description: total emissions allowed }\n", + "\n", + "variables:\n", + " p:\n", + " foreach: [snapshot, generator]\n", + "\n", + "constraints:\n", + " co2_limit:\n", + " foreach: []\n", + " expression: sum(p * emission_factor) <= co2_cap\n", + "\n", + "expressions:\n", + " emissions: sum(p * emission_factor, over=generator)\n", + "\"\"\"\n", + "\n", + "base.add_spec(\n", + " CO2,\n", + " {\n", + " \"snapshot\": snapshot,\n", + " \"generator\": generator,\n", + " \"emission_factor\": pd.Series([0.0, 0.4], index=generator),\n", + " \"co2_cap\": 30.0,\n", + " \"p\": base.variables[\"p\"], # a Variable binds; everything else is data\n", + " },\n", + " name=\"co2\",\n", + ")\n", + "base" + ] + }, + { + "cell_type": "markdown", + "id": "56", + "metadata": {}, + "source": [ + "The repr now says the model is *extended* by a layer and tags what the layer\n", + "owns. `p` is not built twice: the constraint and the named expression read the\n", + "hand-built variable, and the layer sits under `m.spec[\"co2\"]`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "57", + "metadata": {}, + "outputs": [], + "source": [ + "base.solve(solver_name=\"highs\", output_flag=False)\n", + "\n", + "print(base.spec[\"co2\"].expressions[\"emissions\"].solution.to_pandas())\n", + "print(\"\\nunspecified:\", base.spec.unspecified)" + ] + }, + { + "cell_type": "markdown", + "id": "58", + "metadata": {}, + "source": [ + "`unspecified` lists what no layer declares — here the hand-built balance and\n", + "objective — so the drift report stays honest about which maths the spec covers.\n", + "\n", + "Two things a layer may **not** do. It may not declare an objective on a model\n", + "that already has one; put the extra cost into a named expression and add it by\n", + "hand (`base.objective += base.spec[\"co2\"].expressions[...].expression`). And it\n", + "may not re-declare a name the model already holds without binding it: a second\n", + "`p` without a `Variable` in `sources`, or a constraint named `power_balance`, is\n", + "refused before anything is built." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " base.add_spec(\n", + " CO2,\n", + " {\n", + " \"snapshot\": snapshot,\n", + " \"generator\": generator,\n", + " \"emission_factor\": pd.Series([0.0, 0.4], index=generator),\n", + " \"co2_cap\": 30.0,\n", + " },\n", + " name=\"again\",\n", + " )\n", + "except ValueError as e:\n", + " print(\"ValueError:\", e)" + ] + }, + { + "cell_type": "markdown", + "id": "60", + "metadata": {}, + "source": [ + "Layers persist like whole-model specs: `to_netcdf` writes each layer under its\n", + "own prefix and `read_netcdf` restores them in order, bindings included." + ] + }, + { + "cell_type": "markdown", + "id": "61", + "metadata": {}, "source": [ "## Where the code lives\n", "\n", "The feature is a small package, `linopy/spec/`, imported only when you call\n", "`add_spec`/`from_spec` — `import linopy` never pulls in `math_spec`. Roughly:\n", "\n", - "- `accessor.py` — `model.spec`, the `NamedExpression` views, `evaluate`, and\n", + "- `accessor.py` — `model.spec`, a `ModelSpec` over the named `Layer`s a model\n", + " holds (one per `add_spec`, each with its program, text, data and the\n", + " variables it binds), the `NamedExpression` views, `evaluate`, and\n", " typesetting: the spec (`m.spec.typeset`, with `to_latex` / `to_markdown` /\n", " `to_typst` as its named formats), the drift `m.spec.unspecified` reports,\n", " and any single declaration — a named expression, constraint or\n", " variable — via `m.spec.declaration(name)` and math-spec's\n", " `typeset_declaration`.\n", - "- `attach.py` — the three attachment rules; data onto master coordinates.\n", + "- `attach.py` — the three attachment rules; data onto master coordinates, and\n", + " binding: a linopy `Variable` in `sources` is checked against its declaration\n", + " and read instead of built.\n", "- `builder.py` — emits variables, constraints, objective; folds expressions.\n", "- `operators.py` — `sum`, `by=`, `shift`, `at`, `sum_back`.\n", "- `where.py` — `where:` predicates as boolean masks.\n", "- `coverage.py` / `terms.py` — the absence rule from section 6: a missing row\n", " is refused wherever it is used.\n", "- `curves.py` — the data side of `piecewise:` blocks.\n", - "- `netcdf.py` — the factorize-based persistence from section 10.\n", + "- `netcdf.py` — the factorize-based persistence from section 10, one prefix\n", + " per layer.\n", "- `nodes.py` — walks over expression nodes, and the dimensions a node\n", " spans before any data is bound.\n", "\n", diff --git a/linopy/constraints.py b/linopy/constraints.py index 6cc02cdd..3783620b 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -10,7 +10,15 @@ import warnings import weakref from abc import ABC, abstractmethod -from collections.abc import Callable, Generator, Hashable, ItemsView, Iterator, Sequence +from collections.abc import ( + Callable, + Generator, + Hashable, + ItemsView, + Iterator, + Mapping, + Sequence, +) from dataclasses import dataclass from itertools import product from typing import ( @@ -2119,9 +2127,9 @@ def _formatted_names(self) -> dict[str, str]: return {format_string_as_variable_name(n): n for n in self} def _format_items( - self, exclude: set[str] | None = None, tag: set[str] | None = None + self, exclude: set[str] | None = None, tag: Mapping[str, str] | None = None ) -> str: - """Format constraint items, optionally excluding names in a group.""" + """Format constraint items, optionally excluding names in a group and tagging others.""" r = "" count = 0 for name, ds in self.items(): @@ -2133,7 +2141,7 @@ def _format_items( if ds.coords else "" ) - suffix = " [spec]" if tag and name in tag else "" + suffix = f" [{tag[name]}]" if tag and name in tag else "" r += f" * {name}{coords}{suffix}\n" if count == 0: r += "\n" diff --git a/linopy/io.py b/linopy/io.py index 22781750..337c271c 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -48,6 +48,11 @@ DTYPE_ATTR = "_linopy_dtype" EXPR_TYPE_ATTR = "_linopy_expr_type" SPEC_ATTR = "_linopy_spec" +SPEC_LAYERS_ATTR = "_linopy_spec_layers" +SPEC_WHOLE_ATTR = "_linopy_spec_whole" +SPEC_OBJECTIVE_ATTR = "_linopy_spec_objective" +LAYER_TEXT_ATTR = SPEC_ATTR + "-{}-text" +LAYER_BOUND_ATTR = SPEC_ATTR + "-{}-bound" CONTAINER_ORDER_ATTR = "_linopy_{}_order" @@ -1147,17 +1152,19 @@ def _stamped(data: xr.Dataset, coords: Mapping[str, pd.Index]) -> xr.Dataset: def _restamped(found: pd.Index, master: pd.Index | None) -> pd.Index: """ - *master* where *found* is it as a netcdf type gave it back, else *found* itself. + *master*, or its part, where *found* is it as a netcdf type gave it back, else *found* itself. A narrowed int or a widened bool holds the same labels at another dtype and is the one to replace -- which is what ``Index.equals`` asks, since it - compares labels and not dtypes. An index of another length, or of other - labels entirely, belongs to a container that was never built on *master* - and is left alone. + compares labels and not dtypes. A container spanning some of the master's + labels in its order, a variable bound to a spec over more, takes that part. + An index of other labels belongs to a container that was never built on + *master* and is left alone. """ if master is None or found.dtype == master.dtype: return found - return master if found.equals(master) else found + part = master[master.isin(found)] + return part if part.equals(found) else found def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: @@ -1181,10 +1188,13 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: :func:`linopy.io.read_netcdf`. The insertion order of each container is stored as a JSON list in the ``_linopy__order`` attribute. - A model built with :meth:`Model.add_spec` also persists its spec under a - ``spec-`` prefix of its own: the YAML text, the master coordinates and the - parameters the spec retained, apart from ``m.parameters``. ``read_netcdf`` - lowers the program from the text again, so reading such a file needs + A model built or extended with :meth:`Model.add_spec` also persists each + spec layer under a ``spec--`` prefix of its own: the master + coordinates and the parameters the layer retained, apart from + ``m.parameters``, with its YAML text and its bound names as attributes. + The layer order, whether the layers describe the whole model and the + layer owning the objective are attributes of the file. ``read_netcdf`` + lowers each program from its text again, so reading such a file needs the ``math-spec`` package; a file without a spec does not. The SOS reformulation lifecycle token lives only on the in-memory @@ -1235,7 +1245,7 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: if m._spec is not None: from linopy.spec.netcdf import encode - specs = [encode(m._spec)] + specs = [encode(layer) for layer in m._spec.layers.values()] params = [with_prefix(record_dtypes(m.parameters), "parameters")] scalars = {k: getattr(m, k) for k in m.scalar_attrs} @@ -1250,6 +1260,10 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: ("constraints", m.constraints), ): ds.attrs[CONTAINER_ORDER_ATTR.format(kind)] = json.dumps(list(container)) + if m._spec is not None: + ds.attrs[SPEC_LAYERS_ATTR] = json.dumps(list(m._spec.layers)) + ds.attrs[SPEC_WHOLE_ATTR] = int(m._spec.whole) + ds.attrs[SPEC_OBJECTIVE_ATTR] = json.dumps(m._spec.objective_owner) if m._relaxed_registry: ds.attrs["_relaxed_registry"] = json.dumps(m._relaxed_registry) if m._piecewise_formulations: @@ -1381,10 +1395,10 @@ def container_names(kind: str) -> list[str]: m.parameters = restore_dtypes(get_prefix(ds, "parameters")) - if SPEC_ATTR in ds.attrs: - from linopy.spec.netcdf import decode + if SPEC_LAYERS_ATTR in ds.attrs or SPEC_ATTR in ds.attrs: + from linopy.spec.netcdf import read - m._spec = decode(m, ds, ds.attrs[SPEC_ATTR]) + m._spec = read(m, ds) for k in m.scalar_attrs: if k in ds.attrs: diff --git a/linopy/model.py b/linopy/model.py index e5d117c7..05b76b82 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -443,17 +443,20 @@ def solution(self) -> Dataset: @property def spec(self) -> ModelSpec: """ - The math-spec program this model was built from, see :meth:`add_spec`. + The math-spec layers this model was built from or extended by, see :meth:`add_spec`. + + A :class:`linopy.spec.ModelSpec` over the named layers; + ``model.spec[name]`` is one of them. Raises ------ AttributeError - If the model was not built from a spec. + If no spec was added to the model. """ if self._spec is None: raise AttributeError( - "This model was not built from a spec. Use `Model.add_spec` or " - "`Model.from_spec` to build one." + "This model holds no spec. Use `Model.add_spec` or " + "`Model.from_spec` to add one." ) return self._spec @@ -462,9 +465,10 @@ def add_spec( spec: SpecLike, sources: Mapping[str, Any] | Dataset, retain: Retain = "report", + name: str | None = None, ) -> Model: """ - Build a math-spec program with its data into this empty model. + Build a math-spec program with its data into this model. Requires the ``math-spec`` package and linopy's v1 semantics (``linopy.options["semantics"] = "v1"``). Variables, constraints and @@ -472,6 +476,12 @@ def add_spec( parameters the named expressions read and the lookups are kept on the model, and the named expressions are read back through ``model.spec``. + A spec can extend a model that already holds variables: passing a + model variable under a declared variable's name in ``sources`` binds + it, so the spec reads that variable instead of building one. Its + declaration must then match the model variable in dimensions and + domain and carry no bounds or ``where`` of its own. + Parameters ---------- spec : str, pathlib.Path, dict or math_spec.Spec @@ -479,14 +489,18 @@ def add_spec( ``str`` is a path. A lowered ``math_spec.Program`` is refused, since it has no YAML form to keep on the model. sources : mapping or xarray.Dataset - Data keyed by declared name: dimension labels, parameters and - lookups. Read by key on demand and never iterated. + Data keyed by declared name: dimension labels, parameters, + lookups and the model variables to bind. Read by key on demand + and never iterated. A ``Dataset`` cannot carry a binding. retain : {"report", "all", "none"} Which parameters to keep in ``model.spec.parameters``: those the named expressions read, all of them, or none. ``model.parameters`` stays the caller's and is never written to. This decides what a netcdf file holds, not what this session can read: ``model.spec`` falls back to ``sources`` for a parameter it did not keep. + name : str, optional + The layer's name, ``model.spec[name]``. Defaults to the file's + stem, else ``"spec"``. A name already on the model is refused. Returns ------- @@ -496,10 +510,13 @@ def add_spec( Raises ------ ValueError - If the model already holds variables or constraints, or runs - under legacy semantics. + If the model runs under legacy semantics; if a variable the spec + introduces, a constraint or a named expression collides with a + name the model already holds; or if the spec declares an + objective and the model already has one. linopy.spec.SpecDataError - If the data does not fit the spec. + If the data does not fit the spec, or a binding does not fit its + declaration. Warns ----- @@ -510,7 +527,7 @@ def add_spec( """ from linopy.spec.accessor import attach - self._spec = attach(self, spec, sources, retain) + self._spec = attach(self, spec, sources, retain, name) return self @classmethod @@ -519,6 +536,7 @@ def from_spec( spec: SpecLike, sources: Mapping[str, Any] | Dataset, retain: Retain = "report", + name: str | None = None, **model_kwargs: Any, ) -> Model: """ @@ -526,7 +544,7 @@ def from_spec( ``model_kwargs`` are passed to :class:`Model`. """ - return cls(**model_kwargs).add_spec(spec, sources, retain=retain) + return cls(**model_kwargs).add_spec(spec, sources, retain=retain, name=name) @property def dual(self) -> Dataset: @@ -710,28 +728,36 @@ def __repr__(self) -> str: var_names, con_names = _get_piecewise_groups(self) model_string = f"Linopy {self.type} model" - var_tag: set[str] | None = None - con_tag: set[str] | None = None + var_tag: dict[str, str] | None = None + con_tag: dict[str, str] | None = None expr_string = self.expressions._format_items() + descriptions: list[str] = [] if self._spec is not None: - model_string += ", built from a math-spec" - program = self._spec.program + layers = list(self._spec.layers.values()) + if self._spec.whole: + model_string += ", built from a math-spec" + else: + names = ", ".join(layer.name for layer in layers) + model_string += f", extended by math-spec layer(s) {names}" unspecified = self._spec.unspecified if unspecified.variables: - var_tag = set(program.variables) + var_tag = {n: layer.name for layer in layers for n in layer.variables} if unspecified.constraints: - con_tag = set(program.constraints) + con_tag = { + n: layer.name for layer in layers for n in layer.program.constraints + } eager = expr_string if len(self.expressions) else "" spec = "".join( - f" * {name} ({', '.join(e.dims)}) [spec]\n" - for name, e in self._spec.expressions.items() + f" * {name} ({', '.join(e.dims)}) [{layer.name}]\n" + for layer in layers + for name, e in layer.expressions.items() ) expr_string = eager + spec or "\n" + descriptions = [layer.description for layer in layers if layer.description] var_string = self.variables._format_items(exclude=var_names, tag=var_tag) con_string = self.constraints._format_items(exclude=con_names, tag=con_tag) header = f"{model_string}\n{'=' * len(model_string)}\n" - if self._spec is not None and self._spec.description: - header += f"{self._spec.description}\n" + header += "".join(f"{d}\n" for d in descriptions) return ( f"{header}\n" @@ -1628,9 +1654,7 @@ def add_objective( self.objective.sense = sense self.objective.scaling = scaling if self._spec is not None: - # A spec sets its objective through here during its own build, - # while `_spec` is still unset, so only a later call reaches this. - self._spec._objective_replaced = True + self._spec.objective_owner = None def remove_variables(self, name: str) -> None: """ diff --git a/linopy/spec/__init__.py b/linopy/spec/__init__.py index 0ca35127..dcd8014c 100644 --- a/linopy/spec/__init__.py +++ b/linopy/spec/__init__.py @@ -18,6 +18,7 @@ from linopy.spec.accessor import ( Declaration, + Layer, ModelSpec, NamedExpression, NamedExpressions, @@ -30,6 +31,7 @@ __all__ = [ "Attached", "Declaration", + "Layer", "ModelSpec", "NamedExpression", "NamedExpressions", diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index 265b1cf8..9a9b8e05 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -1,5 +1,5 @@ """ -``model.spec``: the program a model was built from, and its named expressions as data. +``model.spec``: the programs a model was built from or extended by, and their named expressions as data. The spec owns its data. The spec text, the retained parameters, the lookups and the master coordinates sit on the accessor rather than in @@ -8,6 +8,12 @@ its parameters the spec owns. All of it round trips through a file, written under the ``spec-`` prefix. +A model holds an ordered set of spec *layers*. The first may be the whole +model, built into an empty one; any layer may extend a model that already +holds variables, binding the ones it reads through ``sources``. Each +:class:`Layer` is one program with its data; :class:`ModelSpec` is the +model-level view over all of them. + A parameter is resolved the same way however much of it was retained: from the retained dataset, else from the sources the model was built with, which the accessor keeps for as long as the model lives. So ``retain`` decides what @@ -19,8 +25,8 @@ import functools import warnings -from collections.abc import Iterator, Mapping -from dataclasses import dataclass +from collections.abc import Callable, Collection, Iterable, Iterator, Mapping +from dataclasses import dataclass, field, replace from pathlib import Path from typing import Any, TypeAlias @@ -38,14 +44,14 @@ from math_spec import program as ms from math_spec.typesetting import FormatName -from linopy.constants import warn_evolving_api +from linopy.constants import SOS_TYPE_ATTR, warn_evolving_api from linopy.model import Model from linopy.semantics import is_v1 from linopy.spec import terms from linopy.spec.attach import EVOLVING_MESSAGE, Attached, Retain from linopy.spec.attach import attach as attach_data from linopy.spec.builder import build -from linopy.spec.context import Context +from linopy.spec.context import Context, Views from linopy.spec.errors import SpecDataError from linopy.spec.evaluate import evaluate_named, fold from linopy.spec.nodes import dims_of @@ -56,6 +62,7 @@ # A note about what is missing, spelled as a comment of the format's own. A # format math-spec grows later renders without one rather than with a wrong one. _DRIFTED = "This model has drifted from the spec typeset here: {}." +_EXTENDS = "This spec extends a model it does not describe: {}." _COMMENT: dict[str, str] = { "latex": "% {}", @@ -89,7 +96,7 @@ class Unspecified: Formulations added by ``add_piecewise_formulation``, named as formulations rather than as the variables and constraints they hold. objective - Whether ``add_objective`` has replaced the spec's objective. The one + Whether the model's objective is one no spec layer declared. The one entry here that a render gets *wrong* rather than leaves out: the typeset objective is the spec's, and the model's is another. """ @@ -135,15 +142,20 @@ def attach( spec: SpecLike, sources: Mapping[str, Any] | xr.Dataset, retain: Retain, + name: str | None = None, ) -> ModelSpec: """ - Build *spec* with *sources* into the empty *model* and return its accessor. + Build *spec* with *sources* into *model* as a layer and return the accessor. + + The layer is named *name*, else the file's stem, else ``"spec"``. Raises ------ ValueError - The model already holds variables or constraints, or runs - under legacy semantics. + The model runs under legacy semantics; a layer of that name is + already on the model; a variable the spec introduces, a constraint + or a named expression collides with a name the model already holds; + or the spec declares an objective and the model already has one. TypeError *spec* is a lowered ``Program``, which has no YAML form to keep on the model. @@ -154,35 +166,105 @@ def attach( "a spec-built model uses linopy's v1 semantics, and the current setting is " "'legacy'. Set linopy.options['semantics'] = 'v1' before building from a spec." ) - if len(model.variables) or len(model.constraints): + text, program = _source(spec) + attached: Attached = attach_data( + program, sources, retain=retain, given=_given(model) + ) + _check_collisions(model, program, attached) + layer_name = _layer_name(spec, name) + if model._spec is not None and layer_name in model._spec.layers: raise ValueError( - "add_spec builds into an empty model, and this one already holds " - f"{len(model.variables)} variable(s) and {len(model.constraints)} constraint(s)." + f"a spec layer named '{layer_name}' is already on this model; pass another name." ) - text, program = _source(spec) - attached: Attached = attach_data(program, sources, retain=retain) # Resolved before the build, so a parameter no declaration reads cannot fail # halfway through one and leave a model too full to build into again. parameters = attached.retained().assign_coords(dict(attached.coords)) + whole = not len(model.variables) and not len(model.constraints) build(model, attached) - return ModelSpec(model, program, text, parameters, attached) + layer = Layer( + model, layer_name, program, text, parameters, attached, attached.names + ) + spec_ = model._spec if model._spec is not None else ModelSpec(model, [], whole) + spec_.layers[layer.name] = layer + if program.objective is not None: + spec_.objective_owner = layer.name + return spec_ + + +def _layers(model: Model) -> list[Layer]: + """The layers already on *model*, in order.""" + return [] if model._spec is None else list(model._spec.layers.values()) + + +def _given(model: Model) -> dict[str, pd.Index]: + """The master coordinates of every layer already on *model*; they agree wherever they meet.""" + return {d: index for layer in _layers(model) for d, index in layer.coords.items()} + + +def _layer_name(spec: SpecLike, name: str | None) -> str: + if name is not None: + return name + if isinstance(spec, str) and "\n" not in spec: + spec = Path(spec) + return spec.stem if isinstance(spec, Path) else "spec" + + +def _check_collisions(model: Model, program: ms.Program, attached: Attached) -> None: + introduced = [ + n for n in program.variables if n not in attached.bound and n in model.variables + ] + if introduced: + raise ValueError( + f"the spec introduces variable(s) {introduced} and the model already holds " + f"them: bind it or rename it. A binding passes the model variable under the " + f"declared name in sources." + ) + constraints = [n for n in program.constraints if n in model.constraints] + if constraints: + raise ValueError( + f"the spec declares constraint(s) {constraints} and the model already holds them." + ) + on = {attached.names.get(s.variable, s.variable) for s in program.sos.values()} + sos = [ + n + for n in on + if n in model.variables and SOS_TYPE_ATTR in model.variables[n].attrs + ] + if sos: + raise ValueError( + f"the spec declares a special-ordered set on variable(s) {sos} and the model " + f"already holds one on them." + ) + if program.objective is not None and not model.objective.expression.empty: + raise ValueError( + "the spec declares an objective and the model already has one. Add extra cost " + "terms through a named expression: " + "`m.objective += m.spec.expressions[name].expression`." + ) + earlier = {n for layer in _layers(model) for n in layer.program.named_expressions} + expressions = [n for n in program.named_expressions if n in earlier] + if expressions: + raise ValueError( + f"the spec declares named expression(s) {expressions} and an earlier spec on " + f"this model already does." + ) -def restore( +def restore_layer( model: Model, + name: str, text: str, parameters: xr.Dataset, - objective_replaced: bool = False, -) -> ModelSpec: + names: Mapping[str, str], +) -> Layer: """ - The accessor for *model*, with the program lowered afresh from *text*. + One layer of *model*, with the program lowered afresh from *text*. Read from a file, so the sources the model was built with are gone and only what ``retain`` kept can be read back. """ - spec = ModelSpec(model, to_program(yaml.safe_load(text)), text, parameters, None) - spec._objective_replaced = objective_replaced - return spec + program = to_program(yaml.safe_load(text)) + return Layer(model, name, program, text, parameters, None, dict(names)) def _source(spec: SpecLike) -> tuple[str, ms.Program]: @@ -215,39 +297,43 @@ def _row(label: str, items: list[str], cap: int = 8) -> str: return f" {label + ':':<13}{', '.join(shown) if shown else '—'}" -class ModelSpec: +@dataclass(frozen=True, eq=False, repr=False) +class Layer: """ - The spec a model was built from. + One spec on a model: its program, its text and its data. Attributes ---------- + name + What the layer was attached as, ``model.spec[name]``. program The lowered spec. text The spec as YAML, verbatim where a file or text was passed. + parameters + The parameters and lookups the layer retained, on the master coordinates. + names + Spec name to model name for every variable the layer reads instead + of building. """ - def __init__( - self, - model: Model, - program: ms.Program, - text: str, - parameters: xr.Dataset, - attached: Attached | None, - ) -> None: - self._model = model - self.program = program - self.text = text - self._parameters = parameters - self._attached = attached - # A build sets the objective through `add_objective` before `_spec` is - # assigned, so only a call after the build ever flips this. - self._objective_replaced = False + model: Model + name: str + program: ms.Program + text: str + parameters: xr.Dataset + attached: Attached | None + names: Mapping[str, str] + _views: Views = field(default_factory=dict) def __repr__(self) -> str: + return "\n".join(self._rows(f"Layer '{self.name}'")) + + def _rows(self, head: str) -> list[str]: p = self.program coords = self.coords - head = f"ModelSpec: {self.description}" if self.description else "ModelSpec" + if self.description: + head = f"{head}: {self.description}" rows = [ head, _row("Dimensions", [_dimension(d, coords) for d in p.dimensions]), @@ -257,24 +343,12 @@ def __repr__(self) -> str: if p.objective is not None: rows.append(_row("Objective", [p.objective.sense])) rows.append(_row("Expressions", list(p.named_expressions))) - return "\n".join(rows) - - def _reattach(self, model: Model, deep: bool = True) -> ModelSpec: - """The same spec, read off *model*, holding its own copy of the parameters.""" - copied = ModelSpec( - model, - self.program, - self.text, - self._parameters.copy(deep=deep), - self._attached, - ) - copied._objective_replaced = self._objective_replaced - return copied + return rows - @property - def parameters(self) -> xr.Dataset: - """The parameters and lookups the spec retained, on the master coordinates.""" - return self._parameters + def _reattach(self, model: Model, deep: bool = True) -> Layer: + """The same layer, read off *model*, holding its own copy of the parameters.""" + parameters = self.parameters.copy(deep=deep) + return replace(self, model=model, parameters=parameters, _views={}) @property def description(self) -> str: @@ -284,9 +358,14 @@ def description(self) -> str: @property def coords(self) -> dict[str, pd.Index]: - """Master coordinates by dimension, as the model was built on them.""" + """Master coordinates by dimension, as the layer was built on them.""" return {str(d): index for d, index in self.parameters.indexes.items()} + @property + def variables(self) -> set[str]: + """The model variables the layer declares, by model name: built as declared, bound as bound.""" + return {self.names.get(n, n) for n in self.program.variables} + @property def lookups(self) -> dict[str, dict[str, xr.DataArray]]: """By dimension, by name, each lookup as an array over its dimension.""" @@ -298,7 +377,7 @@ def lookups(self) -> dict[str, dict[str, xr.DataArray]]: @property def expressions(self) -> NamedExpressions: """Each named expression as a :class:`NamedExpression`: its math, its linopy fold and its solution.""" - return NamedExpressions(self) + return NamedExpressions({n: self for n in self.program.named_expressions}) def declaration(self, name: str) -> Declaration: """ @@ -320,30 +399,248 @@ def _declarations(self) -> list[str]: p = self.program return [*p.named_expressions, *p.constraints, *p.variables] + def typeset(self, fmt: FormatName, **options: Any) -> str: + """ + This layer's spec typeset in *fmt* as a document, the spec alone. + + Parameters + ---------- + fmt : {"latex", "markdown", "typst"} + What spells the math, as ``math_spec.typeset`` takes it. + **options + Passed on to ``math_spec.typeset``: ``symbols``, ``standalone``, + ``legend``, ``numbered``, ``inline_expressions``. + """ + return typeset(self._schema, fmt, **options) + + def to_latex(self, **options: Any) -> str: + """The layer typeset as a LaTeX document, see :meth:`typeset`.""" + return self.typeset("latex", **options) + + def to_markdown(self, **options: Any) -> str: + """The layer typeset as Markdown, its equations in ``$$`` blocks, see :meth:`typeset`.""" + return self.typeset("markdown", **options) + + def to_typst(self, **options: Any) -> str: + """The layer typeset as Typst, see :meth:`typeset`.""" + return self.typeset("typst", **options) + + @property + def _schema(self) -> dict[str, Any]: + """The spec as the mapping the typesetter reads (a bare string it reads as a path).""" + return yaml.safe_load(self.text) + + def evaluate( + self, name: str, sources: Mapping[str, Any] | xr.Dataset + ) -> NamedExpression: + """ + The named expression *name*, with its parameters attached afresh from *sources*. + + For reading the spec against other data than the model was built with, + and for a model read from a file, whose own sources are gone. + ``expressions`` needs neither. *sources* is read the way ``add_spec`` + read it, and must describe the coordinates the model was built on. + + Raises + ------ + SpecDataError + *sources* label a dimension differently than the + model was built on. + """ + attached = attach_data(self.program, sources, retain="none") + coords = self.coords + for dim, index in attached.coords.items(): + if dim in coords and not index.equals(coords[dim]): + raise SpecDataError( + f"sources describe dimension '{dim}' as {index.tolist()[:5]}, and the model " + f"was built on {coords[dim].tolist()[:5]}. evaluate() reads the solution the " + f"model holds, so the data must be attached on the same labels in the same order." + ) + return NamedExpression(self, name, self._context(attached.parameter)) + + def _resolve(self, name: str) -> xr.DataArray: + """The parameter *name*: retained if it was kept, else read from the sources again.""" + if name in self.parameters: + return self.parameters[name] + if self.attached is not None: + return self.attached.parameter(name) + raise SpecDataError( + f"parameter '{name}' was not retained and this model no longer holds the sources " + f"it was built with, which is what a model read from a file looks like. Build with " + f"retain='all' before writing it out, or read the expression with " + f"evaluate(name, sources)." + ) + + def _context(self, resolve: Resolve) -> Context: + return Context( + self.model, + self.program, + self.coords, + self.lookups, + Parameters(self.program, resolve), + solved=True, + names=self.names, + views=self._views, + ) + + +class ModelSpec: + """ + The spec layers of a model, and the model-level view over them. + + ``model.spec[name]`` is one :class:`Layer`. With a single layer the + layer's ``program``, ``text``, ``parameters``, ``coords``, ``lookups``, + ``name`` and ``names`` read through here as well; with several they are + each layer's own. + + Attributes + ---------- + layers + By name, in the order they were attached. + whole + Whether the first layer was built into an empty model, so the layers + together describe the model rather than extend one. + objective_owner + The layer whose objective the model holds, ``None`` where the model's + objective is none of theirs. + """ + + def __init__( + self, + model: Model, + layers: Iterable[Layer], + whole: bool, + objective_owner: str | None = None, + ) -> None: + self._model = model + self.layers: dict[str, Layer] = {layer.name: layer for layer in layers} + self.whole = whole + self.objective_owner = objective_owner + + def __getitem__(self, name: str) -> Layer: + if name not in self.layers: + raise KeyError( + f"unknown spec layer '{name}'. " + did_you_mean(name, self.layers) + ) + return self.layers[name] + + def __repr__(self) -> str: + if len(self.layers) == 1: + return "\n".join(self._only()._rows("ModelSpec")) + head = f"ModelSpec: layers {', '.join(self.layers)}" + return "\n\n".join([head, *map(repr, self.layers.values())]) + + def _reattach(self, model: Model, deep: bool = True) -> ModelSpec: + """The same layers, read off *model*, each holding its own copy of the parameters.""" + layers = [layer._reattach(model, deep) for layer in self.layers.values()] + return ModelSpec(model, layers, self.whole, self.objective_owner) + + def _only(self) -> Layer: + if len(self.layers) != 1: + raise ValueError( + f"this model holds spec layers {list(self.layers)}; read one through " + f"model.spec[name]." + ) + return next(iter(self.layers.values())) + + def _owner(self, name: str, declared: Callable[[Layer], Collection[str]]) -> Layer: + for layer in self.layers.values(): + if name in declared(layer): + return layer + known = [n for layer in self.layers.values() for n in declared(layer)] + raise KeyError(f"unknown declaration '{name}'. " + did_you_mean(name, known)) + + @property + def name(self) -> str: + """The single layer's name, see :attr:`Layer.name`.""" + return self._only().name + + @property + def names(self) -> Mapping[str, str]: + """The single layer's bound names, see :attr:`Layer.names`.""" + return self._only().names + + @property + def program(self) -> ms.Program: + """The single layer's lowered spec.""" + return self._only().program + + @property + def text(self) -> str: + """The single layer's spec as YAML.""" + return self._only().text + + @property + def parameters(self) -> xr.Dataset: + """The single layer's retained parameters and lookups, on the master coordinates.""" + return self._only().parameters + + @property + def description(self) -> str: + """The single layer's description, see :attr:`Layer.description`.""" + return self._only().description + + @property + def coords(self) -> dict[str, pd.Index]: + """The single layer's master coordinates by dimension.""" + return self._only().coords + + @property + def lookups(self) -> dict[str, dict[str, xr.DataArray]]: + """The single layer's lookups, by dimension, by name.""" + return self._only().lookups + + @property + def expressions(self) -> NamedExpressions: + """Every layer's named expressions as :class:`NamedExpression` objects, by name.""" + owners = { + n: layer + for layer in self.layers.values() + for n in layer.program.named_expressions + } + return NamedExpressions(owners) + + def declaration(self, name: str) -> Declaration: + """One declaration typeset on its own, from whichever layer declares it, see :meth:`Layer.declaration`.""" + return self._owner(name, lambda layer: layer._declarations).declaration(name) + + def evaluate( + self, name: str, sources: Mapping[str, Any] | xr.Dataset + ) -> NamedExpression: + """The named expression *name* on fresh *sources*, from the layer that declares it, see :meth:`Layer.evaluate`.""" + owner = self._owner(name, lambda layer: layer.program.named_expressions) + return owner.evaluate(name, sources) + @property def unspecified(self) -> Unspecified: """ - How the model has drifted from this spec, see :class:`Unspecified`. + How the model has drifted from its spec layers, see :class:`Unspecified`. - Falsy for a model that is only what its spec says; everything added - beside the spec lands here, and is what typesetting cannot show. + Falsy for a model that is only what its layers say; everything added + beside them lands here, and is what typesetting cannot show. """ - from linopy.constants import SOS_TYPE_ATTR from linopy.piecewise import _get_piecewise_groups - model, program = self._model, self.program + model = self._model + layers = list(self.layers.values()) pw_variables, pw_constraints = _get_piecewise_groups(model) - declared_sos = {sos.variable for sos in program.sos.values()} + declared_sos = { + layer.names.get(sos.variable, sos.variable) + for layer in layers + for sos in layer.program.sos.values() + } + variables = {n for layer in layers for n in layer.variables} + constraints = {n for layer in layers for n in layer.program.constraints} return Unspecified( variables=tuple( n for n in model.variables - if n not in program.variables and n not in pw_variables + if n not in variables and n not in pw_variables ), constraints=tuple( n for n in model.constraints - if n not in program.constraints and n not in pw_constraints + if n not in constraints and n not in pw_constraints ), expressions=tuple(model.expressions), sos=tuple( @@ -352,12 +649,13 @@ def unspecified(self) -> Unspecified: if SOS_TYPE_ATTR in model.variables[n].attrs and n not in declared_sos ), piecewise=tuple(model._piecewise_formulations), - objective=self._objective_replaced, + objective=self.objective_owner is None + and not model.objective.expression.empty, ) def typeset(self, fmt: FormatName, **options: Any) -> str: """ - The spec this model was built from, typeset in *fmt* as a document. + The spec layers typeset in *fmt* as a document, one rendering after another. The spec, and so not necessarily the whole model: what was added beside the spec carries no declaration to typeset. Where the model @@ -371,7 +669,8 @@ def typeset(self, fmt: FormatName, **options: Any) -> str: What spells the math, as ``math_spec.typeset`` takes it. **options Passed on to ``math_spec.typeset``: ``symbols``, ``standalone``, - ``legend``, ``numbered``, ``inline_expressions``. + ``legend``, ``numbered``, ``inline_expressions``. Several layers + refuse ``standalone``: one document cannot hold two preambles. Warns ----- @@ -397,20 +696,31 @@ def _render( self, fmt: FormatName, options: Mapping[str, Any], stacklevel: int ) -> str: """Typeset in *fmt*, warned and commented where the model holds more than the spec.""" - rendered = typeset(self._schema, fmt, **options) + if len(self.layers) > 1 and options.get("standalone", False): + raise ValueError( + f"a standalone document holds one spec, and this model holds layers " + f"{list(self.layers)}. Typeset one with model.spec[name].typeset(fmt, " + f"standalone=True)." + ) + rendered = "\n\n".join( + layer.typeset(fmt, **options) for layer in self.layers.values() + ) tally = self._tally() if tally is None: return rendered + note = self._note(tally) warnings.warn( - f"this model has drifted from the spec it was built from: {tally}. " - f"What is typeset is the spec, so it is not this model.", + f"{note} What is typeset is the spec, so it is not this model.", UserWarning, stacklevel=stacklevel, ) comment = _COMMENT.get(fmt) if comment is None: return rendered - return f"{comment.format(_DRIFTED.format(tally))}\n{rendered}" + return f"{comment.format(note)}\n{rendered}" + + def _note(self, tally: str) -> str: + return (_DRIFTED if self.whole else _EXTENDS).format(tally) def _tally(self) -> str | None: """How the model has drifted, counted and named; ``None`` when it has not.""" @@ -424,9 +734,12 @@ def _tally(self) -> str | None: _counted(found.sos, "SOS set"), _counted(found.piecewise, "piecewise formulation"), ] - return _joined( - [p for p in parts if p] + ["a replaced objective"] * found.objective + objective = ( + "a replaced objective" + if self.whole + else "an objective this spec does not declare" ) + return _joined([p for p in parts if p] + [objective] * found.objective) def _repr_markdown_(self) -> str: """The spec as Markdown, with a *visible* note where a notebook would swallow the warning.""" @@ -434,87 +747,29 @@ def _repr_markdown_(self) -> str: tally = self._tally() if tally is None: return rendered - return f"{rendered}\n\n*{_DRIFTED.format(tally)}*" - - @property - def _schema(self) -> dict[str, Any]: - """The spec as the mapping the typesetter reads (a bare string it reads as a path).""" - return yaml.safe_load(self.text) - - def evaluate( - self, name: str, sources: Mapping[str, Any] | xr.Dataset - ) -> NamedExpression: - """ - The named expression *name*, with its parameters attached afresh from *sources*. - - For reading the spec against other data than the model was built with, - and for a model read from a file, whose own sources are gone. - ``model.spec.expressions`` needs neither. *sources* is read the way - ``add_spec`` read it, and must describe the coordinates the model was - built on. - - Raises - ------ - SpecDataError - *sources* label a dimension differently than the - model was built on. - """ - attached = attach_data(self.program, sources, retain="none") - coords = self.coords - for dim, index in attached.coords.items(): - if dim in coords and not index.equals(coords[dim]): - raise SpecDataError( - f"sources describe dimension '{dim}' as {index.tolist()[:5]}, and the model " - f"was built on {coords[dim].tolist()[:5]}. evaluate() reads the solution the " - f"model holds, so the data must be attached on the same labels in the same order." - ) - return NamedExpression(self, name, self._context(attached.parameter)) - - def _resolve(self, name: str) -> xr.DataArray: - """The parameter *name*: retained if it was kept, else read from the sources again.""" - if name in self.parameters: - return self.parameters[name] - if self._attached is not None: - return self._attached.parameter(name) - raise SpecDataError( - f"parameter '{name}' was not retained and this model no longer holds the sources " - f"it was built with, which is what a model read from a file looks like. Build with " - f"retain='all' before writing it out, or read the expression with " - f"evaluate(name, sources)." - ) - - def _context(self, resolve: Resolve) -> Context: - return Context( - self._model, - self.program, - self.coords, - self.lookups, - Parameters(self.program, resolve), - solved=True, - ) + return f"{rendered}\n\n*{self._note(tally)}*" class NamedExpressions(Mapping[str, "NamedExpression"]): - """The named expressions of a spec, each a :class:`NamedExpression` on read.""" + """The named expressions of one or more layers, each a :class:`NamedExpression` on read.""" - def __init__(self, spec: ModelSpec) -> None: - self._spec = spec + def __init__(self, owners: Mapping[str, Layer]) -> None: + self._owners = owners def __getitem__(self, name: str) -> NamedExpression: - if name not in self._spec.program.named_expressions: + if name not in self._owners: raise KeyError( f"unknown named expression '{name}'. " - + did_you_mean(name, self._spec.program.named_expressions) + + did_you_mean(name, self._owners) ) - return NamedExpression( - self._spec, name, self._spec._context(self._spec._resolve) - ) + layer = self._owners[name] + return NamedExpression(layer, name, layer._context(layer._resolve)) def __iter__(self) -> Iterator[str]: - return iter(self._spec.program.named_expressions) + return iter(self._owners) def __len__(self) -> int: - return len(self._spec.program.named_expressions) + return len(self._owners) def __repr__(self) -> str: return f"NamedExpressions({list(self)})" @@ -529,8 +784,8 @@ class Declaration: fold and the solution on top of this. """ - def __init__(self, spec: ModelSpec, name: str) -> None: - self._spec = spec + def __init__(self, layer: Layer, name: str) -> None: + self._layer = layer self._name = name def typeset(self, fmt: FormatName, **options: Any) -> str: @@ -541,7 +796,7 @@ def typeset(self, fmt: FormatName, **options: Any) -> str: :meth:`ModelSpec.typeset` can: a declaration is reached by name through the spec, so there is only ever the spec's own math to show. """ - return typeset_declaration(self._spec._schema, self._name, fmt, **options) + return typeset_declaration(self._layer._schema, self._name, fmt, **options) def to_latex(self, **options: Any) -> str: """This declaration typeset as a single LaTeX line, no document around it.""" @@ -574,19 +829,19 @@ class NamedExpression(Declaration): The lowered expression body, math-spec's own AST handle. """ - def __init__(self, spec: ModelSpec, name: str, ctx: Context) -> None: - super().__init__(spec, name) + def __init__(self, layer: Layer, name: str, ctx: Context) -> None: + super().__init__(layer, name) self._ctx = ctx @property def node(self) -> ms.ExpressionNode: """The expression body as lowered, math-spec's own AST handle.""" - return self._spec.program.named_expressions[self._name].expression + return self._layer.program.named_expressions[self._name].expression @property def dims(self) -> tuple[str, ...]: """The dimensions the expression spans, read off the spec without binding data.""" - return dims_of(self.node, self._spec.program) + return dims_of(self.node, self._layer.program) @functools.cached_property def expression(self) -> terms.Value: diff --git a/linopy/spec/attach.py b/linopy/spec/attach.py index e98ac540..e1b41c02 100644 --- a/linopy/spec/attach.py +++ b/linopy/spec/attach.py @@ -2,9 +2,12 @@ Attach user data to a math-spec program. The language fixes three attachment rules and this module enforces them: a -dimension's members come only from the source keyed by the dimension's -name, their order is the source's order and is never sorted, and a -parameter or lookup source is read for values, never for labels. Parameters +dimension's members come from the source keyed by the dimension's name -- +else from an earlier layer or a bound variable that spans it -- their order +is that source's order and is never sorted, and a parameter or lookup source +is read for values, never for labels. One dimension name is one axis: every +claimant to it must agree on the labels, a bound variable at most leaving +some out. Parameters are resolved from ``sources`` on demand and aligned onto the master coordinates without copying an already aligned array. A coordinate a table leaves out becomes NaN (``False`` for a ``bool`` parameter); what that means @@ -15,7 +18,8 @@ from collections.abc import Hashable, Iterable, Mapping, Sequence from dataclasses import dataclass, field -from typing import Any, Literal, get_args +from types import MappingProxyType +from typing import Any, Literal, NoReturn, get_args import numpy as np import pandas as pd @@ -26,10 +30,17 @@ from linopy.constants import warn_evolving_api from linopy.spec.errors import SpecDataError from linopy.spec.nodes import amounts_of, parameters_of, walk +from linopy.variables import Variable Retain = Literal["report", "all", "none"] _RETAIN: tuple[str, ...] = get_args(Retain) +_DEFAULT_BOUNDS: dict[str, tuple[float, float]] = { + "continuous": (-np.inf, np.inf), + "integer": (-np.inf, np.inf), + "binary": (0.0, 1.0), +} + _ACCEPTED_KINDS: dict[str, frozenset[str]] = { "float": frozenset("fiu"), "int": frozenset("iu"), @@ -73,6 +84,7 @@ def attach( sources: Mapping[str, Any] | xr.Dataset, *, retain: Retain = "report", + given: Mapping[str, pd.Index] = MappingProxyType({}), ) -> Attached: """ Attach *sources* to *program*: master coordinates now, parameters on demand. @@ -83,19 +95,26 @@ def attach( The lowered spec. sources Data keyed by declared name. Any mapping works; it is read by - key and never iterated beyond ``sources.keys()``. An ``xr.Dataset`` + key and never iterated beyond ``sources.keys()``. A linopy + ``Variable`` under a declared variable's name binds that variable: + the spec reads it instead of building one. An ``xr.Dataset`` is accepted too: its indexes are dimension sources, its data - variables parameters and lookups. + variables parameters and lookups; it cannot carry a binding. retain Which parameters :meth:`Attached.retained` persists. + given + Master coordinates the model already holds, from the layers built + before this one. A dimension no source keys takes its labels from + here, else from a bound variable that spans it. Raises ------ SpecDataError A ``retain`` outside its three values, a key naming nothing the spec declares, a reached dimension or a lookup with no - source, a duplicated dimension member, or a lookup breaking the - rules a map has. + source, a duplicated dimension member, a lookup breaking the + rules a map has, a binding whose variable does not fit its + declaration, or two claimants labelling one dimension differently. """ warn_evolving_api("spec", EVOLVING_MESSAGE) if retain not in _RETAIN: @@ -106,9 +125,10 @@ def attach( sources = _dataset_sources(sources) keys = frozenset(sources.keys()) _check_keys(program, keys) - coords = _master_coords(program, sources, keys) + bound = _bindings(program, sources, keys) + coords = _master_coords(program, sources, keys, bound, given) lookups = _lookups(program, sources, keys, coords) - return Attached(program, coords, lookups, retain, sources, keys) + return Attached(program, coords, lookups, retain, sources, bound, keys) @dataclass(frozen=True, eq=False) @@ -123,7 +143,9 @@ class Attached: coords Master coordinates by dimension, in source order, each index named after its dimension. A declared dimension nothing reaches - and nothing supplies is absent. + and nothing supplies is absent. A bound variable spanning fewer + labels than the master is read reindexed onto it, absent where it + has none. lookups By dimension, by lookup name, the map as an array over the dimension's master coordinates, NaN where a label is unmapped. @@ -131,6 +153,9 @@ class Attached: Which parameters :meth:`retained` persists. sources The caller's data, read by key on demand. + bound + By declared variable name, the model variable it is bound to and + reads instead of building. """ program: ms.Program @@ -138,8 +163,14 @@ class Attached: lookups: Mapping[str, Mapping[str, xr.DataArray]] retain: Retain sources: Mapping[str, Any] + bound: Mapping[str, Variable] _keys: frozenset[str] = field(repr=False) + @property + def names(self) -> dict[str, str]: + """Spec name to model name for every bound variable.""" + return {name: variable.name for name, variable in self.bound.items()} + def parameter(self, name: str) -> xr.DataArray: """ The parameter *name* resolved from ``sources`` and aligned to ``coords``. @@ -223,6 +254,7 @@ def _attachable(program: ms.Program) -> dict[str, str]: } kinds.update({d: "dimension" for d in program.dimensions}) kinds.update({lk.name: "lookup" for _, lk in program.lookups}) + kinds.update({v: "variable" for v in program.variables}) return kinds @@ -234,9 +266,66 @@ def _check_keys(program: ms.Program, keys: frozenset[str]) -> None: one = len(unknown) == 1 lead = f"source key {unknown[0]!r} names" if one else f"source keys {unknown} name" raise SpecDataError( - f"{lead} neither a parameter, a dimension nor a lookup this spec declares. " - f"{did_you_mean(unknown[0], known)} Pass only what the spec takes." + f"{lead} neither a parameter, a dimension, a lookup nor a variable this spec " + f"declares. {did_you_mean(unknown[0], known)} Pass only what the spec takes." + ) + + +# --------------------------------------------------------------------------- +# bindings +# --------------------------------------------------------------------------- + + +def _bindings( + program: ms.Program, sources: Mapping[str, Any], keys: frozenset[str] +) -> dict[str, Variable]: + bound: dict[str, Variable] = {} + for name in program.variables: + if name not in keys: + continue + variable = sources[name] + if not isinstance(variable, Variable): + raise SpecDataError( + f"the source for variable '{name}' must be a linopy Variable to bind, or " + f"absent to build; it arrived as {type(variable).__name__}." + ) + _check_binding(name, program.variables[name], variable) + bound[name] = variable + return bound + + +def _check_binding( + name: str, declared: ms.VariableDeclaration, variable: Variable +) -> None: + dims = tuple(str(d) for d in variable.dims) + if declared.dims != dims: + raise SpecDataError( + f"variable '{name}' is declared over {list(declared.dims)} and the bound " + f"variable '{variable.name}' spans {list(dims)}. Dimensions match by name, " + f"so the spec must declare the axes the model variable has." + ) + if not _default_bounds(declared) or declared.where is not None: + raise SpecDataError( + f"variable '{name}' is bound to '{variable.name}', and the base model owns this " + f"variable's bounds and mask; the spec only reads it. Declare '{name}' with " + f"no bounds and no where." + ) + attrs = variable.attrs + kind = ( + "binary" if attrs["binary"] else "integer" if attrs["integer"] else "continuous" ) + if declared.variable_type != kind: + raise SpecDataError( + f"variable '{name}' is declared {declared.variable_type} and the bound variable " + f"'{variable.name}' is {kind}." + ) + + +def _default_bounds(declared: ms.VariableDeclaration) -> bool: + lower, upper = declared.lower, declared.upper + if not isinstance(lower, ms.Constant) or not isinstance(upper, ms.Constant): + return False + return (lower.value, upper.value) == _DEFAULT_BOUNDS[declared.variable_type] # --------------------------------------------------------------------------- @@ -257,22 +346,66 @@ def _reached(program: ms.Program) -> set[str]: def _master_coords( - program: ms.Program, sources: Mapping[str, Any], keys: frozenset[str] + program: ms.Program, + sources: Mapping[str, Any], + keys: frozenset[str], + bound: Mapping[str, Variable], + given: Mapping[str, pd.Index], ) -> dict[str, pd.Index]: reached = _reached(program) coords: dict[str, pd.Index] = {} for dim in program.dimensions: + spanning = {n: v.indexes[dim] for n, v in bound.items() if dim in v.dims} if dim in keys: - coords[dim] = _index(dim, sources[dim]) + master = _index(dim, sources[dim]) + if dim in given and not given[dim].equals(master): + _refuse_other_axis(dim, f"sources['{dim}']", master, given[dim]) + elif dim in given: + master = given[dim] + elif spanning: + master = _agreed(dim, spanning) elif dim in reached: raise SpecDataError( f"dimension '{dim}' has no index: pass its labels under key '{dim}' as " - f"{_DIMENSION_SHAPES}. The index is what says which labels exist, and without " - f"one a mistyped label is indistinguishable from a new one." + f"{_DIMENSION_SHAPES}, or bind a variable that spans it. The index is what " + f"says which labels exist, and without one a mistyped label is " + f"indistinguishable from a new one." ) + else: + continue + for name, found in spanning.items(): + _check_axis(name, dim, found, master) + coords[dim] = master return coords +def _agreed(dim: str, spanning: Mapping[str, pd.Index]) -> pd.Index: + """The one axis every bound variable spanning *dim* labels alike; with no master given, none may leave labels out.""" + master = next(iter(spanning.values())) + for name, found in spanning.items(): + if not found.equals(master): + _refuse_other_axis(dim, f"the bound variable '{name}'", found, master) + return master + + +def _check_axis(name: str, dim: str, found: pd.Index, master: pd.Index) -> None: + """*found* is the master, or the master with labels left out, in master order.""" + if master[master.isin(found)].equals(found): + return + _refuse_strangers(name, dim, found, master, kind="variable") + _refuse_other_axis(dim, f"the bound variable '{name}'", found, master) + + +def _refuse_other_axis( + dim: str, claimant: str, found: pd.Index, master: pd.Index +) -> NoReturn: + raise SpecDataError( + f"dimension '{dim}' is {_shown(master.tolist(), 8)}, and {claimant} labels it " + f"{_shown(found.tolist(), 8)}. The same dimension name means the same axis; " + f"a different axis needs a different name." + ) + + def _index(dim: str, obj: Any) -> pd.Index: if isinstance(obj, (pd.Series, xr.DataArray, np.ndarray)): if obj.ndim != 1: @@ -594,12 +727,14 @@ def _refuse_duplicate_coordinates( ) -def _refuse_strangers(name: str, dim: str, labels: pd.Index, known: pd.Index) -> None: +def _refuse_strangers( + name: str, dim: str, labels: pd.Index, known: pd.Index, kind: str = "parameter" +) -> None: strangers = labels[~labels.isin(known)].unique().tolist() if not strangers: return raise SpecDataError( - f"parameter '{name}' has label(s) in dimension '{dim}' that are not coordinates of it: " + f"{kind} '{name}' has label(s) in dimension '{dim}' that are not coordinates of it: " f"{_shown(strangers)}.\n {dim} has: {_shown(known.tolist(), 10)}\n" f"A label that is not a coordinate is a typo: its row joins nothing, so the coordinate it " f"was meant for is left uncovered. Fix the label, or add it to sources['{dim}']." diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py index 6170c55c..6d4456e1 100644 --- a/linopy/spec/builder.py +++ b/linopy/spec/builder.py @@ -44,6 +44,7 @@ def build(model: Model, attached: Attached) -> None: attached.coords, attached.lookups, Parameters(attached.program, attached.parameter), + names=attached.names, ) curves.validate(ctx.program, ctx.parameters) _variables(ctx) @@ -55,7 +56,10 @@ def build(model: Model, attached: Attached) -> None: def _variables(ctx: Context) -> None: + """Every declared variable the layer does not bind, built as its own.""" for name, declared in ctx.program.variables.items(): + if name in ctx.names: + continue rows = evaluate_where(declared.where, ctx) check_bounds_cover(name, declared, ctx, as_linopy_mask(rows)) ctx.model.add_variables( @@ -79,9 +83,15 @@ def _bound(node: ms.ExpressionNode, ctx: Context) -> float | xr.DataArray: def _sos(ctx: Context) -> None: + """ + Special-ordered sets on the model-owned variable object. + + ``add_sos_constraints`` writes attributes onto the variable it is + handed, so only the object ``model.variables`` holds may go in. + """ for sos in ctx.program.sos.values(): ctx.model.add_sos_constraints( - ctx.model.variables[sos.variable], + ctx.model.variables[ctx.names.get(sos.variable, sos.variable)], sos_type=sos.sos_type, sos_dim=sos.over, big_m=sos.big_m, diff --git a/linopy/spec/context.py b/linopy/spec/context.py index ea9d53c6..a7f5c973 100644 --- a/linopy/spec/context.py +++ b/linopy/spec/context.py @@ -4,12 +4,16 @@ from collections.abc import Mapping from dataclasses import dataclass, field, replace +from types import MappingProxyType import pandas as pd import xarray as xr from math_spec import program as ms from linopy.model import Model +from linopy.variables import Variable + +Views = dict[str, tuple[xr.Dataset, Variable]] @dataclass(frozen=True) @@ -20,7 +24,11 @@ class Context: ``solved`` is the fold's switch: a build leaves it false and a variable enters an expression as its linopy term; a fold sets it true and a variable enters as its solved values, so a named expression reads off the - primal. + primal. ``names`` maps a bound spec variable to the model variable it + reads; a variable the spec introduced is absent and keeps its own name. + ``views`` caches each bound variable reindexed onto the master + coordinates, keyed by spec name and good for as long as the model + variable's data is the one it was made from. """ model: Model @@ -29,12 +37,41 @@ class Context: lookups: Mapping[str, Mapping[str, xr.DataArray]] parameters: Mapping[str, xr.DataArray] solved: bool = field(default=False) + names: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + views: Views = field(default_factory=dict) @property def unsolved(self) -> Context: """The same context with the fold's switch off, so a variable enters as its linopy term.""" return replace(self, solved=False) + def variable(self, name: str) -> Variable: + """ + The model variable the spec variable *name* stands for, on the master coordinates. + + A bound variable spanning fewer labels than the master is a reindexed + view, absent where it has none; it is a copy, so nothing written on it + reaches ``model.variables``. + """ + if name not in self.names: + return self.model.variables[name] + owned = self.model.variables[self.names[name]] + cached = self.views.get(name) + if cached is None or cached[0] is not owned.data: + cached = (owned.data, _onto(owned, self.coords)) + self.views[name] = cached + return cached[1] + def lookup(self, name: str, over: str) -> xr.DataArray: """The lookup *name* as an array over *over*, NaN where a label is unmapped.""" return self.lookups[over][name] + + +def _onto(variable: Variable, coords: Mapping[str, pd.Index]) -> Variable: + """*variable* reindexed onto *coords* along every dimension it does not already span whole.""" + partial = { + str(d): coords[str(d)] + for d in variable.dims + if not variable.indexes[d].equals(coords[str(d)]) + } + return variable.reindex(partial) if partial else variable diff --git a/linopy/spec/coverage.py b/linopy/spec/coverage.py index 11da5e04..5b870cdd 100644 --- a/linopy/spec/coverage.py +++ b/linopy/spec/coverage.py @@ -112,7 +112,7 @@ def _divisor_uses(quotient: ms.Divide, ctx: Context, rows: Rows) -> list[Obligat return [] needed = rows for variable in sorted(ms.variables_of(quotient.numerator)): - present = terms.present(ctx.model.variables[variable]) + present = terms.present(ctx.variable(variable)) needed = present if needed is None else needed & present return [(param, needed) for param in sorted(params)] diff --git a/linopy/spec/evaluate.py b/linopy/spec/evaluate.py index 6571e90d..c8db9cbc 100644 --- a/linopy/spec/evaluate.py +++ b/linopy/spec/evaluate.py @@ -126,7 +126,7 @@ def evaluate(node: ms.ExpressionNode, ctx: Context) -> Value: def _variable(name: str, ctx: Context) -> Value: - variable = ctx.model.variables[name] + variable = ctx.variable(name) absence = ctx.program.variable(name).absence if not ctx.solved: return terms.variable_term(variable, absence) diff --git a/linopy/spec/netcdf.py b/linopy/spec/netcdf.py index a5ed7e06..cde138b0 100644 --- a/linopy/spec/netcdf.py +++ b/linopy/spec/netcdf.py @@ -1,10 +1,12 @@ """ -Persist the spec of a spec-built model in its netcdf file. +Persist the spec layers of a model in its netcdf file. Variables, constraints and the solution round trip through :mod:`linopy.io` -already. Besides them a spec-built model carries the spec text, the master -coordinates and the lookups; the program is re-lowered from the text on read, -so no lowered ``Program`` ever reaches the file. +already. Besides them each spec layer carries its text, the names it binds, +its master coordinates and its lookups; the program is re-lowered from the +text on read, so no lowered ``Program`` ever reaches the file. The layer +order, whether the layers describe the whole model and which layer owns the +objective are attributes of the file itself. No netcdf type holds a dtype as written, so every array carries the dtype it had in memory (:func:`linopy.io.record_dtypes`) and is cast back to it on @@ -23,6 +25,8 @@ from __future__ import annotations +import json +from collections.abc import Mapping from typing import Any import numpy as np @@ -31,16 +35,22 @@ from linopy.io import ( DTYPE_ATTR, + LAYER_BOUND_ATTR, + LAYER_TEXT_ATTR, SPEC_ATTR, + SPEC_LAYERS_ATTR, + SPEC_OBJECTIVE_ATTR, + SPEC_WHOLE_ATTR, get_prefix, restamp_coords, with_prefix, ) from linopy.model import Model -from linopy.spec.accessor import ModelSpec, restore +from linopy.spec.accessor import Layer, ModelSpec, restore_layer PREFIX = "spec" -OBJECTIVE_ATTR = "_linopy_spec_objective_replaced" +LEGACY_NAME = "spec" +LEGACY_OBJECTIVE_ATTR = "_linopy_spec_objective_replaced" COORD = "coords__" PARAM = "param__" CODES = "codes__" @@ -50,80 +60,123 @@ HOLES: dict[str, Any] = {"f": np.nan, "O": np.nan, "M": np.datetime64("NaT")} -def encode(spec: ModelSpec) -> xr.Dataset: +def encode(layer: Layer) -> xr.Dataset: """ - The spec's own dataset: its text, its master coordinates and its parameters. - - The spec text is the dataset's one attribute, which the merge lifts to the - file's. Beside it sits one array of labels per master coordinate and, per - parameter, either its values or -- where it is coded -- its codes and its - categories. The dataset carries no coordinates of its own: an index - coordinate is dropped on read together with the dimension it indexes once - no data variable is left over that dimension, and a master coordinate - nothing else reaches has exactly that shape. So a parameter is written - over bare dimensions and put back on the master coordinates on read. + The layer's own dataset: its master coordinates and its parameters, its text and bindings as attributes. + + Everything is written under the prefix ``spec-``, and the two + attributes carry the layer's name too, so the merge of several layers + lifts every one of them to the file's. Beside them sits one array of + labels per master coordinate and, per parameter, either its values or -- + where it is coded -- its codes and its categories. The dataset carries no + coordinates of its own: an index coordinate is dropped on read together + with the dimension it indexes once no data variable is left over that + dimension, and a master coordinate nothing else reaches has exactly that + shape. So a parameter is written over bare dimensions and put back on the + master coordinates on read. + + Raises + ------ + ValueError + An array name holds a ``-``: the prefix is split off at the last + one on read, so such a name would be silently dropped. """ arrays: dict[str, xr.DataArray] = { COORD + dim: _array(index.to_numpy(), (dim,)) - for dim, index in spec.coords.items() + for dim, index in layer.coords.items() } - coded = _coded(spec) - for name, arr in spec.parameters.items(): + coded = _coded(layer) + for name, arr in layer.parameters.items(): if str(name) in coded: arrays.update(_encode(str(name), arr)) else: arrays[PARAM + str(name)] = _array(arr.to_numpy(), arr.dims, str(arr.dtype)) - written = with_prefix(xr.Dataset(arrays), PREFIX).assign_attrs( - {SPEC_ATTR: spec.text} + dashed = sorted(name for name in arrays if "-" in name) + if dashed: + raise ValueError( + f"spec layer '{layer.name}' would write arrays {dashed}, and a netcdf name " + f"is split from its prefix at the last '-'. A dimension or parameter name " + f"cannot hold one." + ) + written = with_prefix(xr.Dataset(arrays), f"{PREFIX}-{layer.name}") + return written.assign_attrs( + { + LAYER_TEXT_ATTR.format(layer.name): layer.text, + LAYER_BOUND_ATTR.format(layer.name): json.dumps(dict(layer.names)), + } ) - if spec.unspecified.objective: - # Only when true, so a file written from an untouched spec is unchanged. - written = written.assign_attrs({OBJECTIVE_ATTR: 1}) - return written -def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: +def read(model: Model, ds: xr.Dataset) -> ModelSpec: """ - Re-lower *text* onto *model* and read back the dataset :func:`encode` wrote. + The spec layers a file holds, restored onto *model* in their order. - The master coordinates, the plainly written parameters and the coded ones - together are the dataset :func:`linopy.spec.accessor.attach` gave the spec - when the model was built. ``model.parameters`` is not touched: it holds + A file written before layers existed holds one spec under the bare + ``spec`` prefix and its text in one attribute; it reads as a single layer + named ``"spec"`` that describes the whole model. + """ + if SPEC_LAYERS_ATTR in ds.attrs: + layers = [ + decode( + model, + get_prefix(ds, f"{PREFIX}-{name}"), + name, + ds.attrs[LAYER_TEXT_ATTR.format(name)], + json.loads(ds.attrs[LAYER_BOUND_ATTR.format(name)]), + ) + for name in json.loads(ds.attrs[SPEC_LAYERS_ATTR]) + ] + whole = bool(ds.attrs[SPEC_WHOLE_ATTR]) + owner = json.loads(ds.attrs[SPEC_OBJECTIVE_ATTR]) + return ModelSpec(model, layers, whole, owner) + layer = decode(model, get_prefix(ds, PREFIX), LEGACY_NAME, ds.attrs[SPEC_ATTR], {}) + replaced = bool(ds.attrs.get(LEGACY_OBJECTIVE_ATTR, 0)) + owned = layer.program.objective is not None and not replaced + return ModelSpec( + model, [layer], whole=True, objective_owner=LEGACY_NAME if owned else None + ) + + +def decode( + model: Model, sub: xr.Dataset, name: str, text: str, names: Mapping[str, str] +) -> Layer: + """ + Re-lower *text* onto *model* as the layer *name* and read back the dataset :func:`encode` wrote. + + *sub* is the layer's part of the file with its prefix given back. The + master coordinates, the plainly written parameters and the coded ones + together are the dataset :func:`linopy.spec.accessor.attach` gave the + layer when it was built. ``model.parameters`` is not touched: it holds what the caller put there and nothing of the spec. """ - sub = get_prefix(ds, PREFIX) coords = { - _stripped(name, COORD): _index(sub[name]) - for name in sub.data_vars - if str(name).startswith(COORD) + _stripped(var, COORD): _index(sub[var]) + for var in sub.data_vars + if str(var).startswith(COORD) } arrays = { - _stripped(name, PARAM): _plain(sub[name], _stripped(name, PARAM), coords) - for name in sub.data_vars - if str(name).startswith(PARAM) + _stripped(var, PARAM): _plain(sub[var], _stripped(var, PARAM), coords) + for var in sub.data_vars + if str(var).startswith(PARAM) } arrays.update( { - _stripped(name, CODES): _decode(sub, _stripped(name, CODES), coords) - for name in sub.data_vars - if str(name).startswith(CODES) + _stripped(var, CODES): _decode(sub, _stripped(var, CODES), coords) + for var in sub.data_vars + if str(var).startswith(CODES) } ) restamp_coords(model, coords) - return restore( - model, - text, - xr.Dataset(arrays).assign_coords(coords), - bool(ds.attrs.get(OBJECTIVE_ATTR, 0)), - ) + parameters = xr.Dataset(arrays).assign_coords(coords) + return restore_layer(model, name, text, parameters, names) -def _coded(spec: ModelSpec) -> set[str]: +def _coded(layer: Layer) -> set[str]: """The parameters written as codes: every lookup and every array of objects.""" - lookups = {name for by_name in spec.lookups.values() for name in by_name} + lookups = {name for by_name in layer.lookups.values() for name in by_name} return { str(name) - for name, arr in spec.parameters.items() + for name, arr in layer.parameters.items() if name in lookups or arr.dtype == object } diff --git a/linopy/spec/where.py b/linopy/spec/where.py index fc301d0c..3c1d10ed 100644 --- a/linopy/spec/where.py +++ b/linopy/spec/where.py @@ -55,7 +55,7 @@ def _node(node: ms.WhereNode, ctx: Context) -> xr.DataArray: ctx.parameters[node.name], ctx.program.parameter(node.name).dtype ) if isinstance(node, ms.VariableDefinedNode): - return terms.present(ctx.model.variables[node.name]) + return terms.present(ctx.variable(node.name)) if isinstance(node, ms.ParameterComparisonNode): arr = ctx.parameters[node.name] result = _PREDICATE_OPS[node.op](arr, _as_the_axis_spells_it(arr, node.value)) diff --git a/linopy/testing.py b/linopy/testing.py index bcd471e5..ec35e326 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -153,9 +153,15 @@ def assert_model_equal(a: Model, b: Model) -> None: assert (a._spec is None) == (b._spec is None) if a._spec is not None and b._spec is not None: - assert a._spec.text == b._spec.text + assert list(a._spec.layers) == list(b._spec.layers) + for name, layer in a._spec.layers.items(): + other = b._spec.layers[name] + assert layer.text == other.text + assert dict(layer.names) == dict(other.names) + assert_datasetequal(layer.parameters, other.parameters) + assert a._spec.whole == b._spec.whole + assert a._spec.objective_owner == b._spec.objective_owner assert a._spec.unspecified == b._spec.unspecified - assert_datasetequal(a._spec.parameters, b._spec.parameters) assert a.status == b.status assert a.termination_condition == b.termination_condition diff --git a/linopy/variables.py b/linopy/variables.py index ea4ada53..935ce9c3 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -1810,9 +1810,9 @@ def __dir__(self) -> list[str]: return base_attributes + formatted_names def _format_items( - self, exclude: set[str] | None = None, tag: set[str] | None = None + self, exclude: set[str] | None = None, tag: Mapping[str, str] | None = None ) -> str: - """Format variable items, optionally excluding names in a group.""" + """Format variable items, optionally excluding names in a group and tagging others.""" r = "" count = 0 for name, ds in self.items(): @@ -1830,7 +1830,7 @@ def _format_items( coords += f" - sos{sos_type} on {sos_dim}" if ds.attrs.get("semi_continuous", False): coords += " - semi-continuous" - suffix = " [spec]" if tag and name in tag else "" + suffix = f" [{tag[name]}]" if tag and name in tag else "" r += f" * {name}{coords}{suffix}\n" if count == 0: r += "\n" diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index 5bca044b..b125da39 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -18,23 +18,38 @@ math_spec = pytest.importorskip("math_spec") yaml = pytest.importorskip("yaml") +from math_spec.typesetting import FormatName # noqa: E402 +from test_spec_builder import ( # noqa: E402 + BASE_MODEL, + EXTRA_DATA, + EXTRA_SPEC, + SECOND_SPEC, + SOS_SPEC, + THREE, + dispatch_p, + extended, + subset_bound, +) + import linopy # noqa: E402 from conftest import ( # noqa: E402 DISPATCH_DATA, DISPATCH_P, EXAMPLE_DISPATCH, GENERATOR, + SNAPSHOT, solved, with_, yaml_dict, ) -from linopy import Model, breakpoints # noqa: E402 +from linopy import LinearExpression, Model, breakpoints # noqa: E402 from linopy.spec import ( # noqa: E402 ModelSpec, NamedExpression, SpecDataError, Unspecified, ) +from linopy.testing import assert_linequal # noqa: E402 pytestmark = [ pytest.mark.v1, @@ -75,13 +90,233 @@ def test_a_lowered_program_is_refused() -> None: Model().add_spec(program, DISPATCH_DATA) -def test_add_spec_needs_an_empty_model() -> None: +def test_a_second_spec_must_bind_or_not_collide() -> None: m = Model() m.add_variables(name="x") - with pytest.raises(ValueError, match="empty model"): + m.add_spec(yaml_dict(), DISPATCH_DATA) + assert list(m.variables) == ["x", "p"] + with pytest.raises(ValueError, match="bind it or rename it"): m.add_spec(yaml_dict(), DISPATCH_DATA) +def test_a_bound_variable_is_read_not_built() -> None: + m = extended() + assert list(m.variables) == ["p"] + assert m.spec.names == {"p": "p"} + total = m.spec.expressions["total"] + assert isinstance(total.expression, LinearExpression) + assert_linequal(total.expression, m.variables["p"].sum()) + m.solve(solver_name="highs", output_flag=False) + assert float(total.solution) == pytest.approx(float(DISPATCH_P.sum())) + + +def test_a_binding_must_be_a_variable() -> None: + with pytest.raises(SpecDataError, match="must be a linopy Variable to bind"): + extended(p=3.0) + + +P_OVER_GENERATOR = with_( + EXTRA_SPEC, + variables={"p": {"foreach": ["generator"]}}, + constraints={"p_cap": {"foreach": ["generator"], "expression": "p <= cap"}}, +) + + +def p_declared(**more: Any) -> dict[str, Any]: + return with_(EXTRA_SPEC, variables={"p": {**EXTRA_SPEC["variables"]["p"], **more}}) + + +@pytest.mark.parametrize( + ("spec", "match"), + [ + pytest.param(P_OVER_GENERATOR, "Dimensions match by name", id="dims"), + pytest.param( + p_declared(bounds={"lower": 0}), + "owns this variable's bounds and mask", + id="bounds", + ), + pytest.param( + p_declared(where="cap > 0"), + "owns this variable's bounds and mask", + id="where", + ), + pytest.param( + p_declared(domain="binary"), + "declared binary and the bound variable 'p' is continuous", + id="domain", + ), + ], +) +def test_a_bound_variable_keeps_its_declared_shape( + spec: dict[str, Any], match: str +) -> None: + with pytest.raises(SpecDataError, match=match): + extended(spec) + + +def test_a_bound_subset_is_reindexed_onto_the_master() -> None: + """The spec spans three generators, the bound ``p`` two: its third column is absent, not a stranger.""" + m = subset_bound() + assert m.spec.coords["generator"].equals(THREE) + assert m.variables["p"].indexes["generator"].equals(GENERATOR) + labels = m.constraints["p_cap"].labels + assert labels.shape == (3, 3) + assert (labels.sel(generator="solar") == -1).all() + assert (labels.sel(generator=GENERATOR) != -1).all() + total = m.spec.expressions["total"] + assert isinstance(total.expression, LinearExpression) + read = total.expression.vars.values + assert set(read[read != -1]) == set(m.variables["p"].labels.values.ravel()) + m.solve(solver_name="highs", output_flag=False) + assert float(total.solution) == pytest.approx(float(DISPATCH_P.sum())) + + wider = Model() + wider.add_variables(coords=[SNAPSHOT, THREE], name="p") + with pytest.raises(SpecDataError, match="variable 'p' has label.*'solar'"): + wider.add_spec(EXTRA_SPEC, {**EXTRA_DATA, "p": wider.variables["p"]}) + + +def test_a_bound_variable_can_supply_a_dimension() -> None: + m = BASE_MODEL() + data = {"cap": EXTRA_DATA["cap"], "p": m.variables["p"]} + m.add_spec(EXTRA_SPEC, data) + assert m.spec.coords["generator"].equals(GENERATOR) + assert m.spec.coords["snapshot"].equals(SNAPSHOT) + with pytest.raises(SpecDataError, match="or bind a variable that spans it"): + Model().add_spec(EXTRA_SPEC, {"cap": EXTRA_DATA["cap"]}) + + +def two_bound_variables(q_generator: pd.Index, first: str) -> None: + """``p`` and ``q`` bound with no generator source, *first* declared before the other.""" + m = BASE_MODEL() + m.add_variables(coords=[SNAPSHOT, q_generator], name="q") + declared = {**EXTRA_SPEC["variables"], "q": {"foreach": ["snapshot", "generator"]}} + ordered = {first: declared[first], **declared} + spec = {**EXTRA_SPEC, "variables": ordered} + data = {"cap": EXTRA_DATA["cap"], "p": m.variables["p"], "q": m.variables["q"]} + m.add_spec(spec, data) + + +def second_layer_disagrees() -> None: + m = extended() + again = { + k: v for k, v in EXTRA_SPEC.items() if k not in ("constraints", "expressions") + } + data = {**EXTRA_DATA, "generator": GENERATOR[::-1], "p": m.variables["p"]} + m.add_spec(again, data, name="again") + + +@pytest.mark.parametrize( + "build", + [ + lambda: extended(generator=GENERATOR[::-1]), + lambda: two_bound_variables(GENERATOR[::-1], "p"), + lambda: two_bound_variables(THREE, "p"), + lambda: two_bound_variables(THREE, "q"), + second_layer_disagrees, + ], + ids=[ + "sources-vs-bound", + "bound-vs-bound", + "narrower-bound-first", + "wider-bound-first", + "layer-vs-layer", + ], +) +def test_dimension_labels_must_agree(build: Callable[[], None]) -> None: + with pytest.raises(SpecDataError, match="same dimension name means the same axis"): + build() + + +def test_a_later_layer_inherits_the_dimensions_it_does_not_key() -> None: + """No generator source and a bound ``p`` over two: the master is the first layer's three.""" + m = subset_bound() + floor = pd.Series([0.0, 0.0, 0.0], index=THREE) + m.add_spec(SECOND_SPEC, {"floor": floor, "p": m.variables["p"]}, name="second") + assert m.spec["second"].coords["generator"].equals(THREE) + assert m.spec["second"].coords["snapshot"].equals(SNAPSHOT) + labels = m.constraints["p_floor"].labels + assert labels.indexes["generator"].equals(THREE) + assert labels.indexes["snapshot"].equals(SNAPSHOT) + assert (labels.sel(generator="solar") == -1).all() + assert (labels.sel(generator=GENERATOR) != -1).all() + + +def with_sos(m: Model) -> None: + m.add_sos_constraints(m.variables["p"], sos_type=1, sos_dim="generator") + + +@pytest.mark.parametrize( + ("spec", "sources", "prepare", "match"), + [ + pytest.param( + EXTRA_SPEC, {"p": None}, None, "bind it or rename it", id="variable" + ), + pytest.param( + with_( + EXTRA_SPEC, + constraints={"power_balance": EXTRA_SPEC["constraints"]["p_cap"]}, + ), + {}, + None, + r"constraint\(s\) \['power_balance'\]", + id="constraint", + ), + pytest.param( + SOS_SPEC, + {}, + with_sos, + r"special-ordered set on variable\(s\) \['p'\]", + id="sos", + ), + ], +) +def test_collisions_are_refused( + spec: dict[str, Any], + sources: dict[str, Any], + prepare: Callable[[Model], None] | None, + match: str, +) -> None: + m = BASE_MODEL() + if prepare is not None: + prepare(m) + data = {**EXTRA_DATA, "p": m.variables["p"], **sources} + data = {k: v for k, v in data.items() if v is not None} + with pytest.raises(ValueError, match=match): + m.add_spec(spec, data) + + +def test_an_expression_name_is_taken_once_across_specs() -> None: + m = extended() + again = {k: v for k, v in EXTRA_SPEC.items() if k != "constraints"} + with pytest.raises(ValueError, match=r"named expression\(s\) \['total'\]"): + m.add_spec(again, {**EXTRA_DATA, "p": m.variables["p"]}) + + +def test_an_objective_on_a_non_empty_model_is_refused() -> None: + spec = with_(EXTRA_SPEC, objective={"sense": "minimize", "expression": "sum(p)"}) + with pytest.raises(ValueError, match="already has one"): + extended(spec) + + m = Model() + dispatch_p(m) + m.add_spec(spec, {**EXTRA_DATA, "p": m.variables["p"]}, name="extra") + assert m.objective.sense == "min" + assert m.spec.name == "extra" + assert m.spec.objective_owner == "extra" + assert m.spec.unspecified.objective is False + + +def test_a_binding_needs_a_mapping_source() -> None: + ds = xr.Dataset( + {"cap": EXTRA_DATA["cap"].to_xarray()}, coords={"snapshot": SNAPSHOT} + ) + with pytest.raises(ValueError, match="bind it or rename it"): + BASE_MODEL().add_spec(EXTRA_SPEC, ds) + m = Model().add_spec(EXTRA_SPEC, ds) + assert "p" in m.variables and m.spec.names == {} + + def test_legacy_semantics_is_refused() -> None: with linopy.options as options: options["semantics"] = "legacy" @@ -90,7 +325,7 @@ def test_legacy_semantics_is_refused() -> None: def test_a_model_without_a_spec_has_no_accessor() -> None: - with pytest.raises(AttributeError, match="not built from a spec"): + with pytest.raises(AttributeError, match="holds no spec"): _ = Model().spec @@ -253,6 +488,15 @@ def test_repr_summarises_every_section() -> None: assert "Expressions: spend, usage" in text +def test_repr_of_several_layers_names_each() -> None: + spec = two_layers().spec + text = repr(spec) + assert text.startswith("ModelSpec: layers spec, extra") + assert "Layer 'spec': Least-cost dispatch" in text + assert "Layer 'extra'\n" in text + assert repr(spec["extra"]).startswith("Layer 'extra'\n Dimensions:") + + def test_repr_caps_long_sections() -> None: spec = with_(yaml_dict(), expressions={f"e{i}": "p / p_max" for i in range(12)}) text = repr(Model.from_spec(spec, DISPATCH_DATA).spec) @@ -271,6 +515,18 @@ def test_model_repr_shows_the_spec_and_tags_only_expressions() -> None: assert "" not in text +def test_model_repr_of_an_extended_model_names_its_layers() -> None: + m = extended() + m.add_variables(lower=0, coords=[GENERATOR], name="reserve") + text = repr(m) + assert "Linopy LP model, extended by math-spec layer(s) extra" in text + assert " * p (snapshot, generator) [extra]" in text + assert " * reserve (generator)\n" in text + assert " * p_cap (snapshot, generator) [extra]" in text + assert " * power_balance (snapshot)\n" in text + assert " * total () [extra]" in text + + def test_model_repr_of_a_spec_without_a_description() -> None: spec = {k: v for k, v in yaml_dict().items() if k != "description"} m = Model.from_spec(spec, DISPATCH_DATA) @@ -293,6 +549,12 @@ def test_hybrid_model_tags_spec_variables_constraints_and_expressions() -> None: assert "" not in text +def two_layers() -> Model: + """The dispatch example built from its spec, then extended by a second layer.""" + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + return m.add_spec(EXTRA_SPEC, {**EXTRA_DATA, "p": m.variables["p"]}, name="extra") + + def test_the_spec_typesets_in_every_format() -> None: spec = Model.from_spec(yaml_dict(), DISPATCH_DATA).spec assert "align" in spec.to_latex() @@ -302,7 +564,52 @@ def test_the_spec_typesets_in_every_format() -> None: @pytest.mark.parametrize("fmt", ["latex", "markdown", "typst"]) -def test_typeset_and_its_named_aliases_agree(fmt: str) -> None: +def test_two_layers_typeset_one_after_the_other(fmt: FormatName) -> None: + spec = two_layers().spec + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + rendered = spec.typeset(fmt) + assert rendered == f"{spec['spec'].typeset(fmt)}\n\n{spec['extra'].typeset(fmt)}" + with pytest.raises(ValueError, match=r"model\.spec\[name\]\.typeset"): + spec.typeset(fmt, standalone=True) + assert spec["extra"].typeset(fmt, standalone=True) + + +def test_model_spec_layers() -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + assert list(m.spec.layers) == ["spec"] + assert m.spec.program is m.spec["spec"].program + assert m.spec.whole and m.spec.objective_owner == "spec" + + m = two_layers() + assert list(m.spec.layers) == ["spec", "extra"] + assert m.spec["extra"].program.constraints.keys() == {"p_cap"} + assert m.spec["extra"].names == {"p": "p"} + with pytest.raises(ValueError, match=r"\['spec', 'extra'\]"): + m.spec.program + assert set(m.spec.expressions) == {"spend", "usage", "total"} + assert m.spec.declaration("p_cap").to_latex() + with pytest.raises(KeyError, match="unknown spec layer 'extr'.*extra"): + m.spec["extr"] + assert m.spec.whole and not extended().spec.whole + + +def test_layer_names(tmp_path: Path) -> None: + path = tmp_path / "dispatch.yaml" + path.write_text(EXAMPLE_DISPATCH) + assert list(Model.from_spec(path, DISPATCH_DATA).spec.layers) == ["dispatch"] + assert list(Model.from_spec(yaml_dict(), DISPATCH_DATA).spec.layers) == ["spec"] + m = extended() + assert list(m.spec.layers) == ["extra"] + again = { + k: v for k, v in EXTRA_SPEC.items() if k not in ("constraints", "expressions") + } + with pytest.raises(ValueError, match="layer named 'extra' is already"): + m.add_spec(again, {**EXTRA_DATA, "p": m.variables["p"]}, name="extra") + + +@pytest.mark.parametrize("fmt", ["latex", "markdown", "typst"]) +def test_typeset_and_its_named_aliases_agree(fmt: FormatName) -> None: """The format is a parameter; the named methods only spell a common one.""" spec = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec declaration = spec.declaration("p") @@ -328,6 +635,26 @@ def test_unspecified_names_what_the_spec_does_not_declare() -> None: piecewise=(), objective=False, ) + found = extended().spec.unspecified + assert found.variables == () + assert found.constraints == ("power_balance",) + + +def test_a_bound_spec_name_does_not_hide_a_hand_variable_of_that_name() -> None: + """A layer declares the model variable it binds, not the spec name it binds under.""" + over = ["snapshot", "generator"] + spec = { + **EXTRA_SPEC, + "variables": {"q": {"foreach": over}}, + "constraints": {"q_cap": {"foreach": over, "expression": "q <= cap"}}, + "expressions": {"total": "sum(q)"}, + } + m = BASE_MODEL() + m.add_variables(lower=0, coords=[GENERATOR], name="q") + m.add_spec(spec, {**EXTRA_DATA, "q": m.variables["p"]}, name="extra") + + assert m.spec["extra"].variables == {"p"} + assert m.spec.unspecified.variables == ("q",) def test_unspecified_sees_what_carries_no_name_of_its_own() -> None: @@ -363,20 +690,41 @@ def test_a_piecewise_formulation_is_named_as_one_and_not_as_its_parts() -> None: assert found.constraints == () +@pytest.mark.parametrize( + ("build", "match", "tallied"), + [ + pytest.param( + hybrid, + "drifted from the spec", + ["1 variable (reserve)", "1 constraint (reserve_cap)"], + id="whole", + ), + pytest.param( + extended, + "extends a model it does not describe", + ["1 constraint (power_balance)"], + id="extended", + ), + ], +) @pytest.mark.parametrize( ("fmt", "opener"), [("latex", "%"), ("markdown", "