Compose spec fragments into a whole model via given

doc/api.rst

@@ -122,6 +122,11 @@ Build a model from a `math-spec
 <https://github.com/energy-models/math-spec>`__ YAML program attached to
 data. Requires the ``spec`` dependency group.
 
+Spec fragments that read variables or constraints they do not build (declared
+under ``given:``) compose into one whole model with ``linopy.spec.merge`` (fold
+fragments together) or ``linopy.spec.override`` (lay a patch over a base). linopy
+builds whole models only: a spec that still reads a ``given:`` name is refused.
+
 A spec's grouped sums and windows build dense by default, which is wasteful on
 a skewed topology (a lookup with a few large groups and many small ones, or a
 wide ``sum_back`` window). Set ``linopy.options["sparse_groupby"] = True`` (v1
@@ -143,6 +148,8 @@ instead of densifying them.
    spec.attach
    spec.Attached
    spec.SpecDataError
+   spec.merge
+   spec.override
 
 
 Variable

doc/release_notes.rst

@@ -29,6 +29,8 @@ Upcoming Version
 
 * ``model.spec.expressions`` evaluates the spec's named expressions against the solved model, and ``model.spec.typeset`` (``to_latex`` / ``to_markdown`` / ``to_typst``) renders the spec, warning where the model has drifted from it.
 
+* Spec fragments compose into a whole model. A fragment declares under ``given:`` the variables and constraints it reads but does not build; ``linopy.spec.merge`` folds those reads into the fragment that introduces them and ``linopy.spec.override`` lays a patch over a base, both producing one whole model. linopy builds whole models only, so a spec that still reads a ``given:`` name is refused.
+
 * Spec models round-trip through ``to_netcdf`` / ``read_netcdf``; a file read without ``math-spec`` installed loads as a plain model with a warning.
 
 

examples/building-models-from-specs.ipynb

@@ -1046,6 +1046,109 @@
    "cell_type": "markdown",
    "id": "53",
    "metadata": {},
+   "source": [
+    "## 11. Composing fragments with `given`\n",
+    "\n",
+    "A large spec need not be one file. A **fragment** builds only part of a model and\n",
+    "declares under `given:` the variables and constraints it *reads but does not\n",
+    "build*. Here the dispatch model splits in two: a `generation` fragment owns the\n",
+    "output variable `p`, and a `balance` fragment writes the power-balance\n",
+    "constraint over a `p` it only reads.\n",
+    "\n",
+    "`linopy.spec.merge` folds every `given` name into the fragment that introduces\n",
+    "it, producing one whole model with an empty `given`. linopy builds whole models,\n",
+    "so you compose first, then build the result."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from linopy.spec import merge, override\n",
+    "\n",
+    "GENERATION = \"\"\"\n",
+    "dimensions:\n",
+    "  snapshot: { dtype: int }\n",
+    "  generator: {}\n",
+    "parameters:\n",
+    "  p_max: { dims: [generator] }\n",
+    "  cost:  { dims: [generator] }\n",
+    "variables:\n",
+    "  p:\n",
+    "    dims: [snapshot, generator]\n",
+    "    where: \"p_max > 0\"\n",
+    "    bounds: { lower: 0, upper: p_max }\n",
+    "objective:\n",
+    "  sense: minimize\n",
+    "  expression: sum(p * cost)\n",
+    "\"\"\"\n",
+    "\n",
+    "BALANCE = \"\"\"\n",
+    "dimensions:\n",
+    "  snapshot: { dtype: int }\n",
+    "  generator: {}\n",
+    "given:\n",
+    "  variables:\n",
+    "    p: { dims: [snapshot, generator] }\n",
+    "parameters:\n",
+    "  load: { dims: [snapshot] }\n",
+    "constraints:\n",
+    "  power_balance:\n",
+    "    dims: [snapshot]\n",
+    "    expression: sum(p, over=generator) == load\n",
+    "\"\"\"\n",
+    "\n",
+    "whole = merge({\"generation\": GENERATION, \"balance\": BALANCE})\n",
+    "print(\"`given` left after merge:\", bool(whole.get(\"given\")))\n",
+    "\n",
+    "m3 = Model.from_spec(whole, dispatch_data)\n",
+    "m3.solve(solver_name=\"highs\", output_flag=False)\n",
+    "print(\"variables  :\", set(m3.variables))\n",
+    "print(\"constraints:\", set(m3.constraints))\n",
+    "print(\"objective  :\", m3.objective.value)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "55",
+   "metadata": {},
+   "source": [
+    "A fragment on its own is **not** a whole model, and building it is refused: its\n",
+    "`p` is read but never built. Compose first. `linopy.spec.override` lays a patch\n",
+    "over a base the same way — here it adds a peak-output constraint on top of the\n",
+    "composed model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "56",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# building the lone fragment is refused - it still reads `p` under `given:`\n",
+    "try:\n",
+    "    Model.from_spec(BALANCE, dispatch_data)\n",
+    "except ValueError as err:\n",
+    "    print(\"refused:\", err)\n",
+    "\n",
+    "PEAK = \"\"\"\n",
+    "constraints:\n",
+    "  peak:\n",
+    "    dims: [generator]\n",
+    "    expression: sum(p, over=snapshot) <= p_max\n",
+    "\"\"\"\n",
+    "m4 = Model.from_spec(override(whole, {\"cap\": PEAK}), dispatch_data)\n",
+    "print(\"after override:\", set(m4.constraints))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "57",
+   "metadata": {},
    "source": [
     "## Where the code lives\n",
     "\n",

linopy/spec/__init__.py

@@ -20,6 +20,8 @@
         message = "linopy.spec needs Python >= 3.12. " + message
     raise ImportError(message)
 
+from math_spec import merge, override
+
 from linopy.spec.accessor import (
     Declaration,
     ModelSpec,
@@ -42,4 +44,6 @@
     "SpecLike",
     "Unspecified",
     "attach",
+    "merge",
+    "override",
 ]

linopy/spec/accessor.py

@@ -148,8 +148,9 @@ def build_into(
     Raises
     ------
     ValueError
-        The model already holds variables or constraints, or runs
-        under legacy semantics.
+        The model already holds variables or constraints, runs under
+        legacy semantics, or the spec is a fragment that still reads a
+        name under ``given:``.
     TypeError
         *spec* is a lowered ``Program`` or an open file, neither of which
         has a YAML form to keep on the model.
@@ -168,6 +169,14 @@ def build_into(
             f"{len(model.variables)} variable(s) and {len(model.constraints)} constraint(s)."
         )
     text, program, stem = normalize_spec(spec)
+    given = (*program.given.variables, *program.given.constraints)
+    if given:
+        raise ValueError(
+            f"the spec reads {len(given)} name(s) under `given:` that it does not build "
+            f"({', '.join(given)}), so it is a fragment rather than a whole model. linopy "
+            f"builds whole models: compose the fragments with linopy.spec.merge, or lay a "
+            f"patch over a base with linopy.spec.override, and build the result."
+        )
     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.

pyproject.toml

@@ -123,7 +123,10 @@ gpu = [
 # keeps the git pin out of the published wheel metadata, which PyPI rejects.
 # Install with `uv sync --group spec` or `uv pip install --group spec`.
 spec = [
-    "math-spec @ git+https://github.com/energy-models/math-spec.git@v0.0.0-alpha.106 ; python_version >= '3.12'",
+    # TODO: repin to the released alpha before merge. The feat/given branch head
+    # (a106 + given: merge/override/Program.given) is used for local iteration
+    # only, until an alpha tag cut from it is available.
+    "math-spec @ git+https://github.com/energy-models/math-spec.git@4afb08cd69bee2358600611ec18a70a261a8f063 ; python_version >= '3.12'",
     "pyyaml ; python_version >= '3.12'",
     "pyarrow ; python_version >= '3.12'",
 ]

test/test_spec_composition.py

@@ -0,0 +1,128 @@
+"""
+Composing spec fragments into one whole model.
+
+A fragment declares under ``given:`` the variables it reads but does not build.
+:func:`linopy.spec.merge` folds those reads into the fragment that introduces
+them and :func:`linopy.spec.override` lays a patch over a base, both producing
+one whole model with no ``given:`` left. A program that still carries a
+``given:`` name is a fragment, not a whole model, and building it is refused.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+import xarray as xr
+
+pytest.importorskip("math_spec")
+
+import linopy  # noqa: E402
+from conftest import DISPATCH_DATA, DISPATCH_P, solved  # noqa: E402
+from linopy import Model, read_netcdf  # noqa: E402
+from linopy.spec import merge, override  # noqa: E402
+from linopy.testing import assert_model_equal  # noqa: E402
+
+pytestmark = [
+    pytest.mark.v1,
+    pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"),
+]
+
+GENERATION: dict[str, Any] = {
+    "dimensions": {"snapshot": {"dtype": "int"}, "generator": {}},
+    "parameters": {"p_max": {"dims": ["generator"]}, "cost": {"dims": ["generator"]}},
+    "variables": {
+        "p": {
+            "dims": ["snapshot", "generator"],
+            "where": "p_max > 0",
+            "bounds": {"lower": 0, "upper": "p_max"},
+        }
+    },
+    "objective": {"sense": "minimize", "expression": "sum(p * cost)"},
+}
+BALANCE: dict[str, Any] = {
+    "dimensions": {"snapshot": {"dtype": "int"}, "generator": {}},
+    "given": {"variables": {"p": {"dims": ["snapshot", "generator"]}}},
+    "parameters": {"load": {"dims": ["snapshot"]}},
+    "constraints": {
+        "power_balance": {
+            "dims": ["snapshot"],
+            "expression": "sum(p, over=generator) == load",
+        }
+    },
+}
+FRAGMENTS = {"generation": GENERATION, "balance": BALANCE}
+
+
+def test_merge_reexport_folds_given_into_a_whole_model_dict() -> None:
+    import linopy.spec as spec
+
+    assert {"merge", "override"} <= set(spec.__all__)
+    merged = spec.merge(FRAGMENTS)
+    assert isinstance(merged, dict) and not merged.get("given")
+    assert {"p"} <= set(merged["variables"])
+    assert "power_balance" in merged["constraints"]
+
+
+def test_override_reexport_lays_a_patch_onto_the_dict() -> None:
+    import linopy.spec as spec
+
+    patched = spec.override(
+        spec.merge(FRAGMENTS), {"cap": {"parameters": {"budget": {"dims": []}}}}
+    )
+    assert isinstance(patched, dict) and "budget" in patched["parameters"]
+    assert not patched.get("given")
+
+
+def test_merged_fragments_build_the_whole_model_and_solve() -> None:
+    m = solved(merge(FRAGMENTS), DISPATCH_DATA)
+    assert set(m.variables) == {"p"} and set(m.constraints) == {"power_balance"}
+    assert m.objective.value == pytest.approx(2500.0)
+    xr.testing.assert_allclose(m.solution["p"], DISPATCH_P)
+
+
+def test_a_merged_model_carries_no_given() -> None:
+    m = Model.from_spec(merge(FRAGMENTS), DISPATCH_DATA)
+    assert not m.spec.program.given.variables
+    assert not m.spec.program.given.constraints
+
+
+def test_override_lays_a_patch_over_a_base() -> None:
+    peak = {
+        "constraints": {
+            "peak": {
+                "dims": ["generator"],
+                "expression": "sum(p, over=snapshot) <= p_max",
+            }
+        }
+    }
+    m = Model.from_spec(override(merge(FRAGMENTS), {"cap": peak}), DISPATCH_DATA)
+    assert set(m.constraints) == {"power_balance", "peak"}
+
+
+@pytest.mark.parametrize(
+    "spec",
+    [
+        pytest.param(BALANCE, id="lone-fragment"),
+        pytest.param(merge({"balance": BALANCE}), id="merge-without-introducer"),
+    ],
+)
+def test_a_program_that_still_reads_a_given_name_is_refused(
+    spec: dict[str, Any],
+) -> None:
+    with pytest.raises(ValueError, match=r"under `given:` that it does not build"):
+        Model.from_spec(spec, DISPATCH_DATA)
+
+
+@pytest.mark.parametrize("engine", ["netcdf4", "scipy"])
+def test_a_composed_model_round_trips_through_netcdf(
+    tmp_path: Path, engine: str
+) -> None:
+    m = Model.from_spec(merge(FRAGMENTS), DISPATCH_DATA)
+    path = tmp_path / f"composed-{engine}.nc"
+    m.to_netcdf(path, engine=engine)
+    p = read_netcdf(path)
+    assert_model_equal(m, p)
+    assert p.spec.text == m.spec.text
+    assert not p.spec.program.given.variables