diff --git a/.github/workflows/test-notebooks.yml b/.github/workflows/test-notebooks.yml index 4050badb2..cfed5914b 100644 --- a/.github/workflows/test-notebooks.yml +++ b/.github/workflows/test-notebooks.yml @@ -30,7 +30,7 @@ jobs: - name: Install package and dependencies run: | python -m pip install uv - uv pip install --system -e ".[docs]" + uv pip install --system -e ".[docs]" --group spec - name: Execute notebooks run: | diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 89c303af4..234c0db11 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -82,7 +82,7 @@ jobs: - name: Install package and dependencies run: | python -m pip install uv - uv pip install --system "$(ls dist/*.whl)[dev,solvers,oetc]" + uv pip install --system "$(ls dist/*.whl)[dev,solvers,oetc]" --group spec - name: Test with pytest env: @@ -120,7 +120,7 @@ jobs: - name: Install package and dependencies run: | python -m pip install uv - uv pip install --system "$(ls dist/*.whl)[dev]" + uv pip install --system "$(ls dist/*.whl)[dev]" --group spec - name: Run type checker (mypy) run: | diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 5eac0ccac..0249aaa7d 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -8,6 +8,12 @@ build: jobs: pre_system_dependencies: - git fetch --unshallow # Needed to get version tags + post_install: + # The spec API is documented via autodoc, which imports linopy.spec and so + # needs math-spec. It lives in the `spec` dependency group, not the docs + # extra; --group needs pip >= 25.1, hence the upgrade. + - python -m pip install --upgrade pip + - python -m pip install --group spec python: install: - method: pip diff --git a/benchmarks/models/__init__.py b/benchmarks/models/__init__.py index 66c9a7c76..2b9f7ecac 100644 --- a/benchmarks/models/__init__.py +++ b/benchmarks/models/__init__.py @@ -21,5 +21,6 @@ qp, sos, sparse_network, + spec_pypsa, storage, ) diff --git a/benchmarks/models/spec_pypsa.py b/benchmarks/models/spec_pypsa.py new file mode 100644 index 000000000..fce1719a2 --- /dev/null +++ b/benchmarks/models/spec_pypsa.py @@ -0,0 +1,55 @@ +""" +Model built from math-spec's ``pypsa.yaml`` example (requires math-spec). + +The subject is :meth:`linopy.Model.from_spec`: lowering a spec of PyPSA's full +statement, binding synthetic data to it and building every variable and +constraint it declares. The example lives outside the wheel, so its directory +comes from ``MATH_SPEC_EXAMPLES`` and the case skips without it. A sweep +value is the number of labels per dimension; 40 of them is about 20k +variables. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import TYPE_CHECKING + +from benchmarks.registry import BUILD, FROM_NETCDF, TO_NETCDF, BenchSpec, register + +if TYPE_CHECKING: + import linopy + +SIZES = (5, 40) + +EXAMPLES = os.environ.get("MATH_SPEC_EXAMPLES") +EXAMPLE = Path(EXAMPLES, "pypsa.yaml") if EXAMPLES else None + + +def build_spec_pypsa(n: int) -> linopy.Model: + """Lower ``pypsa.yaml`` and build it with ``n`` labels per dimension.""" + import pytest + + if EXAMPLE is None or not EXAMPLE.exists(): + pytest.skip("set MATH_SPEC_EXAMPLES to a math-spec examples directory") + import math_spec + + import linopy + from linopy.spec.testing import synthetic_sources + + path = str(EXAMPLE) + sources = synthetic_sources(math_spec.to_program(path), n) + with linopy.options as options: + options["semantics"] = "v1" + return linopy.Model.from_spec(path, sources) + + +SPEC = register( + BenchSpec( + name="spec_pypsa", + build=build_spec_pypsa, + sweep=SIZES, + phases=frozenset({BUILD, TO_NETCDF, FROM_NETCDF}), + requires=("math_spec",), + ) +) diff --git a/conftest.py b/conftest.py new file mode 100644 index 000000000..3fd48ab4c --- /dev/null +++ b/conftest.py @@ -0,0 +1,9 @@ +"""Root pytest configuration for ``--doctest-modules`` collection of ``linopy/``.""" + +from __future__ import annotations + +from importlib.util import find_spec + +collect_ignore: list[str] = [] +if find_spec("math_spec") is None: + collect_ignore.append("linopy/spec") diff --git a/doc/api.rst b/doc/api.rst index 0656a99ae..38daf0480 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -115,6 +115,36 @@ IO model.Model.to_netcdf io.read_netcdf +Building from specs +------------------- + +Build a model from a `math-spec +`__ YAML program attached to +data. Requires the ``spec`` dependency group. + +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 +semantics) to back grouped sums with :mod:`linopy.csr`, and build the model +with ``Model(freeze_constraints=True)`` to keep the constraints CSR-backed +instead of densifying them. + +.. autosummary:: + :toctree: generated/ + + model.Model.add_spec + model.Model.from_spec + model.Model.spec + spec.ModelSpec + spec.Layer + spec.NamedExpressions + spec.NamedExpression + spec.Declaration + spec.Unspecified + spec.attach + spec.Attached + spec.SpecDataError + Variable ======== diff --git a/doc/building-models-from-specs.nblink b/doc/building-models-from-specs.nblink new file mode 100644 index 000000000..f9918a8db --- /dev/null +++ b/doc/building-models-from-specs.nblink @@ -0,0 +1,3 @@ +{ + "path": "../examples/building-models-from-specs.ipynb" +} diff --git a/doc/contributing.rst b/doc/contributing.rst index e0d71cc36..97e47ed65 100644 --- a/doc/contributing.rst +++ b/doc/contributing.rst @@ -45,6 +45,9 @@ To run the test suite: # Install development dependencies uv sync --extra dev --extra solvers + # Also run the math-spec binder tests (needs Python >= 3.12) + uv sync --extra dev --extra solvers --group spec + # Run all tests pytest diff --git a/doc/index.rst b/doc/index.rst index b3c754473..e07d31999 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -116,6 +116,7 @@ This package is published under MIT license. coordinate-alignment migrating-to-v1 manipulating-models + building-models-from-specs .. toctree:: :hidden: diff --git a/doc/release_notes.rst b/doc/release_notes.rst index dae3a6dd4..f1aa24a14 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -21,6 +21,17 @@ Upcoming Version * Every operation whose result changes under v1 emits a ``LinopySemanticsWarning`` under legacy, naming the fix — so a model can be migrated incrementally before opting in. The full rules are specified in :doc:`the arithmetic convention `. +*Build a model from a math-spec program* + +* ``Model.from_spec`` / ``model.add_spec`` build a model from a `math-spec `__ program attached to data, and ``model.spec`` reads it back as named layers, each with its ``program``, ``parameters``, ``coords`` and ``lookups``. Requires the ``spec`` dependency group (``uv sync --group spec``, Python >= 3.12) and v1 semantics. The API emits an :class:`linopy.EvolvingAPIWarning` once per session while it stabilises. See :doc:`building-models-from-specs`. + +* A spec can extend a hand-built model: a ``Variable`` passed in ``sources`` under a declared name binds it instead of building one. Everything a layer builds carries its layer name in ``.spec``; ``model.remove_spec(name)`` unwinds a layer, and removing a spec-owned name any other way is refused. ``model.spec.unspecified`` reports what the model holds beyond its specs. + +* ``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 models round-trip through ``to_netcdf`` / ``read_netcdf``; a file read without ``math-spec`` installed loads as a plain model with a warning. + + *Numerical scaling* * Variables, constraints and the objective accept a ``scaling`` factor that rewrites the problem into better-behaved units for the solver, without changing the answer. Variable scaling is column-like, constraint and objective scaling are row-like, and primal values, duals and the objective are transformed back to the original units after solving. See the :doc:`numerical-scaling` tutorial and the *Numerical scaling* section of the :doc:`user-guide`. diff --git a/doc/user-guide.rst b/doc/user-guide.rst index 92995e3ff..fcf43a075 100644 --- a/doc/user-guide.rst +++ b/doc/user-guide.rst @@ -83,6 +83,20 @@ bound, swap a constraint, or copy it for what-if analysis. variables. +Building a model from a spec +----------------------------- + +Instead of calling ``add_variables`` / ``add_constraints`` directly, +you can declare a model as a `math-spec +`__ YAML program attached to +data, and let linopy build it. + +- :doc:`building-models-from-specs` — ``Model.from_spec`` and + ``model.add_spec``, attaching data to a spec, and reading named + expressions back through ``model.spec`` after solving. Requires the + ``spec`` dependency group and v1 semantics. + + Where to go next ---------------- diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb new file mode 100644 index 000000000..7e2dbfd27 --- /dev/null +++ b/examples/building-models-from-specs.ipynb @@ -0,0 +1,1275 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Building models from math-spec programs\n", + "\n", + "This notebook is a tour of the `linopy.spec` feature: build a full linopy model\n", + "from a **math-spec** program (a YAML description of an optimization problem)\n", + "plus a bag of data, solve it, read named results back as arrays, and round-trip\n", + "the whole thing through netCDF.\n", + "\n", + "The idea in one line: **a spec is the maths, the sources are the numbers.** The\n", + "spec names dimensions, parameters, variables, constraints and an objective over\n", + "labelled axes; you supply the labels and the values separately. `linopy` attaches\n", + "the two together and emits variables, constraints and an objective that align\n", + "and broadcast by dimension, exactly as if you had written them by hand.\n", + "\n", + "We work through, in order:\n", + "\n", + "1. Enabling v1 semantics and the `math-spec` dependency.\n", + "2. The anatomy of a spec, section by section.\n", + "3. Attaching data and building a model with `Model.from_spec`.\n", + "4. Solving, and folding **named expressions** back into arrays.\n", + "5. `retain` modes and `evaluate` — what data stays on the model.\n", + "6. **Absence and coverage** — the rule that decides when a missing row is\n", + " refused. This is the conceptual heart of the feature.\n", + "7. Lookups and grouped sums.\n", + "8. Temporal operators (`shift`).\n", + "9. Synthetic data for any spec.\n", + "10. Persistence: netCDF round-trip and `Model.copy()`.\n", + "\n", + "> This notebook runs headless under `nbconvert`. It needs the `math-spec`\n", + "> package and the HiGHS solver, both pulled in by linopy's `solvers` and `spec`\n", + "> dependency groups." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import math_spec\n", + "import pandas as pd\n", + "import xarray as xr\n", + "import yaml\n", + "\n", + "import linopy\n", + "from linopy import Model, read_netcdf\n", + "from linopy.spec import ModelSpec, SpecDataError\n", + "from linopy.spec.testing import synthetic_sources\n", + "\n", + "# A spec-built model uses linopy's v1 semantics. Set it once, up front.\n", + "linopy.options[\"semantics\"] = \"v1\"\n", + "\n", + "print(\"linopy \", linopy.__version__)\n", + "print(\"math_spec \", math_spec.__version__)\n", + "print(\"solvers \", linopy.available_solvers)\n", + "assert \"highs\" in linopy.available_solvers" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## 1. A worked spec: least-cost dispatch\n", + "\n", + "Here is a complete, self-contained spec. It is the classic **economic\n", + "dispatch** problem: run a fleet of generators as cheaply as possible so that\n", + "supply meets demand in every hour.\n", + "\n", + "Read it top to bottom — every section is explained right after." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "DISPATCH = \"\"\"\n", + "description: Least-cost dispatch of a generator fleet against an hourly load.\n", + "\n", + "dimensions:\n", + " snapshot: { dtype: int, description: dispatch periods }\n", + " generator: { description: generating units }\n", + "\n", + "parameters:\n", + " p_max: { dims: [generator], description: installed capacity }\n", + " load: { dims: [snapshot], description: demand to be met }\n", + " cost: { dims: [generator], description: marginal cost }\n", + "\n", + "variables:\n", + " p:\n", + " description: output of a generator in a snapshot\n", + " dims: [snapshot, generator]\n", + " where: \"p_max > 0\"\n", + " bounds: { lower: 0, upper: p_max }\n", + "\n", + "constraints:\n", + " power_balance:\n", + " dims: [snapshot]\n", + " expression: sum(p, over=generator) == load\n", + "\n", + "objective:\n", + " sense: minimize\n", + " expression: sum(p * cost)\n", + "\n", + "expressions:\n", + " spend: sum(p * cost, over=generator)\n", + " usage: p / p_max\n", + "\"\"\"\n", + "print(DISPATCH)" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "### What each section means\n", + "\n", + "- **`dimensions`** — the labelled axes of the problem. Here `snapshot` (an\n", + " integer time index) and `generator` (unit names). A dimension's `dtype`\n", + " constrains the labels you may supply for it.\n", + "- **`parameters`** — named input data, each declared over some dimensions.\n", + " `p_max` is one number per generator, `load` one per snapshot, `cost` one per\n", + " generator. The spec declares the *shape*; you supply the *values* later.\n", + "- **`variables`** — the unknowns. `p` exists `dims: [snapshot, generator]`,\n", + " so one decision variable per (hour, unit). `where: \"p_max > 0\"` masks the\n", + " variable off wherever a generator has no capacity. `bounds` fixes the feasible\n", + " range: output is non-negative and at most the installed capacity `p_max`.\n", + "- **`constraints`** — `power_balance` holds `dims: [snapshot]`: in every\n", + " hour, the generators' total output must equal the load. `sum(p,\n", + " over=generator)` collapses the generator axis, leaving one equation per\n", + " snapshot.\n", + "- **`objective`** — minimise total spend, `sum(p * cost)` over everything.\n", + "- **`expressions`** — *named* expressions, post-solve read-outs: after solving\n", + " you can ask for `spend` (cost per hour) or `usage` (output as a fraction of\n", + " capacity) and get them back as numeric arrays. Neither sits in the objective\n", + " or a constraint, but since both read the variable `p`, `linopy` builds them\n", + " as ordinary expressions on the model too, tagged as this spec's. More on this\n", + " below.\n", + "\n", + "Notice there are **no numbers** in the spec, except the structural `0`. The\n", + "spec is reusable across any fleet and any set of hours." + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## 2. Supplying the data\n", + "\n", + "Data is a plain mapping keyed by the names the spec declares: one entry per\n", + "dimension (its labels), one per parameter (its values). linopy reads it **by\n", + "key, on demand** — it never iterates your mapping beyond the keys it needs.\n", + "\n", + "Three attachment rules are worth knowing, because they make the result\n", + "predictable:\n", + "\n", + "1. A dimension's members come **only** from the source keyed by that\n", + " dimension's name.\n", + "2. Their **order is your order** — linopy never sorts them.\n", + "3. A parameter source is read for **values, not labels**; it is aligned onto the\n", + " dimension members you gave." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "generator = pd.Index([\"wind\", \"gas\"], name=\"generator\")\n", + "snapshot = pd.Index([0, 1, 2], name=\"snapshot\")\n", + "\n", + "dispatch_data = {\n", + " \"snapshot\": snapshot,\n", + " \"generator\": generator,\n", + " \"p_max\": pd.Series([100.0, 200.0], index=generator),\n", + " \"load\": pd.Series([80.0, 150.0, 50.0], index=snapshot),\n", + " \"cost\": pd.Series([0.0, 50.0], index=generator), # wind free, gas costly\n", + "}\n", + "dispatch_data" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "## 3. Building the model\n", + "\n", + "`Model.from_spec(spec, sources)` lowers the spec, attaches the data and emits a\n", + "normal linopy `Model`. The `spec` argument is flexible: a path, YAML text, a\n", + "`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", + "`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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "m = Model.from_spec(DISPATCH, dispatch_data)\n", + "\n", + "print(\"variables \", list(m.variables))\n", + "print(\"constraints\", list(m.constraints))\n", + "print(\"sense \", m.objective.sense)\n", + "m" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "The variable `p` is a genuine linopy variable over `(snapshot, generator)`, and\n", + "`power_balance` a genuine constraint over `snapshot`. From here everything is\n", + "ordinary linopy — you can inspect, print and manipulate them." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "print(m.variables[\"p\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "print(m.constraints[\"power_balance\"])" + ] + }, + { + "cell_type": "markdown", + "id": "12", + "metadata": {}, + "source": [ + "## 4. Solve, then fold named expressions\n", + "\n", + "Solving is ordinary linopy. Wind is free, so it is used to its 100 MW cap first;\n", + "gas covers the rest. Total spend at the optimum is 2500." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "m.solve(solver_name=\"highs\", output_flag=False)\n", + "print(\"termination:\", m.termination_condition)\n", + "print(\"objective: \", m.objective.value)\n", + "m.solution[\"p\"]" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "### Named expressions become data, and stay maths too\n", + "\n", + "`m.spec` is the accessor onto the program the model was built from. Its\n", + "`expressions` mapping returns a `NamedExpression` for each name — three views of\n", + "the same quantity:\n", + "\n", + "- `.node` — the formula as math-spec's lowered expression: the symbolic handle.\n", + "- `.expression` — the **unsolved** linopy expression, variables still symbolic\n", + " and parameters already attached. When the body reads a variable, as `spend`\n", + " and `usage` both do, this is the very `LinearExpression` the model itself\n", + " holds in `model.expressions`, stamped with the layer's name through its\n", + " `.spec` attribute — a data-only body would instead give back an array or a\n", + " scalar, kept on the spec alone (a named expression is affine, so never\n", + " quadratic).\n", + "- `.solution` — the expression **folded** over the solution: every variable\n", + " replaced by its solved value, every parameter by the data it was attached to, the\n", + " arithmetic run on xarray.\n", + "\n", + "`spend` = `sum(p * cost, over=generator)` folds to the cost incurred each hour;\n", + "`usage` = `p / p_max` folds to each unit's utilisation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "print(repr(m.spec))\n", + "spend = m.spec.expressions[\"spend\"]\n", + "\n", + "print(\"\\nspend.expression (unsolved linopy expression):\")\n", + "print(spend.expression)\n", + "\n", + "print(\"\\nspend.solution (folded over the solution):\")\n", + "print(spend.solution)\n", + "\n", + "print(\"\\nusage.solution:\")\n", + "print(m.spec.expressions[\"usage\"].solution)\n", + "print(\n", + " \"\\nspend.expression is m.expressions['spend']:\",\n", + " spend.expression is m.expressions[\"spend\"],\n", + ")\n", + "print(\"m.expressions['spend'].spec:\", m.expressions[\"spend\"].spec)" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "### The spec as maths\n", + "\n", + "The accessor typesets the spec, delegating to math-spec: `m.spec.typeset(fmt)`\n", + "for any format math-spec knows, with `m.spec.to_latex()`, `.to_markdown()` and\n", + "`.to_typst()` spelling the three it knows today. `to_markdown()` is\n", + "GitHub-flavoured, with math in the verbatim delimiters GitHub renders; in a\n", + "notebook the accessor renders itself as Markdown with the `$$` pairs\n", + "MathJax reads.\n", + "\n", + "The **spec**, which need not be the whole model: a spec-built model goes on\n", + "taking everything linopy can add to it, and none of that carries a math-spec\n", + "declaration to typeset. `m.spec.unspecified` reports the drift — variables,\n", + "constraints and expressions the spec does not declare, special-ordered sets it\n", + "does not declare, piecewise formulations added beside it, and whether\n", + "`add_objective` has replaced its objective, which is the one that makes the\n", + "render *wrong* rather than incomplete. Where there is any, typesetting warns\n", + "and opens the rendered text with a comment of the format's own — `%` in LaTeX,\n", + "gone once compiled, there in the source.\n", + "\n", + "A spec's own `piecewise:` and `sos:` are not drift: math-spec lowers them into\n", + "ordinary declarations, which typeset like any other.\n", + "\n", + "Any single declaration typesets on its own too. `m.spec.declaration(name)`\n", + "takes a named expression, a constraint or a variable and hands back a\n", + "`Declaration` with the same methods; a `NamedExpression` carries them\n", + "directly. A declaration is reached through the spec, so it can never be out of\n", + "step with the model the way the whole-spec render can. These render **one** line — math only, no surrounding document — so\n", + "the string drops straight into a docstring or a table cell, and both a\n", + "`Declaration` and a `NamedExpression` render as their own formula in a notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "m.spec" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"spend.to_latex(): \", spend.to_latex())\n", + "print(\"power_balance.to_latex():\", m.spec.declaration(\"power_balance\").to_latex())\n", + "print(\"p.to_latex(): \", m.spec.declaration(\"p\").to_latex())\n", + "\n", + "# each renders as its own formula in a notebook:\n", + "m.spec.declaration(\"power_balance\")" + ] + }, + { + "cell_type": "markdown", + "id": "19", + "metadata": {}, + "source": [ + "### When the model holds more than the spec\n", + "\n", + "Everything above typeset the *whole* model, because the dispatch model **is** its\n", + "spec. Once you add a variable or constraint by hand, that stops being true: the\n", + "render still shows the spec, so it no longer shows the model. `m.spec.unspecified`\n", + "reports exactly this drift, and typesetting a drifted model warns and opens the\n", + "rendered text with a comment naming what it left out." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "\n", + "# Build the dispatch model again, then add a variable and a constraint by hand.\n", + "# Neither carries a math-spec declaration, so neither can be typeset.\n", + "hybrid = Model.from_spec(DISPATCH, dispatch_data)\n", + "hybrid.add_variables(lower=0, name=\"reserve\")\n", + "hybrid.add_constraints(hybrid.variables[\"reserve\"] <= 10, name=\"reserve_cap\")\n", + "\n", + "print(\"unspecified:\", hybrid.spec.unspecified)\n", + "\n", + "# Typesetting a drifted model warns, and opens the render with a comment of the\n", + "# format's own (here '%' for LaTeX) naming what is missing.\n", + "with warnings.catch_warnings(record=True) as caught:\n", + " warnings.simplefilter(\"always\", UserWarning)\n", + " latex = hybrid.spec.to_latex()\n", + "print(\"\\nUserWarning:\", caught[-1].message)\n", + "print(\"\\nrender opens with:\", latex.splitlines()[0])" + ] + }, + { + "cell_type": "markdown", + "id": "21", + "metadata": {}, + "source": [ + "A named expression that reads only data (no variables) has a `.solution`\n", + "**before** a solve too — it needs a solution only if it actually references a\n", + "variable. Subscripting an unknown name raises a `KeyError` with a suggestion\n", + "(the fold is lazy, so the error is on the subscript, not on a view)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "22", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " m.spec.expressions[\"spent\"]\n", + "except KeyError as e:\n", + " print(\"KeyError:\", e)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "hybrid" + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": {}, + "source": [ + "## 5. `retain`: what data stays on the model\n", + "\n", + "Folding needs the parameters an expression reads. `retain` controls which\n", + "parameters linopy keeps in `model.spec.parameters` after building.\n", + "That is the spec's own dataset — `model.parameters` stays yours, and a\n", + "build never writes to it:\n", + "\n", + "| `retain` | keeps in `model.spec.parameters` |\n", + "|------------|-------------------------------------------------|\n", + "| `\"report\"` | only parameters the named expressions read (default) |\n", + "| `\"all\"` | every parameter |\n", + "| `\"none\"` | nothing |\n", + "\n", + "`spend` reads `cost`, `usage` reads `p_max`, neither reads `load` — so\n", + "`\"report\"` keeps `cost` and `p_max` but drops `load`.\n", + "\n", + "Dropping is about *storage*, not about what you can read. A parameter\n", + "`retain` left out is resolved from the `sources` you built with, which the\n", + "model keeps hold of — so every `retain` folds the same in this session.\n", + "It is writing the model to netCDF that leaves the sources behind: read that\n", + "file back and only what `retain` kept is still there, with\n", + "`m.spec.evaluate(name, sources)` as the way in for the rest." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "for retain in [\"report\", \"all\", \"none\"]:\n", + " mm = Model.from_spec(DISPATCH, dispatch_data, retain=retain)\n", + " print(\n", + " f\"retain={retain!r:9} -> parameters kept: {sorted(mm.spec.parameters.data_vars)}\"\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "### `evaluate`: fold against fresh data\n", + "\n", + "With `retain=\"none\"` nothing is kept, so `expressions[name].solution` cannot\n", + "fold. For that case (or any expression whose parameters were not retained) there\n", + "is `spec.evaluate(name, sources)`: it returns a `NamedExpression` whose\n", + "parameters are reattached from a **fresh** bag of data, folding against the model's\n", + "solution.\n", + "\n", + "The catch: `evaluate` reads the solution the model already holds, so the fresh\n", + "sources must describe the **same dimension labels in the same order**.\n", + "Mislabelling a dimension is refused with a `SpecDataError`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "lean = Model.from_spec(DISPATCH, dispatch_data, retain=\"none\")\n", + "lean.solve(solver_name=\"highs\", output_flag=False)\n", + "\n", + "# .solution cannot fold: no parameters were retained.\n", + "try:\n", + " lean.spec.expressions[\"spend\"].solution\n", + "except SpecDataError as e:\n", + " print(\"SpecDataError:\", str(e)[:90], \"...\\n\")\n", + "\n", + "# evaluate reattaches from fresh sources; .solution folds:\n", + "print(lean.spec.evaluate(\"spend\", dispatch_data).solution)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28", + "metadata": {}, + "outputs": [], + "source": [ + "# Relabelling a dimension is refused: evaluate reads the held solution.\n", + "wrong = {**dispatch_data, \"generator\": pd.Index([\"solar\", \"coal\"], name=\"generator\")}\n", + "try:\n", + " lean.spec.evaluate(\"spend\", wrong)\n", + "except SpecDataError as e:\n", + " print(\"SpecDataError:\", e)" + ] + }, + { + "cell_type": "markdown", + "id": "29", + "metadata": {}, + "source": [ + "## 6. Absence and coverage — one rule, every position\n", + "\n", + "This is the concept that makes spec-built models predictable on **sparse** data.\n", + "Real data has holes: a parameter table may simply not list a value for some\n", + "member. math-spec's answer is **uniform** — a missing row is **refused\n", + "wherever it is used**, no matter which position in the maths it sits in:\n", + "\n", + "- **As a coefficient**, a missing row is refused. It would otherwise read as a\n", + " silent zero and drop the term while the row stays — that's exactly the\n", + " ambiguity the rule closes.\n", + "- **As a variable bound**, a missing row is refused. Zero is a bound, not the\n", + " absence of one, so linopy refuses to guess.\n", + "- **As a constant side** of a constraint, a missing row is refused. It would\n", + " bind the constraint, so it must be present.\n", + "- **As a divisor**, a missing row is refused. Zero is not a divisor.\n", + "- A shift `offset` or window `width` given by a parameter *name* is a\n", + " coefficient too, so a hole there is refused the same way.\n", + "\n", + "Crucially, each rule is checked against the rows the declaration **actually\n", + "builds** — a `where:` that removed a coordinate has already answered, so a slot\n", + "you masked off is never demanded. There is no silent zero-fill anywhere; if\n", + "zero is what you mean, you say so, either by masking the coordinate out or by\n", + "filling the data yourself.\n", + "\n", + "Let's see all four positions refuse the same kind of hole." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "T = pd.Index([0, 1, 2], name=\"t\")\n", + "\n", + "SPARSE = {\n", + " \"dimensions\": {\"t\": {\"dtype\": \"int\"}},\n", + " \"parameters\": {\"c\": {\"dims\": [\"t\"]}, \"w\": {\"dims\": [\"t\"]}},\n", + " \"variables\": {\"x\": {\"dims\": [\"t\"], \"bounds\": {\"lower\": 0, \"upper\": 10}}},\n", + " \"constraints\": {\"cap\": {\"dims\": [\"t\"], \"expression\": \"w * x <= c\"}},\n", + " \"objective\": {\"sense\": \"maximize\", \"expression\": \"sum(x, over=t)\"},\n", + "}\n", + "\n", + "# w has no value at t=0. As the COEFFICIENT of x, the missing row would\n", + "# otherwise be read as 0 and the term dropped -- that's refused, not guessed.\n", + "w_hole = pd.Series([1.0, 1.0], index=T[1:]) # missing t=0\n", + "c_full = pd.Series([0.0, 4.0, 5.0], index=T)\n", + "\n", + "\n", + "def refuse(spec, data, label):\n", + " try:\n", + " Model.from_spec(spec, {\"t\": T, **data})\n", + " except SpecDataError as e:\n", + " print(f\"[{label}]\\n {e}\\n\")\n", + "\n", + "\n", + "refuse(SPARSE, {\"w\": w_hole, \"c\": c_full}, \"coefficient\")" + ] + }, + { + "cell_type": "markdown", + "id": "31", + "metadata": {}, + "source": [ + "The other three positions refuse the same kind of hole, joining the\n", + "coefficient. Each `SpecDataError` names the position and how many rows are\n", + "short." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32", + "metadata": {}, + "outputs": [], + "source": [ + "c_hole = pd.Series([4.0, 5.0], index=T[1:]) # missing t=0\n", + "\n", + "# (a) a hole in a variable bound\n", + "bound_spec = {\n", + " **SPARSE,\n", + " \"variables\": {\"x\": {\"dims\": [\"t\"], \"bounds\": {\"lower\": 0, \"upper\": \"c\"}}},\n", + "}\n", + "refuse(bound_spec, {\"w\": pd.Series([1.0, 1.0, 1.0], index=T), \"c\": c_hole}, \"bound\")\n", + "\n", + "# (b) a hole in a constant side (right-hand side that binds the constraint)\n", + "refuse(SPARSE, {\"w\": pd.Series([1.0, 1.0, 1.0], index=T), \"c\": c_hole}, \"constant side\")\n", + "\n", + "# (c) a hole in a divisor\n", + "div_spec = {\n", + " **SPARSE,\n", + " \"constraints\": {\"cap\": {\"dims\": [\"t\"], \"expression\": \"x / w <= c\"}},\n", + "}\n", + "refuse(div_spec, {\"w\": w_hole, \"c\": c_full}, \"divisor\")" + ] + }, + { + "cell_type": "markdown", + "id": "33", + "metadata": {}, + "source": [ + "Two escape hatches fix the coefficient hole above, and both build and solve.\n", + "\n", + "**(a) `where:`** — the coordinate does not exist there, so there is no row to\n", + "cover. Add `where: \"w\"` to the `cap` constraint and t=0 drops out entirely.\n", + "\n", + "**(b) Fill the data** — if zero really is what you mean, say so:\n", + "`w.fillna(0.0)` (or any dense series) supplies the row instead of leaving a\n", + "hole for linopy to guess at." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "34", + "metadata": {}, + "outputs": [], + "source": [ + "where_spec = {\n", + " **SPARSE,\n", + " \"constraints\": {\"cap\": {\"dims\": [\"t\"], \"where\": \"w\", \"expression\": \"w * x <= c\"}},\n", + "}\n", + "wm = Model.from_spec(where_spec, {\"t\": T, \"w\": w_hole, \"c\": c_full})\n", + "wm.solve(solver_name=\"highs\", output_flag=False)\n", + "print(\"where: t=0 has no cap row ->\", wm.objective.value)\n", + "\n", + "fm2 = Model.from_spec(SPARSE, {\"t\": T, \"w\": w_hole.reindex(T).fillna(0.0), \"c\": c_full})\n", + "fm2.solve(solver_name=\"highs\", output_flag=False)\n", + "print(\"fillna(0.0): t=0's cap is 0*x <= 0 ->\", fm2.objective.value)" + ] + }, + { + "cell_type": "markdown", + "id": "35", + "metadata": {}, + "source": [ + "And the same masking escape hatch on the variable and constraint together:\n", + "`x` and its cap only exist where `live` is true, so the hole in `c` at the\n", + "masked position is fine." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36", + "metadata": {}, + "outputs": [], + "source": [ + "masked_spec = {\n", + " **SPARSE,\n", + " \"parameters\": {**SPARSE[\"parameters\"], \"live\": {\"dims\": [\"t\"], \"dtype\": \"bool\"}},\n", + " \"variables\": {\n", + " \"x\": {\"dims\": [\"t\"], \"where\": \"live\", \"bounds\": {\"lower\": 0, \"upper\": \"c\"}}\n", + " },\n", + " \"constraints\": {\n", + " \"cap\": {\"dims\": [\"t\"], \"where\": \"live\", \"expression\": \"w * x <= c\"}\n", + " },\n", + "}\n", + "live = pd.Series([True, True], index=T[1:]) # off at t=0, where c is missing\n", + "mm = Model.from_spec(\n", + " masked_spec,\n", + " {\"t\": T, \"w\": pd.Series([1.0, 1.0, 1.0], index=T), \"c\": c_hole, \"live\": live},\n", + ")\n", + "built = int((mm.variables[\"x\"].labels != -1).sum())\n", + "print(f\"x occupies {built} of 3 slots; the masked t=0 needed no data.\")" + ] + }, + { + "cell_type": "markdown", + "id": "37", + "metadata": {}, + "source": [ + "## 7. Lookups and grouped sums\n", + "\n", + "A **lookup** maps each member of one dimension to a member of another — think\n", + "\"which bus is this generator on\". The spec declares it under `lookups:`, and an\n", + "expression can then sum a per-generator quantity **into** per-bus totals with\n", + "`sum(..., by=)`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "38", + "metadata": {}, + "outputs": [], + "source": [ + "GROUPED = {\n", + " \"dimensions\": {\"generator\": {}, \"bus\": {\"dtype\": \"str\"}},\n", + " \"lookups\": {\"gen_bus\": {\"over\": \"generator\", \"into\": \"bus\"}},\n", + " \"parameters\": {\"capacity\": {\"dims\": [\"generator\"]}},\n", + " \"variables\": {\"imports\": {\"dims\": [\"bus\"], \"bounds\": {\"lower\": 0, \"upper\": 100}}},\n", + " \"constraints\": {\n", + " \"import_limit\": {\n", + " \"dims\": [\"bus\"],\n", + " \"expression\": \"imports <= sum(capacity, by=gen_bus)\",\n", + " }\n", + " },\n", + " \"objective\": {\"sense\": \"maximize\", \"expression\": \"sum(imports, over=bus)\"},\n", + "}\n", + "gens = pd.Index([\"g1\", \"g2\"], name=\"generator\")\n", + "grouped_data = {\n", + " \"bus\": [\"north\", \"south\"],\n", + " \"generator\": gens,\n", + " \"gen_bus\": pd.Series([\"north\", \"north\"], index=gens), # both gens on north\n", + " \"capacity\": pd.Series([3.0, 4.0], index=gens),\n", + "}\n", + "gm = Model.from_spec(GROUPED, grouped_data)\n", + "gm.solve(solver_name=\"highs\", output_flag=False)\n", + "print(gm.solution[\"imports\"].to_series())\n", + "print(\"south has no generators -> its grouped capacity is 0, not a gap.\")" + ] + }, + { + "cell_type": "markdown", + "id": "39", + "metadata": {}, + "source": [ + "Note `south` has no generators mapped to it. Its group is **empty**, and an\n", + "empty group on a constant side sums to a clean **zero**, not a missing-data gap.\n", + "An empty group is a legitimate answer; a member with no value is still refused." + ] + }, + { + "cell_type": "markdown", + "id": "40", + "metadata": {}, + "source": [ + "### Skewed topologies\n", + "\n", + "`sum(..., by=lookup)` and `sum_back` build dense by default: a group-by\n", + "allocates the full (group, member) rectangle even when most groups are tiny,\n", + "and a trailing window allocates the full window width at every row. On a\n", + "skewed topology — a few large groups among many singletons, or a wide window\n", + "over a long axis — that rectangle is mostly padding. Set\n", + "`linopy.options[\"sparse_groupby\"] = True` (v1 semantics) to back grouped sums\n", + "with the CSR representation in `linopy.csr` instead, and build the model with\n", + "`Model(freeze_constraints=True)` to keep constraints CSR-backed rather than\n", + "densified." + ] + }, + { + "cell_type": "markdown", + "id": "41", + "metadata": {}, + "source": [ + "## 8. Temporal operators: `shift`\n", + "\n", + "For time-coupled problems the language provides operators that walk an axis:\n", + "`shift` (offset a series along a dimension), `at` (index through a lookup),\n", + "`sum_back` (a trailing window). `shift(expr, over=snapshot, offset=1,\n", + "edge='wrap')` gives \"the value one step earlier, wrapping at the ends\" — exactly\n", + "what a storage balance needs.\n", + "\n", + "Below, a battery links consecutive hours: its state of charge equals the\n", + "previous hour's charge, plus what it stored, minus what it released. With a\n", + "cheap-then-expensive price profile, the optimizer buys extra cheap energy, banks\n", + "it, and discharges when power is dear." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "42", + "metadata": {}, + "outputs": [], + "source": [ + "STORAGE = \"\"\"\n", + "description: A battery shifts cheap energy into expensive hours.\n", + "dimensions:\n", + " snapshot: { dtype: int }\n", + "parameters:\n", + " load: { dims: [snapshot] }\n", + " price: { dims: [snapshot] }\n", + " soc_max: { dims: [] }\n", + "variables:\n", + " gen: { dims: [snapshot], bounds: { lower: 0, upper: 1000 } }\n", + " charge: { dims: [snapshot], bounds: { lower: 0, upper: soc_max } }\n", + " discharge: { dims: [snapshot], bounds: { lower: 0, upper: soc_max } }\n", + " soc: { dims: [snapshot], bounds: { lower: 0, upper: soc_max } }\n", + "constraints:\n", + " balance:\n", + " dims: [snapshot]\n", + " expression: gen + discharge - charge == load\n", + " storage:\n", + " dims: [snapshot]\n", + " expression: soc == shift(soc, over=snapshot, offset=1, edge='wrap') + charge - discharge\n", + "objective:\n", + " sense: minimize\n", + " expression: sum(gen * price)\n", + "expressions:\n", + " cost: sum(gen * price, over=snapshot)\n", + "\"\"\"\n", + "snap = pd.Index(range(6), name=\"snapshot\")\n", + "storage_data = {\n", + " \"snapshot\": snap,\n", + " \"load\": pd.Series([10, 10, 10, 10, 10, 10], index=snap, dtype=float),\n", + " \"price\": pd.Series([1, 1, 1, 9, 9, 9], index=snap, dtype=float),\n", + " \"soc_max\": 20.0,\n", + "}\n", + "bm = Model.from_spec(STORAGE, storage_data, retain=\"all\")\n", + "bm.solve(solver_name=\"highs\", output_flag=False)\n", + "print(\"objective:\", bm.objective.value)\n", + "print(\n", + " pd.DataFrame(\n", + " {\n", + " \"price\": storage_data[\"price\"],\n", + " \"gen\": bm.solution[\"gen\"].to_series(),\n", + " \"charge\": bm.solution[\"charge\"].to_series(),\n", + " \"discharge\": bm.solution[\"discharge\"].to_series(),\n", + " \"soc\": bm.solution[\"soc\"].to_series(),\n", + " }\n", + " ).round(1)\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "43", + "metadata": {}, + "source": [ + "The generator over-produces while power is cheap (hour 2 runs at 30 to fill the\n", + "battery), the battery discharges through the expensive hours, and the folded\n", + "`cost` expression reports total generation spend." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "44", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"folded cost:\", float(bm.spec.expressions[\"cost\"].solution))" + ] + }, + { + "cell_type": "markdown", + "id": "45", + "metadata": {}, + "source": [ + "## 9. Synthetic data for any spec\n", + "\n", + "A spec declares exactly what data it needs, which is enough to invent some. The\n", + "`synthetic_sources` helper reads a lowered program and fabricates dense data of\n", + "the right shapes — labels numbered per dimension, parameters a linear ramp. The\n", + "result builds and solves, and tells you nothing about a real system. It is what\n", + "the test suite and benchmarks use to exercise any spec." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "46", + "metadata": {}, + "outputs": [], + "source": [ + "program = math_spec.to_program(yaml.safe_load(DISPATCH))\n", + "fake = synthetic_sources(program, n=4)\n", + "print(\"keys:\", sorted(fake))\n", + "print(\"\\ngenerated 'generator' labels:\", list(fake[\"generator\"]))\n", + "print(\"generated 'load':\")\n", + "print(fake[\"load\"])\n", + "\n", + "fm = Model.from_spec(DISPATCH, fake, retain=\"all\")\n", + "fm.solve(solver_name=\"highs\", output_flag=False)\n", + "print(\"\\nsynthetic model solves:\", fm.termination_condition)" + ] + }, + { + "cell_type": "markdown", + "id": "47", + "metadata": {}, + "source": [ + "## 10. Persistence: netCDF and copy\n", + "\n", + "A spec-built model round-trips through netCDF and through `Model.copy()`. The\n", + "spec travels as its **YAML text**, stored as a top-level attribute and lowered\n", + "again on read. Everything else that must survive is data: the master\n", + "coordinates, the lookups and the retained parameters.\n", + "\n", + "Labels are the delicate part — a partial lookup can hold a `NaN` inside an array\n", + "of strings, and no netCDF type carries that. linopy stores lookups and\n", + "object-dtype parameters as `pandas.factorize` output (integer codes plus a\n", + "category table) and records each parameter's in-memory dtype, so the exact\n", + "dtypes come back on read on both the `netcdf4` and `scipy` engines." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "48", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import tempfile\n", + "\n", + "from linopy.testing import assert_model_equal\n", + "\n", + "m2 = Model.from_spec(DISPATCH, dispatch_data, retain=\"report\")\n", + "m2.solve(solver_name=\"highs\", output_flag=False)\n", + "\n", + "with tempfile.TemporaryDirectory() as d:\n", + " path = os.path.join(d, \"dispatch.nc\")\n", + " m2.to_netcdf(path)\n", + " restored = read_netcdf(path)\n", + "\n", + "# the models are equal, including the spec text and the retained parameters:\n", + "assert_model_equal(m2, restored)\n", + "print(\"round-trip equal:\", True)\n", + "print(\"spec text preserved:\", restored.spec.text == m2.spec.text)\n", + "\n", + "# and the named expressions fold identically after the round-trip:\n", + "for name in restored.spec.expressions:\n", + " xr.testing.assert_equal(\n", + " m2.spec.expressions[name].solution, restored.spec.expressions[name].solution\n", + " )\n", + " print(f\" {name}: identical\")" + ] + }, + { + "cell_type": "markdown", + "id": "49", + "metadata": {}, + "source": [ + "Even a `retain=\"none\"` model round-trips: the spec text and coordinates survive,\n", + "so after loading you can still `evaluate` against fresh data." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50", + "metadata": {}, + "outputs": [], + "source": [ + "with tempfile.TemporaryDirectory() as d:\n", + " path = os.path.join(d, \"lean.nc\")\n", + " lean.to_netcdf(path)\n", + " lean_back = read_netcdf(path)\n", + "\n", + "print(\"no parameters retained:\", list(lean_back.spec.parameters.data_vars) == [])\n", + "print(lean_back.spec.evaluate(\"spend\", dispatch_data))" + ] + }, + { + "cell_type": "markdown", + "id": "51", + "metadata": {}, + "source": [ + "`Model.copy()` carries the spec too, with the accessor reattached to the copy. The\n", + "copy is a fresh, unsolved model (like any linopy copy), so solve it before\n", + "folding an expression that reads a variable — the folded result then matches the\n", + "original." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "52", + "metadata": {}, + "outputs": [], + "source": [ + "clone = m2.copy()\n", + "print(\"copy has spec:\", isinstance(clone.spec, ModelSpec))\n", + "print(\"copy carries a solution:\", \"solution\" in clone.variables[\"p\"].data)\n", + "\n", + "clone.solve(solver_name=\"highs\", output_flag=False)\n", + "xr.testing.assert_equal(\n", + " clone.spec.expressions[\"spend\"].solution, m2.spec.expressions[\"spend\"].solution\n", + ")\n", + "print(\"after solving the copy, folded expressions match the original\")" + ] + }, + { + "cell_type": "markdown", + "id": "53", + "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": "54", + "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": "55", + "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": "56", + "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", + " dims: [snapshot, generator]\n", + "\n", + "constraints:\n", + " co2_limit:\n", + " dims: []\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": "57", + "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": "58", + "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": "59", + "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": "60", + "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": "61", + "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": "62", + "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`, 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, 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, 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", + "### Summary\n", + "\n", + "A spec is the maths over labelled axes; the sources are the numbers. `linopy`\n", + "attaches them into an ordinary model, hands each named expression back as three\n", + "views — its formula, its unsolved linopy expression and its solution — refuses a\n", + "missing parameter row wherever it is used (as a coefficient, bound, constant\n", + "side or divisor alike, with `where:` and filling the data as the escape\n", + "hatches), and round-trips the lot through netCDF by keeping the spec as text\n", + "beside factorized labels." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/linopy/constants.py b/linopy/constants.py index 7936ef1c2..b7496d507 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -4,6 +4,7 @@ """ import logging +import warnings from dataclasses import dataclass, field from enum import StrEnum from typing import Any, Literal, Self, TypeAlias, get_args @@ -93,6 +94,9 @@ class PerformanceWarning(UserWarning): SOS_DIM_ATTR = "sos_dim" SOS_BIG_M_ATTR = "big_m_upper" +# The spec layer that built a variable, constraint or expression +SPEC_LAYER_ATTR = "spec" + # Indicator constraint attribute keys INDICATOR_BINARY_VAR_ATTR = "indicator_binary_var" INDICATOR_BINARY_VAL_ATTR = "indicator_binary_val" @@ -124,6 +128,22 @@ class EvolvingAPIWarning(FutureWarning): """ +_emitted_evolving_warnings: set[str] = set() + + +def warn_evolving_api(key: str, message: str, stacklevel: int = 3) -> None: + """ + Emit an :class:`EvolvingAPIWarning` at most once per session per ``key``. + + ``stacklevel`` counts from the ``warnings.warn`` call: 3 points at the + caller of the function that calls this helper. + """ + if key in _emitted_evolving_warnings: + return + _emitted_evolving_warnings.add(key) + warnings.warn(message, category=EvolvingAPIWarning, stacklevel=stacklevel) + + class ModelStatus(StrEnum): """ Model status. diff --git a/linopy/constraints.py b/linopy/constraints.py index f3b301cce..d35d1c8b1 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -10,7 +10,14 @@ 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, + Sequence, +) from dataclasses import dataclass from itertools import product from typing import ( @@ -72,6 +79,7 @@ GREATER_EQUAL, HELPER_DIMS, LESS_EQUAL, + SPEC_LAYER_ATTR, TERM_DIM, PerformanceWarning, SIGNS_pretty, @@ -152,6 +160,15 @@ def model(self) -> Model: def name(self) -> str: """Get the constraint name.""" + @property + @abstractmethod + def spec(self) -> str | None: + """The spec layer that built this constraint; ``None`` for one built by hand.""" + + @spec.setter + @abstractmethod + def spec(self, layer: str) -> None: ... + @property @abstractmethod def is_assigned(self) -> bool: @@ -625,6 +642,7 @@ class CSRConstraint(ConstraintBase): "_dual", "_binvar_labels", "_binval", + "_spec", "_positional_cache", ) @@ -642,6 +660,7 @@ def __init__( binvar_labels: np.ndarray | None = None, binval: int | np.ndarray | None = None, scaling: np.ndarray | None = None, + spec: str | None = None, ) -> None: self._csr = csr self._active_positions = active_positions @@ -659,6 +678,7 @@ def __init__( self._dual = dual self._binvar_labels = binvar_labels self._binval = binval + self._spec = spec self._positional_cache: _PositionalCache | None = None @property @@ -697,8 +717,18 @@ def attrs(self) -> dict[str, Any]: d: dict[str, Any] = {"name": self._name} if self._cindex is not None: d["label_range"] = (self._cindex, self._cindex + self.full_size) + if self._spec is not None: + d[SPEC_LAYER_ATTR] = self._spec return d + @property + def spec(self) -> str | None: + return self._spec + + @spec.setter + def spec(self, layer: str) -> None: + self._spec = layer + @property def coords(self) -> DatasetCoordinates: return Dataset(coords=self._grid.indexes).coords @@ -750,6 +780,7 @@ def _init_kwargs(self) -> dict[str, Any]: binvar_labels=self._binvar_labels, binval=self._binval, scaling=self._scaling, + spec=self._spec, ) def _replace(self, **changes: Any) -> CSRConstraint: @@ -782,6 +813,7 @@ def assign_labels( sign=self._sign if isinstance(self._sign, str) else self._sign[keep], cindex=cindex, name=name, + spec=None, ) if scaling is not None: changes["scaling"] = scaling[positions] @@ -1083,6 +1115,8 @@ def to_netcdf_ds(self) -> Dataset: } if isinstance(self._sign, str): attrs["sign"] = self._sign + if self._spec is not None: + attrs[SPEC_LAYER_ATTR] = self._spec if self._binvar_labels is not None: attrs["is_indicator"] = True data_vars["_binvar_labels"] = DataArray(self._binvar_labels, dims=["_flat"]) @@ -1151,6 +1185,7 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: binvar_labels=binvar_labels, binval=binval, scaling=scaling, + spec=attrs.get(SPEC_LAYER_ATTR), ) def has_labels(self, labels: np.ndarray) -> bool: @@ -1364,6 +1399,7 @@ def from_dense( binvar_labels=binvar_labels, binval=binval, scaling=scaling, + spec=con.data.attrs.get(SPEC_LAYER_ATTR), ) @classmethod @@ -1482,6 +1518,14 @@ def model(self) -> Model: def name(self) -> str: return self.attrs["name"] + @property + def spec(self) -> str | None: + return self.attrs.get(SPEC_LAYER_ATTR) + + @spec.setter + def spec(self, layer: str) -> None: + self.attrs[SPEC_LAYER_ATTR] = layer + @property def is_assigned(self) -> bool: return self._assigned @@ -2118,8 +2162,10 @@ 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) -> str: - """Format constraint items, optionally excluding names in a group.""" + def _format_items( + self, exclude: set[str] | None = None, tagged: bool = False + ) -> str: + """Format constraint items, optionally excluding names in a group and, if *tagged*, naming each one's spec layer.""" r = "" count = 0 for name, ds in self.items(): @@ -2131,7 +2177,8 @@ def _format_items(self, exclude: set[str] | None = None) -> str: if ds.coords else "" ) - r += f" * {name}{coords}\n" + suffix = f" [{ds.spec}]" if tagged and ds.spec is not None else "" + r += f" * {name}{coords}{suffix}\n" if count == 0: r += "\n" return r @@ -2218,7 +2265,11 @@ def add(self, constraint: ConstraintBase, freeze: bool = False) -> ConstraintBas def remove(self, name: str) -> None: """ Remove constraint `name` from the constraints. + + Refused where a spec layer builds or binds it. """ + if self.model._ownership is not None: + self.model._ownership.refuse_removal("constraint", [name]) self.data.pop(name) self._invalidate_label_position_index() diff --git a/linopy/expressions.py b/linopy/expressions.py index 65472232d..a14a8fbab 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -103,6 +103,7 @@ GROUP_STACK_DIM, HELPER_DIMS, LESS_EQUAL, + SPEC_LAYER_ATTR, STACKED_TERM_DIM, TERM_DIM, ) @@ -1587,6 +1588,15 @@ def name(self) -> str: """ return str(self.attrs["name"]) + @property + def spec(self) -> str | None: + """The spec layer whose named expression this is; ``None`` for an expression built by hand.""" + return self.attrs.get(SPEC_LAYER_ATTR) + + @spec.setter + def spec(self, layer: str) -> None: + self.attrs[SPEC_LAYER_ATTR] = layer + @property def data(self) -> Dataset: return self._data @@ -3536,8 +3546,10 @@ def __dir__(self) -> list[str]: ] return base_attributes + formatted_names - def _format_items(self, exclude: set[str] | None = None) -> str: - """Format expression items, optionally excluding names in a group.""" + def _format_items( + self, exclude: set[str] | None = None, tagged: bool = False + ) -> str: + """Format expression items, optionally excluding names in a group and, if *tagged*, naming each one's spec layer.""" r = "" count = 0 for name, ds in self.items(): @@ -3549,7 +3561,8 @@ def _format_items(self, exclude: set[str] | None = None) -> str: if ds.coords else "" ) - r += f" * {name}{coords}\n" + suffix = f" [{ds.spec}]" if tagged and ds.spec is not None else "" + r += f" * {name}{coords}{suffix}\n" if count == 0: r += "\n" return r @@ -3591,8 +3604,12 @@ def add(self, expression: LinearExpression | QuadraticExpression) -> None: def remove(self, name: str) -> None: """ - Remove variable `name` from the variables. + Remove expression `name` from the expressions. + + Refused where a spec layer builds or binds it. """ + if self.model._ownership is not None: + self.model._ownership.refuse_removal("expression", [name]) self.data.pop(name) @property diff --git a/linopy/io.py b/linopy/io.py index 08701ed03..b5d416df6 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -11,8 +11,9 @@ import shutil import time import warnings -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from importlib.metadata import version +from importlib.util import find_spec from io import BufferedWriter from pathlib import Path from tempfile import TemporaryDirectory @@ -45,7 +46,15 @@ logger = logging.getLogger(__name__) NETCDF_VERSION_ATTR = "_linopy_version" +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" +SPEC_VERSION_ATTR = SPEC_ATTR + "-version" CONTAINER_ORDER_ATTR = "_linopy_{}_order" @@ -1019,6 +1028,147 @@ def non_bool_dict( return {k: int(v) if isinstance(v, bool) else v for k, v in d.items()} +def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: + """*ds* with every dim, coordinate, variable and attribute of it namespaced under *prefix*.""" + to_rename = set([*ds.dims, *ds.coords, *ds]) + ds = ds.rename({d: f"{prefix}-{d}" for d in to_rename}) + ds.attrs = {f"{prefix}-{k}": v for k, v in ds.attrs.items()} + + # Flatten multiindexes + for dim in ds.dims: + if isinstance(ds[dim].to_index(), pd.MultiIndex): + prefix_len = len(prefix) + 1 # leave original index level name + names = [n[prefix_len:] for n in ds[dim].to_index().names] + ds = ds.reset_index(dim) + # scipy netCDF3 backend cannot write unicode-array attrs. + ds.attrs[f"{dim}_multiindex"] = json.dumps(list(names)) + + return ds + + +def has_prefix(k: str, prefix: str) -> bool: + return k.rsplit("-", 1)[0] == prefix + + +def remove_prefix(k: str, prefix: str) -> str: + return k[len(prefix) + 1 :] + + +def parse_multiindex_attr(value: str | Iterable[str]) -> list[str]: + # str = JSON (new); iterable = legacy list from older linopy. + if isinstance(value, str): + return [str(n) for n in json.loads(value)] + return [str(n) for n in value] + + +def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: + """The part of *ds* :func:`with_prefix` wrote under *prefix*, its names given back.""" + ds = ds[[k for k in ds if has_prefix(str(k), prefix)]] + multiindexes = [] + for dim in ds.dims: + attr = ds.attrs.get(f"{dim}_multiindex") + if attr is None: + continue + for name in parse_multiindex_attr(attr): + multiindexes.append(prefix + "-" + name) + ds = ds.drop_vars(set(ds.coords) - set(ds.dims) - set(multiindexes)) + to_rename = set([*ds.dims, *ds.coords, *ds]) + ds = ds.rename({d: remove_prefix(d, prefix) for d in to_rename}) + ds.attrs = { + remove_prefix(k, prefix): v + for k, v in ds.attrs.items() + if has_prefix(k, prefix) + } + + for dim in ds.dims: + if f"{dim}_multiindex" in ds.attrs: + names = parse_multiindex_attr(ds.attrs.pop(f"{dim}_multiindex")) + ds = ds.set_index({dim: names}) # type: ignore[dict-item] + + return ds + + +def record_dtypes(ds: xr.Dataset) -> xr.Dataset: + """ + *ds* with each array's in-memory dtype written as an attribute. + + No netcdf type holds a dtype as written: an engine narrows an int64 to + int32 and hands a bool back as int8, so the dtype travels beside the + values and :func:`restore_dtypes` puts it back. + """ + typed = { + str(name): arr.assign_attrs({DTYPE_ATTR: str(arr.dtype)}) + for name, arr in ds.items() + } + return ds.assign(typed) + + +def restore_dtypes(ds: xr.Dataset) -> xr.Dataset: + """*ds* with each array back at the dtype :func:`record_dtypes` recorded; one written without is left as it is.""" + cast = { + str(name): arr.astype(np.dtype(arr.attrs.pop(DTYPE_ATTR))) + for name, arr in ds.items() + if DTYPE_ATTR in arr.attrs + } + return ds.assign(cast) + + +def restamp_coords(m: Model, coords: Mapping[str, pd.Index]) -> None: + """ + Put *coords* on every container of *m* that was built on them. + + Only on those: a container may carry a dimension of that name and its own + labels -- a hand-added variable beside a spec-built one -- and restamping + it would rewrite labels it never had, or fail outright over a length the + master coordinate does not share. + """ + from linopy.constraints import Constraint, CSRConstraint + from linopy.csr import Grid + + for _, variable in m.variables.items(): + variable._data = _stamped(variable.data, coords) + for _, expression in m.expressions.items(): + expression._data = _stamped(expression.data, coords) + m.objective.expression._data = _stamped(m.objective.expression.data, coords) + for _, constraint in m.constraints.items(): + if isinstance(constraint, Constraint): + constraint._data = _stamped(constraint.data, coords) + elif isinstance(constraint, CSRConstraint): + constraint._grid = Grid( + { + d: _restamped(index, coords.get(str(d))) + for d, index in constraint._grid.indexes.items() + } + ) + + +def _stamped(data: xr.Dataset, coords: Mapping[str, pd.Index]) -> xr.Dataset: + """*data* with *coords* in place of the ones a dtype narrowed.""" + stale = { + str(dim): restamped + for dim, index in data.indexes.items() + if (restamped := _restamped(index, coords.get(str(dim)))) is not index + } + return data.assign_coords(stale) if stale else data + + +def _restamped(found: pd.Index, master: pd.Index | None) -> pd.Index: + """ + *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. 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 + part = master[master.isin(found)] + return part if part.equals(found) else found + + def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: """ Write out the model to a netcdf file. @@ -1040,6 +1190,15 @@ 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 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 Model and is not persisted. If the model has an active SOS reformulation at serialization time, the netcdf contains the @@ -1061,22 +1220,6 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: stacklevel=2, ) - def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: - to_rename = set([*ds.dims, *ds.coords, *ds]) - ds = ds.rename({d: f"{prefix}-{d}" for d in to_rename}) - ds.attrs = {f"{prefix}-{k}": v for k, v in ds.attrs.items()} - - # Flatten multiindexes - for dim in ds.dims: - if isinstance(ds[dim].to_index(), pd.MultiIndex): - prefix_len = len(prefix) + 1 # leave original index level name - names = [n[prefix_len:] for n in ds[dim].to_index().names] - ds = ds.reset_index(dim) - # scipy netCDF3 backend cannot write unicode-array attrs. - ds.attrs[f"{dim}_multiindex"] = json.dumps(list(names)) - - return ds - vars = [ with_prefix(var.data, f"variables-{name}") for name, var in m.variables.items() ] @@ -1100,10 +1243,17 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: if m.objective.value is not None: objective = objective.assign_attrs(value=m.objective.value) obj = [with_prefix(objective, "objective")] - params = [with_prefix(m.parameters, "parameters")] + specs: list[xr.Dataset] = [] + if m._spec is not None: + from linopy.spec.netcdf import encode + + 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} - ds = xr.merge(vars + cons + exprs + obj + params, combine_attrs="drop_conflicts") + ds = xr.merge( + vars + cons + exprs + obj + params + specs, combine_attrs="drop_conflicts" + ) ds = ds.assign_attrs(scalars) ds.attrs[NETCDF_VERSION_ATTR] = version("linopy") for kind, container in ( @@ -1112,6 +1262,10 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ("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: @@ -1134,6 +1288,11 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ds.to_netcdf(*args, **kwargs) +def spec_available() -> bool: + """Whether the ``math-spec`` package is importable, so a file's spec layers can be read.""" + return find_spec("math_spec") is not None + + def read_netcdf(path: Path | str, **kwargs: Any) -> Model: """ Read in a model from a netcdf file. @@ -1176,43 +1335,6 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: m = Model() ds = xr.load_dataset(path, **kwargs) - def has_prefix(k: str, prefix: str) -> bool: - return k.rsplit("-", 1)[0] == prefix - - def remove_prefix(k: str, prefix: str) -> str: - return k[len(prefix) + 1 :] - - def parse_multiindex_attr(value: str | Iterable[str]) -> list[str]: - # str = JSON (new); iterable = legacy list from older linopy. - if isinstance(value, str): - return [str(n) for n in json.loads(value)] - return [str(n) for n in value] - - def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: - ds = ds[[k for k in ds if has_prefix(str(k), prefix)]] - multiindexes = [] - for dim in ds.dims: - attr = ds.attrs.get(f"{dim}_multiindex") - if attr is None: - continue - for name in parse_multiindex_attr(attr): - multiindexes.append(prefix + "-" + name) - ds = ds.drop_vars(set(ds.coords) - set(ds.dims) - set(multiindexes)) - to_rename = set([*ds.dims, *ds.coords, *ds]) - ds = ds.rename({d: remove_prefix(d, prefix) for d in to_rename}) - ds.attrs = { - remove_prefix(k, prefix): v - for k, v in ds.attrs.items() - if has_prefix(k, prefix) - } - - for dim in ds.dims: - if f"{dim}_multiindex" in ds.attrs: - names = parse_multiindex_attr(ds.attrs.pop(f"{dim}_multiindex")) - ds = ds.set_index({dim: names}) # type: ignore[dict-item] - - return ds - def container_names(kind: str) -> list[str]: found = {str(k).rsplit("-", 1)[0] for k in ds if str(k).startswith(kind)} order_attr = ds.attrs.get(CONTAINER_ORDER_ATTR.format(kind)) @@ -1278,7 +1400,20 @@ def container_names(kind: str) -> list[str]: ) m.objective._value = objective.attrs.pop("value", None) - m.parameters = get_prefix(ds, "parameters") + m.parameters = restore_dtypes(get_prefix(ds, "parameters")) + + if SPEC_LAYERS_ATTR in ds.attrs or SPEC_ATTR in ds.attrs: + if spec_available(): + from linopy.spec.netcdf import read + + m._spec = read(m, ds) + else: + warnings.warn( + f"'{path}' holds spec layers and math-spec is not installed; loaded as " + f"a plain model, without model.spec. Install math-spec to read the layers.", + UserWarning, + stacklevel=2, + ) for k in m.scalar_attrs: if k in ds.attrs: @@ -1412,6 +1547,9 @@ def _copy_con_data(con: ConstraintBase) -> xr.Dataset: ) new_model._parameters = m._parameters.copy(deep=deep) + if m._spec is not None: + new_model._spec = m._spec._reattach(new_model, deep=deep) + new_model._ownership = _copy.deepcopy(m._ownership) new_model._blocks = m._blocks.copy(deep=deep) if m._blocks is not None else None for attr in m.scalar_attrs: diff --git a/linopy/model.py b/linopy/model.py index 614adf387..82a035e99 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -47,6 +47,7 @@ SOS_BIG_M_ATTR, SOS_DIM_ATTR, SOS_TYPE_ATTR, + SPEC_LAYER_ATTR, TERM_DIM, ModelStatus, Result, @@ -117,6 +118,8 @@ if TYPE_CHECKING: from linopy.piecewise import PiecewiseFormulation + from linopy.spec import ModelSpec, Retain, SpecLike + from linopy.spec.ownership import Ownership logger = logging.getLogger(__name__) @@ -202,6 +205,8 @@ class Model: "_piecewise_formulations", "_solver", "_sos_reformulation_state", + "_spec", + "_ownership", "__weakref__", ) @@ -305,6 +310,8 @@ def __init__( ) self._solver: solvers.Solver | None = None self._sos_reformulation_state: SOSReformulationResult | None = None + self._spec: ModelSpec | None = None + self._ownership: Ownership | None = None @property def solver(self) -> solvers.Solver | None: @@ -396,6 +403,8 @@ def objective( obj = Objective(obj, self) self._objective = obj + if self._ownership is not None: + self._ownership.objective = None @property def sense(self) -> str: @@ -437,6 +446,168 @@ def solution(self) -> Dataset: """ return self.variables.solution + @property + def spec(self) -> ModelSpec: + """ + 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 no spec was added to the model. + """ + if self._spec is None: + raise AttributeError( + "This model holds no spec. Use `Model.add_spec` or " + "`Model.from_spec` to add one." + ) + return self._spec + + def add_spec( + self, + spec: SpecLike, + sources: Mapping[str, Any] | Dataset, + retain: Retain = "report", + name: str | None = None, + build_expressions: bool = True, + ) -> 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 + the objective are added as the spec declares them; the spec text, the + parameters the named expressions read and the lookups are kept on the + model, and the named expressions are read back through ``model.spec``. + + Everything the layer builds carries the layer's name in its ``attrs`` + under ``"spec"``, read through the ``spec`` property of a + ``Variable``, ``Constraint`` or expression. A bound variable carries + none: the layer reads it and the model owns it. + + A named expression whose body holds a variable term is built as well + and added to ``model.expressions`` under its declared name, stamped + like the rest, so ``model.expressions[name]`` and + ``model.spec.expressions[name].expression`` are one object. A + data-only body (parameters and constants alone) and one reading a + constraint's ``dual`` stay on the spec, read lazily. + + 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. Layers share + variables only this way: one layer's named expression is not visible + to another, since math-spec lowers each layer standalone. + + Parameters + ---------- + spec : str, pathlib.Path, dict or math_spec.Spec + The spec. A ``str`` is YAML text if it holds a newline, opens a + mapping or a sequence, or holds a ``:`` and names no file; any + other ``str`` is a path. A lowered ``math_spec.Program`` and an + open file are refused, having no YAML form to keep on the model. + sources : mapping or xarray.Dataset + Data keyed by declared name: dimension labels, parameters, + lookups and the model variables to bind. ``keys()`` is called + once, and everything after that is read by key on demand. A key + naming nothing the spec declares is ignored, so one mapping can + feed several specs; one close to a declared name is warned about + as a likely typo. 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, as + is one holding ``/`` or ``-``, which a netcdf file cannot carry. + build_expressions : bool, default True + Whether to build the layer's variable-bearing named expressions + into ``model.expressions``. ``False`` stores nothing and + ``model.spec.expressions[name].expression`` folds on read, the + escape hatch for a spec whose named expressions are too large to + hold. Not persisted: a model read from a file folds whatever the + file does not hold. + + Returns + ------- + linopy.Model + This model, for chaining. + + Raises + ------ + ValueError + 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, or a binding does not fit its + declaration. + + Warns + ----- + EvolvingAPIWarning + Once per session: the spec API is newly added and may change in + minor releases. Silence with ``warnings.filterwarnings("ignore", + category=linopy.EvolvingAPIWarning)``. + """ + from linopy.spec.accessor import attach + + self._spec = attach(self, spec, sources, retain, name, build_expressions) + return self + + def remove_spec(self, name: str) -> None: + """ + Take the spec layer *name* off the model with everything it built. + + Its constraints, its named expressions and the variables it built + are removed, in that order, a special-ordered set it declared comes + off the variable carrying it, and the objective goes where the + layer's is the one the model holds. A variable the layer bound + stays: it is the model's. A hand-added constraint reading a variable + the layer built goes with the variable, as :meth:`remove_variables` + takes it. Once the last layer is off, the model holds no spec. + + Raises + ------ + AttributeError + The model holds no spec. + KeyError + No layer of that name. + ValueError + Another layer binds a variable this one built; remove that + layer first. + """ + from linopy.spec.accessor import remove_layer + + remove_layer(self, name) + + @classmethod + def from_spec( + cls, + spec: SpecLike, + sources: Mapping[str, Any] | Dataset, + retain: Retain = "report", + name: str | None = None, + build_expressions: bool = True, + **model_kwargs: Any, + ) -> Model: + """ + A new model built from a math-spec program, see :meth:`add_spec`. + + ``model_kwargs`` are passed to :class:`Model`. + """ + return cls(**model_kwargs).add_spec( + spec, sources, retain=retain, name=name, build_expressions=build_expressions + ) + @property def dual(self) -> Dataset: """ @@ -618,13 +789,29 @@ def __repr__(self) -> str: from linopy.piecewise import _repr_summary as pwl_repr_summary var_names, con_names = _get_piecewise_groups(self) - var_string = self.variables._format_items(exclude=var_names) - con_string = self.constraints._format_items(exclude=con_names) - expr_string = self.expressions._format_items() model_string = f"Linopy {self.type} model" + tag_variables = tag_constraints = tag_expressions = False + descriptions: list[str] = [] + if self._spec is not None: + 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 + tag_variables = bool(unspecified.variables) + tag_constraints = bool(unspecified.constraints) + tag_expressions = bool(unspecified.expressions) + descriptions = [layer.description for layer in layers if layer.description] + var_string = self.variables._format_items(var_names, tag_variables) + expr_string = self.expressions._format_items(tagged=tag_expressions) + con_string = self.constraints._format_items(con_names, tag_constraints) + header = f"{model_string}\n{'=' * len(model_string)}\n" + header += "".join(f"{d}\n" for d in descriptions) return ( - f"{model_string}\n{'=' * len(model_string)}\n\n" + f"{header}\n" f"Variables:\n----------\n{var_string}\n" f"Expressions:\n------------\n{expr_string}\n" f"Constraints:\n------------\n{con_string}" @@ -871,6 +1058,8 @@ def add_variables( name = f"var{self._varnameCounter}" self._varnameCounter += 1 + if self._ownership is not None: + self._ownership.refuse_addition("variable", name) if name in self.variables: raise ValueError(f"Variable '{name}' already assigned to model") @@ -1016,6 +1205,8 @@ def add_expressions( name = f"expr{self._exprnameCounter}" self._exprnameCounter += 1 + if self._ownership is not None: + self._ownership.refuse_addition("expression", name) if name in self.expressions: raise ValueError(f"Expression '{name}' already assigned to model") @@ -1035,6 +1226,7 @@ def add_expressions( if self.chunk: expr = expr.chunk(self.chunk) + expr.attrs.pop(SPEC_LAYER_ATTR, None) expr.attrs["name"] = name self.expressions.add(expr) return expr @@ -1105,6 +1297,8 @@ def add_sos_constraints( def _resolve_constraint_name(self, name: str | None, prefix: str = "con") -> str: """Validate a constraint name or generate one from ``prefix``.""" + if name is not None and self._ownership is not None: + self._ownership.refuse_addition("constraint", name) if name in list(self.constraints): raise ValueError(f"Constraint '{name}' already assigned to model") if name is None: @@ -1377,6 +1571,7 @@ def add_constraints( self.check_force_dim_names(data) data = self._allocate_constraint_labels(data, name, mask) + data.attrs.pop(SPEC_LAYER_ATTR, None) if self.chunk: data = data.chunk(self.chunk) @@ -1517,6 +1712,8 @@ def add_objective( self.objective.expression = expr self.objective.sense = sense self.objective.scaling = scaling + if self._ownership is not None: + self._ownership.objective = None def remove_variables(self, name: str) -> None: """ @@ -1542,6 +1739,8 @@ def remove_variables(self, name: str) -> None: to_remove = [k for k, con in self.constraints.items() if con.has_labels(labels)] + self.variables.remove(name) + if to_remove: warnings.warn( f"Removing variable '{name}' also removes constraints {to_remove} " @@ -1552,13 +1751,12 @@ def remove_variables(self, name: str) -> None: for k in to_remove: self.constraints.remove(k) - self.variables.remove(name) - referenced = self.objective.vars.isin(labels) if FACTOR_DIM in referenced.dims: referenced = referenced.any(FACTOR_DIM) - self.objective = self.objective.sel({TERM_DIM: ~referenced}) + if referenced.any(): + self.objective = self.objective.sel({TERM_DIM: ~referenced}) def remove_constraints(self, name: str | list[str]) -> None: """ @@ -1576,13 +1774,12 @@ def remove_constraints(self, name: str | list[str]) -> None: ------- None. """ - if isinstance(name, list): - for n in name: - logger.debug(f"Removed constraint: {n}") - self.constraints.remove(n) - else: - logger.debug(f"Removed constraint: {name}") - self.constraints.remove(name) + names = [name] if isinstance(name, str) else name + if self._ownership is not None: + self._ownership.refuse_removal("constraint", names) + for n in names: + logger.debug(f"Removed constraint: {n}") + self.constraints.remove(n) def remove_expressions(self, name: str | list[str]) -> None: """ @@ -1601,6 +1798,8 @@ def remove_expressions(self, name: str | list[str]) -> None: None. """ names = [name] if isinstance(name, str) else name + if self._ownership is not None: + self._ownership.refuse_removal("expression", names) for n in names: logger.debug(f"Removed expression: {n}") self.expressions.remove(n) @@ -1628,6 +1827,8 @@ def remove_sos_constraints(self, variable: Variable) -> None: del variable.attrs[SOS_TYPE_ATTR], variable.attrs[SOS_DIM_ATTR] variable.attrs.pop(SOS_BIG_M_ATTR, None) + if self._ownership is not None: + self._ownership.sos.pop(variable.name, None) logger.debug( f"Removed sos{sos_type} constraint on {sos_dim} from {variable.name}" diff --git a/linopy/piecewise.py b/linopy/piecewise.py index 5a07dba43..5349e7789 100644 --- a/linopy/piecewise.py +++ b/linopy/piecewise.py @@ -8,7 +8,6 @@ from __future__ import annotations import logging -import warnings from collections.abc import Sequence from dataclasses import dataclass from numbers import Real @@ -47,8 +46,8 @@ PWL_SELECT_SUFFIX, SEGMENT_DIM, SIGNS, - EvolvingAPIWarning, sign_replace_dict, + warn_evolving_api, ) from linopy.semantics import check_user_nan_breakpoints @@ -61,30 +60,6 @@ logger = logging.getLogger(__name__) -# Each user-facing piecewise entry point fires its EvolvingAPIWarning at -# most once per process. Without dedup, a single model build emits the -# verbose warning hundreds of times and drowns out other output. -_EvolvingApiKey: TypeAlias = Literal[ - "tangent_lines", "add_piecewise_formulation", "Slopes" -] -_emitted_evolving_warnings: set[_EvolvingApiKey] = set() - - -def _warn_evolving_api(key: _EvolvingApiKey, message: str, stacklevel: int = 3) -> None: - """ - Emit an :class:`EvolvingAPIWarning` at most once per session per ``key``. - - ``stacklevel`` defaults to 3 (helper → entry-point function → user - code). Pass a larger value when called from one frame deeper than - a function — e.g. from a dataclass ``__post_init__``, which is - itself invoked by an auto-generated ``__init__``. - """ - if key in _emitted_evolving_warnings: - return - _emitted_evolving_warnings.add(key) - warnings.warn(message, category=EvolvingAPIWarning, stacklevel=stacklevel) - - # Accepted input types for breakpoint-like data BreaksLike: TypeAlias = ( Sequence[float] @@ -172,7 +147,7 @@ class Slopes: def __post_init__(self) -> None: # ``stacklevel=4``: warn → _warn_evolving_api → __post_init__ → # dataclass-generated ``__init__`` → user code. - _warn_evolving_api( + warn_evolving_api( "Slopes", "piecewise: Slopes is a new API; the constructor signature and " "the dispatch rules for inheriting an x grid from sibling tuples " @@ -826,7 +801,7 @@ def tangent_lines( Silence with ``warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning)``. """ - _warn_evolving_api( + warn_evolving_api( "tangent_lines", "piecewise: tangent_lines is a new API; the returned expression " "shape and the piece-dim name may be refined in minor releases. " @@ -1272,7 +1247,7 @@ def add_piecewise_formulation( with ``warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning)``. """ - _warn_evolving_api( + warn_evolving_api( "add_piecewise_formulation", "piecewise: add_piecewise_formulation is a new API; some details " "(e.g. the per-tuple sign convention, active+sign semantics) " diff --git a/linopy/spec/__init__.py b/linopy/spec/__init__.py new file mode 100644 index 000000000..b3dc29f83 --- /dev/null +++ b/linopy/spec/__init__.py @@ -0,0 +1,47 @@ +""" +Build linopy models from math-spec programs. + +The package needs the ``math-spec`` distribution (import name ``math_spec``, +Python >= 3.12). It is imported here and nowhere else in linopy, so +``import linopy`` never pulls it in. +""" + +from __future__ import annotations + +import sys +from importlib.util import find_spec + +if find_spec("math_spec") is None: + message = ( + "linopy.spec needs the math-spec package. Install it from a checkout " + "with `uv sync --group spec` or `pip install --group spec`." + ) + if sys.version_info < (3, 12): + message = "linopy.spec needs Python >= 3.12. " + message + raise ImportError(message) + +from linopy.spec.accessor import ( + Declaration, + Layer, + ModelSpec, + NamedExpression, + NamedExpressions, + SpecLike, + Unspecified, +) +from linopy.spec.attach import Attached, Retain, attach +from linopy.spec.errors import SpecDataError + +__all__ = [ + "Attached", + "Declaration", + "Layer", + "ModelSpec", + "NamedExpression", + "NamedExpressions", + "Retain", + "SpecDataError", + "SpecLike", + "Unspecified", + "attach", +] diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py new file mode 100644 index 000000000..3d1d97ad0 --- /dev/null +++ b/linopy/spec/accessor.py @@ -0,0 +1,1087 @@ +""" +``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 +``model.parameters``, which stays the caller's: a spec never overwrites what +was put there, and nothing reading a spec-built model has to guess which of +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 +a *file* holds, not what a session can read, and it is only after a round trip +that a parameter can be out of reach. +""" + +from __future__ import annotations + +import functools +import io +import re +import warnings +from collections.abc import Callable, Collection, Iterable, Iterator, Mapping +from dataclasses import dataclass, field, replace +from pathlib import Path +from types import MappingProxyType +from typing import Any, TypeAlias + +import pandas as pd +import xarray as xr +import yaml +from math_spec import ( + Spec, + did_you_mean, + to_program, + to_spec, + typeset, + typeset_declaration, +) +from math_spec import program as ms +from math_spec.typesetting import FormatName + +from linopy.constants import SOS_TYPE_ATTR, warn_evolving_api +from linopy.expressions import LinearExpression, QuadraticExpression +from linopy.model import Model +from linopy.semantics import is_v1 +from linopy.spec import terms +from linopy.spec.attach import Attached, Retain +from linopy.spec.attach import attach as attach_data +from linopy.spec.builder import build +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 +from linopy.spec.ownership import Ownership, joined +from linopy.spec.parameters import Parameters, Resolve + +SpecLike: TypeAlias = str | Path | Mapping[str, Any] | Spec + +# 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: {}." + +_ILLEGAL_IN_NAME = re.compile(r"[/-]") + +EVOLVING_MESSAGE = ( + "spec: Model.add_spec, Model.from_spec, model.spec and linopy.spec.attach are " + "newly added and their details may change in minor releases. Silence with " + '`warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning)`.' +) + +_COMMENT: dict[str, str] = { + "latex": "% {}", + "markdown": "", + "typst": "// {}", +} + + +@dataclass(frozen=True) +class Unspecified: + """ + How a spec-built model has drifted from the spec it was built from. + + A model goes on taking everything linopy can add to it, and none of that + carries a math-spec declaration to typeset. A spec's own ``piecewise:`` + and ``sos:`` are not drift: math-spec lowers them into ordinary + declarations, so they sit in the program like any other. + + Drift is tracked by name: a name a layer owns counts as the layer's + whatever it holds now, so a bound or an ``rhs`` edited in place is not + seen here. + + Attributes + ---------- + variables, constraints + Added beside the spec, a piecewise formulation's own aside. + expressions + Added beside the spec. A layer's named expressions that hold a + variable term live in ``model.expressions`` stamped with the layer's + name, and are not drift. + sos + Variables given a special-ordered set no layer owns. A set a layer + declared and ``remove_sos_constraints`` dropped is the layer's no + longer, so one added in its place lands here. + piecewise + Formulations added by ``add_piecewise_formulation``, named as + formulations rather than as the variables and constraints they hold. + objective + 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. + bound + Model variables a layer binds rather than builds. Not drift and not + falsiness: the layer accounts for them, but the model owns their bounds + and mask, which the layer declares without and so the render omits. + """ + + variables: tuple[str, ...] + constraints: tuple[str, ...] + expressions: tuple[str, ...] + sos: tuple[str, ...] + piecewise: tuple[str, ...] + objective: bool + bound: tuple[str, ...] = () + + def __bool__(self) -> bool: + return bool( + self.variables + or self.constraints + or self.expressions + or self.sos + or self.piecewise + or self.objective + ) + + +def _counted(names: tuple[str, ...], kind: str, cap: int = 5) -> str: + """``2 constraints (a, b)``, capped with a ``+N more`` tail; empty for no names.""" + if not names: + return "" + shown = list(names[:cap]) + if len(names) > cap: + shown.append(f"+{len(names) - cap} more") + plural = kind if len(names) == 1 else f"{kind}s" + return f"{len(names)} {plural} ({', '.join(shown)})" + + +def attach( + model: Model, + spec: SpecLike, + sources: Mapping[str, Any] | xr.Dataset, + retain: Retain, + name: str | None = None, + build_expressions: bool = True, +) -> ModelSpec: + """ + 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"``. With + *build_expressions* the named expressions holding a variable term are + built into ``model.expressions``; without, they fold on read. + + Raises + ------ + ValueError + The model runs under legacy semantics; the layer name is empty or + holds a character a netcdf file cannot carry; 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`` or an open file, neither of which + has a YAML form to keep on the model. + FileNotFoundError + *spec* reads as a path and there is no file there. + """ + warn_evolving_api("spec", EVOLVING_MESSAGE, stacklevel=4) + if not is_v1(): + raise ValueError( + "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." + ) + text, program, stem = normalize_spec(spec) + attached: Attached = attach_data( + program, sources, retain=retain, given=_given(model) + ) + foreign = [n for n, v in attached.bound.items() if v.model is not model] + if foreign: + raise SpecDataError( + f"variable(s) {foreign} are bound to a variable of another model. A layer reads " + f"the variables of the model it is added to; pass the variables of this model." + ) + _check_collisions(model, program, attached) + layer_name = _check_layer_name(name if name is not None else stem or "spec") + if model._spec is not None and layer_name in model._spec.layers: + raise ValueError( + f"a spec layer named '{layer_name}' is already on this model; pass another name." + ) + # 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, layer_name, build_expressions) + layer = Layer( + model, layer_name, program, text, parameters, attached, attached.names + ) + if model._spec is None: + model._ownership = Ownership() + model._spec = ModelSpec(model, [], whole) + model._spec._layers[layer.name] = layer + register(model, layer) + return model._spec + + +def remove_layer(model: Model, name: str) -> None: + """ + Take the layer *name* off *model*: its constraints, its expressions, then the variables it built. + + A special-ordered set the layer declared comes off the variable that + carries it, bound or built. The objective goes where the layer's is the + one the model holds. A hand-added constraint reading a variable the + layer built goes with the variable, as ``remove_variables`` takes it. + Once the last layer is off, the model holds no spec. + + Raises + ------ + KeyError + No layer of that name. + ValueError + Another layer binds a variable this one built; that layer comes off + first. + """ + spec = model.spec + spec[name] + owned = spec._ownership + built = owned.held("variable", name) + read = [n for n in built if n in owned.bound] + if read: + binders = sorted({by for n in read for by in owned.bound[n]}) + raise ValueError( + f"spec layer '{name}' built variable(s) {read} that layer(s) {binders} " + f"bind; remove those layers first." + ) + for variable in owned.held("sos", name): + model.remove_sos_constraints(model.variables[variable]) + constraints = owned.held("constraint", name) + expressions = [n for n in owned.held("expression", name) if n in model.expressions] + owns_objective = owned.objective == name + owned.release(name) + model.remove_constraints(constraints) + model.remove_expressions(expressions) + for variable in built: + model.remove_variables(variable) + if owns_objective: + model.remove_objective() + del spec._layers[name] + if not spec._layers: + model._spec = None + model._ownership = None + + +def register(model: Model, layer: Layer) -> None: + """ + Record in the model's registry every name *layer* owns. + + Called once the layer is on the model: after a build, or after a read, + where the constraints and special-ordered sets a file gave back say what + the layer still holds. + """ + if model._ownership is None: + raise RuntimeError("a model holding spec layers has no ownership registry") + p = layer.program + declared_sos = {layer.names.get(s.variable, s.variable) for s in p.sos.values()} + model._ownership.claim( + layer.name, + variables=[n for n in p.variables if n not in layer.names], + bound=layer.names.values(), + constraints=[n for n in p.constraints if n in model.constraints], + expressions=p.named_expressions, + sos=[n for n in declared_sos if SOS_TYPE_ATTR in model.variables[n].attrs], + objective=p.objective is not None, + ) + + +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 _check_layer_name(name: str) -> str: + """*name*, if a netcdf file can carry it: its arrays are named ``spec--``.""" + if not name or _ILLEGAL_IN_NAME.search(name): + raise ValueError( + f"a spec layer cannot be named '{name}': a name is written into a netcdf file " + f"as the prefix of the layer's arrays, which rules out '/' and '-' and the " + f"empty name. Pass name= with another one." + ) + return name + + +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`." + ) + owned = model._ownership.expressions if model._ownership is not None else {} + expressions = [ + n for n in program.named_expressions if n in owned or n in model.expressions + ] + if expressions: + raise ValueError( + f"the spec declares named expression(s) {expressions} and the model or an " + f"earlier spec on it already holds them." + ) + + +def restore_layer( + model: Model, + name: str, + text: str, + parameters: xr.Dataset, + names: Mapping[str, str], +) -> Layer: + """ + 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. + """ + program = to_program(text) + return Layer(model, name, program, text, parameters, None, dict(names)) + + +def _is_yaml_text(spec: str) -> bool: + """A ``str`` is YAML rather than a path if it looks like YAML and names no file.""" + if "\n" in spec or spec.lstrip()[:1] in ("{", "-"): + return True + return ":" in spec and not Path(spec).is_file() + + +def normalize_spec(spec: SpecLike) -> tuple[str, ms.Program, str | None]: + """ + *spec* as the YAML text kept on the model, lowered, and the name a file lends the layer. + + A ``str`` is YAML text if it holds a newline, opens a mapping or a + sequence, or holds a ``:`` and names no file; every other ``str`` is a + path, and so lends the layer its stem. + + Raises + ------ + TypeError + A lowered ``Program`` or an open file: neither has a YAML form to + keep on the model. + FileNotFoundError + *spec* reads as a path and there is no file there. + SpecDataError + The spec declares no dimension, parameter or variable. + """ + if isinstance(spec, ms.Program): + raise TypeError( + "add_spec takes the spec as a path, YAML text, a mapping or a math_spec.Spec, " + "not a lowered Program: a Program has no YAML form to keep on the model." + ) + if isinstance(spec, io.IOBase): + raise TypeError( + "add_spec takes the spec as a path, YAML text, a mapping or a math_spec.Spec, " + "not an open file: pass the path it was opened on, or spec.read()." + ) + stem: str | None = None + if isinstance(spec, str) and not _is_yaml_text(spec): + spec = Path(spec) + if isinstance(spec, Path): + if not spec.is_file(): + raise FileNotFoundError( + f"no spec file at '{spec}'. A str is read as a path unless it holds a " + f"newline, opens a mapping or a sequence, or holds a ':' and names no " + f"file, so YAML text written on one line arrives here as a path." + ) + text, stem = spec.read_text(), spec.stem + elif isinstance(spec, str): + text = spec + else: + text = (to_spec(dict(spec)) if isinstance(spec, Mapping) else spec).to_yaml() + sections = yaml.safe_load(text) + if not isinstance(sections, Mapping): + raise SpecDataError( + f"a spec is a mapping of sections, and this one reads as " + f"{type(sections).__name__}: {text[:80]!r}." + ) + program = to_program(dict(sections)) + if not (program.dimensions or program.parameters or program.variables): + raise SpecDataError( + "the spec declares nothing: no dimension, no parameter and no variable. " + "There is nothing to attach data to and nothing to build." + ) + return text, program, stem + + +def _dimension(dim: str, coords: Mapping[str, pd.Index]) -> str: + """A dimension and how many labels it holds; a declared one nothing reaches holds none.""" + return f"{dim} ({len(coords[dim])})" if dim in coords else f"{dim} (unreached)" + + +def _row(label: str, items: list[str], cap: int = 8) -> str: + """One aligned summary line, capped with a ``(+N more)`` tail.""" + shown = items[:cap] + if len(items) > cap: + shown = shown + [f"(+{len(items) - cap} more)"] + return f" {label + ':':<13}{', '.join(shown) if shown else '—'}" + + +@dataclass(frozen=True, eq=False, repr=False) +class Layer: + """ + 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. + """ + + 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 + if self.description: + head = f"{head}: {self.description}" + rows = [ + head, + _row("Dimensions", [_dimension(d, coords) for d in p.dimensions]), + _row("Variables", list(p.variables)), + _row("Constraints", list(p.constraints)), + ] + if p.objective is not None: + rows.append(_row("Objective", [p.objective.sense])) + rows.append(_row("Expressions", list(p.named_expressions))) + return rows + + 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: + """The spec's own description, its first line, or an empty string.""" + lines = str(self._schema.get("description", "")).strip().splitlines() + return lines[0] if lines else "" + + @property + def coords(self) -> dict[str, pd.Index]: + """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.""" + out: dict[str, dict[str, xr.DataArray]] = {} + for over, lk in self.program.lookups: + out.setdefault(over, {})[lk.name] = self.parameters[lk.name] + return out + + @property + def expressions(self) -> NamedExpressions: + """Each named expression as a :class:`NamedExpression`: its math, its linopy fold and its solution.""" + return NamedExpressions({n: self for n in self.program.named_expressions}) + + def declaration(self, name: str) -> Declaration: + """ + One declaration typeset on its own: a named expression, constraint or variable. + + Its math as a single line, no document around it. A named expression + also carries its linopy fold and solution through :attr:`expressions`; + this handle is the typesetting one every declaration shares. + """ + if name not in self._declarations: + raise KeyError( + f"unknown declaration '{name}'. " + + did_you_mean(name, self._declarations) + ) + return Declaration(self, name) + + @property + 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), + self.name, + 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. Read-only: a layer is + added with :meth:`Model.add_spec` and taken off with + :meth:`Model.remove_spec`. + 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) -> None: + self._model = model + self._layers: dict[str, Layer] = {layer.name: layer for layer in layers} + self.layers: Mapping[str, Layer] = MappingProxyType(self._layers) + self.whole = whole + + @property + def _ownership(self) -> Ownership: + owned = self._model._ownership + if owned is None: + raise RuntimeError("a model holding spec layers has no ownership registry") + return owned + + @property + def objective_owner(self) -> str | None: + """The layer whose objective the model holds, ``None`` once it was replaced or edited.""" + return self._ownership.objective + + 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) + + 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 its spec layers, see :class:`Unspecified`. + + 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.piecewise import _get_piecewise_groups + + model = self._model + owned = self._ownership + pw_variables, pw_constraints = _get_piecewise_groups(model) + return Unspecified( + variables=tuple( + n + for n in model.variables + if owned.owner("variable", n) is None and n not in pw_variables + ), + constraints=tuple( + n + for n in model.constraints + if owned.owner("constraint", n) is None and n not in pw_constraints + ), + expressions=tuple( + n for n in model.expressions if owned.owner("expression", n) is None + ), + sos=tuple( + n + for n, v in model.variables.items() + if SOS_TYPE_ATTR in v.attrs and owned.owner("sos", n) is None + ), + piecewise=tuple(model._piecewise_formulations), + objective=owned.objective is None and not model.objective.expression.empty, + bound=tuple(n for n in model.variables if n in owned.bound), + ) + + def typeset(self, fmt: FormatName, **options: Any) -> str: + """ + 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 + holds such a thing, :attr:`unspecified` names it, a warning says so, + and the rendered text opens with the same tally as a comment of + *fmt*'s own -- gone once compiled, there in the source. + + 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``. Several layers + refuse ``standalone``: one document cannot hold two preambles. + + Warns + ----- + UserWarning + The model holds variables or constraints the spec does not + declare, which are not in the rendered text. + """ + return self._render(fmt, options, 3) + + def to_latex(self, **options: Any) -> str: + """The spec typeset as a LaTeX document, see :meth:`typeset`.""" + return self._render("latex", options, 3) + + def to_markdown(self, **options: Any) -> str: + """The spec typeset as Markdown, its equations in ``$$`` blocks, see :meth:`typeset`.""" + return self._render("markdown", options, 3) + + def to_typst(self, **options: Any) -> str: + """The spec typeset as Typst, see :meth:`typeset`.""" + return self._render("typst", options, 3) + + 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.""" + 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() + ) + notes = [] + tally = self._tally() + if tally is not None: + note = self._note(tally) + warnings.warn( + f"{note} What is typeset is the spec, so it is not this model.", + UserWarning, + stacklevel=stacklevel, + ) + notes.append(note) + bound = self._bound_note() + if bound is not None: + notes.append(bound) + comment = _COMMENT.get(fmt) + if not notes or comment is None: + return rendered + header = "\n".join(comment.format(n) for n in notes) + return f"{header}\n{rendered}" + + def _note(self, tally: str) -> str: + return (_DRIFTED if self.whole else _EXTENDS).format(tally) + + def _bound_note(self) -> str | None: + """A layer reads these model variables; their bounds and mask are not in the render.""" + names = self.unspecified.bound + if not names: + return None + return ( + f"This spec reads {_counted(names, 'variable')} it binds from the host " + "model, whose bounds and mask the render leaves out." + ) + + def _tally(self) -> str | None: + """How the model has drifted, counted and named; ``None`` when it has not.""" + found = self.unspecified + if not found: + return None + parts = [ + _counted(found.variables, "variable"), + _counted(found.constraints, "constraint"), + _counted(found.expressions, "expression"), + _counted(found.sos, "SOS set"), + _counted(found.piecewise, "piecewise formulation"), + ] + 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.""" + rendered = _notebook_math(self._render("markdown", {}, 3)) + notes = [] + tally = self._tally() + if tally is not None: + notes.append(self._note(tally)) + bound = self._bound_note() + if bound is not None: + notes.append(bound) + if not notes: + return rendered + footer = "\n\n".join(f"*{n}*" for n in notes) + return f"{rendered}\n\n{footer}" + + +class NamedExpressions(Mapping[str, "NamedExpression"]): + """The named expressions of one or more layers, each a :class:`NamedExpression` on read.""" + + def __init__(self, owners: Mapping[str, Layer]) -> None: + self._owners = owners + + def __getitem__(self, name: str) -> NamedExpression: + if name not in self._owners: + raise KeyError( + f"unknown named expression '{name}'. " + + did_you_mean(name, self._owners) + ) + layer = self._owners[name] + held = layer.model.expressions.data.get(name) + stored = held if held is not None and held.spec == layer.name else None + return NamedExpression(layer, name, layer._context(layer._resolve), stored) + + def __iter__(self) -> Iterator[str]: + return iter(self._owners) + + def __len__(self) -> int: + return len(self._owners) + + def __repr__(self) -> str: + return f"NamedExpressions({list(self)})" + + +class Declaration: + """ + One declaration of a spec, typeset on its own: math only, no document. + + A named expression, a constraint or a variable, reached by name through + :meth:`ModelSpec.declaration`. :class:`NamedExpression` adds the linopy + fold and the solution on top of this. + """ + + def __init__(self, layer: Layer, name: str) -> None: + self._layer = layer + self._name = name + + def typeset(self, fmt: FormatName, **options: Any) -> str: + """ + This declaration typeset in *fmt* as a single line, no document around it. + + Nothing here can be out of step with the model the way + :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._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.""" + return self.typeset("latex", **options) + + def to_markdown(self, **options: Any) -> str: + """This declaration typeset as a single Markdown math line, no ``$$`` around it.""" + return self.typeset("markdown", **options) + + def to_typst(self, **options: Any) -> str: + """This declaration typeset as a single Typst line, no document around it.""" + return self.typeset("typst", **options) + + def _repr_markdown_(self) -> str: + return f"$$\n{self.to_markdown()}\n$$" + + +def _notebook_math(markdown: str) -> str: + r""" + GitHub's verbatim math delimiters as the ``$``-pairs a notebook's MathJax reads. + + math-spec prints ``$\`...\`$`` and ```` ```math ```` fences because GitHub + runs Markdown's escape pass inside ``$...$``; Jupyter does not, and renders + only the classic pair. + """ + fenced = re.sub(r"```math\n(.*?)\n```", r"$$\n\1\n$$", markdown, flags=re.S) + return re.sub(r"\$`(.*?)`\$", r"$\1$", fenced) + + +class NamedExpression(Declaration): + """ + One named expression, in three views: its math, its linopy fold and its solution. + + The object pins the data sources it was made with for its lifetime, so the + three views agree. ``expressions[name]`` reads the model's own data -- + what ``retain`` kept, and the sources behind it for the rest; + ``evaluate(name, sources)`` attaches fresh data instead. + + Where the layer built the expression into ``model.expressions``, + ``expressions[name].expression`` is that stored object and needs no + parameters at all; ``solution`` folds afresh and does. + + Attributes + ---------- + node + The lowered expression body, math-spec's own AST handle. + """ + + def __init__( + self, + layer: Layer, + name: str, + ctx: Context, + stored: LinearExpression | QuadraticExpression | None = None, + ) -> None: + super().__init__(layer, name) + self._ctx = ctx + self._stored = stored + + @property + def node(self) -> ms.ExpressionNode: + """The expression body as lowered, math-spec's own AST handle.""" + 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._layer.program) + + @functools.cached_property + def expression(self) -> terms.Value: + """ + The linopy symbolic expression, its variables unsolved. + + The entry ``model.expressions`` holds where the layer built one, + else folded here: a ``LinearExpression`` where the body carries + variables, a bare ``Variable``, a ``DataArray`` for a data-only body + or a ``float`` for a constant. Not wrapped: a degree-0 array can hold + holes that ``from_constant`` would refuse. + """ + if self._stored is not None: + return self._stored + return evaluate_named(self._name, self._ctx.unsolved) + + @property + def solution(self) -> xr.DataArray: + """ + The expression folded over the model's solution, as data. + + Folded afresh on every read, so it follows the model: a body over + data alone reads without a solve at all. + + Raises + ------ + RuntimeError + The body reads a variable the model holds no solution for, or a + constraint it holds no dual for. + SpecDataError + A parameter the body reads was neither retained nor + still reachable through the model's sources. + """ + return fold(self._name, self._ctx) + + def __repr__(self) -> str: + value = self.__dict__.get("expression") + if isinstance(value, xr.DataArray): + return f"NamedExpression('{self._name}', dims={tuple(value.dims)})" + return f"NamedExpression('{self._name}')" diff --git a/linopy/spec/attach.py b/linopy/spec/attach.py new file mode 100644 index 000000000..a0e21833c --- /dev/null +++ b/linopy/spec/attach.py @@ -0,0 +1,866 @@ +""" +Attach user data to a math-spec program. + +The language fixes three attachment rules and this module enforces them: a +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 +is the builder's question, not this module's. +""" + +from __future__ import annotations + +import difflib +import warnings +from collections.abc import Hashable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Literal, NoReturn, get_args + +import numpy as np +import pandas as pd +import xarray as xr +from math_spec import did_you_mean +from math_spec import program as ms + +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"), + "bool": frozenset("b"), + "str": frozenset("OUS"), + "datetime": frozenset("M"), +} +_KIND_NAMES: dict[str, str] = { + "f": "float", + "i": "int", + "u": "int", + "b": "bool", + "O": "str", + "U": "str", + "S": "str", + "M": "datetime", +} +_EMPTY_DTYPES: dict[str, Any] = { + "float": float, + "int": int, + "bool": bool, + "str": object, +} +_SCALARS = (bool, int, float, str, np.number, np.bool_) +_DIMENSION_SHAPES = "a pandas Index, a list, a tuple, a 1-D numpy array, a pandas Series or a 1-D DataArray" +_LOOKUP_SHAPES = "a pandas Series indexed by '{over}', a dict keyed by '{over}' labels, or a 1-D DataArray over '{over}'" +_PARAMETER_SHAPES = ( + "a DataArray over {dims}, a pandas Series whose (Multi)Index levels are {dims}, " + "a DataFrame with columns {columns} or in wide form, a dict keyed by label, or one number" +) + + +def attach( + program: ms.Program, + sources: Mapping[str, Any] | xr.Dataset, + *, + retain: Retain = "report", + given: Mapping[str, pd.Index] = MappingProxyType({}), + strict: bool = False, +) -> Attached: + """ + Attach *sources* to *program*: master coordinates now, parameters on demand. + + Parameters + ---------- + program + The lowered spec. + sources + Data keyed by declared name. Any mapping works; ``keys()`` is + called once, and everything after that is read by key. 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; 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. + strict + Whether a key naming nothing the spec declares is refused. By + default it is ignored and reported by :attr:`Attached.unused`, so + one mapping can feed several specs; a key close to a declared name + is warned about either way. + + Raises + ------ + SpecDataError + A ``retain`` outside its three values, a reached dimension or a + lookup with no source, a duplicated dimension member, a lookup + breaking the rules a map has, a binding whose variable does not fit + its declaration, two claimants labelling one dimension differently, + or, under *strict*, a key naming nothing the spec declares. + TypeError + *sources* offers no ``keys()``. + """ + if retain not in _RETAIN: + raise SpecDataError( + f"retain={retain!r} is not one of {_shown(_RETAIN)}. {did_you_mean(retain, _RETAIN)}" + ) + if isinstance(sources, xr.Dataset): + sources = _dataset_sources(sources) + read = Sources(sources, _attachable(program), strict) + bound = _bindings(program, read) + coords = _master_coords(program, read, bound, given) + lookups = _lookups(program, read, coords) + return Attached(program, coords, lookups, retain, read, bound) + + +@dataclass(frozen=True, eq=False) +class Attached: + """ + A program attached to its data. + + Attributes + ---------- + program + The lowered spec the data is attached to. + coords + Master coordinates by dimension, in source order, each index + named after its dimension. A declared dimension nothing reaches + 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. + retain + 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 + coords: Mapping[str, pd.Index] + lookups: Mapping[str, Mapping[str, xr.DataArray]] + retain: Retain + sources: Sources + bound: Mapping[str, Variable] + + @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()} + + @property + def unused(self) -> frozenset[str]: + """ + The source keys nothing has read. + + A key naming nothing the spec declares, and a declared name whose + data no declaration reaches, both land here. It shrinks as the + parameters are read, so it only tells the whole story once the + build is done. + """ + return self.sources.unused + + def parameter(self, name: str) -> xr.DataArray: + """ + The parameter *name* resolved from ``sources`` and aligned to ``coords``. + + Resolved on every call and never cached. An already aligned array is + returned without a copy; a mismatching one is reindexed onto the + master coordinates, leaving NaN (``False`` for ``bool``) where no row + was supplied. + + Raises + ------ + SpecDataError + No data, a shape no reader accepts, a rank other + than declared, a label its dimension lacks, two rows for one + coordinate, a null value in a row, or values of another type + than declared. + """ + declared = self._declaration(name) + if name not in self.sources: + raise SpecDataError(f"no data provided for parameter '{name}'") + arr = _numpy(_as_array(name, declared, self.sources[name], self.coords)) + onto = {d: self.coords[d] for d in declared.dims} + return _aligned(name, arr, onto, _fill(declared)) + + def retained(self) -> xr.Dataset: + """The lookups plus the parameters ``retain`` keeps, as one dataset.""" + arrays = {n: self.parameter(n) for n in self._retained_names()} + for by_name in self.lookups.values(): + arrays.update(by_name) + return xr.Dataset(arrays) + + def _declaration(self, name: str) -> ms.ParameterDeclaration: + if name not in self.program.parameters: + raise SpecDataError( + f"unknown parameter '{name}'. {did_you_mean(name, self.program.parameters)}" + ) + declared = self.program.parameters[name] + if declared.derivation is not None: + raise SpecDataError( + f"parameter '{name}' is emitted by piecewise block '{declared.derivation.block}' " + f"and is filled from the block's own breakpoints, not attached from sources." + ) + return declared + + def _retained_names(self) -> list[str]: + if self.retain == "none": + return [] + parameters = self.program.parameters + keep = ( + set(parameters) if self.retain == "all" else _report_closure(self.program) + ) + return [n for n, p in parameters.items() if p.derivation is None and n in keep] + + +def _report_closure(program: ms.Program) -> set[str]: + """Every parameter a named expression reads, by node or by name.""" + bodies = tuple(d.expression for d in program.named_expressions.values()) + names = set(parameters_of(*bodies)) + for node in walk(*bodies): + names.update(amounts_of(node)) + if isinstance(node, ms.Cases): + for region in node.regions: + names |= region.when.names_read + return names & set(program.parameters) + + +# --------------------------------------------------------------------------- +# sources and keys +# --------------------------------------------------------------------------- + + +def _dataset_sources(ds: xr.Dataset) -> dict[str, Any]: + sources: dict[str, Any] = {str(d): index for d, index in ds.indexes.items()} + sources.update({str(n): ds[n] for n in ds.data_vars}) + return sources + + +def _attachable(program: ms.Program) -> dict[str, str]: + kinds = { + n: "parameter" for n, p in program.parameters.items() if p.derivation is None + } + 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 + + +class Sources: + """ + The caller's data, read by key and remembering what was read. + + ``keys()`` is called once, to see which declared names have data; + everything after that goes through ``__getitem__``, so a source is + read when a declaration reaches it and never otherwise. + + A key naming nothing the spec declares is ignored, so one mapping can + feed several specs, and :attr:`unused` reports it; under *strict* it is + refused instead. A key close to a declared name is a typo either way, + and is warned about. + """ + + def __init__( + self, mapping: Mapping[str, Any], declared: Mapping[str, str], strict: bool + ) -> None: + self._mapping = mapping + try: + self._keys = frozenset(mapping.keys()) + except AttributeError: + raise TypeError( + f"sources must offer keys(): it is called once, to see which declared names " + f"have data, and everything after that is read by key. " + f"{type(mapping).__name__} offers no keys()." + ) from None + self.used: set[str] = set() + _check_extras(self._keys, declared, strict) + + def __contains__(self, key: str) -> bool: + return key in self._keys + + def __getitem__(self, key: str) -> Any: + self.used.add(key) + return self._mapping[key] + + @property + def unused(self) -> frozenset[str]: + """The keys nothing has read yet.""" + return self._keys - self.used + + +def _check_extras( + keys: frozenset[str], declared: Mapping[str, str], strict: bool +) -> None: + unknown = sorted(keys - set(declared)) + if not unknown: + return + one = len(unknown) == 1 + lead = f"source key {unknown[0]!r} names" if one else f"source keys {unknown} name" + body = ( + f"{lead} neither a parameter, a dimension, a lookup nor a variable this spec " + f"declares." + ) + if strict: + raise SpecDataError( + f"{body} {did_you_mean(unknown[0], declared)} Pass only what the spec takes." + ) + typos = {k: near for k in unknown if (near := _near(k, declared)) is not None} + if typos: + warnings.warn( + f"{body} {_shown(sorted(typos))} read like a typo: " + f"{_shown([f'{k} -> {near}' for k, near in sorted(typos.items())])}. " + f"An unknown key is ignored, so a mistyped one leaves its declaration " + f"without data.", + UserWarning, + stacklevel=4, + ) + + +def _near(key: str, declared: Mapping[str, str]) -> str | None: + """The one declared name *key* is close enough to be a typo of.""" + near = difflib.get_close_matches(key, sorted(declared), n=1, cutoff=0.6) + return near[0] if near else None + + +# --------------------------------------------------------------------------- +# bindings +# --------------------------------------------------------------------------- + + +def _bindings(program: ms.Program, sources: Sources) -> dict[str, Variable]: + bound: dict[str, Variable] = {} + for name in program.variables: + if name not in sources: + 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.domain != kind: + raise SpecDataError( + f"variable '{name}' is declared {declared.domain} 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.domain] + + +# --------------------------------------------------------------------------- +# dimensions +# --------------------------------------------------------------------------- + + +def _reached(program: ms.Program) -> set[str]: + dims: set[str] = set() + for declared in (program.parameters, program.variables, program.constraints): + dims.update(d for decl in declared.values() for d in decl.dims) + dims.update(pw.over for pw in program.piecewise.values()) + for over, lk in program.lookups: + dims.add(over) + dims.add(lk.target) + return dims + + +def _master_coords( + program: ms.Program, + sources: Sources, + 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 sources: + master = _index(dim, sources[dim], program.dimensions[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}, 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, declared: ms.DimensionDeclaration) -> pd.Index: + if isinstance(obj, (pd.Series, xr.DataArray, np.ndarray)): + if obj.ndim != 1: + raise SpecDataError( + f"index for dimension '{dim}' is {obj.ndim}-dimensional; pass {_DIMENSION_SHAPES}." + ) + values: Any = np.asarray(obj) + elif isinstance(obj, (pd.Index, list, tuple)): + values = obj + else: + raise SpecDataError( + f"index for dimension '{dim}': cannot read labels out of {type(obj).__name__}; pass {_DIMENSION_SHAPES}." + ) + flat = values if isinstance(values, pd.Index) else pd.Index(values) + if isinstance(flat, pd.MultiIndex): + raise SpecDataError( + f"index for dimension '{dim}' is a MultiIndex; a spec dimension is one flat axis, " + f"flatten or split it." + ) + index = flat.rename(dim) + _check_value_dtype(dim, declared.dtype, index.dtype, what="dimension") + if index.has_duplicates: + twice = index[index.duplicated()].unique().tolist() + raise SpecDataError( + f"dimension '{dim}' lists {_shown(twice)} more than once. A dimension's members are a set: " + f"each label appears once, in the order the source gives it." + ) + return index + + +# --------------------------------------------------------------------------- +# lookups +# --------------------------------------------------------------------------- + + +def _lookups( + program: ms.Program, + sources: Sources, + coords: Mapping[str, pd.Index], +) -> dict[str, dict[str, xr.DataArray]]: + out: dict[str, dict[str, xr.DataArray]] = {} + for over, lk in program.lookups: + space = lk.target or lk.name + if lk.name not in sources: + raise SpecDataError( + f"no data provided for lookup '{lk.name}'. Pass it under key '{lk.name}' as " + f"{_LOOKUP_SHAPES.format(over=over)}, holding a '{space}' value for each " + f"'{over}' label it maps and nothing for a label it does not." + ) + series = _lookup_series(lk.name, over, sources[lk.name]) + _check_lookup(series, lk, over, coords) + padded = series.reindex(coords[over]) + out.setdefault(over, {})[lk.name] = _numpy(xr.DataArray(padded, name=lk.name)) + return out + + +def _numpy(arr: xr.DataArray) -> xr.DataArray: + """ + *arr* backed by a numpy array. + + xarray keeps a pandas extension array as it arrives, and pandas 3 hands + strings over as one. Its ``dtype`` is no ``np.dtype``, so nothing + downstream that records or restores a dtype can name it, and xarray's + positional indexing refuses the Arrow-backed variant. + """ + if isinstance(arr.dtype, np.dtype): + return arr + return arr.copy(data=arr.to_numpy()) + + +def _lookup_series(name: str, over: str, obj: Any) -> pd.Series: + if isinstance(obj, xr.DataArray): + if obj.dims != (over,) or over not in obj.indexes: + raise SpecDataError( + f"lookup '{name}' arrived as a DataArray over {list(obj.dims)}, and it is a map " + f"out of '{over}': pass a 1-D DataArray with '{over}' as its labelled dimension." + ) + return obj.to_series() + if isinstance(obj, Mapping): + return pd.Series(dict(obj)).rename_axis(over) + if isinstance(obj, pd.Series): + if obj.index.name not in (None, over): + raise SpecDataError( + f"lookup '{name}' is a Series indexed by '{obj.index.name}', and it is a map out of " + f"'{over}': index it by '{over}' labels." + ) + return obj.rename_axis(over) + raise SpecDataError( + f"lookup '{name}': cannot adapt {type(obj).__name__} to a map; pass {_LOOKUP_SHAPES.format(over=over)}." + ) + + +def _check_lookup( + series: pd.Series, + lk: ms.LookupDeclaration, + over: str, + coords: Mapping[str, pd.Index], +) -> None: + space = lk.target or lk.name + holes = series.isna() + if holes.any(): + at = _coordinates_shown((over,), series.index[holes][:5]) + raise SpecDataError( + f"lookup '{lk.name}' carries {int(holes.sum())} row(s) with a null in '{space}': {at}. A map is " + f"partial by leaving a label out, not by mapping it to nothing: drop the row and the " + f"label is unmapped, which is what every operator reading the lookup already means by it." + ) + if series.index.has_duplicates: + twice = series.index[series.index.duplicated()].unique().tolist() + raise SpecDataError( + f"lookup '{lk.name}' maps {len(twice)} '{over}' label(s) more than once: {_shown(twice)}. " + f"A lookup is single-valued, so each label it maps takes exactly one row." + ) + strays = series.index[~series.index.isin(coords[over])].tolist() + if strays: + raise SpecDataError( + f"lookup '{lk.name}' maps {_shown(strays)}, which are not labels of '{over}'. " + f"'{over}' takes its labels from sources['{over}'], and they are " + f"{_shown(coords[over].tolist(), 8)}. A map maps the labels that exist: a key matching " + f"none of them would place its terms nowhere, so it is a typo on one side or a label " + f"missing from the other." + ) + values = pd.Index(series.to_numpy()) + foreign = values[~values.isin(coords[lk.target])].unique().tolist() + if foreign: + raise SpecDataError( + f"dimension '{over}' lookup '{lk.name}' has value(s) that are not '{lk.target}' labels: " + f"{_shown(foreign)}. Every value must be a declared '{lk.target}' label, otherwise " + f"sum(by={lk.name}) drops those terms in the join that places them, and the model " + f"builds and solves without them." + ) + + +# --------------------------------------------------------------------------- +# parameters +# --------------------------------------------------------------------------- + + +def _fill(declared: ms.ParameterDeclaration) -> Any: + return False if declared.dtype == "bool" else np.nan + + +def _as_array( + name: str, + declared: ms.ParameterDeclaration, + obj: Any, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + if isinstance(obj, xr.DataArray): + return _from_dense(name, declared, obj) + if isinstance(obj, pd.DataFrame): + return _from_frame(name, declared, obj, coords) + if isinstance(obj, pd.Series): + return _from_rows(name, declared, obj, coords) + if isinstance(obj, Mapping): + return _from_rows(name, declared, pd.Series(dict(obj)), coords) + if isinstance(obj, _SCALARS): + return _from_scalar(name, declared, obj, coords) + dims = declared.dims + raise SpecDataError( + f"parameter '{name}': cannot adapt {type(obj).__name__} to an array over {list(dims)}; " + f"pass {_parameter_shapes(dims)}." + ) + + +def _parameter_shapes(dims: Sequence[str]) -> str: + return _PARAMETER_SHAPES.format(dims=list(dims), columns=[*dims, "value"]) + + +def _from_scalar( + name: str, + declared: ms.ParameterDeclaration, + obj: Any, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + if pd.isna(obj): + raise SpecDataError( + f"parameter '{name}' is one value and that value is a hole (null or NaN). " + f"A number was meant, or the parameter has no data and should not be passed." + ) + value = np.asarray(obj) + _check_value_dtype(name, declared.dtype, value.dtype) + if declared.dtype == "float": + value = value.astype(float) + arr = xr.DataArray(value, name=name) + if declared.dims: + arr = arr.expand_dims({d: coords[d] for d in declared.dims}) + return arr + + +def _from_dense( + name: str, declared: ms.ParameterDeclaration, arr: xr.DataArray +) -> xr.DataArray: + dims = declared.dims + _check_value_dtype(name, declared.dtype, arr.dtype) + if set(arr.dims) != set(dims) or len(arr.dims) != len(dims): + raise SpecDataError( + f"parameter '{name}' arrived as a DataArray over {list(arr.dims)}, and '{name}' is over " + f"{list(dims)}. The dims must be the declared ones, in any order." + ) + for d in dims: + if d not in arr.indexes: + raise SpecDataError( + f"parameter '{name}' has no coordinate labels along '{d}'. A parameter is read for " + f"values against its labels, so every dimension needs an index coordinate." + ) + _refuse_duplicate_coordinates(name, (d,), arr.indexes[d]) + return arr.transpose(*dims) + + +def _from_frame( + name: str, + declared: ms.ParameterDeclaration, + df: pd.DataFrame, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + dims = declared.dims + tidy = df + if not set(dims) <= set(df.columns) and set(dims) <= _headers(df): + tidy = df.reset_index() + if "value" in tidy.columns and set(dims) <= set(tidy.columns): + indexed = tidy.set_index(list(dims)) if dims else tidy + return _from_rows(name, declared, indexed["value"], coords) + if len(dims) == 2: + return _from_dense(name, declared, xr.DataArray(_wide(name, dims, df))) + raise SpecDataError( + f"parameter '{name}' arrived as a DataFrame with columns {list(df.columns)}; a table for " + f"'{name}' carries columns {[*dims, 'value']}." + ) + + +def _headers(df: pd.DataFrame) -> set[Any]: + return set(df.columns) | set(df.index.names) + + +def _wide(name: str, dims: tuple[str, ...], df: pd.DataFrame) -> pd.DataFrame: + names = (df.index.name, df.columns.name) + if names == (None, None): + if len(df.index) == len(df.columns): + raise SpecDataError( + f"parameter '{name}' arrived as a {len(df.index)}x{len(df.columns)} wide DataFrame " + f"with neither axis named, and '{name}' is over {list(dims)}. A square frame does " + f"not say which axis is which. Name the index and columns after the two dims, or " + f"pass a table with columns {[*dims, 'value']}." + ) + return df.rename_axis(index=dims[0], columns=dims[1]) + if set(names) == set(dims): + return df + raise SpecDataError( + f"parameter '{name}' arrived as a wide DataFrame with index '{names[0]}' and columns " + f"'{names[1]}', and '{name}' is over {list(dims)}. Name the index and columns after the " + f"two dims, or pass a table with columns {[*dims, 'value']}." + ) + + +def _from_rows( + name: str, + declared: ms.ParameterDeclaration, + series: pd.Series, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + dims = declared.dims + if not dims: + if len(series) != 1: + raise SpecDataError( + f"parameter '{name}' is declared with no dims, which means one value broadcast " + f"everywhere, but its source has {len(series)} rows. Declare the dims it is indexed " + f"by, or pass one number." + ) + return _from_scalar(name, declared, series.iloc[0], coords) + series = _with_dims(name, dims, series) + if series.empty: + series = series.astype(_EMPTY_DTYPES[declared.dtype]) + holes = series.isna() + if holes.any(): + raise SpecDataError( + f"parameter '{name}' carries {int(holes.sum())} row(s) with no value, null or NaN: " + f"{_coordinates_shown(dims, series.index[holes][:3])}. In a table the absence of a " + f"value is the absence of the row, and such a row says the coordinate exists and denies " + f"it in the same breath. Drop those rows, or supply the values." + ) + _check_value_dtype(name, declared.dtype, series.dtype) + _refuse_duplicate_coordinates(name, dims, series.index) + for d in dims: + _refuse_strangers(name, d, series.index.get_level_values(d), coords[d]) + onto = [coords[d] for d in dims] + full = onto[0] if len(dims) == 1 else pd.MultiIndex.from_product(onto, names=dims) + values = series.reindex(full, fill_value=_fill(declared)).to_numpy() + dense = values.reshape(tuple(len(index) for index in onto)) + return xr.DataArray(dense, dims=dims, coords=dict(zip(dims, onto)), name=name) + + +def _with_dims(name: str, dims: tuple[str, ...], series: pd.Series) -> pd.Series: + index = series.index + if index.nlevels != len(dims): + raise SpecDataError( + f"parameter '{name}': a Series or dict carries one label per level, and its index has " + f"{index.nlevels} level(s) where '{name}' is over {list(dims)}. " + f"Pass {_parameter_shapes(dims)}." + ) + names = list(index.names) + if all(n is None for n in names): + return series.set_axis(index.set_names(list(dims))) + if set(names) != set(dims): + raise SpecDataError( + f"parameter '{name}' is indexed by {names}, and '{name}' is over {list(dims)}. " + f"Name the index levels after the declared dims." + ) + if tuple(names) != dims: + series = series.reorder_levels(list(dims)) + return series + + +def _refuse_duplicate_coordinates( + name: str, dims: tuple[str, ...], index: pd.Index +) -> None: + duplicated = index.duplicated() + if not duplicated.any(): + return + counts = index[duplicated].value_counts() + shown = "; ".join( + f"{_coordinate(dims, key)} ({n + 1} rows)" for key, n in counts.iloc[:3].items() + ) + raise SpecDataError( + f"parameter '{name}' has more than one row for a coordinate: {shown}. A parameter is a " + f"function of its dims, so which value applies is undefined; aggregate the source to one " + f"row per {list(dims)} before attaching it." + ) + + +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"{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}']." + ) + + +def _aligned( + name: str, arr: xr.DataArray, onto: Mapping[str, pd.Index], fill: Any +) -> xr.DataArray: + if all(arr.indexes[d].equals(index) for d, index in onto.items()): + stale = {d: i for d, i in onto.items() if arr.indexes[d].dtype != i.dtype} + return arr.assign_coords(stale) if stale else arr + for d, index in onto.items(): + _refuse_strangers(name, d, arr.indexes[d], index) + return arr.reindex(onto, fill_value=fill) + + +def _check_value_dtype( + name: str, declared: str, dtype: Any, what: str = "parameter" +) -> None: + if str(dtype.kind) in _ACCEPTED_KINDS[declared]: + return + arrived = _KIND_NAMES.get(str(dtype.kind), str(dtype)) + raise SpecDataError( + f"{what} '{name}' is declared '{declared}' and its values arrived as '{arrived}'. " + f"A declared dtype is a claim about the values, and it is checked here: the file says what " + f"the values are, or the values are not attached.\n" + f" Cast the values to {declared}, if the declaration is what you meant\n" + f" Or declare what the data has: {{dtype: {arrived}}}" + ) + + +# --------------------------------------------------------------------------- +# wording +# --------------------------------------------------------------------------- + + +def _shown(labels: Sequence[Any], limit: int = 5) -> str: + head = ", ".join(repr(x) for x in labels[:limit]) + return head + (f" (and {len(labels) - limit} more)" if len(labels) > limit else "") + + +def _coordinate(dims: Sequence[str], key: Hashable) -> str: + row = key if isinstance(key, tuple) else (key,) + return ", ".join(f"{d}={v!r}" for d, v in zip(dims, row)) + + +def _coordinates_shown(dims: Sequence[str], rows: Iterable[Hashable]) -> str: + return "; ".join(_coordinate(dims, row) for row in rows) diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py new file mode 100644 index 000000000..a6b76584c --- /dev/null +++ b/linopy/spec/builder.py @@ -0,0 +1,257 @@ +""" +Program plus attached data to linopy declarations. + +A build hands every variable to linopy as its term, then adds special-ordered +sets, constraints, the objective and the named expressions that carry a +variable term; which linopy call each construct becomes is one branch of +:func:`linopy.spec.evaluate.evaluate`. Everything built is stamped with the +layer's name, a bound variable excepted: the layer reads it, the model owns it. +""" + +from __future__ import annotations + +import warnings + +import xarray as xr +from math_spec import program as ms + +from linopy.expressions import LinearExpression, QuadraticExpression +from linopy.model import Model +from linopy.spec import curves +from linopy.spec.attach import Attached, _coordinates_shown +from linopy.spec.context import Context +from linopy.spec.coverage import check_bounds_cover, check_coverage +from linopy.spec.errors import SpecDataError +from linopy.spec.evaluate import carried, evaluate +from linopy.spec.nodes import walk +from linopy.spec.parameters import Parameters +from linopy.spec.terms import Term, Value, live_rows +from linopy.spec.where import as_linopy_mask, evaluate_where +from linopy.variables import Variable + +_SIGN = {"==": "=", "<=": "<=", ">=": ">="} +_FLIPPED = {"==": "==", "<=": ">=", ">=": "<="} +_SENSE = {"minimize": "min", "maximize": "max"} + + +def build( + model: Model, attached: Attached, layer: str, build_expressions: bool = True +) -> None: + """ + Add every declaration of the attached program to *model* as the layer *layer*. + + Variables, special-ordered sets, constraints, the objective and, with + *build_expressions*, the named expressions holding a variable term, in + that order. Every named expression is checked for divisor and + coefficient coverage either way, so a body that cannot be folded is + refused at build rather than at read. + """ + check_supported(attached.program) + ctx = Context( + model, + attached.program, + attached.coords, + attached.lookups, + Parameters(attached.program, attached.parameter), + layer, + names=attached.names, + ) + curves.validate(ctx.program, ctx.parameters) + _variables(ctx) + _sos(ctx) + _constraints(ctx) + _objective(ctx) + _expressions(ctx, build_expressions) + + +def check_supported(program: ms.Program) -> None: + """ + Refuse the constructs of *program* linopy cannot build, before any of it is built. + + A product of two variable-carrying operands is a quadratic term, and + linopy carries one in the objective only: a constraint holding one has no + linopy form to be built into. + """ + if "constraint" in program.footprint.quadratic: + raise NotImplementedError( + "a constraint of the spec multiplies two variable-carrying operands, and linopy " + "carries a quadratic term in the objective only. Move the product into the " + "objective, or write the constraint so that at most one side of each product " + "holds a variable." + ) + + +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)) + variable = ctx.model.add_variables( + lower=_bound(declared.lower, ctx), + upper=_bound(declared.upper, ctx), + coords={d: ctx.coords[d] for d in declared.dims}, + name=name, + mask=as_linopy_mask(rows), + binary=declared.domain == "binary", + integer=declared.domain == "integer", + ) + variable.spec = ctx.layer + + +def _bound(node: ms.ExpressionNode, ctx: Context) -> float | xr.DataArray: + """A bound as linopy takes it, read raw: an uncovered slot stays NaN for :func:`check_bounds_cover`.""" + if isinstance(node, ms.Constant): + return node.value + if isinstance(node, ms.Parameter): + return ctx.parameters[node.name] + raise TypeError(f"a bound is a number or a parameter, not {type(node).__name__}") + + +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[ctx.names.get(sos.variable, sos.variable)], + sos_type=sos.sos_type, + sos_dim=sos.over, + big_m=sos.big_m, + ) + + +def _constraints(ctx: Context) -> None: + for name, row in ctx.program.constraints.items(): + rows = evaluate_where(row.where, ctx) + mask = as_linopy_mask(rows) + check_coverage( + f"constraint '{name}'", (row.lhs, row.rhs), ctx, mask, comparison=True + ) + lhs, rhs = evaluate(row.lhs, ctx), evaluate(row.rhs, ctx) + if _term_free(lhs) and _term_free(rhs): + continue + term, other, sense = _sides(lhs, rhs, row.sense) + _check_live(name, row, term, other, rows, ctx) + if isinstance(other, xr.DataArray): + term, other = carried(term, other) + built = ctx.model.add_constraints( + term, _SIGN[sense], other, name=name, mask=mask + ) + built.spec = ctx.layer + + +def _check_live( + name: str, + declared: ms.ConstraintDeclaration, + term: Term, + other: Value, + rows: xr.DataArray, + ctx: Context, +) -> None: + """ + Refuse a row the data emptied of every variable term. + + Such a row reads as ``0 sense rhs``: linopy carries no column there, the + row leaves the problem and the constraint silently stops binding. Where + every variable of the row declares ``absence: zero`` the zero is what the + math says, so a zero other side is warned about rather than refused. + """ + dead = rows & ~live_rows(term) + if not bool(dead.any()): + return + said = ( + f"constraint '{name}': {int(dead.sum())} row(s) hold no variable term once the data " + f"is attached, the first at {_dead_at(dead)}. Nothing is left to constrain there, so " + f"the row leaves the problem without saying so." + ) + zeroed = all( + ctx.program.variable(v).absence == "zero" + for v in ms.variables_of(declared.lhs, declared.rhs) + ) + if zeroed and not _binds(other, dead): + warnings.warn( + f"{said} Every variable there is absence: zero and the other side is 0, " + f"so the row is trivially true.", + UserWarning, + stacklevel=2, + ) + return + supply = ( + " Supply the rows of the variables, if the row is meant to bind." + if zeroed + else " Declare absence: zero on the variables, if an absent term is a zero there." + ) + raise SpecDataError( + f"{said}\n" + f" Mask them out with a where on the constraint, if the row should not exist there.\n" + f"{supply}" + ) + + +def _binds(other: Value, dead: xr.DataArray) -> bool: + """Whether the side without the variable term is anything but 0 on a row *dead* names.""" + if isinstance(other, xr.DataArray): + return bool((other.where(dead, 0.0) != 0).any()) + return other != 0 + + +def _dead_at(dead: xr.DataArray) -> str: + """The first coordinates *dead* marks, spelled the way every other refusal spells them.""" + dims = tuple(str(d) for d in dead.dims) + if not dims: + return "the only row" + stacked = dead.stack(_dead=dims) + return _coordinates_shown(dims, stacked.indexes["_dead"][stacked.values][:3]) + + +def _sides(lhs: Value, rhs: Value, sense: str) -> tuple[Term, Value, str]: + """The comparison with a term on the left, as linopy takes it; a swap flips the sense.""" + if isinstance(lhs, Variable | LinearExpression | QuadraticExpression): + return lhs, rhs, sense + if isinstance(rhs, Variable | LinearExpression | QuadraticExpression): + return rhs, lhs, _FLIPPED[sense] + raise TypeError("a constraint needs a variable term on one side") + + +def _term_free(side: Value) -> bool: + """Whether *side* has nowhere for a variable term to sit: data, or an expression the data emptied.""" + if isinstance(side, Variable): + return False + if isinstance(side, LinearExpression | QuadraticExpression): + return side.nterm == 0 + return True + + +def _objective(ctx: Context) -> None: + declared = ctx.program.objective + if declared is None: + return + check_coverage("the objective", (declared.expression,), ctx, None) + expr = evaluate(declared.expression, ctx) + if not isinstance(expr, Variable | LinearExpression | QuadraticExpression): + raise SpecDataError( + "the objective carries no variable term once the data is attached, so there is nothing to optimize" + ) + ctx.model.add_objective(expr, overwrite=True, sense=_SENSE[declared.sense]) + + +def _expressions(ctx: Context, build: bool) -> None: + """ + Every named expression coverage-checked; with *build*, the ones holding a variable term added to the model. + + A data-only body has no linopy term to hold and stays on the spec; so + does one reading a ``dual``, which needs a solved model. + """ + for name, declared in ctx.program.named_expressions.items(): + body = declared.expression + check_coverage(f"expression '{name}'", (body,), ctx, None) + if not build or any(isinstance(n, ms.Dual) for n in walk(body)): + continue + value = evaluate(body, ctx) + if isinstance(value, Variable | LinearExpression | QuadraticExpression): + ctx.model.add_expressions(value, name=name).spec = ctx.layer diff --git a/linopy/spec/context.py b/linopy/spec/context.py new file mode 100644 index 000000000..ecb4f7bb2 --- /dev/null +++ b/linopy/spec/context.py @@ -0,0 +1,80 @@ +"""The data an evaluation reads: the parameters, and the model, coordinates and lookups beside them.""" + +from __future__ import annotations + +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) +class Context: + """ + Everything evaluating a node needs beyond the node. + + ``layer`` is the name of the layer the program is attached as, the stamp + everything a build adds to the model carries. ``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. ``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 + program: ms.Program + coords: Mapping[str, pd.Index] + lookups: Mapping[str, Mapping[str, xr.DataArray]] + parameters: Mapping[str, xr.DataArray] + layer: str + 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 new file mode 100644 index 000000000..1c9b6c359 --- /dev/null +++ b/linopy/spec/coverage.py @@ -0,0 +1,195 @@ +""" +Is the data there where a declaration needs it? Every position asks. + +A parameter row that no source supplies is a hole, and the spec refuses it +wherever the row is used: as a coefficient, where the missing row would +silently drop its term; as a bound, where zero is a bound rather than the +absence of one; as a constant side, where it binds; and as a divisor, where +zero is not a divisor at all. Each is decided against the rows the declaration +actually builds, so a ``where`` that removed the coordinate has already +answered. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field + +import xarray as xr +from math_spec import program as ms + +from linopy.spec import terms +from linopy.spec.context import Context +from linopy.spec.errors import SpecDataError +from linopy.spec.nodes import amounts_of, parameters_of +from linopy.spec.where import evaluate_where + +Rows = xr.DataArray | None +Obligation = tuple[str, Rows] + + +@dataclass +class Obligations: + """Every parameter use under a declaration, each with the rows it has to cover, gathered in one walk.""" + + divisors: list[Obligation] = field(default_factory=list) + constants: list[Obligation] = field(default_factory=list) + coefficients: list[Obligation] = field(default_factory=list) + + +def gaps_under(array: xr.DataArray, rows: Rows) -> int: + """How many slots of *array* are null where *rows* still admits the row; ``None`` narrows nothing.""" + missing = array.isnull() + if rows is not None: + missing = missing & rows + return int(missing.sum()) + + +def check_coverage( + subject: str, + expressions: Sequence[ms.ExpressionNode], + ctx: Context, + rows: Rows, + *, + comparison: bool = False, +) -> None: + """ + Refuse *subject* if a parameter it reads leaves a row it builds uncovered. + + One walk over *expressions* collects what every parameter has to cover, + narrowed at each ``cases:`` region; divisors are judged first, then, for a + *comparison*, the side without a variable term, then every coefficient. + """ + found = obligations_of(expressions, ctx, rows, comparison=comparison) + check_divisors(subject, found.divisors, ctx) + check_constant_sides(subject, found.constants, ctx) + check_coefficients(subject, found.coefficients, ctx) + + +def obligations_of( + expressions: Sequence[ms.ExpressionNode], + ctx: Context, + rows: Rows, + *, + comparison: bool = False, +) -> Obligations: + """What the parameters under *expressions* have to cover, a side of a *comparison* without a variable being its constant side.""" + found = Obligations() + for expression in expressions: + constant = comparison and not ms.carries_variable(expression) + _collect(expression, ctx, rows, constant, found) + return found + + +def _collect( + node: ms.ExpressionNode, + ctx: Context, + rows: Rows, + constant: bool, + into: Obligations, +) -> None: + if isinstance(node, ms.Multiply): + rows = _where_present(rows, ms.variables_of(node), ctx) + if isinstance(node, ms.Divide): + into.divisors.extend(_divisor_uses(node, ctx, rows)) + if isinstance(node, ms.Parameter): + if constant: + into.constants.append((node.name, rows)) + into.coefficients.append((node.name, rows)) + into.coefficients.extend((name, None) for name in amounts_of(node)) + if isinstance(node, ms.Cases): + for region in node.regions: + inside = evaluate_where(region.when, ctx) + narrowed = inside if rows is None else rows & inside + _collect(region.value, ctx, narrowed, constant, into) + return + for child in ms.children(node): + _collect(child, ctx, rows, constant, into) + + +def _divisor_uses(quotient: ms.Divide, ctx: Context, rows: Rows) -> list[Obligation]: + """Each parameter in the divisor, with the rows the quotient is divided over: the region, narrowed by the presence of every numerator variable.""" + params = parameters_of(quotient.divisor) + if not params: + return [] + needed = _where_present(rows, ms.variables_of(quotient.numerator), ctx) + return [(param, needed) for param in sorted(params)] + + +def _where_present(rows: Rows, variables: Iterable[str], ctx: Context) -> Rows: + """*rows* narrowed to the coordinates every one of *variables* occupies, a term absent there carrying no parameter with it.""" + for variable in sorted(variables): + present = terms.present(ctx.variable(variable)) + rows = present if rows is None else rows & present + return rows + + +def check_divisors(subject: str, found: Sequence[Obligation], ctx: Context) -> None: + """ + A divisor must have a value wherever *subject* divides by it. + + Reached before evaluation, the last moment the gap is visible: the + coefficient fill would turn it into a division by zero. + """ + for param, needed in found: + missing = gaps_under(ctx.parameters[param], needed) + if missing: + raise SpecDataError( + f"{subject}: parameter '{param}' is used as a divisor but covers {missing} " + f"fewer coordinates than it is divided over. A missing row means a zero " + f"coefficient everywhere else, and zero is not a divisor: the term would drop " + f"and the row would silently stop constraining.\n" + f" Supply the missing rows, or mask the coordinates out with a where." + ) + + +def check_constant_sides( + subject: str, found: Sequence[Obligation], ctx: Context +) -> None: + """A comparison's constant side must have values wherever the row is built, or the zero is the bound.""" + for param, needed in sorted(found, key=lambda pair: pair[0]): + missing = gaps_under(ctx.parameters[param], needed) + if missing: + raise SpecDataError( + f"{subject}: parameter '{param}' covers {missing} fewer coordinates " + f"than the rows built here. A missing row is read as 0, and on the constant side " + f"that zero is a bound rather than an absence: the row still exists, and it binds.\n" + f" Supply the missing rows, if the value is what was meant.\n" + f" Mask them out with a where, if the row should not exist there." + ) + + +def check_coefficients(subject: str, found: Sequence[Obligation], ctx: Context) -> None: + """ + A coefficient parameter must reach every row it is built over. + + A missing coefficient row would otherwise read as a zero, dropping its term + while the row stays. A shift offset or window width given by name is a + coefficient too, and stands or falls over its own coordinates. + """ + for param, needed in found: + missing = gaps_under(ctx.parameters[param], needed) + if missing: + raise SpecDataError( + f"{subject}: parameter '{param}' is used as a coefficient but leaves " + f"{missing} of the rows built here uncovered. A missing row reads as a zero " + f"coefficient, dropping the term while the row stays.\n" + f" Supply the missing rows, if a value other than 0 was meant.\n" + f" Mask them out with a where, if the row should not exist there." + ) + + +def check_bounds_cover( + name: str, declared: ms.VariableDeclaration, ctx: Context, rows: Rows +) -> None: + """A bound parameter must have a value at every coordinate the variable occupies.""" + names = sorted(parameters_of(declared.lower, declared.upper)) + missing = sum(gaps_under(ctx.parameters[p], rows) for p in names) + if missing: + raise SpecDataError( + f"variable '{name}': {missing} rows have NULL bounds, a bound parameter is missing " + f"values for some coordinates. The two ways out build different models, so neither " + f"is picked:\n" + f" supply the value the variable exists there, bounded (`inf` is a value)\n" + f' where: "" the variable does not exist there at all' + ) diff --git a/linopy/spec/curves.py b/linopy/spec/curves.py new file mode 100644 index 000000000..ae9c8428a --- /dev/null +++ b/linopy/spec/curves.py @@ -0,0 +1,184 @@ +""" +The data-time side of a ``piecewise:`` block. + +The language decides a curve's shape and can decide nothing about its +numbers. This module fills the parameters an expansion emitted from the +block's own breakpoints, and checks that the numbers hold what the block's +method rests on: the conditions are the program's :data:`~math_spec.program.Check` +values and :func:`~math_spec.program.check_message` words each refusal. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TypeVar + +import numpy as np +import xarray as xr +from math_spec import program as ms + +from linopy.spec.errors import SpecDataError + +_C = TypeVar("_C", bound=ms.Check) + + +def derive( + derivation: ms.Derivation, + parameters: Mapping[str, xr.DataArray], + program: ms.Program, +) -> xr.DataArray: + """ + An emitted ``bool`` parameter, built from the parameters it hangs off. + + A :class:`~math_spec.program.MaskOf` is true wherever the nominated + breakpoints have a row; :class:`~math_spec.program.FirstOf` and + :class:`~math_spec.program.LastOf` mark, per curve, the first and last + breakpoint the mask admits. + """ + if isinstance(derivation, ms.MaskOf): + return parameters[derivation.values].notnull() + mask = parameters[derivation.mask] + over = program.piecewise[derivation.block].over + ordinal = xr.DataArray(np.arange(mask.sizes[over]), dims=[over]) + if isinstance(derivation, ms.FirstOf): + edge = ordinal.where(mask, np.inf).min(over) + else: + edge = ordinal.where(mask, -np.inf).max(over) + return (mask & (ordinal == edge)).transpose(*mask.dims) + + +def validate(program: ms.Program, parameters: Mapping[str, xr.DataArray]) -> None: + """ + Refuse curves the data does not supply everywhere they are built, or that bend against their method. + + Raises + ------ + SpecDataError + A breakpoint parameter with a hole where the block + builds a weight, a ``points:`` mask that is not one run per curve, + breakpoints that do not increase, a one-point curve under + ``method: lp``, or a curve of the curvature the method is not + exact for. + """ + for block, decl in program.piecewise.items(): + run = _one(decl.checks, ms.Contiguous) + mask = None + if run is not None: + mask = parameters[run.mask] + _check_one_run(block, decl, run, mask) + for values in decl.breakpoints: + _check_extent(block, values, parameters[values], mask, run) + curved = _one(decl.checks, ms.Curved) + if curved is not None: + _check_curves(block, decl, curved, parameters, mask) + + +def _one(checks: tuple[ms.Check, ...], kind: type[_C]) -> _C | None: + return next((check for check in checks if isinstance(check, kind)), None) + + +def _check_extent( + block: str, + name: str, + values: xr.DataArray, + mask: xr.DataArray | None, + run: ms.Contiguous | None, +) -> None: + needed = ( + xr.ones_like(values, dtype=bool) + if mask is None + else mask.any([d for d in mask.dims if d not in values.dims]) + ) + holes = needed & values.isnull() + if not bool(holes.any()): + return + points = None if run is None else (run.values or run.mask) + remedy = ( + f" Shorten it '{points}' claims this breakpoint, so either it is one row too long " + f"or the value is missing\n" + f" Or supply it a value everywhere the mask says the curve runs" + if points + else ( + " Say how far points: a mask over the curve, true up to each one's last " + "breakpoint\n" + " Or supply it a value at every coordinate of the axis" + ) + ) + raise SpecDataError( + f"piecewise '{block}': parameter '{name}' has no value at ({_first(holes)}), and every " + f"breakpoint the block builds gets a weight, so a missing row is not a shorter " + f"curve: read as a zero coefficient it is a breakpoint at the origin.\n{remedy}" + ) + + +def _check_one_run( + block: str, decl: ms.PiecewiseDeclaration, run: ms.Contiguous, mask: xr.DataArray +) -> None: + over = decl.over + ordinal = xr.DataArray(np.arange(mask.sizes[over]), dims=[over]) + marked = mask.sum(over) + span = ( + ordinal.where(mask, -np.inf).max(over) + - ordinal.where(mask, np.inf).min(over) + + 1 + ) + broken = (marked == 0) | (span != marked) + if not bool(broken.any()): + return + message = ms.check_message(block, decl, run) + if not broken.dims: + raise SpecDataError(message) + raise SpecDataError(f"{message}\n Not so at {_first(broken)}") + + +def _first(flags: xr.DataArray) -> str: + """The first coordinate *flags* is true at, written as the reader would look for it.""" + stacked = flags.stack(_at=flags.dims) + at = stacked["_at"].to_index()[stacked.to_numpy()].tolist()[0] + return ", ".join(f"{d}={v!r}" for d, v in zip(flags.dims, at)) + + +def _check_curves( + block: str, + decl: ms.PiecewiseDeclaration, + curved: ms.Curved, + parameters: Mapping[str, xr.DataArray], + mask: xr.DataArray | None, +) -> None: + over = decl.over + xs, ys = xr.broadcast(parameters[curved.x], parameters[curved.y]) + on_curve = xs.notnull() & ys.notnull() + if mask is not None: + on_curve = on_curve & mask + xs, ys, on_curve = xr.broadcast(xs, ys, on_curve) + frame = [d for d in xs.dims if d != over] + x = xs.transpose(*frame, over).to_numpy().reshape(-1, xs.sizes[over]) + y = ys.transpose(*frame, over).to_numpy().reshape(-1, xs.sizes[over]) + keep = on_curve.transpose(*frame, over).to_numpy().reshape(-1, xs.sizes[over]) + increasing = _one(decl.checks, ms.Increasing) + segment = _one(decl.checks, ms.AtLeastTwo) + for row_x, row_y, row_keep in zip(x, y, keep): + px, py = row_x[row_keep].astype(float), row_y[row_keep].astype(float) + if segment is not None and px.size < 2: + raise SpecDataError( + f"{ms.check_message(block, decl, segment)}\n This curve carries {px.size}" + ) + dx = np.diff(px) + if increasing is not None and not bool((dx > 0).all()): + raise SpecDataError( + f"{ms.check_message(block, decl, increasing)} (got {px.tolist()})" + ) + if _bends_wrong(dx, np.diff(py), curved.curvature): + raise SpecDataError( + f"{ms.check_message(block, decl, curved)} (got {py.tolist()})" + ) + + +def _bends_wrong(dx: np.ndarray, dy: np.ndarray, curvature: str) -> bool: + slopes = dy / dx + bend = np.diff(slopes) + tol = 1e-9 * float(np.abs(slopes).max(initial=0.0)) + rises, falls = bool((bend > tol).any()), bool((bend < -tol).any()) + if curvature == "either": + return rises and falls + return falls if curvature == "convex" else rises diff --git a/linopy/spec/errors.py b/linopy/spec/errors.py new file mode 100644 index 000000000..88f9fa577 --- /dev/null +++ b/linopy/spec/errors.py @@ -0,0 +1,12 @@ +"""Errors raised while attaching data to a math-spec program.""" + +from __future__ import annotations + + +class SpecDataError(ValueError): + """ + Data attached to a valid spec is missing, malformed or the wrong shape. + + Every refusal names the symbol, the dimension(s) and the offending labels, + so the message points back at the ``sources`` entry to fix. + """ diff --git a/linopy/spec/evaluate.py b/linopy/spec/evaluate.py new file mode 100644 index 000000000..523bd5738 --- /dev/null +++ b/linopy/spec/evaluate.py @@ -0,0 +1,214 @@ +"""The recursive evaluator: one expression node to its linopy term, array or number.""" + +from __future__ import annotations + +import functools +import operator +from collections.abc import Callable +from typing import assert_never + +import xarray as xr +from math_spec import did_you_mean +from math_spec import program as ms + +from linopy.expressions import LinearExpression, QuadraticExpression +from linopy.spec import operators, terms +from linopy.spec.context import Context +from linopy.spec.coverage import check_divisors, obligations_of +from linopy.spec.errors import SpecDataError +from linopy.spec.terms import Array, Term, Value +from linopy.spec.where import evaluate_where +from linopy.variables import Variable + + +def evaluate_named(name: str, ctx: Context) -> Value: + """The named expression *name* as its linopy term, array or number over *ctx*, its divisors checked first.""" + if name not in ctx.program.named_expressions: + raise KeyError( + f"unknown named expression '{name}'. " + + did_you_mean(name, ctx.program.named_expressions) + ) + body = ctx.program.named_expressions[name].expression + found = obligations_of((body,), ctx, None) + check_divisors(f"expression '{name}'", found.divisors, ctx) + value = evaluate(body, ctx) + return _named(value, name) if isinstance(value, xr.DataArray) else value + + +def fold(name: str, ctx: Context) -> xr.DataArray: + """The named expression *name* as data, folded over the solution and the parameters *ctx* holds.""" + value = evaluate_named(name, ctx) + if isinstance(value, xr.DataArray): + return value + if isinstance(value, float | int): + return xr.DataArray(float(value), name=name) + raise TypeError( + f"expression '{name}' folded to a {type(value).__name__}, not to data" + ) + + +def _named(value: xr.DataArray, name: str) -> xr.DataArray: + """*value* with its stray non-dimension coordinates dropped and renamed to *name*.""" + stray = [c for c in value.coords if c not in value.dims] + return value.drop_vars(stray).rename(name) + + +def evaluate(node: ms.ExpressionNode, ctx: Context) -> Value: + """One node as a linopy term, an array or a number.""" + if isinstance(node, ms.Constant): + return node.value + if isinstance(node, ms.Variable): + return _variable(node.name, ctx) + if isinstance(node, ms.Dual): + return _dual(node.constraint, ctx) + if isinstance(node, ms.Parameter): + return terms.coefficient(ctx.parameters[node.name]) + if isinstance(node, ms.Negate): + return -evaluate(node.operand, ctx) + if isinstance(node, ms.Add): + return _combine( + operator.add, evaluate(node.left, ctx), evaluate(node.right, ctx) + ) + if isinstance(node, ms.Multiply): + return _combine( + operator.mul, evaluate(node.left, ctx), evaluate(node.right, ctx) + ) + if isinstance(node, ms.Divide): + return _combine( + operator.truediv, evaluate(node.numerator, ctx), evaluate(node.divisor, ctx) + ) + if isinstance(node, ms.Power): + return _combine( + operator.pow, evaluate(node.base, ctx), evaluate(node.exponent, ctx) + ) + if isinstance(node, ms.Sum): + summed = _array(evaluate(node.operand, ctx)) + for dimension in node.over: + summed = operators.sum_over(summed, dimension) + return summed + if isinstance(node, ms.GroupSum): + return operators.grouped_sum( + _array(evaluate(node.operand, ctx)), + _lookup_arrays(node.over, node.coordinate, ctx), + into=node.into, + labels=ctx.coords, + ) + if isinstance(node, ms.At): + return operators.at( + _array(evaluate(node.operand, ctx)), + _lookup_arrays(node.over, node.coordinate, ctx), + into=node.into, + ) + if isinstance(node, ms.Translate): + return operators.shift( + _array(evaluate(node.operand, ctx)), + over=node.dimension, + offset=_amount(node.offset, ctx), + wrap=node.wrap, + fill=node.fill, + by=_partition(node, ctx), + ) + if isinstance(node, ms.Window): + return operators.sum_back( + _array(evaluate(node.operand, ctx)), + over=node.dimension, + within=_amount(node.width, ctx), + wrap=node.wrap, + by=_partition(node, ctx), + ) + if isinstance(node, ms.Cases): + regions = ( + _in_region(evaluate(region.value, ctx), evaluate_where(region.when, ctx)) + for region in node.regions + ) + return functools.reduce(lambda a, b: _combine(operator.add, a, b), regions) + assert_never(node) + + +def _variable(name: str, ctx: Context) -> Value: + variable = ctx.variable(name) + absence = ctx.program.variable(name).absence + if not ctx.solved: + return terms.variable_term(variable, absence) + if "solution" not in variable.data: + raise RuntimeError( + f"variable '{name}' has no solution yet: solve the model before reading a named expression" + ) + return terms.solution(variable, absence) + + +def _dual(constraint: str, ctx: Context) -> xr.DataArray: + if not ctx.solved: + raise TypeError( + f"the dual of constraint '{constraint}' has no symbolic form, read `.solution`" + ) + data = ctx.model.constraints[constraint].data + if "dual" not in data: + raise RuntimeError( + f"constraint '{constraint}' has no dual yet: solve the model, with a solver " + f"and a problem that report duals, before reading one" + ) + return data["dual"] + + +def _combine(op: Callable[[Value, Value], Value], left: Value, right: Value) -> Value: + """*left* and *right* combined by *op*, once two arrays agree on their shared coordinates and a hole beside a term has become its absence.""" + if isinstance(left, xr.DataArray) and isinstance(right, xr.DataArray): + for dim in set(left.dims) & set(right.dims): + if not left.indexes[dim].equals(right.indexes[dim]): + raise SpecDataError( + f"operands are not aligned on '{dim}': {left.indexes[dim].tolist()[:5]} against " + f"{right.indexes[dim].tolist()[:5]}. Every operand is read on the master " + f"coordinates, so the data was attached against other labels than the model was built on." + ) + elif isinstance(left, xr.DataArray) and isinstance( + right, Variable | LinearExpression | QuadraticExpression + ): + right, left = carried(right, left) + elif isinstance(right, xr.DataArray) and isinstance( + left, Variable | LinearExpression | QuadraticExpression + ): + left, right = carried(left, right) + return op(left, right) + + +def carried(term: Term, data: xr.DataArray) -> tuple[Term, xr.DataArray]: + """A hole an operator left in *data* is an absence the term takes: the slot leaves the row, and the hole reads as a harmless one.""" + if not bool(data.isnull().any()): + return term, data + return term.where(data.notnull()), data.fillna(1.0) + + +def _array(value: Value) -> Array: + if isinstance(value, float | int): + raise TypeError("a shape operator takes an array or a term, not a bare number") + return value + + +def _in_region(value: Value, rows: xr.DataArray) -> Value: + """*value* where the region holds and a hard zero everywhere else: a fill, so absence inside the region stands.""" + if isinstance(value, float | int): + return rows * value + if isinstance(value, Variable): + value = value.to_linexpr() + return value.where(rows, 0) + + +def _amount(amount: int | str, ctx: Context) -> operators.Amount: + if isinstance(amount, str): + return terms.coefficient(ctx.parameters[amount]) + return amount + + +def _partition(node: ms.Translate | ms.Window, ctx: Context) -> xr.DataArray | None: + """The lookup a windowed operator stays inside, named for the dimension its values are labels of.""" + if node.partition is None: + return None + array = ctx.lookup(node.partition, node.dimension) + return array.rename(ctx.program.dimension(node.dimension).targets[node.partition]) + + +def _lookup_arrays( + over: str, names: tuple[str, ...], ctx: Context +) -> tuple[xr.DataArray, ...]: + return tuple(ctx.lookup(name, over) for name in names) diff --git a/linopy/spec/groups.py b/linopy/spec/groups.py new file mode 100644 index 000000000..c39a4c6bc --- /dev/null +++ b/linopy/spec/groups.py @@ -0,0 +1,68 @@ +"""How a lookup partitions an axis: the shape every group-wise operator reads.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import xarray as xr + + +def unmapped(key: object) -> bool: + """Whether a lookup left this member in no group: ``None``, or the NaN that never equals itself.""" + return key is None or key != key + + +@dataclass(frozen=True) +class Groups: + labels: np.ndarray + grouped: xr.DataArray + belongs: xr.DataArray + within: xr.DataArray + size: xr.DataArray + roster: np.ndarray + names: tuple[object, ...] + counts: tuple[int, ...] + + +def grouped(over: str, labels: np.ndarray, groups: xr.DataArray) -> Groups: + """ + How the lookup *groups* partitions the axis *over*. + + A coordinate the lookup sends nowhere belongs to no group: its ``within`` + is 0, its ``size`` 1 and its ``grouped`` False. + """ + keys = np.asarray(groups.sel({over: labels}).values, dtype=object) + peers: dict[object, list[int]] = {} + within = np.zeros(len(labels), dtype=int) + held = np.zeros(len(labels), dtype=bool) + for k, key in enumerate(keys): + if unmapped(key): + continue + held[k] = True + beside = peers.setdefault(key, []) + within[k] = len(beside) + beside.append(k) + order = {key: g for g, key in enumerate(peers)} + widest = max((len(beside) for beside in peers.values()), default=1) + roster = np.zeros((max(len(peers), 1), widest), dtype=int) + for key, beside in peers.items(): + roster[order[key], : len(beside)] = beside + belongs = np.array([order.get(key, 0) for key in keys], dtype=int) + span = np.array( + [len(peers[key]) if inside else 1 for key, inside in zip(keys, held)], dtype=int + ) + + def on_axis(values: np.ndarray) -> xr.DataArray: + return xr.DataArray(values, coords={over: labels}, dims=[over]) + + return Groups( + labels, + on_axis(held), + on_axis(belongs), + on_axis(within), + on_axis(span), + roster, + tuple(peers), + tuple(len(beside) for beside in peers.values()), + ) diff --git a/linopy/spec/netcdf.py b/linopy/spec/netcdf.py new file mode 100644 index 000000000..6380d0ec3 --- /dev/null +++ b/linopy/spec/netcdf.py @@ -0,0 +1,303 @@ +""" +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 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, and so is a header naming the +math-spec version that lowered the text and the number of this layout, +``FORMAT``; a read under another of either warns. + +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 +read. That is enough for a parameter, but not for a partial lookup, which +holds NaN in an array of labels: a hole in a string array comes back as an +empty string, indistinguishable from a label. So a lookup, and any array of +objects, is written instead as integer codes into its own table of +categories, ``-1`` where a label is missing. Decoding indexes the table and +fills the holes back in, which reproduces what attach built, values and +dtype alike. + +The master coordinates are canonical: a container's coordinates for a +dimension are re-stamped from them on read, so the whole model agrees on one +dtype per dimension however the engine returned it. +""" + +from __future__ import annotations + +import json +import warnings +from collections.abc import Mapping +from typing import Any + +import numpy as np +import pandas as pd +import xarray as xr +from math_spec import __version__ as MATH_SPEC_VERSION + +from linopy.io import ( + DTYPE_ATTR, + LAYER_BOUND_ATTR, + LAYER_TEXT_ATTR, + SPEC_ATTR, + SPEC_LAYERS_ATTR, + SPEC_OBJECTIVE_ATTR, + SPEC_VERSION_ATTR, + SPEC_WHOLE_ATTR, + get_prefix, + restamp_coords, + with_prefix, +) +from linopy.model import Model +from linopy.spec.accessor import Layer, ModelSpec, register, restore_layer +from linopy.spec.ownership import Ownership + +PREFIX = "spec" +FORMAT = 1 +LEGACY_NAME = "spec" +LEGACY_OBJECTIVE_ATTR = "_linopy_spec_objective_replaced" +COORD = "coords__" +PARAM = "param__" +CODES = "codes__" +CATEGORIES = "cats__" +CATEGORY_DIM = "category__" + +HOLES: dict[str, Any] = {"f": np.nan, "O": np.nan, "M": np.datetime64("NaT")} + + +def encode(layer: Layer) -> xr.Dataset: + """ + 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 layer.coords.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)) + 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)), + SPEC_VERSION_ATTR: json.dumps( + {"math_spec": MATH_SPEC_VERSION, "format": FORMAT} + ), + } + ) + + +def read(model: Model, ds: xr.Dataset) -> ModelSpec: + """ + The spec layers a file holds, restored onto *model* in their order. + + 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. + + Warns + ----- + UserWarning + The file names another math-spec version or another layout number + than this reader's. + """ + _check_header(ds) + 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 _restored(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 _restored(model, [layer], True, LEGACY_NAME if owned else None) + + +def _check_header(ds: xr.Dataset) -> None: + """Warn where the file was written under another math-spec version or layout; a file without a header is older than both.""" + if SPEC_VERSION_ATTR not in ds.attrs: + return + header = json.loads(ds.attrs[SPEC_VERSION_ATTR]) + if header["format"] != FORMAT: + warnings.warn( + f"the file writes its spec layers in layout {header['format']} and this " + f"linopy reads layout {FORMAT}; what the layers hold may not come back as written.", + UserWarning, + stacklevel=4, + ) + if header["math_spec"] != MATH_SPEC_VERSION: + warnings.warn( + f"the file's spec layers were lowered by math-spec {header['math_spec']} and " + f"are re-lowered by {MATH_SPEC_VERSION}; the same text may lower differently.", + UserWarning, + stacklevel=4, + ) + + +def _restored( + model: Model, layers: list[Layer], whole: bool, owner: str | None +) -> ModelSpec: + """The accessor over *layers*, the registry rebuilt from what the file gave back.""" + model._ownership = Ownership() + spec = ModelSpec(model, layers, whole) + model._spec = spec + for layer in layers: + register(model, layer) + model._ownership.objective = owner + return spec + + +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. + """ + coords = { + _stripped(var, COORD): _index(sub[var]) + for var in sub.data_vars + if str(var).startswith(COORD) + } + arrays = { + _stripped(var, PARAM): _plain(sub[var], _stripped(var, PARAM), coords) + for var in sub.data_vars + if str(var).startswith(PARAM) + } + arrays.update( + { + _stripped(var, CODES): _decode(sub, _stripped(var, CODES), coords) + for var in sub.data_vars + if str(var).startswith(CODES) + } + ) + restamp_coords(model, coords) + parameters = xr.Dataset(arrays).assign_coords(coords) + return restore_layer(model, name, text, parameters, names) + + +def _coded(layer: Layer) -> set[str]: + """The parameters written as codes: every lookup and every array of objects.""" + lookups = {name for by_name in layer.lookups.values() for name in by_name} + return { + str(name) + for name, arr in layer.parameters.items() + if name in lookups or arr.dtype == object + } + + +def _encode(name: str, arr: xr.DataArray) -> dict[str, xr.DataArray]: + codes, categories = pd.factorize(arr.to_numpy().ravel()) + written = { + CODES + name: _array( + codes.astype(np.int32).reshape(arr.shape), arr.dims, str(arr.dtype) + ) + } + if len(categories): + written[CATEGORIES + name] = _array( + np.asarray(categories), (CATEGORY_DIM + name,) + ) + return written + + +def _plain(arr: xr.DataArray, name: str, coords: dict[str, pd.Index]) -> xr.DataArray: + """A parameter written as its own values, back on the master coordinates at its own dtype.""" + dims = tuple(str(d) for d in arr.dims) + return xr.DataArray( + _values(arr), coords={d: coords[d] for d in dims}, dims=dims, name=name + ) + + +def _decode(sub: xr.Dataset, name: str, coords: dict[str, pd.Index]) -> xr.DataArray: + codes = sub[CODES + name] + dtype = np.dtype(codes.attrs[DTYPE_ATTR]) + categories = _categories(sub, name, dtype) + positions = codes.to_numpy().astype(int) + mapped = positions >= 0 + if mapped.all(): + values = categories[positions] + else: + values = np.full(positions.shape, HOLES[dtype.kind], dtype=dtype) + values[mapped] = categories[positions[mapped]] + dims = tuple(str(d) for d in codes.dims) + return xr.DataArray( + values, coords={d: coords[d] for d in dims}, dims=dims, name=name + ) + + +def _categories(sub: xr.Dataset, name: str, dtype: np.dtype) -> np.ndarray: + """ + The table a coded array indexes. + + A map that leaves every label unmapped has no table: netCDF3 writes a + zero-length dimension as the unlimited one, of which a file holds one. + """ + written = CATEGORIES + name + if written in sub.data_vars: + return _values(sub[written]) + return np.empty(0, dtype=dtype) + + +def _array( + values: np.ndarray, dims: tuple[Any, ...], dtype: str | None = None +) -> xr.DataArray: + return xr.DataArray( + values, dims=dims, attrs={DTYPE_ATTR: dtype or str(values.dtype)} + ) + + +def _stripped(name: Any, prefix: str) -> str: + return str(name)[len(prefix) :] + + +def _values(arr: xr.DataArray) -> np.ndarray: + """The array as it was in memory, undoing what the netcdf type could not hold.""" + return arr.to_numpy().astype(np.dtype(arr.attrs[DTYPE_ATTR])) + + +def _index(arr: xr.DataArray) -> pd.Index: + return pd.Index(_values(arr), name=_stripped(arr.name, COORD)) diff --git a/linopy/spec/nodes.py b/linopy/spec/nodes.py new file mode 100644 index 000000000..d73ff9355 --- /dev/null +++ b/linopy/spec/nodes.py @@ -0,0 +1,45 @@ +"""The parameters a node names, and the dimensions a node spans.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from math_spec import program as ms +from math_spec.program import walk + + +def amounts_of(node: ms.ExpressionNode) -> Iterator[str]: + """The parameters *node* names as an amount: a translation's offset or a window's width.""" + if isinstance(node, ms.Translate) and isinstance(node.offset, str): + yield node.offset + elif isinstance(node, ms.Window) and isinstance(node.width, str): + yield node.width + + +def parameters_of(*nodes: ms.ExpressionNode) -> frozenset[str]: + """Every parameter named anywhere under *nodes*.""" + return frozenset(n.name for n in walk(*nodes) if isinstance(n, ms.Parameter)) + + +def dims_of(node: ms.ExpressionNode, program: ms.Program) -> tuple[str, ...]: + """The dimensions *node* spans, in the program's dimension order, before any data is bound.""" + spanned = _dims(node, program) + return tuple(d for d in program.dimensions if d in spanned) + + +def _dims(node: ms.ExpressionNode, program: ms.Program) -> frozenset[str]: + if isinstance(node, ms.Constant): + return frozenset() + if isinstance(node, ms.Variable): + return frozenset(program.variables[node.name].dims) + if isinstance(node, ms.Parameter): + return frozenset(program.parameters[node.name].dims) + if isinstance(node, ms.Dual): + return frozenset(program.constraints[node.constraint].dims) + if isinstance(node, ms.Sum): + return _dims(node.operand, program) - set(node.over) + if isinstance(node, ms.GroupSum | ms.At): + return (_dims(node.operand, program) - {node.over}) | set(node.into) + if isinstance(node, ms.Cases): + return frozenset().union(*(_dims(r.value, program) for r in node.regions)) + return frozenset().union(*(_dims(c, program) for c in ms.children(node))) diff --git a/linopy/spec/operators.py b/linopy/spec/operators.py new file mode 100644 index 000000000..ee32c3ed0 --- /dev/null +++ b/linopy/spec/operators.py @@ -0,0 +1,356 @@ +""" +The language's built-in operators, evaluated on xarray and linopy values. + +Each entry point takes an operand that is already a value, a ``DataArray`` +for data or a linopy term for anything carrying a variable, and returns the +same kind. Nothing here reads the program or the model: the builder +evaluates the operands and the keywords and calls in. +""" + +from __future__ import annotations + +import operator +from collections.abc import Hashable, Mapping +from dataclasses import dataclass +from functools import reduce +from typing import cast, overload + +import numpy as np +import pandas as pd +import xarray as xr + +from linopy.expressions import LinearExpression +from linopy.spec.errors import SpecDataError +from linopy.spec.groups import Groups, grouped +from linopy.spec.terms import Array, Term +from linopy.variables import Variable + +Amount = int | xr.DataArray + + +def filled(expression: Array, fill: float) -> Array: + """*expression* with every absence in it standing as *fill*.""" + if isinstance(expression, Variable): + expression = expression.to_linexpr() + return expression.fillna(fill) + + +def vacated( + shifted: Array, operand: Array, over: str, vacated: xr.DataArray, fill: float +) -> Array: + """ + *shifted*, with the positions the shift vacated filled, and only those. + + The fill lands where the shift vacated and the operand carries the + coordinate; every other slot keeps the absence it arrived with, so no row + is invented at a coordinate the operand never had. + """ + carried = (~operand.isnull()).any(over) + keep = carried & (~shifted.isnull() | vacated) + return filled(shifted, fill).where(keep) + + +def sum_over(array: Array, over: str) -> Array: + """Sum *array* over *over*; a term beside an empty dimension is built as the constant zero.""" + if not isinstance(array, xr.DataArray) and any( + not array.sizes[dim] for dim in array.coord_dims if dim != over + ): + kept = [dim for dim in array.coord_dims if dim != over] + zeros = xr.DataArray( + np.zeros([array.sizes[dim] for dim in kept]), + coords={dim: array.indexes[dim] for dim in kept}, + dims=kept, + ) + return LinearExpression.from_constant(array.model, zeros) + return array.sum(over) + + +def grouped_sum( + array: Array, + mappings: tuple[xr.DataArray, ...], + *, + into: tuple[str, ...], + labels: Mapping[str, pd.Index], +) -> Array: + """ + Sum *array* through the lookups *mappings*, replacing their dimension by *into*. + + A member a lookup sends nowhere contributes nowhere. The result is put + onto every declared label of *into*: a group no member reaches holds the + empty sum, which is 0 and not an absence. + """ + mappings = _renamed(mappings, into) + present = _present(mappings) + dim = str(mappings[0].dims[0]) + if not bool(present.any()): + return _empty_groups(array, dim, into=into, labels=labels) + if not bool(present.all()): + keep = present.to_numpy() + mappings = tuple(m.isel({dim: keep}) for m in mappings) + array = array.isel({dim: keep}) + attached = array.assign_coords( + {target: (dim, m.to_numpy()) for target, m in zip(into, mappings)} + ) + summed = attached.groupby(list(into)).sum() + return summed.reindex({d: labels[d] for d in into}).fillna(0.0) + + +def _empty_groups( + array: Array, + dim: str, + *, + into: tuple[str, ...], + labels: Mapping[str, pd.Index], +) -> Array: + """ + The grouped sum of an operand no member of *dim* is mapped out of. + + Every declared group holds the empty sum, which is 0. Grouping cannot say + so itself: filtering the operand down to its mapped members leaves nothing, + and an empty dimension is one xarray refuses to group over. + """ + kept = [d for d in _coord_dims(array) if d != dim] + zeros = xr.DataArray( + np.zeros([array.sizes[d] for d in kept] + [len(labels[d]) for d in into]), + coords={ + **{d: array.indexes[d] for d in kept}, + **{d: labels[d] for d in into}, + }, + dims=kept + list(into), + ) + if isinstance(array, xr.DataArray): + return zeros + return LinearExpression.from_constant(array.model, zeros) + + +def _coord_dims(array: Array) -> list[str]: + """The dimensions the operand is labelled over, without a term's own ``_term``.""" + dims = array.dims if isinstance(array, xr.DataArray) else array.coord_dims + return [str(d) for d in dims] + + +@overload +def at( + array: xr.DataArray, mappings: tuple[xr.DataArray, ...], *, into: tuple[str, ...] +) -> xr.DataArray: ... + + +@overload +def at( + array: Term, mappings: tuple[xr.DataArray, ...], *, into: tuple[str, ...] +) -> Term: ... + + +def at( + array: Array, mappings: tuple[xr.DataArray, ...], *, into: tuple[str, ...] +) -> Array: + """ + Read *array* through the lookups *mappings*: the adjoint of :func:`grouped_sum`. + + A member a lookup sends nowhere reads nothing, and its row keeps the + operand's own absence rather than a zero. + """ + mappings = _renamed(mappings, into) + present = _present(mappings) + if bool(present.all()): + return array.sel(dict(zip(into, mappings))) + dim = str(mappings[0].dims[0]) + keep = present.to_numpy() + picked = array.sel(dict(zip(into, (m.isel({dim: keep}) for m in mappings)))) + return picked.reindex({dim: mappings[0][dim]}) + + +@dataclass(frozen=True) +class _Edge: + wrap: bool + fill: float | None + + +def shift( + array: Array, + *, + over: str, + offset: Amount, + wrap: bool, + fill: float | None, + by: xr.DataArray | None = None, +) -> Array: + """ + Translate *array* along *over*: the value at ``t - offset``. + + *wrap* is cyclic and vacates nothing, *fill* is what the vacated + positions contribute, and neither leaves them absent. An *offset* that + is an array differs per entity and is a gather. *by* is the lookup whose + groups the translation stays inside. + """ + edge = _Edge(wrap, fill) + if by is not None: + partition = grouped(over, np.asarray(array.indexes[over]), by) + return _gather_in_groups(array, over, _per_group(offset, by), partition, edge) + if isinstance(offset, xr.DataArray) and offset.ndim: + return _gather_by_offset(array, over, offset, edge) + amount: dict[Hashable, int] = {over: int(offset)} + if wrap: + if isinstance(array, xr.DataArray): + return array.roll(amount, roll_coords=False) + return array.roll(amount) + if isinstance(array, xr.DataArray): + return array.shift(amount, fill_value=np.nan if fill is None else fill) + shifted = array.shift(amount) + if fill is None: + return shifted + return vacated(shifted, array, over, _off_the_axis(array, over, amount[over]), fill) + + +def sum_back( + array: Array, + *, + over: str, + within: Amount, + wrap: bool, + by: xr.DataArray | None = None, +) -> Array: + """ + Sum *array* over a trailing window along *over*: positions ``t - within + 1`` through ``t``. + + A position the window cannot reach contributes a zero; a window that + reaches nothing keeps no row. *by* stops the window at each group's edge. + + With *wrap* the window continues from the other end of *over*, so a width + above the axis size still reads every position and is taken as the whole + axis. Without it such a width asks for positions that do not exist and is + refused. + """ + if by is not None: + within = _per_group(within, by) + asked = _widest(within) + size = int(array.sizes[over]) + if asked > size and not wrap: + raise SpecDataError( + f"a trailing sum over '{over}' asks for a window of {asked} position(s) where " + f"'{over}' holds {size}. Narrow 'within' to the axis, or write edge='wrap' so the " + f"window continues from the other end." + ) + widest = max(1, min(asked, size)) + probe = _Edge(wrap=wrap, fill=None) + partition = ( + None if by is None else grouped(over, np.asarray(array.indexes[over]), by) + ) + lagged_terms: list[Array] = [] + reached: list[xr.DataArray] = [] + for lag in range(widest): + lagged = ( + _gather_by_offset(array, over, lag, probe) + if partition is None + else _gather_in_groups(array, over, lag, partition, probe) + ) + live, term = ~lagged.isnull(), filled(lagged, 0.0) + if isinstance(within, xr.DataArray): + live, term = live & (within > lag), term * (within > lag).astype(float) + lagged_terms.append(term) + reached.append(live) + return _merged(lagged_terms).where(reduce(operator.or_, reached)) + + +def _widest(within: Amount) -> int: + """The widest window the data asks for; a width no member carries is a window of nothing.""" + if not isinstance(within, xr.DataArray): + return int(within) + widths = np.asarray(within, dtype=float) + return 0 if np.isnan(widths).all() else int(np.nanmax(widths)) + + +def _merged(values: list[Array]) -> Array: + """The sum of *values* in one step: a running sum would re-concatenate the term axis once per lag.""" + if isinstance(values[0], xr.DataArray): + return reduce(operator.add, values) + from linopy import merge + + return cast(LinearExpression, merge(cast(list[Term], values))) + + +def _renamed( + mappings: tuple[xr.DataArray, ...], into: tuple[str, ...] +) -> tuple[xr.DataArray, ...]: + return tuple(mapping.rename(target) for mapping, target in zip(mappings, into)) + + +def _present(mappings: tuple[xr.DataArray, ...]) -> xr.DataArray: + return reduce(operator.and_, (m.notnull() for m in mappings)) + + +def _gather_by_offset(array: Array, over: str, offset: Amount, edge: _Edge) -> Array: + """ + Translate *array* along *over* by an offset that may differ per entity. + + Selection is by label, so a non-integer axis works. Out-of-range + positions are clipped onto the axis and emptied again, so an edge means + what it does for a scalar shift. + """ + card = int(array.sizes[over]) + labels = np.asarray(array.indexes[over]) + ordinal = xr.DataArray(np.arange(card), coords={over: labels}, dims=[over]) + source = (ordinal - offset).astype(int) + + def gathered(ordinals: xr.DataArray) -> Array: + picked = array.sel({over: _labelled(labels, ordinals)}) + return picked.assign_coords({over: labels}) + + if edge.wrap: + return gathered(source % card) + inside = ((source >= 0) & (source < card)).assign_coords({over: labels}) + moved = gathered(source.clip(0, card - 1)).where(inside) + if edge.fill is None: + return moved + return vacated(moved, array, over, ~inside, edge.fill) + + +def _per_group(offset: Amount, groups: xr.DataArray) -> Amount: + """*offset* at every coordinate where it is declared over the group's own dimension.""" + target = groups.name + if not isinstance(offset, xr.DataArray) or target not in offset.dims: + return offset + return at(offset, (groups,), into=(str(target),)).drop_vars(str(target)) + + +def _gather_in_groups( + array: Array, over: str, offset: Amount, groups: Groups, edge: _Edge +) -> Array: + """ + Translate *array* inside each group rather than along the axis. + + A coordinate in no group reaches nothing, which is not the same as + reaching off a group's edge: only the second is what a fill speaks for. + """ + reached = groups.within - offset + if edge.wrap: + reached = reached % groups.size + inside = groups.grouped & (reached >= 0) & (reached < groups.size) + + def peer(group: np.ndarray, position: np.ndarray) -> np.ndarray: + return groups.roster[group, position] + + source = xr.apply_ufunc(peer, groups.belongs, reached.where(inside, 0).astype(int)) + labels = groups.labels + gathered = ( + array.sel({over: _labelled(labels, source)}) + .assign_coords({over: labels}) + .where(inside) + ) + if edge.fill is None: + return gathered + return vacated(gathered, array, over, groups.grouped & ~inside, edge.fill) + + +def _off_the_axis(array: Array, over: str, offset: int) -> xr.DataArray: + labels = np.asarray(array.indexes[over]) + source = xr.DataArray(np.arange(len(labels)), coords={over: labels}, dims=[over]) + source = source - offset + return (source < 0) | (source >= len(labels)) + + +def _labelled(labels: np.ndarray, ordinals: xr.DataArray) -> xr.DataArray: + """*ordinals* as the labels they stand for, carrying no coordinates of their own.""" + return xr.DataArray( + labels[ordinals.transpose(*ordinals.dims).values], dims=ordinals.dims + ) diff --git a/linopy/spec/ownership.py b/linopy/spec/ownership.py new file mode 100644 index 000000000..7ee6f91e0 --- /dev/null +++ b/linopy/spec/ownership.py @@ -0,0 +1,130 @@ +""" +Which spec layer owns which name of a model. + +A layer builds variables, constraints and expressions, binds model variables +it reads instead of building, declares special-ordered sets and may own the +objective. The model's containers consult this registry before they drop or +shadow a name, so a layer is never left referencing something the model no +longer holds. Names are the unit: what a name holds can drift without the +registry seeing it. + +This module imports nothing of ``math_spec``, so ``linopy.model`` can name +the registry without pulling the spec package in. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass, field +from typing import Literal + +Kind = Literal["variable", "constraint", "expression", "sos"] + + +def joined(parts: list[str]) -> str: + """``a``, ``a and b``, ``a, b and c``.""" + if len(parts) < 3: + return " and ".join(parts) + return f"{', '.join(parts[:-1])} and {parts[-1]}" + + +@dataclass +class Ownership: + """ + Model names by the spec layer that owns them. + + Attributes + ---------- + variables + Built variable name to layer. + bound + Model variable name to the layers that bind it, in the order they + were attached; several layers may read one variable. + constraints, expressions, sos + Constraint, named expression and special-ordered-set variable name to + layer. An expression is recorded whether it was built into the model + or folds lazily on read. + objective + The layer whose objective the model holds, ``None`` where the + objective is none of theirs. + """ + + variables: dict[str, str] = field(default_factory=dict) + bound: dict[str, list[str]] = field(default_factory=dict) + constraints: dict[str, str] = field(default_factory=dict) + expressions: dict[str, str] = field(default_factory=dict) + sos: dict[str, str] = field(default_factory=dict) + objective: str | None = None + + def _records(self, kind: Kind) -> dict[str, str]: + return { + "variable": self.variables, + "constraint": self.constraints, + "expression": self.expressions, + "sos": self.sos, + }[kind] + + def owner(self, kind: Kind, name: str) -> str | None: + """The layer owning *name* as a *kind*; for a variable, the one binding it counts too.""" + layer = self._records(kind).get(name) + if layer is None and kind == "variable" and name in self.bound: + return self.bound[name][0] + return layer + + def claim( + self, + layer: str, + *, + variables: Iterable[str] = (), + bound: Iterable[str] = (), + constraints: Iterable[str] = (), + expressions: Iterable[str] = (), + sos: Iterable[str] = (), + objective: bool = False, + ) -> None: + """Record every name *layer* owns.""" + self.variables.update(dict.fromkeys(variables, layer)) + for name in bound: + self.bound.setdefault(name, []).append(layer) + self.constraints.update(dict.fromkeys(constraints, layer)) + self.expressions.update(dict.fromkeys(expressions, layer)) + self.sos.update(dict.fromkeys(sos, layer)) + if objective: + self.objective = layer + + def held(self, kind: Kind, layer: str) -> list[str]: + """The names *layer* owns as a *kind*.""" + return [n for n, by in self._records(kind).items() if by == layer] + + def release(self, layer: str) -> None: + """Drop every record of *layer*, the objective included where it was the layer's.""" + for records in (self.variables, self.constraints, self.expressions, self.sos): + for name in [n for n, by in records.items() if by == layer]: + del records[name] + for name, binders in list(self.bound.items()): + if layer in binders: + binders.remove(layer) + if not binders: + del self.bound[name] + if self.objective == layer: + self.objective = None + + def refuse_removal(self, kind: Kind, names: Iterable[str]) -> None: + """Refuse to drop a name a layer builds or binds, which would strand its layer.""" + hit = sorted(n for n in names if self.owner(kind, n) is not None) + if hit: + verb = "is" if len(hit) == 1 else "are" + raise ValueError( + f"{joined(hit)} {verb} declared or bound by a spec layer; a layer " + "cannot be left referencing a name the model no longer holds. " + "Model.remove_spec(name) takes a layer off with everything it built." + ) + + def refuse_addition(self, kind: Kind, name: str) -> None: + """Refuse a name a layer already owns as a *kind*, built, bound or declared lazily.""" + layer = self.owner(kind, name) + if layer is not None: + raise ValueError( + f"{kind} '{name}' is declared or bound by spec layer '{layer}'; " + f"a layer's names are its own. Pick another name." + ) diff --git a/linopy/spec/parameters.py b/linopy/spec/parameters.py new file mode 100644 index 000000000..afc7fd1b6 --- /dev/null +++ b/linopy/spec/parameters.py @@ -0,0 +1,50 @@ +""" +Every parameter of a program by name, resolved once. + +A declared parameter is resolved from the caller's data and aligned by the +binder; one a ``piecewise:`` expansion emitted is derived from the block's own +breakpoints. Which of the two a name is, is the declaration's answer, and this +module is where it is asked. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping + +import xarray as xr +from math_spec import program as ms + +from linopy.spec import curves + +Resolve = Callable[[str], xr.DataArray] + + +class Parameters(Mapping[str, xr.DataArray]): + """ + Every parameter of a program by name, each resolved on first read and then held. + + A declared parameter comes from *resolve*; one a ``piecewise:`` expansion + emitted is derived from the block's own breakpoints the way its + derivation says, so a caller never supplies it. + """ + + def __init__(self, program: ms.Program, resolve: Resolve) -> None: + self._program = program + self._resolve = resolve + self._arrays: dict[str, xr.DataArray] = {} + + def __getitem__(self, name: str) -> xr.DataArray: + if name not in self._arrays: + derivation = self._program.parameter(name).derivation + self._arrays[name] = ( + self._resolve(name) + if derivation is None + else curves.derive(derivation, self, self._program) + ) + return self._arrays[name] + + def __iter__(self) -> Iterator[str]: + return iter(self._program.parameters) + + def __len__(self) -> int: + return len(self._program.parameters) diff --git a/linopy/spec/terms.py b/linopy/spec/terms.py new file mode 100644 index 000000000..be429d482 --- /dev/null +++ b/linopy/spec/terms.py @@ -0,0 +1,48 @@ +""" +What an expression node evaluates to, and how absence is spelled at each position. + +Absence is positional: one missing parameter row is a zero in a coefficient, +a refusal in ``bounds:`` and false in a ``where`` operand, so there is no +single fill applied once and each position states its own answer. The +convention underneath is linopy v1's, which a spec-built model requires. +""" + +from __future__ import annotations + +import xarray as xr + +from linopy.constants import FACTOR_DIM, TERM_DIM +from linopy.expressions import LinearExpression, QuadraticExpression +from linopy.variables import Variable + +Term = Variable | LinearExpression | QuadraticExpression +Array = xr.DataArray | Term +Value = float | Array + + +def present(variable: Variable) -> xr.DataArray: + """The coordinates the variable occupies; ``-1`` is linopy's marker for an absent slot.""" + return variable.labels != -1 + + +def live_rows(term: Term) -> xr.DataArray: + """The rows *term* still holds a variable in, once the data has emptied the slots it left.""" + if isinstance(term, Variable): + return present(term) + helpers = [d for d in (TERM_DIM, FACTOR_DIM) if d in term.vars.dims] + return (term.vars != -1).any(helpers) + + +def variable_term(variable: Variable, absence: str) -> Term: + """The variable as it enters a built expression, carrying its declared ``absence:``.""" + return variable.fillna(0) if absence == "zero" else variable + + +def solution(variable: Variable, absence: str) -> xr.DataArray: + """The solved variable as it enters a fold, carrying its declared ``absence:``.""" + return variable.solution.fillna(0) if absence == "zero" else variable.solution + + +def coefficient(parameter: xr.DataArray) -> xr.DataArray: + """A parameter in a coefficient position, its uncovered slots at zero.""" + return parameter.fillna(0.0) diff --git a/linopy/spec/testing.py b/linopy/spec/testing.py new file mode 100644 index 000000000..2771d4afe --- /dev/null +++ b/linopy/spec/testing.py @@ -0,0 +1,68 @@ +""" +Synthetic data for a spec, for tests and benchmarks. + +A spec says what data it takes, which is enough to make some up: the shape is +the declaration's, only the values are invented. What comes out builds and +solves, and says nothing about a real system. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pandas as pd +import xarray as xr +from math_spec import program as ms + +_START = "2030-01-01" + + +def synthetic_sources(program: ms.Program, n: int = 3) -> dict[str, Any]: + """ + Dense data for every declaration of *program*, *n* labels per dimension. + + Labels are numbered after their dimension, parameters are a linear ramp, + and each lookup cycles through the labels it maps into. + """ + sources: dict[str, Any] = { + dim: _labels(dim, decl.dtype, n) for dim, decl in program.dimensions.items() + } + for over, lookup in program.lookups: + into = sources[lookup.target] + sources[lookup.name] = pd.Series( + [into[i % len(into)] for i in range(n)], index=sources[over] + ) + for name, parameter in program.parameters.items(): + if parameter.derivation is None: + sources[name] = _parameter(name, parameter, sources, n) + return sources + + +def _labels(name: str, dtype: str | None, n: int) -> pd.Index: + """*n* labels of the declared dtype, named after what they label.""" + if dtype == "int": + return pd.Index(range(n), name=name) + if dtype == "datetime": + return pd.date_range(_START, periods=n, freq="h", name=name) + return pd.Index([f"{name}{i}" for i in range(n)], name=name) + + +def _parameter( + name: str, declared: ms.ParameterDeclaration, sources: dict[str, Any], n: int +) -> Any: + dims = declared.dims + shape = [n] * len(dims) + if declared.dtype == "bool": + values: Any = np.ones(shape, dtype=bool) + elif declared.dtype == "int": + values = np.ones(shape, dtype=int) + elif declared.dtype == "str": + values = np.full(shape, "a", dtype=object) + elif dims: + values = np.broadcast_to(1.0 + np.arange(n), shape).copy() + else: + values = np.array(1.0) + if not dims: + return values.item() + return xr.DataArray(values, coords={d: sources[d] for d in dims}, dims=list(dims)) diff --git a/linopy/spec/where.py b/linopy/spec/where.py new file mode 100644 index 000000000..3c1d10ed9 --- /dev/null +++ b/linopy/spec/where.py @@ -0,0 +1,144 @@ +"""A ``where:`` predicate as a boolean array over the coordinates it masks.""" + +from __future__ import annotations + +import operator +from collections.abc import Callable +from typing import assert_never + +import numpy as np +import xarray as xr +from math_spec import program as ms + +from linopy.spec import terms +from linopy.spec.context import Context +from linopy.spec.errors import SpecDataError +from linopy.spec.groups import grouped + +_PREDICATE_OPS: dict[str, Callable[..., xr.DataArray]] = { + "==": operator.eq, + "!=": operator.ne, + "<": operator.lt, + ">": operator.gt, + "<=": operator.le, + ">=": operator.ge, +} + + +def evaluate_where(mask: ms.Mask | None, ctx: Context) -> xr.DataArray: + """The rows *mask* admits, as a boolean array; no mask is a 0-d ``True``.""" + if mask is None: + return xr.DataArray(True) + return _node(mask.root, ctx) + + +def as_linopy_mask(mask: xr.DataArray) -> xr.DataArray | None: + """*mask* as linopy's ``mask=`` takes it: ``None`` where nothing is masked.""" + if mask.ndim == 0 and bool(mask): + return None + return mask + + +def _node(node: ms.WhereNode, ctx: Context) -> xr.DataArray: + """ + One predicate node as a boolean array. + + A masked-out variable coordinate and a comparison over NaN both read as + exclusion. A null lookup value is excluded explicitly: numpy answers + ``None != 'north'`` with True, so a ``!=`` would otherwise keep exactly + the labels that map nowhere. + """ + if isinstance(node, ms.BooleanLiteralNode): + return xr.DataArray(node.value) + if isinstance(node, ms.ParameterDefinedNode): + return _defined( + ctx.parameters[node.name], ctx.program.parameter(node.name).dtype + ) + if isinstance(node, ms.VariableDefinedNode): + 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)) + return result.fillna(False).astype(bool) + if isinstance(node, ms.DimensionComparisonNode): + labels = ctx.coords[node.name] + arr = xr.DataArray(labels, coords={node.name: labels}, dims=[node.name]) + result = _PREDICATE_OPS[node.op](arr, _as_the_axis_spells_it(arr, node.value)) + return result.fillna(False).astype(bool) + if isinstance(node, ms.DimensionPositionNode): + return _position(node, ctx) + if isinstance(node, ms.LookupComparisonNode): + arr = ctx.lookup(node.name, node.over) + compared = _PREDICATE_OPS[node.op](arr, node.value) & arr.notnull() + return compared.fillna(False).astype(bool) + if isinstance(node, ms.LookupPairComparisonNode): + left = ctx.lookup(node.name, node.over) + right = ctx.lookup(node.other, node.over) + compared = ( + _PREDICATE_OPS[node.op](left, right) & left.notnull() & right.notnull() + ) + return compared.fillna(False).astype(bool) + if isinstance(node, ms.LookupDefinedNode): + return ctx.lookup(node.name, node.over).notnull() + if isinstance(node, ms.NotNode): + return ~_node(node.operand, ctx) + if isinstance(node, ms.AndNode): + return _node(node.left, ctx) & _node(node.right, ctx) + if isinstance(node, ms.OrNode): + return _node(node.left, ctx) | _node(node.right, ctx) + assert_never(node) + + +def _defined(arr: xr.DataArray, dtype: str) -> xr.DataArray: + """What a bare parameter name asks: a bool is its own answer, a str is defined where it has a row, a number must be finite too.""" + if dtype == "bool": + return arr.fillna(False).astype(bool) + if dtype == "str": + return arr.notnull() + return arr.notnull() & np.isfinite(arr) + + +def _position(node: ms.DimensionPositionNode, ctx: Context) -> xr.DataArray: + labels = ctx.coords[node.name] + if node.by is not None: + groups = ctx.lookup(node.by, node.name) + arr = _group_offsets(node, groups, np.asarray(labels)) + compared = _PREDICATE_OPS[node.op](arr, 0) & arr.notnull() + return compared.fillna(False).astype(bool) + at = node.position + len(labels) if node.position < 0 else node.position + if not 0 <= at < len(labels): + raise SpecDataError( + f"where: position({node.name}) {node.op} {node.position} names position {at} of " + f"'{node.name}', which has {len(labels)} coordinate(s). A boundary that names no " + f"coordinate leaves the rows it was to seed unseeded." + ) + arr = xr.DataArray( + np.arange(len(labels)), coords={node.name: labels}, dims=[node.name] + ) + return _PREDICATE_OPS[node.op](arr, at).astype(bool) + + +def _group_offsets( + node: ms.DimensionPositionNode, groups: xr.DataArray, labels: np.ndarray +) -> xr.DataArray: + """Each coordinate's distance from the boundary of its own group; NaN where it is in no group.""" + partition = grouped(node.name, labels, groups) + needed = node.position + 1 if node.position >= 0 else -node.position + short = sorted( + str(g) for g, n in zip(partition.names, partition.counts) if n < needed + ) + if short: + raise SpecDataError( + f"where: position({node.name}, by={node.by}) {node.op} {node.position} names position " + f"{node.position} within each group, and {len(short)} of them are shorter than that: " + f"{short[:5]}. A boundary that names no coordinate leaves the rows it was to seed unseeded." + ) + target = node.position if node.position >= 0 else partition.size + node.position + return partition.within.where(partition.grouped) - target + + +def _as_the_axis_spells_it(arr: xr.DataArray, value: object) -> object: + """A ``where`` literal in the spelling of the axis it is compared against: a date on a datetime axis is a ``datetime64``.""" + if arr.dtype.kind == "M": + return np.datetime64(str(value)) + return value diff --git a/linopy/testing.py b/linopy/testing.py index e914f7b8d..2534cc6cf 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -111,29 +111,61 @@ def assert_conequal(a: ConstraintBase, b: ConstraintBase, strict: bool = True) - assert_equal(a.rhs, b.rhs) +def _dtypes(ds: xr.Dataset) -> dict[str, str]: + """The dtype of every variable and coordinate, which assert_equal ignores.""" + return {str(name): str(arr.dtype) for name, arr in {**ds.variables}.items()} + + +def assert_datasetequal(a: xr.Dataset, b: xr.Dataset) -> None: + """ + Assert that two datasets hold the same values at the same dtypes. + + xarray's ``assert_equal`` compares values and labels but not dtypes, and a + netcdf engine is free to narrow an int64 or widen a bool, so the dtypes + are compared here on top of it. + """ + assert_equal(a, b) + assert _dtypes(a) == _dtypes(b), f"dtypes differ: {_dtypes(a)} != {_dtypes(b)}" + + def assert_model_equal(a: Model, b: Model) -> None: """Assert that two models are equal.""" for k in a.dataset_attrs: - assert_equal(getattr(a, k), getattr(b, k)) + assert_datasetequal(getattr(a, k), getattr(b, k)) assert list(a.variables) == list(b.variables) assert list(a.constraints) == list(b.constraints) for v in a.variables: assert_varequal(a.variables[v], b.variables[v]) + assert a.variables[v].spec == b.variables[v].spec for c in a.constraints: assert_conequal(a.constraints[c], b.constraints[c]) + assert a.constraints[c].spec == b.constraints[c].spec assert list(a.expressions) == list(b.expressions) for e in a.expressions: assert_exprequal(a.expressions[e], b.expressions[e]) + assert a.expressions[e].spec == b.expressions[e].spec assert_exprequal(a.objective.expression, b.objective.expression, check_name=False) assert a.objective.sense == b.objective.sense assert a.objective.value == b.objective.value + assert (a._spec is None) == (b._spec is None) + if a._spec is not None and b._spec is not None: + 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 a.status == b.status assert a.termination_condition == b.termination_condition diff --git a/linopy/variables.py b/linopy/variables.py index e0ad70bc9..cebeb461e 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -59,6 +59,7 @@ HELPER_DIMS, SOS_DIM_ATTR, SOS_TYPE_ATTR, + SPEC_LAYER_ATTR, STASHED_ATTRS, STASHED_LOWER, STASHED_UPPER, @@ -920,6 +921,16 @@ def name(self) -> str: """ return str(self.attrs["name"]) + @property + def spec(self) -> str | None: + """The spec layer that built this variable; ``None`` for one built by hand, a layer's binding included.""" + layer = self.attrs.get(SPEC_LAYER_ATTR) + return None if layer is None else str(layer) + + @spec.setter + def spec(self, layer: str) -> None: + self.attrs[SPEC_LAYER_ATTR] = layer + @property def labels(self) -> DataArray: """ @@ -1809,8 +1820,10 @@ def __dir__(self) -> list[str]: ] return base_attributes + formatted_names - def _format_items(self, exclude: set[str] | None = None) -> str: - """Format variable items, optionally excluding names in a group.""" + def _format_items( + self, exclude: set[str] | None = None, tagged: bool = False + ) -> str: + """Format variable items, optionally excluding names in a group and, if *tagged*, naming each one's spec layer.""" r = "" count = 0 for name, ds in self.items(): @@ -1828,7 +1841,8 @@ def _format_items(self, exclude: set[str] | None = None) -> str: coords += f" - sos{sos_type} on {sos_dim}" if ds.attrs.get("semi_continuous", False): coords += " - semi-continuous" - r += f" * {name}{coords}\n" + suffix = f" [{ds.spec}]" if tagged and ds.spec is not None else "" + r += f" * {name}{coords}{suffix}\n" if count == 0: r += "\n" return r @@ -1872,7 +1886,11 @@ def add(self, variable: Variable) -> None: def remove(self, name: str) -> None: """ Remove variable `name` from the variables. + + Refused where a spec layer builds or binds it. """ + if self.model._ownership is not None: + self.model._ownership.refuse_removal("variable", [name]) self.data.pop(name) self._invalidate_label_position_index() diff --git a/pyproject.toml b/pyproject.toml index b9c0c051f..2079a5b39 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -118,6 +118,32 @@ gpu = [ # "cupdlpx>=0.1.2", pip package currently unstable, install manually ] +[dependency-groups] +# math-spec is not on PyPI yet and needs Python >= 3.12. A dependency group +# 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.88 ; python_version >= '3.12'", + "pyyaml ; python_version >= '3.12'", + "pyarrow ; python_version >= '3.12'", +] +# datarecord feeds model data as sources into Model.from_spec (adapter showcased +# in dev-scripts). Pre-1.0 git dep, gated on 3.12 like spec. narwhals 2.21.0 has +# a join regression that breaks datarecord's name-uniqueness check, so it is +# excluded until a fix ships. Install with `uv sync --group datarecord`. +datarecord = [ + { include-group = "spec" }, + "datarecord @ git+https://github.com/energy-models/datarecord.git@3b1dd503ace7c9ae12a9adf38f8222e876f09311 ; python_version >= '3.12'", + "narwhals!=2.21.0 ; python_version >= '3.12'", + "pyarrow ; python_version >= '3.12'", +] +# Runs dev-scripts/spec/pypsa_spec_lowering.py: a PyPSA example network lowered +# through math-spec's examples/pypsa.yaml. +pypsa = [ + { include-group = "spec" }, + "pypsa>=1.3 ; python_version >= '3.12'", +] + [tool.uv] # cuopt-cu12 pulls cudf-cu12, which pins pandas<3.0.4, while benchmarks pins # pandas==3.0.5. Resolve the two extras in separate forks instead of together. @@ -151,6 +177,7 @@ filterwarnings = [ # collection of ``linopy/variables.py`` in the source tree on # Windows CI. "ignore:piecewise:FutureWarning", + "ignore:spec:FutureWarning", ] [tool.coverage.run] @@ -161,7 +188,7 @@ omit = ["test/*"] exclude_also = ["if TYPE_CHECKING:"] [tool.mypy] -exclude = ['dev/*', 'examples/*', '^benchmark/', 'doc/*'] +exclude = ['dev/*', 'examples/*', '^benchmark/', 'doc/*', '^conftest\.py$'] ignore_missing_imports = true no_implicit_optional = true warn_unused_ignores = true diff --git a/test/conftest.py b/test/conftest.py index d636778d1..df87d3265 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -5,7 +5,8 @@ import os import warnings from collections.abc import Generator -from typing import TYPE_CHECKING +from importlib.util import find_spec +from typing import TYPE_CHECKING, Any import pandas as pd import pytest @@ -157,3 +158,186 @@ def u(m: Model) -> Variable: idx.name = "dim_3" m.add_variables(coords=[idx], name="u") return m.variables["u"] + + +if find_spec("math_spec") is not None: + import math_spec + import xarray as xr + import yaml + + from linopy import Model + + EXAMPLE_DISPATCH = """ +description: Least-cost dispatch of a generator fleet against an hourly load. + +dimensions: + snapshot: { dtype: int, description: dispatch periods } + generator: { description: generating units } + +parameters: + p_max: { dims: [generator], description: installed capacity } + load: { dims: [snapshot], description: demand to be met } + cost: { dims: [generator], description: marginal cost } + +variables: + p: + description: output of a generator in a snapshot + dims: [snapshot, generator] + where: "p_max > 0" + bounds: { lower: 0, upper: p_max } + +constraints: + power_balance: + dims: [snapshot] + expression: sum(p, over=generator) == load + +objective: + sense: minimize + expression: sum(p * cost) + +expressions: + spend: sum(p * cost, over=generator) + usage: p / p_max +""" + + GENERATOR = pd.Index(["wind", "gas"], name="generator") + SNAPSHOT = pd.Index([0, 1, 2], name="snapshot") + DISPATCH_DATA: dict[str, Any] = { + "snapshot": SNAPSHOT, + "generator": GENERATOR, + "p_max": pd.Series([100.0, 200.0], index=GENERATOR), + "load": pd.Series([80.0, 150.0, 50.0], index=SNAPSHOT), + "cost": pd.Series([0.0, 50.0], index=GENERATOR), + } + DISPATCH_P = xr.DataArray( + [[80.0, 0.0], [100.0, 50.0], [50.0, 0.0]], + coords={"snapshot": SNAPSHOT, "generator": GENERATOR}, + ) + + def solved(spec: Any, sources: Any, **kwargs: Any) -> Model: + m = Model.from_spec(spec, sources, **kwargs) + m.solve(solver_name="highs", output_flag=False, reformulate_sos=True) + return m + + def yaml_dict() -> dict[str, Any]: + return math_spec.to_spec(yaml.safe_load(EXAMPLE_DISPATCH)).to_dict() + + def with_(spec: dict[str, Any], **sections: dict[str, Any]) -> dict[str, Any]: + out = dict(spec) + for section, entries in sections.items(): + out[section] = {**spec.get(section, {}), **entries} + return out + + TT = pd.Index([0, 1, 2, 3], name="t") + S = pd.Index(["a", "b"], name="s") + DAYS = pd.date_range("2030-01-01", periods=4, freq="D", name="d") + + WHERE_SPEC: dict[str, Any] = { + "dimensions": { + "t": {"dtype": "int"}, + "s": {"dtype": "str"}, + "d": {"dtype": "datetime"}, + }, + "lookups": { + "season_of": {"over": "t", "into": "s"}, + "other_of": {"over": "t", "into": "s"}, + }, + "parameters": { + "flag": {"dims": ["t"], "dtype": "bool"}, + "cost": {"dims": ["t"]}, + "label": {"dims": ["t"], "dtype": "str"}, + "day_cost": {"dims": ["d"]}, + }, + "variables": { + "x": {"dims": ["t"], "bounds": {"lower": 0, "upper": 1}}, + "y": {"dims": ["d"], "bounds": {"lower": 0, "upper": 1}}, + }, + "objective": {"sense": "minimize", "expression": "sum(x) + sum(y)"}, + } + WHERE_DATA: dict[str, Any] = { + "t": TT, + "s": S, + "d": DAYS, + "season_of": pd.Series(["a", "a", "b"], index=TT[:3]), + "other_of": pd.Series(["a", "b", "b", "a"], index=TT), + "flag": pd.Series([True, False], index=TT[:2]), + "cost": pd.Series([1.0, float("inf"), 3.0], index=TT[:3]), + "label": pd.Series(["u", "v"], index=TT[1:3]), + "day_cost": pd.Series([1.0, 2.0, 3.0, 4.0], index=DAYS), + } + + BP = pd.Index([0, 1, 2, 3], name="bp") + UNITS = pd.Index(["hydro", "gas"], name="generator") + CURVE_SPEC: dict[str, Any] = { + "dimensions": { + "snapshot": {"dtype": "int"}, + "generator": {"dtype": "str"}, + "bp": {"dtype": "int"}, + }, + "parameters": { + "p_max": {"dims": ["generator"]}, + "load": {"dims": ["snapshot"]}, + "bp_x": {"dims": ["generator", "bp"]}, + "bp_y": {"dims": ["generator", "bp"]}, + }, + "variables": { + "p": { + "dims": ["snapshot", "generator"], + "bounds": {"lower": 0, "upper": "p_max"}, + }, + "op_cost": {"dims": ["snapshot", "generator"], "bounds": {"lower": 0}}, + }, + "piecewise": { + "cost_curve": { + "over": "bp", + "links": [["p", "bp_x"], ["op_cost", "bp_y", ">="]], + "method": "lp", + } + }, + "expressions": {"spend": "sum(op_cost, over=generator)"}, + "constraints": { + "balance": { + "dims": ["snapshot"], + "expression": "sum(p, over=generator) == load", + } + }, + "objective": {"sense": "minimize", "expression": "sum(op_cost)"}, + } + MASKED_CURVE_SPEC = with_( + CURVE_SPEC, + piecewise={ + "cost_curve": {**CURVE_SPEC["piecewise"]["cost_curve"], "points": "bp_x"} + }, + ) + + def curve(points: dict[tuple[str, int], float]) -> pd.Series: + index = pd.MultiIndex.from_tuples(list(points), names=["generator", "bp"]) + return pd.Series(list(points.values()), index=index) + + FULL_X = curve( + {(g, k): x for g in UNITS for k, x in enumerate([0.0, 20.0, 50.0, 80.0])} + ) + FULL_Y = curve( + {(g, k): y for g in UNITS for k, y in enumerate([0.0, 150.0, 450.0, 900.0])} + ) + RAGGED_X = curve( + { + ("hydro", 0): 0.0, + ("hydro", 1): 40.0, + **{("gas", k): x for k, x in enumerate([0.0, 20.0, 50.0, 80.0])}, + } + ) + RAGGED_Y = curve( + { + ("hydro", 0): 0.0, + ("hydro", 1): 200.0, + **{("gas", k): y for k, y in enumerate([0.0, 150.0, 450.0, 900.0])}, + } + ) + CURVE_DATA: dict[str, Any] = { + "snapshot": [0], + "generator": UNITS, + "bp": BP, + "p_max": pd.Series([40.0, 80.0], index=UNITS), + "load": pd.Series([50.0], index=pd.Index([0], name="snapshot")), + } diff --git a/test/test_io.py b/test/test_io.py index 825ca16a9..0317cba1e 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -120,6 +120,25 @@ def test_model_to_netcdf(model: Model, tmp_path: Path) -> None: assert_model_equal(m, p) +@pytest.mark.parametrize("engine", ["netcdf4", "scipy"]) +def test_model_to_netcdf_keeps_parameter_dtypes( + model: Model, tmp_path: Path, engine: str +) -> None: + if engine == "netcdf4" and not HAS_NETCDF4: + pytest.skip("needs the netCDF4 backend") + model.parameters["count"] = xr.DataArray( + np.array([1, 2, 3, 4], dtype=np.int64), dims=["x"] + ) + model.parameters["flag"] = xr.DataArray(np.array([True, False]), dims=["y"]) + fn = tmp_path / f"dtypes-{engine}.nc" + model.to_netcdf(fn, engine=engine) + p = read_netcdf(fn) + + for name in ("count", "flag"): + assert p.parameters[name].dtype == model.parameters[name].dtype + assert p.parameters[name].equals(model.parameters[name]) + + @pytest.fixture def unsorted_model() -> Model: m = Model() diff --git a/test/test_piecewise_constraints.py b/test/test_piecewise_constraints.py index 788a0674e..6c4745e09 100644 --- a/test/test_piecewise_constraints.py +++ b/test/test_piecewise_constraints.py @@ -3266,7 +3266,7 @@ def _reset_dedup(self) -> Generator[None, None, None]: Warnings dedup is module-global so order between tests would otherwise matter. Clear before each test. """ - from linopy.piecewise import _emitted_evolving_warnings + from linopy.constants import _emitted_evolving_warnings _emitted_evolving_warnings.clear() yield diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py new file mode 100644 index 000000000..a3f4b31bd --- /dev/null +++ b/test/test_spec_accessor.py @@ -0,0 +1,1155 @@ +""" +``model.spec``, ``ModelSpec``, ``NamedExpression``, ``evaluate``, typesetting, +and the ``add_spec``/``from_spec`` argument handling that builds them. +""" + +from __future__ import annotations + +import io +import warnings +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +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 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, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +# --------------------------------------------------------------------------- +# inputs and model integration +# --------------------------------------------------------------------------- + + +SPEC_FORMS: dict[str, Callable[[Path], Any]] = { + "path": lambda path: path, + "path-string": str, + "yaml-text": lambda path: path.read_text(), + "flow-yaml": lambda path: yaml.safe_dump( + math_spec.to_spec(path).to_dict(), default_flow_style=True, width=10**6 + ).strip(), + "dict": lambda path: math_spec.to_spec(path).to_dict(), + "spec": lambda path: math_spec.to_spec(path), +} + + +@pytest.mark.parametrize("form", SPEC_FORMS.values(), ids=SPEC_FORMS.keys()) +def test_spec_forms_build_the_same_model( + tmp_path: Path, form: Callable[[Path], Any] +) -> None: + path = tmp_path / "dispatch.yaml" + path.write_text(EXAMPLE_DISPATCH) + m = Model.from_spec(form(path), DISPATCH_DATA) + assert list(m.variables) == ["p"] + assert list(m.constraints) == ["power_balance"] + reread = math_spec.to_program(yaml.safe_load(m.spec.text)) + assert reread.constraints == m.spec.program.constraints + assert isinstance(m.spec, ModelSpec) + + +def test_a_lowered_program_is_refused() -> None: + program = math_spec.to_program(yaml_dict()) + with pytest.raises(TypeError, match="not a lowered Program"): + Model().add_spec(program, DISPATCH_DATA) + + +def _opened(tmp_path: Path) -> Any: + path = tmp_path / "dispatch.yaml" + path.write_text(EXAMPLE_DISPATCH) + with path.open() as handle: + return handle + + +UNREADABLE_SPECS: list[Any] = [ + pytest.param(_opened, TypeError, "not an open file", id="file-object"), + pytest.param( + lambda tmp_path: io.StringIO(EXAMPLE_DISPATCH), + TypeError, + "not an open file", + id="string-io", + ), + pytest.param( + lambda tmp_path: tmp_path / "absent.yaml", + FileNotFoundError, + "no spec file at", + id="missing-path", + ), + pytest.param( + lambda tmp_path: str(tmp_path / "absent.yaml"), + FileNotFoundError, + "no spec file at", + id="missing-path-string", + ), + pytest.param( + lambda tmp_path: {"description": "a spec that declares nothing"}, + SpecDataError, + "the spec declares nothing", + id="empty-spec", + ), + pytest.param( + lambda tmp_path: "- p_max\n- load\n", + SpecDataError, + "a spec is a mapping of sections", + id="sequence", + ), +] + + +@pytest.mark.parametrize(("form", "error", "match"), UNREADABLE_SPECS) +def test_an_unreadable_spec_is_refused( + tmp_path: Path, form: Callable[[Path], Any], error: type[Exception], match: str +) -> None: + with pytest.raises(error, match=match): + Model().add_spec(form(tmp_path), DISPATCH_DATA) + + +@pytest.mark.parametrize("name", ["dispatch-v2", "a/b", ""]) +def test_a_layer_name_a_file_cannot_carry_is_refused(name: str) -> None: + with pytest.raises(ValueError, match="cannot be named"): + Model().add_spec(yaml_dict(), DISPATCH_DATA, name=name) + + +def test_a_dashed_file_stem_is_refused_before_it_reaches_a_file(tmp_path: Path) -> None: + path = tmp_path / "dispatch-v2.yaml" + path.write_text(EXAMPLE_DISPATCH) + with pytest.raises(ValueError, match="cannot be named 'dispatch-v2'"): + Model.from_spec(path, DISPATCH_DATA) + assert Model.from_spec(path, DISPATCH_DATA, name="dispatch").spec.layers.keys() == { + "dispatch" + } + + +def test_a_second_spec_must_bind_or_not_collide() -> None: + m = Model() + m.add_variables(name="x") + 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()) + assert float(m.variables["p"].upper.max()) == 200.0 + m.solve(solver_name="highs", output_flag=False) + assert float(total.solution) == pytest.approx(float(DISPATCH_P.sum())) + assert float(m.solution["p"].sel(generator="wind").max()) == pytest.approx(90.0) + + +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) + + +def test_a_binding_must_be_a_variable_of_this_model() -> None: + with pytest.raises(SpecDataError, match="bound to a variable of another model"): + extended(p=BASE_MODEL().variables["p"]) + + +P_OVER_GENERATOR = with_( + EXTRA_SPEC, + variables={"p": {"dims": ["generator"]}}, + constraints={"p_cap": {"dims": ["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": {"dims": ["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" + with pytest.raises(ValueError, match="v1"): + Model.from_spec(yaml_dict(), DISPATCH_DATA) + + +def test_a_model_without_a_spec_has_no_accessor() -> None: + with pytest.raises(AttributeError, match="holds no spec"): + _ = Model().spec + + +def test_from_spec_passes_model_kwargs_and_chains() -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA, force_dim_names=True) + assert m.force_dim_names + assert Model().add_spec( + yaml_dict(), DISPATCH_DATA + ).spec.program.variables.keys() == {"p"} + + +# --------------------------------------------------------------------------- +# retain and evaluate +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("retain", "kept"), + [ + ("report", {"cost", "p_max"}), + ("all", {"cost", "load", "p_max"}), + ("none", set()), + ], +) +def test_retain_decides_what_is_kept_and_not_what_can_be_read( + retain: str, kept: set[str] +) -> None: + """A parameter retain dropped is read from the sources the model still holds.""" + m = solved(yaml_dict(), DISPATCH_DATA, retain=retain) + assert set(m.spec.parameters.data_vars) == kept + assert not m.parameters.data_vars + want = (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") + xr.testing.assert_allclose(m.spec.expressions["spend"].solution, want) + xr.testing.assert_allclose(m.spec.evaluate("spend", DISPATCH_DATA).solution, want) + + +def test_the_spec_keeps_its_parameters_off_the_model() -> None: + """``model.parameters`` is the caller's: a build neither reads nor writes it.""" + own = xr.DataArray(np.array(["a", "b", "c"], dtype=object), dims=["own"]) + m = Model() + m.parameters["cost"] = own + m.add_spec(yaml_dict(), DISPATCH_DATA, retain="all") + + assert m.parameters["cost"].equals(own) + assert m.spec.parameters["cost"].dims == ("generator",) + + +def test_a_build_that_cannot_retain_leaves_the_model_buildable() -> None: + """retain='all' reaches parameters no declaration does, and must not half-build on one.""" + spec = with_(yaml_dict(), parameters={"spare": {"dims": ["generator"]}}) + m = Model() + with pytest.raises(SpecDataError, match="no data provided for parameter 'spare'"): + m.add_spec(spec, DISPATCH_DATA, retain="all") + assert not len(m.variables) and not len(m.constraints) + + spare = pd.Series([1.0, 2.0], index=GENERATOR) + m.add_spec(spec, {**DISPATCH_DATA, "spare": spare}, retain="all") + assert "spare" in m.spec.parameters + + +def test_a_declared_dimension_with_no_source_still_reprs() -> None: + """A dimension nothing reaches needs no source, so the repr must do without its labels.""" + spec = with_( + yaml_dict(), dimensions={"spare": {"dtype": "int", "description": "unreached"}} + ) + m = Model.from_spec(spec, DISPATCH_DATA) + + assert "spare (unreached)" in repr(m.spec) + assert "snapshot (3)" in repr(m.spec) + + +def test_an_unknown_expression_is_a_key_error_with_a_hint() -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + with pytest.raises(KeyError, match="unknown named expression 'spent'.*spend"): + m.spec.expressions["spent"] + + +def test_a_fold_over_variables_needs_a_solution_and_one_over_data_does_not() -> None: + spec = { + **yaml_dict(), + "parameters": { + **yaml_dict()["parameters"], + "rate": {"dims": []}, + "years": {"dims": []}, + }, + "expressions": { + "spend": "sum(p * cost, over=generator)", + "growth": "rate ** years", + }, + } + m = Model.from_spec(spec, {**DISPATCH_DATA, "rate": 1.05, "years": 3.0}) + assert float(m.spec.expressions["growth"].solution) == pytest.approx(1.05**3) + with pytest.raises(RuntimeError, match="no solution yet"): + m.spec.expressions["spend"].solution + + +# --------------------------------------------------------------------------- +# three views: math, the linopy expression and the solution +# --------------------------------------------------------------------------- + +VIEWS_SPEC: dict[str, Any] = { + **math_spec.to_spec(yaml.safe_load(EXAMPLE_DISPATCH)).to_dict(), + "expressions": { + "spend": "sum(p * cost, over=generator)", + "bare": "p", + "levels": "cost * 2", + "answer": "6 * 7", + }, +} + + +@pytest.mark.parametrize( + ("name", "kind"), + [ + ("spend", linopy.LinearExpression), + ("bare", linopy.LinearExpression), + ("levels", xr.DataArray), + ("answer", float), + ], +) +def test_expression_is_the_unsolved_linopy_term(name: str, kind: type) -> None: + m = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA) + assert isinstance(m.spec.expressions[name].expression, kind) + + +@pytest.mark.parametrize("name", ["spend", "bare"]) +def test_a_variable_bearing_named_expression_is_the_model_expression(name: str) -> None: + """The container holds one object per name; the data-only bodies stay on the spec.""" + m = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA) + assert set(m.expressions) == {"spend", "bare"} + assert m.spec.expressions[name].expression is m.expressions[name] + assert m.expressions[name].spec == "spec" + + +def test_expression_reads_unsolved_but_solution_waits_for_a_solve() -> None: + m = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA) + e = m.spec.expressions["spend"] + assert isinstance(e.expression, linopy.LinearExpression) + with pytest.raises(RuntimeError, match="no solution yet"): + e.solution + m.solve("highs") + xr.testing.assert_allclose( + e.solution, (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") + ) + + +def test_the_named_expression_bundles_the_three_views() -> None: + m = solved(VIEWS_SPEC, DISPATCH_DATA) + e = m.spec.expressions["spend"] + assert e.node is m.spec.program.named_expressions["spend"].expression + assert isinstance(e.expression, linopy.LinearExpression) + xr.testing.assert_allclose( + e.solution, (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") + ) + + +def test_evaluate_returns_a_named_expression() -> None: + m = solved(VIEWS_SPEC, DISPATCH_DATA, retain="none") + e = m.spec.evaluate("spend", DISPATCH_DATA) + assert isinstance(e, NamedExpression) + assert isinstance(e.expression, linopy.LinearExpression) + xr.testing.assert_allclose( + e.solution, (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") + ) + + +def test_repr_summarises_every_section() -> None: + text = repr(Model.from_spec(yaml_dict(), DISPATCH_DATA).spec) + assert text.startswith("ModelSpec: Least-cost dispatch") + assert "Dimensions: snapshot (3), generator (2)" in text + assert "Variables: p" in text + assert "Constraints: power_balance" in text + assert "Objective: minimize" in text + 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) + assert "(+6 more)" in text + assert "e11" not in text + + +def test_model_repr_of_a_whole_spec_model_carries_no_tags() -> None: + """Everything is the spec's, so naming the layer on every line would say nothing.""" + text = repr(Model.from_spec(yaml_dict(), DISPATCH_DATA)) + assert "Linopy LP model, built from a math-spec" in text + assert "Least-cost dispatch of a generator fleet against an hourly load." in text + assert " * spend (snapshot)\n" in text + assert " * usage (snapshot, generator)\n" in text + assert " * p (snapshot, generator)\n" in text + assert " * power_balance (snapshot)\n" in text + assert "[spec]" not in text + assert "" not in text + + +def test_model_repr_of_an_extended_model_names_its_layers() -> None: + """A bound variable is the model's, so it stays untagged; a container the layer fills alone needs no tag either.""" + 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)\n" 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\n" in text + + m.add_expressions(m.variables["reserve"] * 2.0, name="reserve_cost") + text = repr(m) + assert " * total [extra]" in text + assert " * reserve_cost (generator)\n" 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) + assert m.spec.description == "" + assert repr(m).startswith("Linopy LP model, built from a math-spec\n=") + + +def test_hybrid_model_tags_spec_variables_constraints_and_expressions() -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + v = m.add_variables(lower=0, coords=[GENERATOR], name="reserve") + m.add_expressions(v * 2.0, name="reserve_cost") + m.add_constraints(v <= 10.0, name="reserve_cap") + text = repr(m) + assert " * p (snapshot, generator) [spec]" in text + assert " * reserve (generator)\n" in text + assert " * power_balance (snapshot) [spec]" in text + assert " * reserve_cap (generator)\n" in text + assert " * reserve_cost (generator)\n" in text + assert " * spend (snapshot) [spec]" in text + 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() + assert "```math" in spec.to_markdown() + assert spec.to_typst() + shown = spec._repr_markdown_() + assert "```math" not in shown and "$`" not in shown + assert "$$\n" in shown and " $t$ " in shown + + +@pytest.mark.parametrize("fmt", ["latex", "markdown", "typst"]) +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) + body = f"{spec['spec'].typeset(fmt)}\n\n{spec['extra'].typeset(fmt)}" + assert rendered.endswith(body) + assert "binds from the host model" in rendered + 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") + assert spec.typeset(fmt) == getattr(spec, f"to_{fmt}")() + assert declaration.typeset(fmt) == getattr(declaration, f"to_{fmt}")() + + +def hybrid() -> Model: + """A spec-built model grown past its spec by hand.""" + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + m.add_variables(lower=0, coords=[GENERATOR], name="reserve") + m.add_constraints(m.variables["reserve"] <= 10.0, name="reserve_cap") + return m + + +def test_unspecified_names_what_the_spec_does_not_declare() -> None: + assert not Model.from_spec(yaml_dict(), DISPATCH_DATA).spec.unspecified + assert hybrid().spec.unspecified == Unspecified( + variables=("reserve",), + constraints=("reserve_cap",), + expressions=(), + sos=(), + piecewise=(), + objective=False, + ) + m = extended() + found = m.spec.unspecified + assert found.variables == () + assert found.constraints == ("power_balance",) + assert found.expressions == () and "total" in m.expressions + assert found.bound == ("p",) + assert not Unspecified((), (), (), (), (), False, ("p",)) + + +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": {"dims": over}}, + "constraints": {"q_cap": {"dims": 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_a_layer_refuses_removal_of_a_name_it_owns() -> None: + """A bound variable and a built constraint each strand a layer, so removal raises; a hand name still goes.""" + m = extended() + with pytest.raises(ValueError, match="p is declared or bound by a spec layer"): + m.remove_variables("p") + with pytest.raises(ValueError, match="p_cap is declared or bound"): + m.remove_constraints("p_cap") + with pytest.raises(ValueError, match="total is declared or bound"): + m.remove_expressions("total") + assert {"power_balance", "p_cap"} <= set(m.constraints) + m.add_variables(lower=0, coords=[GENERATOR], name="free") + m.remove_variables("free") + assert "free" not in m.variables + m.add_expressions(m.variables["p"].sum(), name="free_expr") + m.remove_expressions("free_expr") + assert "free_expr" not in m.expressions + + +def test_a_refused_list_removal_drops_nothing() -> None: + """A list with one owned name is refused as a whole, so the hand names before it stay.""" + m = extended() + m.add_constraints(m.variables["p"].sum() <= 1, name="hand") + m.add_expressions(m.variables["p"].sum(), name="hand_expr") + with pytest.raises(ValueError, match="p_cap is declared or bound"): + m.remove_constraints(["hand", "p_cap"]) + with pytest.raises(ValueError, match="total is declared or bound"): + m.remove_expressions(["hand_expr", "total"]) + assert "hand" in m.constraints + assert "hand_expr" in m.expressions + + +@pytest.mark.parametrize( + "remove", + [ + lambda m: m.variables.remove("p"), + lambda m: m.constraints.remove("p_cap"), + lambda m: m.expressions.remove("total"), + ], + ids=["variable", "constraint", "expression"], +) +def test_the_containers_refuse_removal_of_what_a_layer_owns( + remove: Callable[[Model], None], +) -> None: + m = extended() + with pytest.raises(ValueError, match="declared or bound by a spec layer"): + remove(m) + + +@pytest.mark.parametrize("build_expressions", [True, False], ids=["built", "lazy"]) +@pytest.mark.parametrize( + ("add", "match"), + [ + (lambda m: m.add_variables(name="p"), "variable 'p'"), + ( + lambda m: m.add_constraints(m.variables["p"] >= 0, name="power_balance"), + "constraint 'power_balance'", + ), + ( + lambda m: m.add_expressions(m.variables["p"].sum(), name="spend"), + "expression 'spend'", + ), + ], + ids=["variable", "constraint", "expression"], +) +def test_a_name_a_layer_declares_cannot_be_added_by_hand( + build_expressions: bool, add: Callable[[Model], Any], match: str +) -> None: + """A lazy named expression is in no container, and is the layer's all the same.""" + m = Model.from_spec(yaml_dict(), DISPATCH_DATA, build_expressions=build_expressions) + with pytest.raises(ValueError, match=f"{match} is declared or bound by spec layer"): + add(m) + + +@pytest.mark.parametrize("build_expressions", [True, False], ids=["built", "lazy"]) +def test_a_hand_expression_collides_with_a_declared_one_however_it_is_held( + build_expressions: bool, +) -> None: + m = BASE_MODEL() + m.add_expressions(m.variables["p"].sum(), name="total") + data = {**EXTRA_DATA, "p": m.variables["p"]} + with pytest.raises(ValueError, match=r"named expression\(s\) \['total'\]"): + m.add_spec(EXTRA_SPEC, data, build_expressions=build_expressions) + + +def _assign(m: Model) -> None: + m.objective = m.variables["p"].sum() * 2.0 + + +def _augment(m: Model) -> None: + m.objective += m.variables["p"].sum() + + +@pytest.mark.parametrize( + "edit", + [ + _assign, + _augment, + Model.remove_objective, + lambda m: m.add_objective(m.variables["p"].sum(), overwrite=True), + ], + ids=["assign", "augment", "remove", "overwrite"], +) +def test_an_edited_objective_is_no_layers(edit: Callable[[Model], None]) -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + assert m.spec.objective_owner == "spec" + edit(m) + assert m.spec.objective_owner is None + assert m.spec.unspecified.objective is not m.objective.expression.empty + + +def test_removing_a_hand_variable_the_objective_never_read_keeps_its_owner() -> None: + m = hybrid() + with pytest.warns(UserWarning, match="also removes constraints"): + m.remove_variables("reserve") + assert m.spec.objective_owner == "spec" + assert not m.spec.unspecified + + +def test_a_dropped_sos_is_the_layers_no_longer() -> None: + m = extended(SOS_SPEC) + assert m.spec.unspecified.sos == () + m.remove_sos_constraints(m.variables["p"]) + assert m.spec.unspecified.sos == () + m.add_sos_constraints(m.variables["p"], sos_type=2, sos_dim="generator") + assert m.spec.unspecified.sos == ("p",) + + +def test_remove_spec_takes_a_layer_off_with_what_it_built() -> None: + """The second layer binds what the first built, so the first goes last; the objective goes with its layer.""" + m = two_layers() + m.add_constraints(m.variables["p"].sum() <= 1e6, name="hand") + with pytest.raises( + ValueError, + match=r"built variable\(s\) \['p'\] that layer\(s\) \['extra'\] bind", + ): + m.remove_spec("spec") + with pytest.raises(KeyError, match="unknown spec layer 'extr'"): + m.remove_spec("extr") + + m.remove_spec("extra") + assert list(m.spec.layers) == ["spec"] + assert "p_cap" not in m.constraints and "total" not in m.expressions + assert m.spec.objective_owner == "spec" + assert m.spec.unspecified == Unspecified((), ("hand",), (), (), (), False) + + with pytest.warns(UserWarning, match=r"also removes constraints \['hand'\]"): + m.remove_spec("spec") + assert not len(m.variables) and not len(m.constraints) + assert not len(m.expressions) and m.objective.expression.empty + with pytest.raises(AttributeError, match="holds no spec"): + _ = m.spec + assert list(Model.from_spec(yaml_dict(), DISPATCH_DATA).variables) == ["p"] + + +def test_remove_spec_leaves_a_bound_variable_and_takes_its_sos() -> None: + m = extended(SOS_SPEC) + m.remove_spec("extra") + assert list(m.variables) == ["p"] and list(m.constraints) == ["power_balance"] + assert "sos_type" not in m.variables["p"].attrs + assert not m.objective.expression.empty + with pytest.raises(AttributeError, match="holds no spec"): + _ = m.spec + m.add_spec(SOS_SPEC, {**EXTRA_DATA, "p": m.variables["p"]}, name="extra") + assert m.spec.unspecified.sos == () and "p_cap" in m.constraints + + +def test_layers_are_read_only() -> None: + m = extended() + with pytest.raises(TypeError): + m.spec.layers["other"] = m.spec["extra"] # type: ignore[index] + with pytest.raises(TypeError): + del m.spec.layers["extra"] # type: ignore[attr-defined] + + +def test_unspecified_sees_what_carries_no_name_of_its_own() -> None: + """An SOS is attributes on a variable, and a replaced objective is no name at all.""" + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + m.add_expressions(m.variables["p"].sum("generator"), name="hand_expr") + m.add_sos_constraints(m.variables["p"], sos_type=2, sos_dim="generator") + m.add_objective(m.variables["p"].sum() * 3.0, overwrite=True) + + found = m.spec.unspecified + assert found.expressions == ("hand_expr",) + assert found.sos == ("p",) + assert found.objective + assert found.variables == () and found.constraints == () + + +def test_a_piecewise_formulation_is_named_as_one_and_not_as_its_parts() -> None: + """Its own variables and constraints are the formulation's business, not the tally's.""" + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + k = pd.Index([0, 1], name="k") + pts = {"k": k, "_breakpoint": [0, 1, 2]} + x = m.add_variables(lower=0, upper=10, coords=[k], name="pw_x") + y = m.add_variables(lower=0, upper=10, coords=[k], name="pw_y") + m.add_piecewise_formulation( + (x, breakpoints(xr.DataArray([[0.0, 5.0, 10.0]] * 2, coords=pts))), + (y, breakpoints(xr.DataArray([[0.0, 1.0, 4.0]] * 2, coords=pts))), + name="curve", + ) + + found = m.spec.unspecified + assert found.piecewise == ("curve",) + assert found.variables == ("pw_x", "pw_y") + 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", "