Extend a built model with spec layers that bind existing variables

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

doc/release_notes.rst

@@ -25,6 +25,12 @@ Upcoming Version
 
 * ``Model.from_spec`` / ``model.add_spec`` build a model from a `math-spec <https://github.com/energy-models/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-<name>`` 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.

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",

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

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_<kind>_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-<name>-`` 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: