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..f74c7c5c8 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -115,6 +115,26 @@ 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. + +.. autosummary:: + :toctree: generated/ + + model.Model.add_spec + model.Model.from_spec + model.Model.spec + spec.ModelSpec + spec.NamedExpressions + spec.NamedExpression + 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..9210ff807 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -21,6 +21,15 @@ 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 `__ YAML program attached to data, and ``model.spec`` (a ``linopy.spec.ModelSpec``) reads it back. Requires the ``spec`` dependency group (``uv sync --group spec`` / ``uv pip install --group spec``, Python >= 3.12) and v1 semantics. Data is attached onto the spec's dimensions and parameters with ``linopy.spec.attach``, raising a ``linopy.spec.SpecDataError`` on mismatched or missing data; ``linopy.spec.Attached`` carries the attached result. The parameters the spec retains live in ``model.spec.parameters``, its own dataset; ``model.parameters`` stays the caller's and a build never writes to it. ``retain`` decides what a netcdf file holds, not what a session can read: a parameter it dropped is resolved from the sources the model was built with, and only a model read back from a file can run out of data. The spec API emits an :class:`linopy.EvolvingAPIWarning` once per session while it stabilises. See :doc:`building-models-from-specs` for a worked example. + +* ``model.spec.expressions`` (a ``linopy.spec.NamedExpressions`` mapping) returns a ``linopy.spec.NamedExpression`` for each declared name, with three views: ``.node`` (the lowered formula), ``.expression`` (the unsolved linopy expression — a ``LinearExpression``, bare ``Variable``, array or scalar) and ``.solution`` (the expression folded over the solved model). ``model.spec.evaluate(name, sources)`` returns the same object with its parameters attached afresh. + +* ``model.spec.to_latex`` / ``.to_markdown`` / ``.to_typst`` typeset the whole model, and ``model.spec.declaration(name)`` returns a ``linopy.spec.Declaration`` whose same three methods typeset one named expression, constraint or variable as a single line (math only, no document); a ``NamedExpression`` carries those methods too. A ``ModelSpec``, a ``Declaration`` and a ``NamedExpression`` all render as Markdown in a notebook. + + *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`. @@ -31,6 +40,14 @@ Upcoming Version *Other* +* ``Model.add_spec`` resolves the parameters it retains before it builds. A ``retain="all"`` build that could not read a parameter no declaration uses raised after the variables and constraints were already added, leaving a model that the "builds into an empty model" guard then refused to build into again. + +* ``repr(model.spec)`` no longer raises ``KeyError`` for a dimension the spec declares but nothing reaches, which needs no source and so has no coordinates; it is shown as ``unreached``. + +* A grouped sum through a lookup that maps no member at all now holds the empty sum, ``0``, on every declared group, as its documented rule says. It raised xarray's ``ValueError: must not be empty`` instead. + +* ``read_netcdf`` no longer rewrites the coordinates of a container that merely shares a dimension's *name* with a spec-built model's master coordinates. A hand-added variable on its own labels kept them; before, it was silently relabelled onto the master ones, or the read failed outright when the two lengths differed. + * ``add_piecewise_formulation`` gained a ``mask`` parameter declaring which breakpoint slots hold a real breakpoint. It is needed for **ragged** curves — entities with different numbers of breakpoints — which are stored densely with the surplus slots left absent. Under v1 that absence must be declared (``mask=x_pts.notnull()``) rather than read off the NaN padding. (https://github.com/PyPSA/linopy/issues/884) *Internal* 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..cfb2ec048 --- /dev/null +++ b/examples/building-models-from-specs.ipynb @@ -0,0 +1,970 @@ +{ + "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", + " foreach: [snapshot, generator]\n", + " where: \"p_max > 0\"\n", + " bounds: { lower: 0, upper: p_max }\n", + "\n", + "constraints:\n", + " power_balance:\n", + " foreach: [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 `foreach: [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 `foreach: [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. These are **not** part of the\n", + " optimization. They are post-solve read-outs: after solving you can ask for\n", + " `spend` (cost per hour) or `usage` (output as a fraction of capacity) and get\n", + " them back as numeric arrays. More on this 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", + "`add_spec` builds into an *empty* model; `from_spec` is sugar that makes the\n", + "model for you and forwards any `Model(...)` keyword arguments." + ] + }, + { + "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. A `LinearExpression`, a bare `Variable`, an\n", + " array, or a scalar (a named expression is affine, so never 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)" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": "### The model as maths\n\nThe accessor typesets the whole model, delegating to math-spec:\n`m.spec.to_latex()`, `.to_markdown()` and `.to_typst()`. In a notebook the\naccessor renders as Markdown on its own; here we show it explicitly.\n\nAny single declaration typesets on its own too. `m.spec.declaration(name)`\ntakes a named expression, a constraint or a variable and hands back a\n`Declaration` with the same three methods; a `NamedExpression` carries them\ndirectly. These render **one** line — math only, no surrounding document — so\nthe 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": [ + "from IPython.display import Markdown\n", + "\n", + "Markdown(m.spec.to_markdown())" + ] + }, + { + "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", + "Markdown(f\"$$\\n{m.spec.declaration('power_balance').to_markdown()}\\n$$\")" + ] + }, + { + "cell_type": "markdown", + "id": "19", + "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": "20", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " m.spec.expressions[\"spent\"]\n", + "except KeyError as e:\n", + " print(\"KeyError:\", e)" + ] + }, + { + "cell_type": "markdown", + "id": "21", + "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": "22", + "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": "23", + "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": "24", + "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": "25", + "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": "26", + "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": "27", + "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\": {\"foreach\": [\"t\"], \"bounds\": {\"lower\": 0, \"upper\": 10}}},\n", + " \"constraints\": {\"cap\": {\"foreach\": [\"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": "28", + "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": "29", + "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\": {\"foreach\": [\"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\": {\"foreach\": [\"t\"], \"expression\": \"x / w <= c\"}},\n", + "}\n", + "refuse(div_spec, {\"w\": w_hole, \"c\": c_full}, \"divisor\")" + ] + }, + { + "cell_type": "markdown", + "id": "30", + "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": "31", + "metadata": {}, + "outputs": [], + "source": [ + "where_spec = {\n", + " **SPARSE,\n", + " \"constraints\": {\n", + " \"cap\": {\"foreach\": [\"t\"], \"where\": \"w\", \"expression\": \"w * x <= c\"}\n", + " },\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": "32", + "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": "33", + "metadata": {}, + "outputs": [], + "source": [ + "masked_spec = {\n", + " **SPARSE,\n", + " \"parameters\": {**SPARSE[\"parameters\"], \"live\": {\"dims\": [\"t\"], \"dtype\": \"bool\"}},\n", + " \"variables\": {\n", + " \"x\": {\"foreach\": [\"t\"], \"where\": \"live\", \"bounds\": {\"lower\": 0, \"upper\": \"c\"}}\n", + " },\n", + " \"constraints\": {\n", + " \"cap\": {\"foreach\": [\"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": "34", + "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": "35", + "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\": {\n", + " \"imports\": {\"foreach\": [\"bus\"], \"bounds\": {\"lower\": 0, \"upper\": 100}}\n", + " },\n", + " \"constraints\": {\n", + " \"import_limit\": {\n", + " \"foreach\": [\"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": "36", + "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": "37", + "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": "38", + "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: { foreach: [snapshot], bounds: { lower: 0, upper: 1000 } }\n", + " charge: { foreach: [snapshot], bounds: { lower: 0, upper: soc_max } }\n", + " discharge: { foreach: [snapshot], bounds: { lower: 0, upper: soc_max } }\n", + " soc: { foreach: [snapshot], bounds: { lower: 0, upper: soc_max } }\n", + "constraints:\n", + " balance:\n", + " foreach: [snapshot]\n", + " expression: gen + discharge - charge == load\n", + " storage:\n", + " foreach: [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": "39", + "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": "40", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"folded cost:\", float(bm.spec.expressions[\"cost\"].solution))" + ] + }, + { + "cell_type": "markdown", + "id": "41", + "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": "42", + "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": "43", + "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": "44", + "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": "45", + "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": "46", + "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.parameters.data_vars) == [])\n", + "print(lean_back.spec.evaluate(\"spend\", dispatch_data))" + ] + }, + { + "cell_type": "markdown", + "id": "47", + "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": "48", + "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": "49", + "metadata": {}, + "source": "## Where the code lives\n\nThe feature is a small package, `linopy/spec/`, imported only when you call\n`add_spec`/`from_spec` — `import linopy` never pulls in `math_spec`. Roughly:\n\n- `accessor.py` — `model.spec`, the `NamedExpression` views, `evaluate`, and\n typesetting: the whole model (`m.spec.to_latex` / `.to_markdown` /\n `.to_typst`) and any single declaration — a named expression, constraint or\n variable — via `m.spec.declaration(name)` and math-spec's\n `typeset_declaration`.\n- `attach.py` — the three attachment rules; data onto master coordinates.\n- `builder.py` — emits variables, constraints, objective; folds expressions.\n- `operators.py` — `sum`, `by=`, `shift`, `at`, `sum_back`.\n- `where.py` — `where:` predicates as boolean masks.\n- `coverage.py` / `terms.py` — the absence rule from section 6: a missing row\n is refused wherever it is used.\n- `curves.py` — the data side of `piecewise:` blocks.\n- `netcdf.py` — the factorize-based persistence from section 10.\n- `nodes.py` — walks over expression nodes, and the dimensions a node\n spans before any data is bound.\n\n### Summary\n\nA spec is the maths over labelled axes; the sources are the numbers. `linopy`\nattaches them into an ordinary model, hands each named expression back as three\nviews — its formula, its unsolved linopy expression and its solution — refuses a\nmissing parameter row wherever it is used (as a coefficient, bound, constant\nside or divisor alike, with `where:` and filling the data as the escape\nhatches), and round-trips the lot through netCDF by keeping the spec as text\nbeside factorized labels." + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "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..91211e61d 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 @@ -124,6 +125,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..6cc02cdd9 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -2118,7 +2118,9 @@ def _formatted_names(self) -> dict[str, str]: """ return {format_string_as_variable_name(n): n for n in self} - def _format_items(self, exclude: set[str] | None = None) -> str: + def _format_items( + self, exclude: set[str] | None = None, tag: set[str] | None = None + ) -> str: """Format constraint items, optionally excluding names in a group.""" r = "" count = 0 @@ -2131,7 +2133,8 @@ def _format_items(self, exclude: set[str] | None = None) -> str: if ds.coords else "" ) - r += f" * {name}{coords}\n" + suffix = " [spec]" if tag and name in tag else "" + r += f" * {name}{coords}{suffix}\n" if count == 0: r += "\n" return r diff --git a/linopy/io.py b/linopy/io.py index 08701ed03..22781750e 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -11,7 +11,7 @@ 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 io import BufferedWriter from pathlib import Path @@ -45,7 +45,9 @@ logger = logging.getLogger(__name__) NETCDF_VERSION_ATTR = "_linopy_version" +DTYPE_ATTR = "_linopy_dtype" EXPR_TYPE_ATTR = "_linopy_expr_type" +SPEC_ATTR = "_linopy_spec" CONTAINER_ORDER_ATTR = "_linopy_{}_order" @@ -1019,6 +1021,145 @@ 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* where *found* is it as a netcdf type gave it back, else *found* itself. + + A narrowed int or a widened bool holds the same labels at another dtype + and is the one to replace -- which is what ``Index.equals`` asks, since it + compares labels and not dtypes. An index of another length, or of other + labels entirely, belongs to a container that was never built on *master* + and is left alone. + """ + if master is None or found.dtype == master.dtype: + return found + return master if found.equals(master) else found + + def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: """ Write out the model to a netcdf file. @@ -1040,6 +1181,12 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: :func:`linopy.io.read_netcdf`. The insertion order of each container is stored as a JSON list in the ``_linopy__order`` attribute. + A model built with :meth:`Model.add_spec` also persists its spec under a + ``spec-`` prefix of its own: the YAML text, the master coordinates and the + parameters the spec retained, apart from ``m.parameters``. ``read_netcdf`` + lowers the program from the text again, so reading such a file needs + 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 +1208,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 +1231,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(m._spec)] + 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 ( @@ -1176,43 +1314,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 +1379,12 @@ 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_ATTR in ds.attrs: + from linopy.spec.netcdf import decode + + m._spec = decode(m, ds, ds.attrs[SPEC_ATTR]) for k in m.scalar_attrs: if k in ds.attrs: @@ -1412,6 +1518,8 @@ 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._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..41c1e2bd2 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -117,6 +117,7 @@ if TYPE_CHECKING: from linopy.piecewise import PiecewiseFormulation + from linopy.spec import ModelSpec, Retain, SpecLike logger = logging.getLogger(__name__) @@ -202,6 +203,7 @@ class Model: "_piecewise_formulations", "_solver", "_sos_reformulation_state", + "_spec", "__weakref__", ) @@ -305,6 +307,7 @@ def __init__( ) self._solver: solvers.Solver | None = None self._sos_reformulation_state: SOSReformulationResult | None = None + self._spec: ModelSpec | None = None @property def solver(self) -> solvers.Solver | None: @@ -437,6 +440,94 @@ def solution(self) -> Dataset: """ return self.variables.solution + @property + def spec(self) -> ModelSpec: + """ + The math-spec program this model was built from, see :meth:`add_spec`. + + Raises + ------ + AttributeError + If the model was not built from a spec. + """ + if self._spec is None: + raise AttributeError( + "This model was not built from a spec. Use `Model.add_spec` or " + "`Model.from_spec` to build one." + ) + return self._spec + + def add_spec( + self, + spec: SpecLike, + sources: Mapping[str, Any] | Dataset, + retain: Retain = "report", + ) -> Model: + """ + Build a math-spec program with its data into this empty 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``. + + Parameters + ---------- + spec : str, pathlib.Path, dict or math_spec.Spec + The spec. A ``str`` containing a newline is YAML text, any other + ``str`` is a path. A lowered ``math_spec.Program`` is refused, + since it has no YAML form to keep on the model. + sources : mapping or xarray.Dataset + Data keyed by declared name: dimension labels, parameters and + lookups. Read by key on demand and never iterated. + 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. + + Returns + ------- + linopy.Model + This model, for chaining. + + Raises + ------ + ValueError + If the model already holds variables or constraints, or runs + under legacy semantics. + linopy.spec.SpecDataError + If the data does not fit the spec. + + 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) + return self + + @classmethod + def from_spec( + cls, + spec: SpecLike, + sources: Mapping[str, Any] | Dataset, + retain: Retain = "report", + **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) + @property def dual(self) -> Dataset: """ @@ -618,13 +709,33 @@ 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" + var_tag: set[str] | None = None + con_tag: set[str] | None = None + expr_string = self.expressions._format_items() + if self._spec is not None: + model_string += ", built from a math-spec" + program = self._spec.program + spec_vars = set(program.variables) + spec_cons = set(program.constraints) + if any(v not in spec_vars for v in self.variables): + var_tag = spec_vars + if any(c not in spec_cons for c in self.constraints): + con_tag = spec_cons + eager = expr_string if len(self.expressions) else "" + spec = "".join( + f" * {name} ({', '.join(e.dims)}) [spec]\n" + for name, e in self._spec.expressions.items() + ) + expr_string = eager + spec or "\n" + var_string = self.variables._format_items(exclude=var_names, tag=var_tag) + con_string = self.constraints._format_items(exclude=con_names, tag=con_tag) + header = f"{model_string}\n{'=' * len(model_string)}\n" + if self._spec is not None and self._spec.description: + header += f"{self._spec.description}\n" 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}" 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..d3d76bbb9 --- /dev/null +++ b/linopy/spec/__init__.py @@ -0,0 +1,39 @@ +""" +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 + +from importlib.util import find_spec + +if find_spec("math_spec") is None: + raise ImportError( + "linopy.spec needs the math-spec package. Install it with " + "`pip install math-spec` (Python >= 3.12) and try again." + ) + +from linopy.spec.accessor import ( + Declaration, + ModelSpec, + NamedExpression, + NamedExpressions, + SpecLike, +) +from linopy.spec.attach import Attached, Retain, attach +from linopy.spec.errors import SpecDataError + +__all__ = [ + "Attached", + "Declaration", + "ModelSpec", + "NamedExpression", + "NamedExpressions", + "Retain", + "SpecDataError", + "SpecLike", + "attach", +] diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py new file mode 100644 index 000000000..b8cec5795 --- /dev/null +++ b/linopy/spec/accessor.py @@ -0,0 +1,424 @@ +""" +``model.spec``: the program a model was built from, and its 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 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 +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Any, TypeAlias + +import pandas as pd +import xarray as xr +import yaml +from math_spec import ( + Spec, + did_you_mean, + to_latex, + to_markdown, + to_program, + to_spec, + to_typst, + typeset_declaration, +) +from math_spec import program as ms + +from linopy.constants import warn_evolving_api +from linopy.model import Model +from linopy.semantics import is_v1 +from linopy.spec import terms +from linopy.spec.attach import EVOLVING_MESSAGE, Attached, Retain +from linopy.spec.attach import attach as attach_data +from linopy.spec.builder import build +from linopy.spec.context import Context +from linopy.spec.errors import SpecDataError +from linopy.spec.evaluate import evaluate_named, fold +from linopy.spec.nodes import dims_of +from linopy.spec.parameters import Parameters, Resolve + +SpecLike: TypeAlias = str | Path | Mapping[str, Any] | Spec + + +def attach( + model: Model, + spec: SpecLike, + sources: Mapping[str, Any] | xr.Dataset, + retain: Retain, +) -> ModelSpec: + """ + Build *spec* with *sources* into the empty *model* and return its accessor. + + Raises + ------ + ValueError + The model already holds variables or constraints, or runs + under legacy semantics. + TypeError + *spec* is a lowered ``Program``, which has no YAML form to + keep on the model. + """ + 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." + ) + if len(model.variables) or len(model.constraints): + raise ValueError( + "add_spec builds into an empty model, and this one already holds " + f"{len(model.variables)} variable(s) and {len(model.constraints)} constraint(s)." + ) + text, program = _source(spec) + attached: Attached = attach_data(program, sources, retain=retain) + # Resolved before the build, so a parameter no declaration reads cannot fail + # halfway through one and leave a model too full to build into again. + parameters = attached.retained().assign_coords(dict(attached.coords)) + build(model, attached) + return ModelSpec(model, program, text, parameters, attached) + + +def restore(model: Model, text: str, parameters: xr.Dataset) -> ModelSpec: + """ + The accessor for *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. + """ + return ModelSpec(model, to_program(yaml.safe_load(text)), text, parameters, None) + + +def _source(spec: SpecLike) -> tuple[str, ms.Program]: + """The spec as the YAML text kept on the model, and lowered.""" + 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, str) and "\n" not in spec: + spec = Path(spec) + if isinstance(spec, Path): + return spec.read_text(), to_program(spec) + if isinstance(spec, str): + return spec, to_program(yaml.safe_load(spec)) + loaded = to_spec(dict(spec)) if isinstance(spec, Mapping) else spec + return loaded.to_yaml(), to_program(loaded) + + +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 '—'}" + + +class ModelSpec: + """ + The spec a model was built from. + + Attributes + ---------- + program + The lowered spec. + text + The spec as YAML, verbatim where a file or text was passed. + """ + + def __init__( + self, + model: Model, + program: ms.Program, + text: str, + parameters: xr.Dataset, + attached: Attached | None, + ) -> None: + self._model = model + self.program = program + self.text = text + self._parameters = parameters + self._attached = attached + + def __repr__(self) -> str: + p = self.program + coords = self.coords + head = f"ModelSpec: {self.description}" if self.description else "ModelSpec" + 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 "\n".join(rows) + + def _reattach(self, model: Model, deep: bool = True) -> ModelSpec: + """The same spec, read off *model*, holding its own copy of the parameters.""" + return ModelSpec( + model, + self.program, + self.text, + self._parameters.copy(deep=deep), + self._attached, + ) + + @property + def parameters(self) -> xr.Dataset: + """The parameters and lookups the spec retained, on the master coordinates.""" + return self._parameters + + @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 model was built on them.""" + return {str(d): index for d, index in self.parameters.indexes.items()} + + @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(self) + + 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 to_latex(self, **options: Any) -> str: + """The whole model typeset as a LaTeX document.""" + return to_latex(self._schema, **options) + + def to_markdown(self, **options: Any) -> str: + """The whole model typeset as Markdown, its equations in ``$$`` blocks.""" + return to_markdown(self._schema, **options) + + def to_typst(self, **options: Any) -> str: + """The whole model typeset as Typst.""" + return to_typst(self._schema, **options) + + def _repr_markdown_(self) -> str: + return self.to_markdown() + + @property + def _schema(self) -> dict[str, Any]: + """The spec as the mapping the typesetter reads (a bare string it reads as a path).""" + return yaml.safe_load(self.text) + + def evaluate( + self, name: str, sources: Mapping[str, Any] | xr.Dataset + ) -> NamedExpression: + """ + The named expression *name*, with its parameters attached afresh from *sources*. + + For reading the spec against other data than the model was built with, + and for a model read from a file, whose own sources are gone. + ``model.spec.expressions`` needs neither. *sources* is read the way + ``add_spec`` read it, and must describe the coordinates the model was + built on. + + Raises + ------ + SpecDataError + *sources* label a dimension differently than the + model was built on. + """ + attached = attach_data(self.program, sources, retain="none") + coords = self.coords + for dim, index in attached.coords.items(): + if dim in coords and not index.equals(coords[dim]): + raise SpecDataError( + f"sources describe dimension '{dim}' as {index.tolist()[:5]}, and the model " + f"was built on {coords[dim].tolist()[:5]}. evaluate() reads the solution the " + f"model holds, so the data must be attached on the same labels in the same order." + ) + return NamedExpression(self, name, self._context(attached.parameter)) + + def _resolve(self, name: str) -> xr.DataArray: + """The parameter *name*: retained if it was kept, else read from the sources again.""" + if name in self.parameters: + return self.parameters[name] + if self._attached is not None: + return self._attached.parameter(name) + raise SpecDataError( + f"parameter '{name}' was not retained and this model no longer holds the sources " + f"it was built with, which is what a model read from a file looks like. Build with " + f"retain='all' before writing it out, or read the expression with " + f"evaluate(name, sources)." + ) + + def _context(self, resolve: Resolve) -> Context: + return Context( + self._model, + self.program, + self.coords, + self.lookups, + Parameters(self.program, resolve), + solved=True, + ) + + +class NamedExpressions(Mapping[str, "NamedExpression"]): + """The named expressions of a spec, each a :class:`NamedExpression` on read.""" + + def __init__(self, spec: ModelSpec) -> None: + self._spec = spec + + def __getitem__(self, name: str) -> NamedExpression: + if name not in self._spec.program.named_expressions: + raise KeyError( + f"unknown named expression '{name}'. " + + did_you_mean(name, self._spec.program.named_expressions) + ) + return NamedExpression( + self._spec, name, self._spec._context(self._spec._resolve) + ) + + def __iter__(self) -> Iterator[str]: + return iter(self._spec.program.named_expressions) + + def __len__(self) -> int: + return len(self._spec.program.named_expressions) + + 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, spec: ModelSpec, name: str) -> None: + self._spec = spec + self._name = name + + def to_latex(self, **options: Any) -> str: + """This declaration typeset as a single LaTeX line, no document around it.""" + return typeset_declaration(self._spec._schema, self._name, "latex", **options) + + def to_markdown(self, **options: Any) -> str: + """This declaration typeset as a single Markdown math line, no ``$$`` around it.""" + return typeset_declaration( + self._spec._schema, self._name, "markdown", **options + ) + + def to_typst(self, **options: Any) -> str: + """This declaration typeset as a single Typst line, no document around it.""" + return typeset_declaration(self._spec._schema, self._name, "typst", **options) + + def _repr_markdown_(self) -> str: + return f"$$\n{self.to_markdown()}\n$$" + + +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. + + Attributes + ---------- + node + The lowered expression body, math-spec's own AST handle. + """ + + def __init__(self, spec: ModelSpec, name: str, ctx: Context) -> None: + super().__init__(spec, name) + self._ctx = ctx + + @property + def node(self) -> ms.ExpressionNode: + """The expression body as lowered, math-spec's own AST handle.""" + return self._spec.program.named_expressions[self._name].expression + + @property + def dims(self) -> tuple[str, ...]: + """The dimensions the expression spans, read off the spec without binding data.""" + return dims_of(self.node, self._spec.program) + + @functools.cached_property + def expression(self) -> terms.Value: + """ + The linopy symbolic expression, its variables unsolved. + + A named expression is read affinely, so this is 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. + """ + return evaluate_named(self._name, self._ctx.unsolved) + + @functools.cached_property + def solution(self) -> xr.DataArray: + """ + The expression folded over the model's solution, as data. + + Raises + ------ + RuntimeError + The model reads a variable but holds no solution yet. + 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("solution", 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..e98ac540d --- /dev/null +++ b/linopy/spec/attach.py @@ -0,0 +1,651 @@ +""" +Attach user data to a math-spec program. + +The language fixes three attachment rules and this module enforces them: a +dimension's members come only from the source keyed by the dimension's +name, their order is the source's order and is never sorted, and a +parameter or lookup source is read for values, never for labels. Parameters +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 + +from collections.abc import Hashable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Literal, 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.constants import warn_evolving_api +from linopy.spec.errors import SpecDataError +from linopy.spec.nodes import amounts_of, parameters_of, walk + +Retain = Literal["report", "all", "none"] +_RETAIN: tuple[str, ...] = get_args(Retain) + +_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" +) + +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)`.' +) + + +def attach( + program: ms.Program, + sources: Mapping[str, Any] | xr.Dataset, + *, + retain: Retain = "report", +) -> 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; it is read by + key and never iterated beyond ``sources.keys()``. An ``xr.Dataset`` + is accepted too: its indexes are dimension sources, its data + variables parameters and lookups. + retain + Which parameters :meth:`Attached.retained` persists. + + Raises + ------ + SpecDataError + A ``retain`` outside its three values, a key naming + nothing the spec declares, a reached dimension or a lookup with no + source, a duplicated dimension member, or a lookup breaking the + rules a map has. + """ + warn_evolving_api("spec", EVOLVING_MESSAGE) + 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) + keys = frozenset(sources.keys()) + _check_keys(program, keys) + coords = _master_coords(program, sources, keys) + lookups = _lookups(program, sources, keys, coords) + return Attached(program, coords, lookups, retain, sources, keys) + + +@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. + 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. + """ + + program: ms.Program + coords: Mapping[str, pd.Index] + lookups: Mapping[str, Mapping[str, xr.DataArray]] + retain: Retain + sources: Mapping[str, Any] + _keys: frozenset[str] = field(repr=False) + + 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._keys: + 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}) + return kinds + + +def _check_keys(program: ms.Program, keys: frozenset[str]) -> None: + known = _attachable(program) + unknown = sorted(keys - set(known)) + if not unknown: + return + one = len(unknown) == 1 + lead = f"source key {unknown[0]!r} names" if one else f"source keys {unknown} name" + raise SpecDataError( + f"{lead} neither a parameter, a dimension nor a lookup this spec declares. " + f"{did_you_mean(unknown[0], known)} Pass only what the spec takes." + ) + + +# --------------------------------------------------------------------------- +# 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) + if lk.target is not None: + dims.add(lk.target) + return dims + + +def _master_coords( + program: ms.Program, sources: Mapping[str, Any], keys: frozenset[str] +) -> dict[str, pd.Index]: + reached = _reached(program) + coords: dict[str, pd.Index] = {} + for dim in program.dimensions: + if dim in keys: + coords[dim] = _index(dim, sources[dim]) + elif dim in reached: + raise SpecDataError( + f"dimension '{dim}' has no index: pass its labels under key '{dim}' as " + f"{_DIMENSION_SHAPES}. The index is what says which labels exist, and without " + f"one a mistyped label is indistinguishable from a new one." + ) + return coords + + +def _index(dim: str, obj: Any) -> 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}." + ) + index = pd.Index(values, name=dim) + 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: Mapping[str, Any], + keys: frozenset[str], + 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 keys: + 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." + ) + if lk.dtype is not None: + _check_value_dtype(lk.name, lk.dtype, series.dtype, kind="lookup") + if lk.target is None: + return + 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): + 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) -> None: + strangers = labels[~labels.isin(known)].unique().tolist() + if not strangers: + return + raise SpecDataError( + f"parameter '{name}' has label(s) in dimension '{dim}' that are not coordinates of it: " + f"{_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, kind: str = "parameter" +) -> None: + if str(dtype.kind) in _ACCEPTED_KINDS[declared]: + return + arrived = _KIND_NAMES.get(str(dtype.kind), str(dtype)) + raise SpecDataError( + f"{kind} '{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..6170c55cf --- /dev/null +++ b/linopy/spec/builder.py @@ -0,0 +1,135 @@ +""" +Program plus attached data to linopy declarations. + +A build hands every variable to linopy as its term, then adds special-ordered +sets, constraints and the objective; which linopy call each construct becomes +is one branch of :func:`linopy.spec.evaluate.evaluate`. +""" + +from __future__ import annotations + +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 +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.parameters import Parameters +from linopy.spec.terms import Term, Value +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) -> None: + """ + Add every declaration of the attached program to *model*. + + Variables, special-ordered sets, constraints and the objective, in that + order; then every named expression is checked for divisor and coefficient + coverage, so a body that cannot be folded is refused at build rather than + at read. + """ + ctx = Context( + model, + attached.program, + attached.coords, + attached.lookups, + Parameters(attached.program, attached.parameter), + ) + curves.validate(ctx.program, ctx.parameters) + _variables(ctx) + _sos(ctx) + _constraints(ctx) + _objective(ctx) + for name, declared in ctx.program.named_expressions.items(): + check_coverage(f"expression '{name}'", (declared.expression,), ctx, None) + + +def _variables(ctx: Context) -> None: + for name, declared in ctx.program.variables.items(): + rows = evaluate_where(declared.where, ctx) + check_bounds_cover(name, declared, ctx, as_linopy_mask(rows)) + 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.variable_type == "binary", + integer=declared.variable_type == "integer", + ) + + +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: + for sos in ctx.program.sos.values(): + ctx.model.add_sos_constraints( + ctx.model.variables[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) + if isinstance(other, xr.DataArray): + term, other = carried(term, other) + ctx.model.add_constraints(term, _SIGN[sense], other, name=name, mask=mask) + + +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]) diff --git a/linopy/spec/context.py b/linopy/spec/context.py new file mode 100644 index 000000000..ea9d53c6d --- /dev/null +++ b/linopy/spec/context.py @@ -0,0 +1,40 @@ +"""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 + +import pandas as pd +import xarray as xr +from math_spec import program as ms + +from linopy.model import Model + + +@dataclass(frozen=True) +class Context: + """ + Everything evaluating a node needs beyond the node. + + ``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. + """ + + model: Model + program: ms.Program + coords: Mapping[str, pd.Index] + lookups: Mapping[str, Mapping[str, xr.DataArray]] + parameters: Mapping[str, xr.DataArray] + solved: bool = field(default=False) + + @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 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] diff --git a/linopy/spec/coverage.py b/linopy/spec/coverage.py new file mode 100644 index 000000000..11da5e045 --- /dev/null +++ b/linopy/spec/coverage.py @@ -0,0 +1,188 @@ +""" +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 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.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 = rows + for variable in sorted(ms.variables_of(quotient.numerator)): + present = terms.present(ctx.model.variables[variable]) + needed = present if needed is None else needed & present + return [(param, needed) for param in sorted(params)] + + +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..6571e90d0 --- /dev/null +++ b/linopy/spec/evaluate.py @@ -0,0 +1,208 @@ +"""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.model.variables[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 RuntimeError( + f"constraint '{constraint}' has no dual yet: solve the model before reading a dual" + ) + return ctx.model.constraints[constraint].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..7f1a28132 --- /dev/null +++ b/linopy/spec/netcdf.py @@ -0,0 +1,189 @@ +""" +Persist the spec of a spec-built model in its netcdf file. + +Variables, constraints and the solution round trip through :mod:`linopy.io` +already. Besides them a spec-built model carries the spec text, the master +coordinates and the lookups; the program is re-lowered from the text on read, +so no lowered ``Program`` ever reaches the file. + +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 + +from typing import Any + +import numpy as np +import pandas as pd +import xarray as xr + +from linopy.io import ( + DTYPE_ATTR, + SPEC_ATTR, + get_prefix, + restamp_coords, + with_prefix, +) +from linopy.model import Model +from linopy.spec.accessor import ModelSpec, restore + +PREFIX = "spec" +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(spec: ModelSpec) -> xr.Dataset: + """ + The spec's own dataset: its text, its master coordinates and its parameters. + + The spec text is the dataset's one attribute, which the merge lifts to the + file's. Beside it sits one array of labels per master coordinate and, per + parameter, either its values or -- where it is coded -- its codes and its + categories. The dataset carries no coordinates of its own: an index + coordinate is dropped on read together with the dimension it indexes once + no data variable is left over that dimension, and a master coordinate + nothing else reaches has exactly that shape. So a parameter is written + over bare dimensions and put back on the master coordinates on read. + """ + arrays: dict[str, xr.DataArray] = { + COORD + dim: _array(index.to_numpy(), (dim,)) + for dim, index in spec.coords.items() + } + coded = _coded(spec) + for name, arr in spec.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)) + return with_prefix(xr.Dataset(arrays), PREFIX).assign_attrs({SPEC_ATTR: spec.text}) + + +def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: + """ + Re-lower *text* onto *model* and read back the dataset :func:`encode` wrote. + + The master coordinates, the plainly written parameters and the coded ones + together are the dataset :func:`linopy.spec.accessor.attach` gave the spec + when the model was built. ``model.parameters`` is not touched: it holds + what the caller put there and nothing of the spec. + """ + sub = get_prefix(ds, PREFIX) + coords = { + _stripped(name, COORD): _index(sub[name]) + for name in sub.data_vars + if str(name).startswith(COORD) + } + arrays = { + _stripped(name, PARAM): _plain(sub[name], _stripped(name, PARAM), coords) + for name in sub.data_vars + if str(name).startswith(PARAM) + } + arrays.update( + { + _stripped(name, CODES): _decode(sub, _stripped(name, CODES), coords) + for name in sub.data_vars + if str(name).startswith(CODES) + } + ) + restamp_coords(model, coords) + return restore(model, text, xr.Dataset(arrays).assign_coords(coords)) + + +def _coded(spec: ModelSpec) -> set[str]: + """The parameters written as codes: every lookup and every array of objects.""" + lookups = {name for by_name in spec.lookups.values() for name in by_name} + return { + str(name) + for name, arr in spec.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..9c95c9680 --- /dev/null +++ b/linopy/spec/nodes.py @@ -0,0 +1,51 @@ +"""Walks over a program's expression nodes, and the dimensions a node spans.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from math_spec import program as ms + + +def walk(*nodes: ms.ExpressionNode) -> Iterator[ms.ExpressionNode]: + """Every node under *nodes*, each of them included, parents first.""" + for node in nodes: + yield node + yield from walk(*ms.children(node)) + + +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..7b86565a2 --- /dev/null +++ b/linopy/spec/operators.py @@ -0,0 +1,343 @@ +""" +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.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. + """ + if by is not None: + within = _per_group(within, by) + asked = _widest(within) + widest = max(1, min(asked, int(array.sizes[over]))) + 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/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..2432122e0 --- /dev/null +++ b/linopy/spec/terms.py @@ -0,0 +1,39 @@ +""" +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.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 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..5bec8301f --- /dev/null +++ b/linopy/spec/testing.py @@ -0,0 +1,72 @@ +""" +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] + if lookup.target is not None + else _labels(lookup.name, lookup.dtype, n) + ) + 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..fc301d0cb --- /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.model.variables[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..2cf5d28c4 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -111,10 +111,27 @@ 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) @@ -134,6 +151,11 @@ def assert_model_equal(a: Model, b: Model) -> None: 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 a._spec.text == b._spec.text + assert_datasetequal(a._spec.parameters, b._spec.parameters) + assert a.status == b.status assert a.termination_condition == b.termination_condition diff --git a/linopy/variables.py b/linopy/variables.py index e0ad70bc9..ea4ada53d 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -1809,7 +1809,9 @@ def __dir__(self) -> list[str]: ] return base_attributes + formatted_names - def _format_items(self, exclude: set[str] | None = None) -> str: + def _format_items( + self, exclude: set[str] | None = None, tag: set[str] | None = None + ) -> str: """Format variable items, optionally excluding names in a group.""" r = "" count = 0 @@ -1828,7 +1830,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 = " [spec]" if tag and name in tag else "" + r += f" * {name}{coords}{suffix}\n" if count == 0: r += "\n" return r diff --git a/pyproject.toml b/pyproject.toml index b9c0c051f..f8b3ec145 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.76 ; 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..cbc576737 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,188 @@ 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 + foreach: [snapshot, generator] + where: "p_max > 0" + bounds: { lower: 0, upper: p_max } + +constraints: + power_balance: + foreach: [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"}, + "tag": {"over": "t", "dtype": "str"}, + }, + "parameters": { + "flag": {"dims": ["t"], "dtype": "bool"}, + "cost": {"dims": ["t"]}, + "label": {"dims": ["t"], "dtype": "str"}, + "day_cost": {"dims": ["d"]}, + }, + "variables": { + "x": {"foreach": ["t"], "bounds": {"lower": 0, "upper": 1}}, + "y": {"foreach": ["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), + "tag": pd.Series(["p", "q"], index=TT[:2]), + "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": { + "foreach": ["snapshot", "generator"], + "bounds": {"lower": 0, "upper": "p_max"}, + }, + "op_cost": {"foreach": ["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": { + "foreach": ["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..91605cecb --- /dev/null +++ b/test/test_spec_accessor.py @@ -0,0 +1,402 @@ +""" +``model.spec``, ``ModelSpec``, ``NamedExpression``, ``evaluate``, typesetting, +and the ``add_spec``/``from_spec`` argument handling that builds them. +""" + +from __future__ import annotations + +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") + +import linopy # noqa: E402 +from conftest import ( # noqa: E402 + DISPATCH_DATA, + DISPATCH_P, + EXAMPLE_DISPATCH, + GENERATOR, + solved, + with_, + yaml_dict, +) +from linopy import Model # noqa: E402 +from linopy.spec import ModelSpec, NamedExpression, SpecDataError # 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(), + "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 test_add_spec_needs_an_empty_model() -> None: + m = Model() + m.add_variables(name="x") + with pytest.raises(ValueError, match="empty model"): + m.add_spec(yaml_dict(), DISPATCH_DATA) + + +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="not built from a 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.Variable), + ("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) + + +def test_expression_reads_unsolved_but_solution_waits_for_a_solve() -> None: + e = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.expressions["spend"] + assert isinstance(e.expression, linopy.LinearExpression) + with pytest.raises(RuntimeError, match="no solution yet"): + e.solution + + +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_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_shows_the_spec_and_tags_only_expressions() -> None: + 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) [spec]" in text + assert " * usage (snapshot, generator) [spec]" in text + assert " * p (snapshot, generator)\n" in text + assert " * power_balance (snapshot)\n" in text + assert "" not 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 test_the_whole_model_typesets() -> None: + spec = Model.from_spec(yaml_dict(), DISPATCH_DATA).spec + assert "align" in spec.to_latex() + assert "$$" in spec.to_markdown() + assert spec.to_typst() + assert spec._repr_markdown_() == spec.to_markdown() + + +@pytest.mark.parametrize("fmt", ["to_latex", "to_markdown", "to_typst"]) +def test_a_named_expression_typesets_to_one_line(fmt: str) -> None: + e = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.expressions["spend"] + line = getattr(e, fmt)() + assert "spend" in line + assert "\n" not in line + assert "align" not in line and "$$" not in line + + +def test_a_named_expression_repr_markdown_wraps_only_itself() -> None: + e = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.expressions["spend"] + assert e._repr_markdown_() == f"$$\n{e.to_markdown()}\n$$" + + +def test_a_named_expression_typeset_passes_options() -> None: + e = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.expressions["spend"] + assert e.to_latex( + symbols={"notation": "latex", "names": {"spend": "S"}} + ).startswith("S") + + +@pytest.mark.parametrize("name", ["power_balance", "p"]) +@pytest.mark.parametrize("fmt", ["to_latex", "to_markdown", "to_typst"]) +def test_a_constraint_or_variable_typesets_to_one_line(name: str, fmt: str) -> None: + d = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.declaration(name) + line = getattr(d, fmt)() + assert line + assert "\n" not in line + assert "align" not in line and "$$" not in line + + +def test_declaration_reaches_every_kind_and_an_unknown_name_is_a_key_error() -> None: + spec = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec + assert spec.declaration("spend").to_latex() == spec.expressions["spend"].to_latex() + with pytest.raises(KeyError, match="unknown declaration 'spent'.*spend"): + spec.declaration("spent") + + +def test_a_constant_expression_folds_to_a_scalar() -> None: + spec = {**yaml_dict(), "expressions": {"answer": "6 * 7"}} + got = Model.from_spec(spec, DISPATCH_DATA).spec.expressions["answer"].solution + assert got.ndim == 0 and float(got) == 42.0 + + +OTHER = pd.Index(["x", "y"], name="generator") + + +@pytest.mark.parametrize( + ("generator", "match"), + [ + pytest.param(GENERATOR[::-1], "as \\['gas', 'wind'\\]", id="reordered"), + pytest.param(OTHER, "as \\['x', 'y'\\]", id="relabelled"), + ], +) +def test_evaluate_refuses_sources_on_other_labels_than_the_model( + generator: pd.Index, match: str +) -> None: + m = solved({**yaml_dict(), "expressions": {"twice": "cost * 2"}}, DISPATCH_DATA) + sources = { + **DISPATCH_DATA, + "generator": generator, + "p_max": pd.Series([100.0, 200.0], index=generator), + "cost": pd.Series([0.0, 50.0], index=generator), + } + with pytest.raises(SpecDataError, match=f"dimension 'generator' {match}"): + m.spec.evaluate("twice", sources) + + +def test_a_reported_dual_folds_to_the_constraint_dual() -> None: + spec = {**yaml_dict(), "expressions": {"price": "dual(power_balance)"}} + m = solved(spec, DISPATCH_DATA) + xr.testing.assert_allclose( + m.spec.expressions["price"].solution, + m.constraints["power_balance"].dual.rename("price"), + ) + + +def test_a_dual_needs_a_solution() -> None: + spec = {**yaml_dict(), "expressions": {"price": "dual(power_balance)"}} + m = Model.from_spec(spec, DISPATCH_DATA) + with pytest.raises(RuntimeError, match="no dual yet"): + m.spec.expressions["price"].expression + + +def test_spec_api_warns_once_per_session() -> None: + from linopy import EvolvingAPIWarning + from linopy.constants import _emitted_evolving_warnings + + _emitted_evolving_warnings.discard("spec") + with pytest.warns(EvolvingAPIWarning, match="spec: Model.add_spec"): + Model.from_spec(EXAMPLE_DISPATCH, DISPATCH_DATA) + with warnings.catch_warnings(): + warnings.simplefilter("error", EvolvingAPIWarning) + Model.from_spec(EXAMPLE_DISPATCH, DISPATCH_DATA) + + +@pytest.mark.parametrize( + ("name", "dims"), [("spend", ("snapshot",)), ("usage", ("snapshot", "generator"))] +) +def test_named_expression_dims_are_static(name: str, dims: tuple[str, ...]) -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + expr = m.spec.expressions[name] + assert expr.dims == dims + assert set(expr.expression.coord_dims) == set(dims) diff --git a/test/test_spec_attach.py b/test/test_spec_attach.py new file mode 100644 index 000000000..56b0175dd --- /dev/null +++ b/test/test_spec_attach.py @@ -0,0 +1,827 @@ +"""Binding user data to a math-spec program.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +math_spec = pytest.importorskip("math_spec") + +from linopy.spec import SpecDataError, attach # noqa: E402 + +SPEC: dict[str, Any] = { + "dimensions": {"f": {"dtype": "str"}, "t": {"dtype": "int"}, "g": {"dtype": "str"}}, + "lookups": {"grp": {"over": "f", "into": "g"}}, + "parameters": { + "cost": {"dims": ["f"]}, + "cap": {"dims": ["f", "t"]}, + "flag": {"dims": ["f"], "dtype": "bool"}, + "rate": {"dims": []}, + "lead": {"dims": ["f"], "dtype": "int"}, + }, + "variables": { + "x": { + "foreach": ["f", "t"], + "where": "flag", + "bounds": {"lower": 0, "upper": "cap"}, + } + }, + "constraints": { + "k": {"foreach": ["g", "t"], "expression": "sum(x, by=grp) <= 10"}, + "s": { + "foreach": ["f", "t"], + "expression": "shift(x, over=t, offset=lead, edge=0) >= 0", + }, + }, + "objective": {"sense": "maximize", "expression": "sum(x * cost)"}, + "expressions": { + "spend": "sum(x * cost, over=t)", + "total": "sum(spend, over=f) * rate", + }, +} + +F = pd.Index(["b", "a", "c"], name="f") +T = pd.Index([0, 1, 2], name="t") +CAP = xr.DataArray(np.arange(9.0).reshape(3, 3), coords={"f": F, "t": T}, name="cap") +COST = pd.Series([1.0, 2.0, 3.0], index=F) + +DUP_ROWS = pd.Series([1.0, 9.0, 2.0], index=pd.Index(["a", "a", "b"], name="f")) +NULL_ROW = pd.Series({"a": 1.0, "b": None}) +NAN_ROW = pd.Series({"a": 1.0, "b": np.nan}) +NULL_FRAME = pd.DataFrame({"f": ["a", "b"], "value": [1.0, None]}) +STRAY_ROW = pd.Series({"a": 1.0, "zz": 2.0}) +DEEP_INDEX = pd.MultiIndex.from_tuples([("a", 0), ("b", 0)], names=["f", "k"]) +DEEP_ROWS = pd.Series([5.0, 5.0], index=DEEP_INDEX) + + +def sources_from( + base: Mapping[str, Any], override: Mapping[str, Any] +) -> dict[str, Any]: + """*base* with *override* applied; a ``None`` value drops the key instead.""" + merged = {**base, **override} + for key, value in override.items(): + if value is None: + merged.pop(key) + return merged + + +@pytest.fixture(scope="module") +def program() -> Any: + return math_spec.to_program(SPEC) + + +@pytest.fixture +def good() -> dict[str, Any]: + return { + "f": list(F), + "t": list(T), + "g": ["n", "e"], + "cost": COST, + "cap": CAP, + "flag": pd.Series([True, False, True], index=F), + "rate": 0.5, + "lead": pd.Series([1, 0, 1], index=F), + "grp": pd.Series(["n", "e", "n"], index=F), + } + + +def read_all(program: Any, sources: Mapping[str, Any]) -> list[xr.DataArray]: + attached = attach(program, sources) + return [attached.parameter(name) for name in program.parameters] + + +CAP_SHAPES = { + "dataarray": CAP, + "dataarray-transposed": CAP.transpose("t", "f"), + "series": CAP.to_series(), + "series-transposed": CAP.transpose("t", "f").to_series(), + "series-unnamed-levels": CAP.to_series().rename_axis([None, None]), + "tidy-frame": CAP.to_series().reset_index(name="value"), + "tidy-frame-extra-column": CAP.to_series() + .reset_index(name="value") + .assign(note="x"), + "wide-frame": CAP.to_pandas(), + "wide-frame-transposed": CAP.to_pandas().T, + "wide-frame-unnamed": CAP.to_pandas().rename_axis(index=None, columns=None), + "dict": CAP.to_series().to_dict(), +} + + +@pytest.mark.parametrize("cap", CAP_SHAPES.values(), ids=CAP_SHAPES.keys()) +def test_rank_two_shapes_attach_alike( + program: Any, good: dict[str, Any], cap: Any +) -> None: + got = attach(program, {**good, "cap": cap}).parameter("cap") + xr.testing.assert_equal(got, CAP) + assert got.dims == ("f", "t") + + +COST_SHAPES = { + "series": COST, + "series-unnamed": COST.rename_axis(None), + "dataarray": xr.DataArray(COST), + "dict": COST.to_dict(), + "tidy-frame": COST.reset_index(name="value"), + "indexed-frame": COST.to_frame("value"), +} + + +@pytest.mark.parametrize("cost", COST_SHAPES.values(), ids=COST_SHAPES.keys()) +def test_rank_one_shapes_attach_alike( + program: Any, good: dict[str, Any], cost: Any +) -> None: + got = attach(program, {**good, "cost": cost}).parameter("cost") + xr.testing.assert_equal(got, xr.DataArray(COST, name="cost")) + + +DIMENSION_SHAPES = { + "index": F, + "list": list(F), + "tuple": tuple(F), + "ndarray": F.to_numpy(), + "series": pd.Series(F), + "dataarray": xr.DataArray(list(F), dims=["f"]), +} + + +@pytest.mark.parametrize("f", DIMENSION_SHAPES.values(), ids=DIMENSION_SHAPES.keys()) +def test_dimension_shapes_keep_source_order( + program: Any, good: dict[str, Any], f: Any +) -> None: + coords = attach(program, {**good, "f": f}).coords + assert coords["f"].tolist() == ["b", "a", "c"] + assert coords["f"].name == "f" + assert list(coords) == ["f", "t", "g"] + + +def test_lookup_is_padded_onto_the_dimension( + program: Any, good: dict[str, Any] +) -> None: + attached = attach(program, {**good, "grp": {"a": "n"}}) + grp = attached.lookups["f"]["grp"] + assert grp.dims == ("f",) + assert grp.sel(f="a").item() == "n" + assert pd.isna(grp.sel(f=["b", "c"])).all() + + +LOOKUP_SHAPES = { + "series": pd.Series(["n", "e", "n"], index=F), + "series-unnamed": pd.Series(["n", "e", "n"], index=F.rename(None)), + "dict": {"b": "n", "a": "e", "c": "n"}, + "dataarray": xr.DataArray(["n", "e", "n"], coords={"f": F}), +} + + +@pytest.mark.parametrize("grp", LOOKUP_SHAPES.values(), ids=LOOKUP_SHAPES.keys()) +def test_lookup_shapes_attach_alike( + program: Any, good: dict[str, Any], grp: Any +) -> None: + got = attach(program, {**good, "grp": grp}).lookups["f"]["grp"] + assert got.values.tolist() == ["n", "e", "n"] + + +@pytest.mark.parametrize("storage", ["python", "pyarrow"]) +@pytest.mark.parametrize("shape", ["series", "dataarray"]) +def test_extension_strings_attach_as_numpy_objects( + program: Any, good: dict[str, Any], storage: str, shape: str +) -> None: + if storage == "pyarrow": + pytest.importorskip("pyarrow") + series = pd.Series(["n", "e"], index=F[:2], dtype=pd.StringDtype(storage)) + grp = xr.DataArray(series) if shape == "dataarray" else series + got = attach(program, {**good, "grp": grp}).lookups["f"]["grp"] + assert got.dtype == np.dtype(object) + assert got.values[:2].tolist() == ["n", "e"] + assert pd.isna(got.values[2]) + assert got.sel(f=["b", "a"]).values.tolist() == ["n", "e"] + + +def test_missing_rows_become_nan_and_false(program: Any, good: dict[str, Any]) -> None: + sparse = { + **good, + "cost": pd.Series({"a": 1.0}), + "lead": pd.Series({"a": 1}), + "flag": pd.Series({"a": True}), + "cap": CAP.sel(t=[0, 1]), + } + attached = attach(program, sparse) + cost = attached.parameter("cost") + assert cost.sel(f="a").item() == 1.0 + assert cost.sel(f=["b", "c"]).isnull().all() + lead = attached.parameter("lead") + assert lead.dtype == np.float64 + assert lead.sel(f="a").item() == 1.0 + assert lead.sel(f=["b", "c"]).isnull().all() + flag = attached.parameter("flag") + assert flag.dtype == bool + assert flag.values.tolist() == [False, True, False] + cap = attached.parameter("cap") + assert cap.dims == ("f", "t") + assert cap.sel(t=2).isnull().all() + + +@pytest.mark.parametrize( + ("name", "value", "expected_dtype"), + [ + ("cost", 2, float), + ("cap", 1.5, float), + ("flag", True, bool), + ("lead", 3, np.int64), + ], +) +def test_scalar_is_broadcast_over_declared_dims( + program: Any, good: dict[str, Any], name: str, value: Any, expected_dtype: Any +) -> None: + got = attach(program, {**good, name: value}).parameter(name) + assert got.dims == tuple(SPEC["parameters"][name]["dims"]) + assert got.dtype == expected_dtype + assert (got == value).all() + + +def test_scalar_parameter_stays_scalar(program: Any, good: dict[str, Any]) -> None: + got = attach(program, good).parameter("rate") + assert got.dims == () + assert got.item() == 0.5 + + +EMPTY_SOURCES = { + "dict": {}, + "object-series": pd.Series(dtype=object), + "float-series": pd.Series(dtype=float), +} + + +@pytest.mark.parametrize("cost", EMPTY_SOURCES.values(), ids=EMPTY_SOURCES.keys()) +def test_empty_source_attaches_as_all_nan( + program: Any, good: dict[str, Any], cost: Any +) -> None: + got = attach(program, {**good, "cost": cost}).parameter("cost") + assert got.dtype == np.float64 + assert got.isnull().all() + assert got.indexes["f"].equals(F) + + +def test_missing_parameter_is_refused_when_read( + program: Any, good: dict[str, Any] +) -> None: + good.pop("cost") + attached = attach(program, good) + with pytest.raises(SpecDataError, match="no data provided for parameter 'cost'"): + attached.parameter("cost") + + +def test_undeclared_parameter_is_refused_with_a_hint( + program: Any, good: dict[str, Any] +) -> None: + with pytest.raises(SpecDataError, match="unknown parameter 'csot'.*'cost'"): + attach(program, good).parameter("csot") + + +def test_retain_is_validated_before_attaching( + program: Any, good: dict[str, Any] +) -> None: + with pytest.raises(SpecDataError, match=r"'report', 'all', 'none'") as error: + attach(program, good, retain="reports") # type: ignore[arg-type] + assert "Did you mean 'report'?" in str(error.value) + + +REFUSALS = [ + pytest.param( + {"f": ["a", "a", "b"]}, + r"dimension 'f' lists 'a' more than once", + id="duplicate-member", + ), + pytest.param( + {"cost": STRAY_ROW}, r"parameter 'cost'.*'f'.*'zz'", id="unknown-label" + ), + pytest.param( + {"cap": CAP.assign_coords(t=[0, 1, 9])}, + r"parameter 'cap'.*'t'.*\b9\b", + id="unknown-label-dense", + ), + pytest.param( + {"cap": CAP.assign_coords(t=[9, 0, 7])}, + r"not coordinates of it: 9, 7\.", + id="unknown-labels-dense-in-source-order", + ), + pytest.param( + {"cap": CAP.to_series().reset_index(name="value").assign(t=[9, 0, 7] * 3)}, + r"not coordinates of it: 9, 7\.", + id="unknown-labels-rows-in-source-order", + ), + pytest.param( + {"cost": DUP_ROWS}, + r"parameter 'cost' has more than one row for a coordinate: f='a' \(2 rows\)", + id="duplicated-coordinate-row", + ), + pytest.param( + {"cap": xr.DataArray([1.0, 2.0], coords={"f": ["a", "a"]})}, + r"parameter 'cap' arrived as a DataArray over \['f'\]", + id="dense-wrong-dims", + ), + pytest.param( + {"cap": xr.DataArray(np.ones((3, 3)), coords={"f": list(F), "t": [0, 0, 1]})}, + r"parameter 'cap' has more than one row for a coordinate: t=0 \(2 rows\)", + id="dense-duplicate-coordinate", + ), + pytest.param( + {"cap": COST}, + r"parameter 'cap'.*1 level\(s\) where 'cap' is over \['f', 't'\]", + id="wrong-rank", + ), + pytest.param( + {"cap": DEEP_ROWS.rename_axis(["f", "q"])}, + r"parameter 'cap' is indexed by \['f', 'q'\]", + id="wrong-level-names", + ), + pytest.param( + {"rate": COST}, + r"parameter 'rate' is declared with no dims.*3 rows", + id="rows-for-scalar", + ), + pytest.param( + {"cost": {"a", "b"}}, + r"parameter 'cost': cannot adapt set", + id="unsupported-shape", + ), + pytest.param( + {"cost": pd.DataFrame({"f": ["a"], "amount": [1.0]})}, + r"parameter 'cost' arrived as a DataFrame with columns \['f', 'amount'\]", + id="frame-without-value-column", + ), + pytest.param( + {"cap": CAP.to_pandas().rename_axis(index="f", columns="q")}, + r"parameter 'cap' arrived as a wide DataFrame with index 'f' and columns 'q'", + id="wide-frame-wrong-axis-names", + ), + pytest.param( + {"cap": xr.DataArray(np.ones((3, 3)), dims=["f", "t"])}, + r"parameter 'cap' has no coordinate labels along 'f'", + id="dense-without-labels", + ), + pytest.param( + {"cost": NULL_ROW}, r"parameter 'cost' carries 1 row.*f='b'", id="null-row" + ), + pytest.param({"cost": NAN_ROW}, r"parameter 'cost' carries 1 row", id="nan-row"), + pytest.param( + {"rate": float("nan")}, + r"parameter 'rate' is one value and that value is a hole", + id="nan-scalar", + ), + pytest.param( + {"rate": pd.DataFrame({"value": [None]})}, + r"parameter 'rate' is one value and that value is a hole", + id="nan-scalar-frame", + ), + pytest.param( + {"lead": pd.Series([1.5, 0.0, 1.0], index=F)}, + r"'lead' is declared 'int'.*'float'", + id="float-for-int", + ), + pytest.param( + {"flag": pd.Series([1, 0, 1], index=F)}, + r"'flag' is declared 'bool'.*'int'", + id="int-for-bool", + ), + pytest.param( + {"flag": 1.0}, r"'flag' is declared 'bool'.*'float'", id="float-scalar-for-bool" + ), + pytest.param( + {"rate": "1.5"}, r"'rate' is declared 'float'.*'str'", id="numeric-str-scalar" + ), + pytest.param( + {"rate": True}, + r"'rate' is declared 'float'.*'bool'", + id="bool-scalar-for-float", + ), + pytest.param( + {"rate": "abc"}, r"'rate' is declared 'float'.*'str'", id="str-scalar-for-float" + ), + pytest.param( + {"cost": pd.Series(["x", "y", "z"], index=F)}, + r"'cost' is declared 'float'.*'str'", + id="str-for-float", + ), + pytest.param( + {"csot": COST}, r"source key 'csot'.*Did you mean 'cost'", id="unknown-key" + ), + pytest.param({"f": None}, r"dimension 'f' has no index", id="missing-dimension"), + pytest.param( + {"f": {"a": 1}}, + r"index for dimension 'f': cannot read labels out of dict", + id="dimension-shape", + ), + pytest.param( + {"f": np.ones((2, 2))}, + r"index for dimension 'f' is 2-dimensional", + id="dimension-rank", + ), + pytest.param( + {"grp": xr.DataArray(["n"], coords={"t": [0]})}, + r"lookup 'grp' arrived as a DataArray over \['t'\]", + id="lookup-wrong-dataarray-dim", + ), + pytest.param( + {"grp": None}, r"no data provided for lookup 'grp'", id="missing-lookup" + ), + pytest.param( + {"grp": {"zz": "n"}}, + r"lookup 'grp' maps 'zz', which are not labels of 'f'", + id="lookup-stray-key", + ), + pytest.param( + {"grp": {"a": "zz"}}, + r"lookup 'grp' has value\(s\) that are not 'g' labels: 'zz'", + id="lookup-stray-value", + ), + pytest.param( + {"grp": pd.Series(["n", "e"], index=pd.Index(["a", "a"], name="f"))}, + r"lookup 'grp' maps 1 'f' label\(s\) more than once: 'a'", + id="lookup-two-values", + ), + pytest.param( + {"grp": {"a": None, "b": "n"}}, + r"lookup 'grp' carries 1 row\(s\) with a null in 'g': f='a'", + id="lookup-null", + ), + pytest.param( + {"grp": pd.DataFrame({"f": ["a"], "g": ["n"]})}, + r"lookup 'grp': cannot adapt DataFrame", + id="lookup-shape", + ), + pytest.param( + {"grp": pd.Series(["n"], index=pd.Index(["a"], name="t"))}, + r"lookup 'grp' is a Series indexed by 't'", + id="lookup-wrong-index", + ), +] + + +@pytest.mark.parametrize(("override", "match"), REFUSALS) +def test_malformed_data_is_refused_naming_the_symbol( + program: Any, good: dict[str, Any], override: dict[str, Any], match: str +) -> None: + with pytest.raises(SpecDataError, match=match): + read_all(program, sources_from(good, override)) + + +def test_int_labels_are_shown_as_written(program: Any, good: dict[str, Any]) -> None: + with pytest.raises(SpecDataError, match=r"\b99\b") as error: + read_all(program, {**good, "cap": CAP.assign_coords(t=[0, 1, 99])}) + assert "int64" not in str(error.value) + + +def test_dataset_is_a_source(program: Any, good: dict[str, Any]) -> None: + dims = {"f": F, "t": T, "g": ["n", "e"]} + values = {k: xr.DataArray(v) for k, v in good.items() if k not in dims} + from_dataset = attach(program, xr.Dataset(values, coords=dims)) + from_mapping = attach(program, good) + assert from_dataset.coords["f"].equals(from_mapping.coords["f"]) + for name in program.parameters: + xr.testing.assert_equal( + from_dataset.parameter(name), from_mapping.parameter(name) + ) + xr.testing.assert_equal( + from_dataset.lookups["f"]["grp"], from_mapping.lookups["f"]["grp"] + ) + + +class Counting(Mapping[str, Any]): + def __init__(self, data: dict[str, Any]) -> None: + self.data = data + self.pulled: list[str] = [] + + def __getitem__(self, key: str) -> Any: + self.pulled.append(key) + return self.data[key] + + def __iter__(self) -> Iterator[str]: + raise AssertionError("sources must not be iterated") + + def __len__(self) -> int: + return len(self.data) + + def keys(self) -> Any: + return self.data.keys() + + +def test_sources_are_pulled_by_key_on_demand( + program: Any, good: dict[str, Any] +) -> None: + sources = Counting(good) + attached = attach(program, sources) + assert set(sources.pulled) == {"f", "t", "g", "grp"} + attached.parameter("cap") + attached.parameter("cap") + assert sources.pulled.count("cap") == 2 + + +@pytest.mark.parametrize( + ("retain", "expected"), + [ + ("report", {"cost", "rate", "grp"}), + ("all", {"cost", "cap", "flag", "rate", "lead", "grp"}), + ("none", {"grp"}), + ], +) +def test_retained_follows_the_named_expressions( + program: Any, good: dict[str, Any], retain: Any, expected: set[str] +) -> None: + retained = attach(program, good, retain=retain).retained() + assert set(retained.data_vars) == expected + assert retained.coords["f"].values.tolist() == ["b", "a", "c"] + + +def test_report_closure_reads_names_and_masks() -> None: + spec = { + "dimensions": {"f": {"dtype": "str"}, "t": {"dtype": "int"}}, + "parameters": { + "cost": {"dims": ["f"]}, + "lag": {"dims": ["f"], "dtype": "int"}, + "span": {"dims": ["f"], "dtype": "int"}, + "on": {"dims": ["f"], "dtype": "bool"}, + "other": {"dims": ["f"]}, + }, + "variables": {"x": {"foreach": ["f", "t"], "bounds": {"lower": 0, "upper": 1}}}, + "objective": {"sense": "maximize", "expression": "sum(x * other)"}, + "expressions": { + "recent": "sum_back(x, over=t, within=span)", + "late": { + "foreach": ["f", "t"], + "cases": { + "active": { + "when": "on", + "expression": "shift(x, over=t, offset=lag, edge=0)", + } + }, + "otherwise": "x * cost", + }, + }, + } + program = math_spec.to_program(spec) + f = pd.Index(["a"], name="f") + sources = { + "f": f, + "t": [0, 1], + "cost": pd.Series([1.0], index=f), + "lag": pd.Series([1], index=f), + "span": pd.Series([2], index=f), + "on": pd.Series([True], index=f), + "other": pd.Series([2.0], index=f), + } + retained = attach(program, sources).retained() + assert set(retained.data_vars) == {"cost", "lag", "span", "on"} + + +def test_unreached_dimension_needs_no_source() -> None: + dimensions = {**PARITY_SPEC["dimensions"], "z": {"dtype": "int"}} + program = math_spec.to_program({**PARITY_SPEC, "dimensions": dimensions}) + assert list(attach(program, GOOD).coords) == ["f"] + + +@pytest.mark.parametrize("shape", ["dataarray", "dataarray-transposed", "wide-frame"]) +def test_aligned_array_is_not_copied( + program: Any, good: dict[str, Any], shape: str +) -> None: + source = CAP_SHAPES[shape] + attached = attach(program, {**good, "cap": source}) + assert np.shares_memory(np.asarray(source), attached.parameter("cap").values) + assert np.shares_memory(np.asarray(source), attached.parameter("cap").values) + + +def test_master_coordinate_dtype_wins_without_a_copy( + program: Any, good: dict[str, Any] +) -> None: + source = CAP.assign_coords(t=T.astype("int32")) + got = attach(program, {**good, "cap": source}).parameter("cap") + assert got.indexes["t"].dtype == np.int64 + assert np.shares_memory(np.asarray(source), got.values) + + +def test_derived_parameter_is_not_bound_from_sources() -> None: + spec = { + "dimensions": {"bp": {"dtype": "int"}}, + "parameters": {"bp_x": {"dims": ["bp"]}, "bp_y": {"dims": ["bp"]}}, + "variables": { + "x": {"foreach": [], "bounds": {"lower": 0, "upper": 10}}, + "y": {"foreach": []}, + }, + "piecewise": { + "curve": { + "over": "bp", + "method": "lp", + "points": "bp_x", + "links": [["x", "bp_x"], ["y", "bp_y", ">="]], + } + }, + "objective": {"sense": "minimize", "expression": "y"}, + } + program = math_spec.to_program(spec) + derived = [n for n, p in program.parameters.items() if p.derivation is not None] + assert derived + bp = pd.Index([0, 1, 2], name="bp") + sources = { + "bp": bp, + "bp_x": pd.Series([0.0, 5.0, 10.0], index=bp), + "bp_y": pd.Series([0.0, 2.0, 8.0], index=bp), + } + attached = attach(program, sources, retain="all") + assert set(attached.retained().data_vars) == {"bp_x", "bp_y"} + with pytest.raises(SpecDataError, match="emitted by piecewise block 'curve'"): + attached.parameter(derived[0]) + with pytest.raises(SpecDataError, match=derived[0]): + attach(program, {**sources, derived[0]: 1.0}) + + +# --------------------------------------------------------------------------- +# lpspec data-parity cases, eager representation +# --------------------------------------------------------------------------- + +PARITY_SPEC: dict[str, Any] = { + "dimensions": {"f": {"dtype": "str"}}, + "parameters": {"cost": {"dims": ["f"]}, "cap": {"dims": ["f"]}}, + "variables": {"x": {"foreach": ["f"], "bounds": {"lower": 0, "upper": "cap"}}}, + "constraints": {"k": {"foreach": ["f"], "expression": "x <= cap"}}, + "objective": {"sense": "maximize", "expression": "sum(x * cost)"}, +} + +GOOD = { + "f": ["a", "b"], + "cost": pd.Series({"a": 1.0, "b": 2.0}), + "cap": pd.Series({"a": 5.0, "b": 5.0}), +} +ACCEPTED = "accepted" + +PARITY_CASES = [ + pytest.param({}, ACCEPTED, id="valid"), + pytest.param({"cap": None}, SpecDataError, id="parameter-missing-entirely"), + pytest.param({"cost": pd.Series({"a": 1.0})}, ACCEPTED, id="coefficient-sparse"), + pytest.param({"cost": DUP_ROWS}, SpecDataError, id="duplicated-coordinate-row"), + pytest.param({"cost": STRAY_ROW}, SpecDataError, id="label-not-in-the-dimension"), + pytest.param({"cost": NULL_ROW}, SpecDataError, id="a-null-value"), + pytest.param({"cost": NAN_ROW}, SpecDataError, id="a-nan-value"), + pytest.param({"cap": NULL_ROW}, SpecDataError, id="a-hole-in-a-bound"), + pytest.param({"cost": float("nan")}, SpecDataError, id="a-hole-as-a-scalar"), + pytest.param({"cost": [1.0, np.nan]}, SpecDataError, id="a-hole-in-a-sequence"), + pytest.param({"cost": {"a": 1.0, "b": None}}, SpecDataError, id="a-hole-in-a-dict"), + pytest.param({"cost": NULL_FRAME}, SpecDataError, id="a-hole-in-a-tidy-frame"), + pytest.param({"cost": pd.Series({"a": 1, "b": 2})}, ACCEPTED, id="whole-numbers"), + pytest.param({"csot": COST}, SpecDataError, id="an-undeclared-source-key"), + pytest.param({"cost": DEEP_ROWS}, SpecDataError, id="a-series-too-deep"), +] + + +@pytest.mark.parametrize(("override", "verdict"), PARITY_CASES) +def test_parity_with_lpspec_data_verdicts( + override: dict[str, Any], verdict: Any +) -> None: + program = math_spec.to_program(PARITY_SPEC) + sources = sources_from(GOOD, override) + if verdict is ACCEPTED: + read_all(program, sources) + return + with pytest.raises(verdict): + read_all(program, sources) + + +def test_a_hole_is_named_where_it_sits() -> None: + program = math_spec.to_program(PARITY_SPEC) + with pytest.raises(SpecDataError, match="parameter 'cost'") as error: + read_all(program, {**GOOD, "cost": NULL_ROW}) + assert "divisor" not in str(error.value) + assert "f='b'" in str(error.value) + + +FLAG_SPEC = { + "dimensions": {"g": {"dtype": "str"}}, + "parameters": {"active": {"dims": ["g"], "dtype": "bool"}}, + "variables": { + "x": {"foreach": ["g"], "where": "active", "bounds": {"lower": 0, "upper": 1}} + }, + "objective": {"sense": "maximize", "expression": "sum(x)"}, +} + + +@pytest.mark.parametrize( + ("column", "verdict"), + [ + pytest.param(pd.Series({"a": True, "b": False}), ACCEPTED, id="a-bool-column"), + pytest.param(pd.Series({"a": 1, "b": 0}), SpecDataError, id="a-1-0-int-column"), + pytest.param( + pd.Series({"a": 1.0, "b": 0.0}), SpecDataError, id="a-1-0-float-column" + ), + ], +) +def test_a_flag_attaches_by_its_declaration(column: pd.Series, verdict: Any) -> None: + program = math_spec.to_program(FLAG_SPEC) + sources = {"g": ["a", "b"], "active": column} + if verdict is ACCEPTED: + assert attach(program, sources).parameter("active").dtype == bool + return + with pytest.raises(SpecDataError, match="declared 'bool'"): + read_all(program, sources) + + +LOOKUP_SPEC = { + "dimensions": {"g": {}, "b": {"dtype": "str"}}, + "lookups": {"gen_bus": {"over": "g", "into": "b"}}, + "parameters": {"p_max": {"dims": ["g"]}}, + "variables": {"x": {"foreach": ["g"], "bounds": {"lower": 0, "upper": "p_max"}}}, + "constraints": {"k": {"foreach": ["b"], "expression": "sum(x, by=gen_bus) <= 10"}}, + "objective": {"sense": "maximize", "expression": "sum(x)"}, +} +G_TWICE = pd.Index(["w", "w", "s"], name="g") +LOOKUP_GOOD = { + "p_max": pd.Series({"w": 5.0, "s": 5.0}), + "g": ["w", "s"], + "b": ["n", "e"], + "gen_bus": pd.Series({"w": "n", "s": "e"}), +} + + +@pytest.mark.parametrize( + ("override", "match"), + [ + pytest.param( + {"g": None, "b": None}, "dimension 'g' has no index", id="a-map-no-labels" + ), + pytest.param( + {"gen_bus": None}, "no data provided for lookup", id="an-index-no-map" + ), + pytest.param( + {"gen_bus": pd.Series({"w": "n", "s": "zz"})}, + "not 'b' labels", + id="a-stray-value", + ), + pytest.param( + {"gen_bus": pd.Series(["n", "e", "e"], index=G_TWICE)}, + "more than once", + id="two-values-for-one-label", + ), + pytest.param( + {"gen_bus": pd.Series({"w": None, "s": "e"})}, + "null in 'b'", + id="mapping-a-label-to-nothing", + ), + pytest.param( + {"gen_bus": pd.Series([None, "n", "n"], index=G_TWICE)}, + "null in 'b'", + id="a-label-held-twice-with-a-null", + ), + ], +) +def test_a_lookup_defect_is_refused(override: dict[str, Any], match: str) -> None: + program = math_spec.to_program(LOOKUP_SPEC) + with pytest.raises(SpecDataError, match=match): + read_all(program, sources_from(LOOKUP_GOOD, override)) + + +TAG_SPEC = { + **LOOKUP_SPEC, + "lookups": {"tag": {"over": "g", "dtype": "int"}}, + "constraints": {"k": {"foreach": ["g"], "expression": "x <= 10"}}, +} +TAG_GOOD = sources_from(LOOKUP_GOOD, {"gen_bus": None, "b": None}) + + +def test_a_label_space_lookup_is_padded_onto_the_dimension() -> None: + program = math_spec.to_program(TAG_SPEC) + attached = attach(program, {**TAG_GOOD, "tag": {"s": 7}}) + tag = attached.lookups["g"]["tag"] + assert tag.dims == ("g",) + assert tag.indexes["g"].tolist() == ["w", "s"] + assert np.isnan(tag.sel(g="w").item()) + assert tag.sel(g="s").item() == 7 + + +@pytest.mark.parametrize( + ("tag", "match"), + [ + pytest.param({"w": None, "s": 7}, "null in 'tag': g='w'", id="a-null"), + pytest.param( + {"w": "x", "s": "y"}, "lookup 'tag' is declared 'int'.*'str'", id="a-str" + ), + ], +) +def test_a_label_space_lookup_defect_is_refused( + tag: dict[str, Any], match: str +) -> None: + program = math_spec.to_program(TAG_SPEC) + with pytest.raises(SpecDataError, match=match): + attach(program, {**TAG_GOOD, "tag": tag}) + + +def test_a_stray_lookup_value_over_an_int_target_is_shown_as_written() -> None: + program = math_spec.to_program( + {**LOOKUP_SPEC, "dimensions": {"g": {}, "b": {"dtype": "int"}}} + ) + numbered = {"b": [1, 2], "gen_bus": pd.Series({"w": 1, "s": 99})} + sources = sources_from(LOOKUP_GOOD, numbered) + with pytest.raises(SpecDataError, match=r"not 'b' labels: 99\b") as error: + read_all(program, sources) + assert "int64" not in str(error.value) diff --git a/test/test_spec_builder.py b/test/test_spec_builder.py new file mode 100644 index 000000000..b7ee54f55 --- /dev/null +++ b/test/test_spec_builder.py @@ -0,0 +1,420 @@ +""" +Building declarations from math-spec programs: variables, constraints, the +objective, SOS-constrained curves, coverage refusals and side-swapped +expressions. +""" + +from __future__ import annotations + +import glob +import os +from pathlib import Path +from typing import Any + +import pandas as pd +import pytest +import xarray as xr + +math_spec = pytest.importorskip("math_spec") +yaml = pytest.importorskip("yaml") + +import linopy # noqa: E402 +from conftest import ( # noqa: E402, F401 + CURVE_DATA, + CURVE_SPEC, + DISPATCH_DATA, + DISPATCH_P, + EXAMPLE_DISPATCH, + FULL_X, + FULL_Y, + GENERATOR, + WHERE_DATA, + WHERE_SPEC, + solved, + with_, + yaml_dict, +) +from linopy import Model # noqa: E402 +from linopy.spec import SpecDataError # noqa: E402 +from linopy.spec.testing import synthetic_sources # noqa: E402 + +pytestmark = [ + pytest.mark.v1, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +EXAMPLES_DIR = os.environ.get("MATH_SPEC_EXAMPLES") +EXAMPLES = ( + sorted(glob.glob(f"{EXAMPLES_DIR}/*.yaml") + glob.glob(f"{EXAMPLES_DIR}/*/*.yaml")) + if EXAMPLES_DIR + else [] +) + + +@pytest.mark.skipif( + not EXAMPLES, reason="set MATH_SPEC_EXAMPLES to a math-spec examples directory" +) +@pytest.mark.parametrize( + "path", EXAMPLES, ids=lambda p: str(Path(p).relative_to(EXAMPLES_DIR or "")) +) +def test_every_math_spec_example_builds_and_solves(path: str) -> None: + if "/symbols/" in path: + pytest.skip("typesetting input, not a spec") + program = math_spec.to_program(path) + m = solved(path, synthetic_sources(program), retain="all") + assert m.nvars == sum(int(m.variables[v].labels.count()) for v in program.variables) + assert m.termination_condition in ("optimal", "infeasible") + + +def test_the_dispatch_example_solves_and_its_expressions_fold() -> None: + m = solved(yaml_dict(), DISPATCH_DATA) + assert m.objective.value == pytest.approx(2500.0) + xr.testing.assert_allclose(m.solution["p"], DISPATCH_P) + spend = m.spec.expressions["spend"].solution + xr.testing.assert_allclose( + spend, (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") + ) + usage = m.spec.expressions["usage"].solution + xr.testing.assert_allclose(usage, (DISPATCH_P / [100.0, 200.0]).rename("usage")) + assert ( + set(m.spec.expressions) == {"spend", "usage"} and len(m.spec.expressions) == 2 + ) + assert set(m.spec.parameters.data_vars) == {"cost", "p_max"} + assert m.spec.coords["generator"].equals(GENERATOR) + + +# --------------------------------------------------------------------------- +# absence: a missing row by position +# --------------------------------------------------------------------------- + +T = pd.Index([0, 1, 2], name="t") +SPARSE_SPEC: dict[str, Any] = { + "dimensions": {"t": {"dtype": "int"}}, + "parameters": {"c": {"dims": ["t"]}, "w": {"dims": ["t"]}}, + "variables": {"x": {"foreach": ["t"], "bounds": {"lower": 0, "upper": 10}}}, + "constraints": {"cap": {"foreach": ["t"], "expression": "w * x <= c"}}, + "objective": {"sense": "maximize", "expression": "sum(x, over=t)"}, +} +FULL_W = pd.Series([1.0, 1.0, 1.0], index=T) +FULL_C = pd.Series([0.0, 4.0, 5.0], index=T) +HOLE_AT_0 = pd.Series([4.0, 5.0], index=T[1:]) +W_HOLE_AT_0 = pd.Series([1.0, 1.0], index=T[1:]) + +NO_W_CONSTRAINT = {"cap": {"foreach": ["t"], "expression": "x <= c"}} + + +@pytest.mark.parametrize( + ("spec", "data", "match"), + [ + pytest.param( + SPARSE_SPEC, + {"w": W_HOLE_AT_0, "c": FULL_C}, + "constraint 'cap'.*parameter 'w' is used as a coefficient", + id="coefficient-in-a-constraint", + ), + pytest.param( + with_( + SPARSE_SPEC, + constraints=NO_W_CONSTRAINT, + objective={"sense": "maximize", "expression": "sum(w * x, over=t)"}, + ), + {"w": W_HOLE_AT_0, "c": FULL_C}, + "the objective.*parameter 'w' is used as a coefficient", + id="coefficient-in-the-objective", + ), + pytest.param( + with_(SPARSE_SPEC, constraints=NO_W_CONSTRAINT, expressions={"e": "w * x"}), + {"w": W_HOLE_AT_0, "c": FULL_C}, + "expression 'e'.*parameter 'w' is used as a coefficient", + id="coefficient-in-a-named-expression", + ), + pytest.param( + SPARSE_SPEC, + {"w": FULL_W, "c": HOLE_AT_0}, + "constraint 'cap'.*covers 1 fewer", + id="constant-side", + ), + pytest.param( + with_( + SPARSE_SPEC, + variables={ + "x": {"foreach": ["t"], "bounds": {"lower": 0, "upper": "c"}} + }, + ), + {"w": FULL_W, "c": HOLE_AT_0}, + "variable 'x': 1 rows have NULL bounds", + id="bound", + ), + pytest.param( + with_( + SPARSE_SPEC, + constraints={"cap": {"foreach": ["t"], "expression": "x / w <= c"}}, + ), + {"w": W_HOLE_AT_0, "c": FULL_C}, + "constraint 'cap'.*divisor", + id="divisor-in-a-constraint", + ), + pytest.param( + with_( + SPARSE_SPEC, + constraints=NO_W_CONSTRAINT, + objective={"sense": "maximize", "expression": "sum(x / w, over=t)"}, + ), + {"w": W_HOLE_AT_0, "c": FULL_C}, + "the objective.*divisor", + id="divisor-in-the-objective", + ), + pytest.param( + with_( + SPARSE_SPEC, + constraints=NO_W_CONSTRAINT, + expressions={"ratio": "x / w"}, + ), + {"w": W_HOLE_AT_0, "c": FULL_C}, + "expression 'ratio'.*divisor", + id="divisor-in-a-named-expression", + ), + ], +) +def test_a_missing_row_is_refused_wherever_it_is_used( + spec: dict[str, Any], data: dict[str, Any], match: str +) -> None: + with pytest.raises(SpecDataError, match=match): + Model.from_spec(spec, {"t": T, **data}) + + +def test_a_masked_variable_bound_needs_no_row_where_it_is_masked() -> None: + spec = with_( + SPARSE_SPEC, + parameters={ + **SPARSE_SPEC["parameters"], + "live": {"dims": ["t"], "dtype": "bool"}, + }, + variables={ + "x": { + "foreach": ["t"], + "where": "live", + "bounds": {"lower": 0, "upper": "c"}, + } + }, + constraints={ + "cap": {"foreach": ["t"], "where": "live", "expression": "w * x <= c"} + }, + ) + live = pd.Series([True, True], index=T[1:]) + m = Model.from_spec(spec, {"t": T, "w": FULL_W, "c": HOLE_AT_0, "live": live}) + assert int(m.variables["x"].labels.count()) == 3 + assert int((m.variables["x"].labels != -1).sum()) == 2 + + +WHERE_MASKS = with_( + SPARSE_SPEC, + constraints={"cap": {**SPARSE_SPEC["constraints"]["cap"], "where": "w"}}, +) + + +@pytest.mark.parametrize( + ("spec", "data", "objective"), + [ + pytest.param(SPARSE_SPEC, {"w": FULL_W, "c": FULL_C}, 9.0, id="fully-covered"), + pytest.param( + WHERE_MASKS, + {"w": W_HOLE_AT_0, "c": FULL_C}, + 19.0, + id="a-where-masks-a-coefficient-hole", + ), + pytest.param( + with_( + SPARSE_SPEC, + constraints={ + "cap": {**SPARSE_SPEC["constraints"]["cap"], "where": "c"} + }, + ), + {"w": FULL_W, "c": HOLE_AT_0}, + 19.0, + id="a-where-masks-a-constant-side-hole", + ), + ], +) +def test_a_covered_or_masked_row_builds( + spec: dict[str, Any], data: dict[str, Any], objective: float +) -> None: + m = solved(spec, {"t": T, **data}) + assert m.objective.value == pytest.approx(objective) + + +F = pd.Index(["a", "b"], name="f") +ENVELOPE_SPEC: dict[str, Any] = { + "dimensions": {"f": {"dtype": "str"}}, + "parameters": {"gate": {"dims": ["f"], "dtype": "bool"}, "relmax": {"dims": ["f"]}}, + "variables": { + "x": {"foreach": ["f"], "bounds": {"lower": 0, "upper": 100}}, + "size": { + "foreach": ["f"], + "where": "gate", + "bounds": {"lower": 0, "upper": 50}, + }, + }, + "constraints": { + "envelope": {"foreach": ["f"], "expression": "x - relmax * size <= 0"} + }, + "objective": {"sense": "maximize", "expression": "sum(x, over=f)"}, +} +ENVELOPE_DATA: dict[str, Any] = { + "f": F, + "gate": pd.Series([True], index=F[:1]), + "relmax": pd.Series([0.5, 0.5], index=F), +} +DEFINED_SPEC = with_( + ENVELOPE_SPEC, + constraints={ + "envelope": { + "foreach": ["f"], + "where": "size", + "expression": "x - relmax * size <= 0", + }, + "pinned": {"foreach": ["f"], "where": "NOT size", "expression": "x <= 0"}, + }, +) + + +@pytest.mark.parametrize( + ("spec", "unsized"), + [ + pytest.param(ENVELOPE_SPEC, 100.0, id="an-absent-term-drops-the-row"), + pytest.param( + DEFINED_SPEC, 0.0, id="a-bare-variable-in-a-where-asks-whether-it-exists" + ), + ], +) +def test_an_absent_variable_takes_its_row_unless_a_where_says_otherwise( + spec: dict[str, Any], unsized: float +) -> None: + m = solved(spec, ENVELOPE_DATA) + x = m.solution["x"] + assert float(x.sel(f="a")) == pytest.approx(25.0) + assert float(x.sel(f="b")) == pytest.approx(unsized) + + +SCALAR_SWITCH: dict[str, Any] = { + "dimensions": {"i": {"dtype": "int"}}, + "parameters": {"on": {"dims": [], "dtype": "bool"}}, + "variables": { + "x": {"foreach": ["i"], "bounds": {"lower": 1, "upper": 5}, "where": "on"}, + "y": {"foreach": ["i"], "bounds": {"lower": 2, "upper": 5}}, + }, + "objective": {"sense": "minimize", "expression": "sum(x) + sum(y)"}, +} + + +@pytest.mark.parametrize(("on", "objective"), [(True, 6.0), (False, 4.0)]) +def test_a_scalar_where_gates_a_whole_variable(on: bool, objective: float) -> None: + m = solved(SCALAR_SWITCH, {"i": [1, 2], "on": on}) + assert m.objective.value == pytest.approx(objective) + + +def test_a_dimension_with_no_members_builds_no_row() -> None: + spec = with_( + SPARSE_SPEC, + constraints={"budget": {"foreach": [], "expression": "sum(x, over=t) <= 10"}}, + ) + empty = pd.Index([], name="t", dtype=int) + m = Model.from_spec( + spec, + { + "t": empty, + "w": pd.Series([], index=empty, dtype=float), + "c": pd.Series([], index=empty, dtype=float), + }, + ) + assert "budget" not in m.constraints + + +@pytest.mark.parametrize( + ("absence", "masked_reads_nan"), + [("undefined", True), ("zero", False)], + ids=["undefined-leaves-a-masked-slot-nan", "zero-fills-a-masked-slot"], +) +def test_a_fold_reads_a_masked_slot_the_way_its_absence_says( + absence: str, masked_reads_nan: bool +) -> None: + spec = yaml_dict() + spec["variables"]["p"]["absence"] = absence + spec["expressions"] = {"spend_by_unit": "p * cost"} + data = {**DISPATCH_DATA, "p_max": pd.Series([200.0, 0.0], index=GENERATOR)} + spend = solved(spec, data).spec.expressions["spend_by_unit"].solution + masked = spend.sel(generator="gas") + assert bool(masked.isnull().all()) is masked_reads_nan + if not masked_reads_nan: + assert float(masked.max()) == pytest.approx(0.0) + assert not bool(spend.sel(generator="wind").isnull().any()) + + +# --------------------------------------------------------------------------- +# piecewise curves as SOS constraints +# --------------------------------------------------------------------------- + + +def test_a_sos2_curve_is_built_as_a_special_ordered_set() -> None: + spec = with_( + CURVE_SPEC, + piecewise={ + "cost_curve": { + "over": "bp", + "links": [["p", "bp_x"], ["op_cost", "bp_y"]], + "method": "sos2", + } + }, + ) + m = Model.from_spec(spec, {**CURVE_DATA, "bp_x": FULL_X, "bp_y": FULL_Y}) + assert m.variables["cost_curve_lam"].attrs["sos_type"] == 2 + + +# --------------------------------------------------------------------------- +# a constant on the left is the same row, a power hides nothing +# --------------------------------------------------------------------------- + + +def test_a_constant_on_the_left_is_the_same_row() -> None: + flipped = with_( + SPARSE_SPEC, constraints={"cap": {"foreach": ["t"], "expression": "c >= w * x"}} + ) + m = solved(flipped, {"t": T, "w": FULL_W, "c": FULL_C}) + assert m.objective.value == pytest.approx(9.0) + + +@pytest.mark.parametrize( + ("expression", "match"), + [ + pytest.param( + "x <= c ** 2", + "constraint 'cap'.*covers 1 fewer", + id="constant-side-under-a-power", + ), + pytest.param( + "x / (c ** 2) <= 1", "constraint 'cap'.*divisor", id="divisor-under-a-power" + ), + ], +) +def test_a_parameter_under_a_power_is_still_checked_for_coverage( + expression: str, match: str +) -> None: + spec = with_( + SPARSE_SPEC, constraints={"cap": {"foreach": ["t"], "expression": expression}} + ) + with pytest.raises(SpecDataError, match=match): + Model.from_spec(spec, {"t": T, "w": FULL_W, "c": HOLE_AT_0}) + + +def test_an_operator_under_a_power_keeps_its_parameters_retained() -> None: + spec = with_( + SPARSE_SPEC, + parameters={**SPARSE_SPEC["parameters"], "lag": {"dims": [], "dtype": "int"}}, + expressions={"e": "shift(c, over=t, offset=lag, edge=0) ** 1"}, + ) + m = Model.from_spec(spec, {"t": T, "w": FULL_W, "c": FULL_C, "lag": 1}) + assert {"c", "lag"} <= set(m.spec.parameters.data_vars) + xr.testing.assert_allclose( + m.spec.expressions["e"].solution, + xr.DataArray([0.0, 0.0, 4.0], coords={"t": T}, name="e"), + ) diff --git a/test/test_spec_curves.py b/test/test_spec_curves.py new file mode 100644 index 000000000..d0c3feb0f --- /dev/null +++ b/test/test_spec_curves.py @@ -0,0 +1,160 @@ +""" +Piecewise curve derivation and validation: whole and ragged breakpoint +tables, the checks that refuse a curve a method cannot build, and the +convex-hull method's single-bend requirement. +""" + +from __future__ import annotations + +from typing import Any + +import pandas as pd +import pytest + +math_spec = pytest.importorskip("math_spec") +yaml = pytest.importorskip("yaml") + +import linopy # noqa: E402 +from conftest import ( # noqa: E402 + CURVE_DATA, + CURVE_SPEC, + FULL_X, + FULL_Y, + MASKED_CURVE_SPEC, + RAGGED_X, + RAGGED_Y, + UNITS, + curve, + solved, + with_, +) +from linopy import Model # noqa: E402 +from linopy.spec import SpecDataError # noqa: E402 + +pytestmark = [ + pytest.mark.v1, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +# --------------------------------------------------------------------------- +# piecewise curves +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("spec", "data", "spend"), + [ + pytest.param( + CURVE_SPEC, {"bp_x": FULL_X, "bp_y": FULL_Y}, 400.0, id="whole-curves" + ), + pytest.param( + MASKED_CURVE_SPEC, + {"bp_x": RAGGED_X, "bp_y": RAGGED_Y}, + 275.0, + id="ragged-curves-under-points", + ), + ], +) +def test_a_piecewise_cost_lands_on_the_curve( + spec: dict[str, Any], data: dict[str, Any], spend: float +) -> None: + m = solved(spec, {**CURVE_DATA, **data}, retain="all") + assert m.spec.expressions["spend"].solution.item() == pytest.approx(spend) + assert m.objective.value == pytest.approx(spend) + + +def without(series: pd.Series, *keys: tuple[str, int]) -> pd.Series: + return series.drop(index=list(keys)) + + +@pytest.mark.parametrize( + ("spec", "data", "match"), + [ + pytest.param( + CURVE_SPEC, + {"bp_x": without(FULL_X, ("gas", 3)), "bp_y": FULL_Y}, + "parameter 'bp_x' has no value at \\(generator='gas', bp=3\\)", + id="a-hole-in-a-whole-curve", + ), + pytest.param( + MASKED_CURVE_SPEC, + {"bp_x": RAGGED_X, "bp_y": without(RAGGED_Y, ("gas", 3))}, + "Shorten it 'bp_x' claims this breakpoint", + id="a-hole-inside-the-mask", + ), + pytest.param( + MASKED_CURVE_SPEC, + {"bp_x": without(FULL_X, ("gas", 1)), "bp_y": FULL_Y}, + "Not so at generator='gas'", + id="a-mask-with-a-gap", + ), + pytest.param( + CURVE_SPEC, + { + "bp_x": curve( + { + (g, k): x + for g in UNITS + for k, x in enumerate([0.0, 20.0, 20.0, 80.0]) + } + ), + "bp_y": FULL_Y, + }, + "strictly increasing", + id="breakpoints-that-do-not-increase", + ), + pytest.param( + CURVE_SPEC, + { + "bp_x": FULL_X, + "bp_y": curve( + { + (g, k): y + for g in UNITS + for k, y in enumerate([0.0, 300.0, 500.0, 600.0]) + } + ), + }, + "exact only for a convex curve", + id="a-concave-curve-under-lp", + ), + pytest.param( + MASKED_CURVE_SPEC, + {"bp_x": without(RAGGED_X, ("hydro", 1)), "bp_y": RAGGED_Y}, + "This curve carries 1", + id="a-one-point-curve-under-lp", + ), + ], +) +def test_a_curve_the_method_cannot_build_is_refused( + spec: dict[str, Any], data: dict[str, Any], match: str +) -> None: + with pytest.raises(SpecDataError, match=match): + Model.from_spec(spec, {**CURVE_DATA, **data}) + + +def test_a_convex_hull_curve_may_bend_either_way_but_not_both() -> None: + spec = with_( + CURVE_SPEC, + piecewise={ + "cost_curve": { + "over": "bp", + "links": [["p", "bp_x"], ["op_cost", "bp_y"]], + "method": "convex", + } + }, + ) + concave = curve( + {(g, k): y for g in UNITS for k, y in enumerate([0.0, 300.0, 500.0, 600.0])} + ) + mixed = curve( + {(g, k): y for g in UNITS for k, y in enumerate([0.0, 300.0, 350.0, 600.0])} + ) + assert ( + "cost_curve_lam" + in Model.from_spec( + spec, {**CURVE_DATA, "bp_x": FULL_X, "bp_y": concave} + ).variables + ) + with pytest.raises(SpecDataError, match="exact only for a single bend"): + Model.from_spec(spec, {**CURVE_DATA, "bp_x": FULL_X, "bp_y": mixed}) diff --git a/test/test_spec_io.py b/test/test_spec_io.py new file mode 100644 index 000000000..4dddf94ee --- /dev/null +++ b/test/test_spec_io.py @@ -0,0 +1,295 @@ +""" +Round trips of a spec-built model through netcdf and through ``copy``. + +The spec itself is persisted as its YAML text and lowered again on read, so +what has to survive besides the model is data: the master coordinates, the +lookups and the retained parameters. Labels are the delicate part — a partial +lookup holds NaN in an array of strings — so every lookup shape is checked +value by value and dtype by dtype, on both netcdf engines ``test_io`` uses. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pandas as pd +import pytest +import xarray as xr + +math_spec = pytest.importorskip("math_spec") + +from test_spec_builder import ( # noqa: E402 + DISPATCH_DATA, + EXAMPLE_DISPATCH, + EXAMPLES_DIR, + WHERE_DATA, + WHERE_SPEC, + solved, +) + +import linopy # noqa: E402 +from linopy import Model, read_netcdf # noqa: E402 +from linopy.io import SPEC_ATTR # noqa: E402 +from linopy.spec import SpecDataError # noqa: E402 +from linopy.spec.testing import synthetic_sources # noqa: E402 +from linopy.testing import assert_model_equal # noqa: E402 + +pytestmark = [ + pytest.mark.v1, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +ENGINES = ["netcdf4", "scipy"] + +S1 = pd.Index(["a", "b", "c"], name="s1") +S2 = pd.Index(["p", "q"], name="s2") +I1 = pd.Index([10, 20, 30], name="i1") +I2 = pd.Index([1, 2], name="i2") + +LOOKUP_SPEC: dict[str, Any] = { + "dimensions": { + "s1": {"dtype": "str"}, + "s2": {"dtype": "str"}, + "i1": {"dtype": "int"}, + "i2": {"dtype": "int"}, + }, + "lookups": { + "str_to_str": {"over": "s1", "into": "s2"}, + "str_to_int": {"over": "s1", "into": "i2"}, + "int_to_str": {"over": "i1", "into": "s2"}, + "int_to_int": {"over": "i1", "into": "i2"}, + }, + "parameters": {"cost": {"dims": ["s1"]}}, + "variables": {"x": {"foreach": ["s1"], "bounds": {"lower": 0, "upper": 1}}}, + "objective": {"sense": "minimize", "expression": "sum(x * cost)"}, +} +LOOKUP_OVER = {"str_to_str": S1, "str_to_int": S1, "int_to_str": I1, "int_to_int": I1} +LOOKUP_INTO = {"str_to_str": S2, "str_to_int": I2, "int_to_str": S2, "int_to_int": I2} + +DTYPE_SPEC: dict[str, Any] = { + "dimensions": {"s1": {"dtype": "str"}}, + "parameters": { + "count": {"dims": ["s1"], "dtype": "int"}, + "flag": {"dims": ["s1"], "dtype": "bool"}, + "cost": {"dims": ["s1"]}, + "tag": {"dims": ["s1"], "dtype": "str"}, + }, + "variables": {"x": {"foreach": ["s1"], "bounds": {"lower": 0, "upper": 1}}}, + "objective": {"sense": "minimize", "expression": "sum(x * cost)"}, +} +DTYPE_DATA: dict[str, Any] = { + "s1": S1, + "count": pd.Series([1, 2, 3], index=S1), + "flag": pd.Series([True, False, True], index=S1), + "cost": pd.Series([1.0, 2.0, 3.0], index=S1), + "tag": pd.Series(["u", "v", "w"], index=S1), +} + + +def lookup_sources(mapped: int) -> dict[str, Any]: + """Data for ``LOOKUP_SPEC``, each lookup mapping only its first *mapped* labels.""" + sources: dict[str, Any] = { + "s1": S1, + "s2": S2, + "i1": I1, + "i2": I2, + "cost": pd.Series([1.0, 2.0, 3.0], index=S1), + } + for name, over in LOOKUP_OVER.items(): + into = LOOKUP_INTO[name] + sources[name] = pd.Series( + [into[i % len(into)] for i in range(mapped)], index=over[:mapped] + ) + return sources + + +def roundtrip(m: Model, tmp_path: Path, engine: str) -> Model: + path = tmp_path / f"model-{engine}.nc" + m.to_netcdf(path, engine=engine) + return read_netcdf(path) + + +def assert_arrayequal(a: xr.DataArray, b: xr.DataArray) -> None: + """Assert equal values and dtype — the dtype is what a netcdf type drops.""" + assert a.dtype == b.dtype, f"dtypes differ: {a.dtype} != {b.dtype}" + xr.testing.assert_equal(a, b) + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize("retain", ["report", "all"]) +def test_a_spec_built_model_round_trips( + tmp_path: Path, engine: str, retain: str +) -> None: + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain=retain) + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert p.spec.text == m.spec.text + assert p.spec.program.constraints == m.spec.program.constraints + assert set(p.spec.expressions) == set(m.spec.expressions) + for name in m.spec.expressions: + assert_arrayequal( + m.spec.expressions[name].solution, p.spec.expressions[name].solution + ) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_a_retain_none_model_evaluates_after_a_round_trip( + tmp_path: Path, engine: str +) -> None: + """A file is where retain bites: the sources the built model still read are gone.""" + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="none") + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert not p.spec.parameters.data_vars + with pytest.raises(SpecDataError, match="no longer holds the sources"): + p.spec.expressions["spend"].solution + assert_arrayequal( + m.spec.expressions["spend"].solution, + p.spec.evaluate("spend", DISPATCH_DATA).solution, + ) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_the_caller_parameters_and_the_spec_ones_stay_apart( + tmp_path: Path, engine: str +) -> None: + """A model parameter of the caller's is written beside the spec's, not into them.""" + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="all") + m.parameters["cost"] = xr.DataArray([1, 2, 3], dims=["own"]) + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert_arrayequal(p.parameters["cost"], m.parameters["cost"]) + assert_arrayequal(p.spec.parameters["cost"], m.spec.parameters["cost"]) + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize("mapped", [3, 2, 0], ids=["full", "partial", "empty"]) +@pytest.mark.parametrize("name", LOOKUP_OVER) +def test_a_lookup_round_trips_exactly( + tmp_path: Path, engine: str, mapped: int, name: str +) -> None: + m = Model.from_spec(LOOKUP_SPEC, lookup_sources(mapped), retain="all") + over = str(LOOKUP_OVER[name].name) + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert name in p.spec.lookups[over] + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize("name", ["count", "flag", "cost", "tag"]) +def test_a_parameter_keeps_its_dtype(tmp_path: Path, engine: str, name: str) -> None: + m = Model.from_spec(DTYPE_SPEC, DTYPE_DATA, retain="all") + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert p.spec.parameters[name].dtype == m.spec.parameters[name].dtype + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize( + "labels", [[10, 11, 12], [0, 1]], ids=["relabelled", "shorter"] +) +def test_a_hand_added_variable_keeps_its_own_labels( + tmp_path: Path, engine: str, labels: list[int] +) -> None: + """A container sharing a master dimension's name but not its labels is left alone.""" + own = pd.Index(labels, name="snapshot") + m = Model.from_spec(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="all") + m.add_variables(coords=[own], name="side") + p = roundtrip(m, tmp_path, engine) + + assert p.variables["side"].indexes["snapshot"].equals(own) + assert p.spec.coords["snapshot"].equals(m.spec.coords["snapshot"]) + assert p.variables["p"].indexes["snapshot"].equals(m.spec.coords["snapshot"]) + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize("frozen", [False, True], ids=["dataset", "csr"]) +def test_every_container_shares_the_master_coordinate_dtypes( + tmp_path: Path, engine: str, frozen: bool +) -> None: + """The master coordinates are canonical: no container may disagree with them.""" + if frozen and engine == "scipy": + pytest.skip( + "netCDF3 holds no unicode-array attr, and a CSR constraint writes one" + ) + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="all", freeze_constraints=frozen) + p = roundtrip(m, tmp_path, engine) + + master = {dim: index.dtype for dim, index in p.spec.coords.items()} + holders = [ + *(v.data for _, v in p.variables.items()), + *(c.data for _, c in p.constraints.items()), + p.objective.expression.data, + ] + assert master == {dim: index.dtype for dim, index in m.spec.coords.items()} + for data in holders: + for dim, index in data.indexes.items(): + if str(dim) in master: + assert index.dtype == master[str(dim)], f"{dim} differs on {data}" + + +@pytest.mark.parametrize("engine", ENGINES) +def test_labelled_parameters_and_unreached_coordinates_round_trip( + tmp_path: Path, engine: str +) -> None: + """A str parameter with holes, and a dimension only a lookup reaches.""" + m = Model.from_spec(WHERE_SPEC, WHERE_DATA, retain="all") + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert set(p.spec.coords) == set(m.spec.coords) + + +@pytest.mark.parametrize("deep", [True, False]) +def test_a_copy_carries_the_spec(deep: bool) -> None: + """The copy's spec reads the copy, and only a deep copy owns its buffers.""" + m = Model.from_spec(WHERE_SPEC, WHERE_DATA, retain="all") + p = m.copy(deep=deep) + p.spec.parameters["label"].values[1] = "changed" + + assert p.spec.text == m.spec.text + assert p.spec.parameters["label"].values[1] == "changed" + assert m.spec.parameters["label"].values[1] == ("u" if deep else "changed") + + +def test_a_copy_can_still_read_what_retain_dropped() -> None: + """A copy keeps the sources, so it folds an unretained parameter like its original.""" + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="none") + + assert_arrayequal( + m.copy(include_solution=True).spec.expressions["spend"].solution, + m.spec.expressions["spend"].solution, + ) + + +def test_a_model_without_a_spec_carries_none(tmp_path: Path) -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") + m.add_objective(x.sum()) + path = tmp_path / "plain.nc" + m.to_netcdf(path) + + assert SPEC_ATTR not in xr.load_dataset(path).attrs + assert read_netcdf(path)._spec is None + assert m.copy()._spec is None + + +@pytest.mark.skipif( + EXAMPLES_DIR is None, reason="set MATH_SPEC_EXAMPLES to a math-spec examples dir" +) +@pytest.mark.parametrize("engine", ENGINES) +def test_the_pypsa_example_round_trips(tmp_path: Path, engine: str) -> None: + """Nine lookups into one dimension, a datetime axis, bool and str parameters.""" + path = Path(EXAMPLES_DIR or "", "pypsa.yaml") + program = math_spec.to_program(str(path)) + m = Model.from_spec(path, synthetic_sources(program, 3), retain="all") + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert set(p.spec.coords) == set(m.spec.coords) diff --git a/test/test_spec_operators.py b/test/test_spec_operators.py new file mode 100644 index 000000000..bd9a1b45b --- /dev/null +++ b/test/test_spec_operators.py @@ -0,0 +1,333 @@ +""" +Operators built as a constraint and folded as a named expression: sum, +grouped sum, ``at``, shift/translate, sum_back windows, and the where/mask +predicates that gate which rows exist. +""" + +from __future__ import annotations + +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") + +import linopy # noqa: E402 +from conftest import TT, WHERE_DATA, WHERE_SPEC, solved, with_ # noqa: E402 +from linopy import Model # noqa: E402 +from linopy.spec import SpecDataError # noqa: E402 + +pytestmark = [ + pytest.mark.v1, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +# --------------------------------------------------------------------------- +# operators, built as a constraint and folded as a named expression +# --------------------------------------------------------------------------- + +S = pd.Index(["a", "b"], name="s") +V = np.array([1.0, 2.0, 4.0, 8.0]) +OPERATORS: dict[str, tuple[str, list[str], list[float]]] = { + "shift-edge-0": ("shift(x, over=t, offset=1, edge=0)", ["t"], [0, 1, 2, 4]), + "shift-ahead-edge-0": ("shift(x, over=t, offset=-1, edge=0)", ["t"], [2, 4, 8, 0]), + "shift-wrap": ("shift(x, over=t, offset=1, edge='wrap')", ["t"], [8, 1, 2, 4]), + "shift-wrap-in-groups": ( + "shift(x, over=t, offset=1, edge='wrap', by=season_of)", + ["t"], + [2, 1, 8, 4], + ), + "shift-by-group-offset": ( + "shift(x, over=t, offset=lag, edge=0, by=season_of)", + ["t"], + [0, 1, 0, 0], + ), + "sum-back": ("sum_back(x, over=t, within=2)", ["t"], [1, 3, 6, 12]), + "sum-back-wrap": ( + "sum_back(x, over=t, within=2, edge='wrap')", + ["t"], + [9, 3, 6, 12], + ), + "sum-back-in-groups": ( + "sum_back(x, over=t, within=2, by=season_of)", + ["t"], + [1, 3, 4, 12], + ), + "sum-back-group-width": ( + "sum_back(x, over=t, within=width, by=season_of)", + ["t"], + [1, 2, 4, 12], + ), + "sum-by": ("sum(x, by=season_of)", ["s"], [3, 12]), + "at": ("x * at(z, by=season_of)", ["t"], [10, 20, 80, 160]), + "cases": ("x_state", ["t"], [100, 1, 2, 4]), +} + + +def operator_spec() -> dict[str, Any]: + spec: dict[str, Any] = { + "dimensions": {"t": {"dtype": "int"}, "s": {"dtype": "str"}}, + "lookups": {"season_of": {"over": "t", "into": "s"}}, + "parameters": { + "v": {"dims": ["t"]}, + "z": {"dims": ["s"]}, + "lag": {"dims": ["s"], "dtype": "int"}, + "width": {"dims": ["s"], "dtype": "int"}, + }, + "variables": {"x": {"foreach": ["t"], "bounds": {"lower": 0, "upper": 100}}}, + "constraints": {"fix": {"foreach": ["t"], "expression": "x == v"}}, + "expressions": { + "x_state": { + "foreach": ["t"], + "cases": {"first": {"when": "position(t) == 0", "expression": 100}}, + "otherwise": "shift(x, over=t, offset=1)", + } + }, + "objective": {"sense": "minimize", "expression": "sum(x)"}, + } + for key, (expression, dims, _) in OPERATORS.items(): + name = key.replace("-", "_") + spec["variables"][f"y_{name}"] = { + "foreach": dims, + "bounds": {"lower": -1000, "upper": 1000}, + } + spec["constraints"][f"link_{name}"] = { + "foreach": dims, + "expression": f"y_{name} == {expression}", + } + spec["expressions"][f"probe_{name}"] = expression + return spec + + +OPERATOR_DATA: dict[str, Any] = { + "t": TT, + "s": S, + "season_of": pd.Series(["a", "a", "b", "b"], index=TT), + "v": pd.Series(V, index=TT), + "z": pd.Series([10.0, 20.0], index=S), + "lag": pd.Series([1, 2], index=S), + "width": pd.Series([1, 2], index=S), +} + + +@pytest.fixture(scope="module") +def operators_model() -> Model: + with linopy.options as options: + options["semantics"] = "v1" + return solved(operator_spec(), OPERATOR_DATA, retain="all") + + +@pytest.mark.parametrize("key", OPERATORS) +def test_an_operator_builds_and_folds_alike(operators_model: Model, key: str) -> None: + _, dims, expected = OPERATORS[key] + name = key.replace("-", "_") + want = xr.DataArray(expected, coords={dims[0]: OPERATOR_DATA[dims[0]]}, dims=dims) + built = operators_model.solution[f"y_{name}"] + folded = operators_model.spec.expressions[f"probe_{name}"].solution + xr.testing.assert_allclose(built, want.rename(f"y_{name}")) + xr.testing.assert_allclose(folded, want.rename(f"probe_{name}")) + + +AMOUNT_SPEC: dict[str, Any] = { + "dimensions": {"t": {"dtype": "int"}, "g": {"dtype": "int"}}, + "lookups": {"grp": {"over": "t", "into": "g"}}, + "parameters": {"v": {"dims": ["t"]}, "lag": {"dims": ["g"], "dtype": "int"}}, + "variables": { + "x": {"foreach": ["t"], "bounds": {"lower": 0, "upper": 100}}, + "y": {"foreach": ["t"], "bounds": {"lower": -100, "upper": 100}}, + }, + "constraints": { + "fix": {"foreach": ["t"], "expression": "x == v"}, + "link": { + "foreach": ["t"], + "expression": "y == shift(x, over=t, offset=lag, edge=0, by=grp)", + }, + }, + "objective": {"sense": "minimize", "expression": "sum(x)"}, +} + + +def test_a_missing_shift_amount_is_refused() -> None: + t = pd.Index([0, 1, 2], name="t") + g = pd.Index([0, 1], name="g") + data = { + "t": t, + "g": g, + "grp": pd.Series([0, 0, 1], index=t), + "v": pd.Series([0.0, 4.0, 5.0], index=t), + "lag": pd.Series([1], index=g[:1]), + } + with pytest.raises(SpecDataError, match="parameter 'lag' is used as a coefficient"): + Model.from_spec(AMOUNT_SPEC, data) + + +# --------------------------------------------------------------------------- +# grouped sum +# --------------------------------------------------------------------------- + +GROUPED_SPEC: dict[str, Any] = { + "dimensions": {"generator": {}, "bus": {"dtype": "str"}}, + "lookups": {"gen_bus": {"over": "generator", "into": "bus"}}, + "parameters": {"capacity": {"dims": ["generator"]}}, + "variables": { + "imports": {"foreach": ["bus"], "bounds": {"lower": 0, "upper": 100}} + }, + "constraints": { + "import_limit": { + "foreach": ["bus"], + "expression": "imports <= sum(capacity, by=gen_bus)", + } + }, + "objective": {"sense": "maximize", "expression": "sum(imports, over=bus)"}, +} +GENS = pd.Index(["g1", "g2"], name="generator") + + +def grouped_sources(capacity: pd.Series) -> dict[str, Any]: + return { + "bus": ["north", "south"], + "generator": GENS, + "gen_bus": pd.Series(["north", "north"], index=GENS), + "capacity": capacity, + } + + +def test_an_empty_group_on_the_constant_side_is_a_zero_and_not_a_gap() -> None: + m = solved(GROUPED_SPEC, grouped_sources(pd.Series([3.0, 4.0], index=GENS))) + assert m.objective.value == pytest.approx(7.0) + assert float(m.solution["imports"].sel(bus="south")) == pytest.approx(0.0) + + +def test_a_lookup_that_maps_nothing_leaves_every_group_at_the_empty_sum() -> None: + """Filtering to the mapped members leaves nothing, and nothing is what xarray will not group.""" + spec = with_( + GROUPED_SPEC, + variables={ + "out": { + "foreach": ["generator"], + "bounds": {"lower": 0, "upper": "capacity"}, + } + }, + expressions={"per_bus": "sum(out, by=gen_bus)"}, + ) + sources = grouped_sources(pd.Series([3.0, 4.0], index=GENS)) + sources["gen_bus"] = pd.Series([], dtype=object) + m = solved(spec, sources) + + assert m.objective.value == pytest.approx(0.0) + per_bus = m.spec.expressions["per_bus"] + assert per_bus.expression.nterm == 0 + assert per_bus.solution.indexes["bus"].tolist() == ["north", "south"] + np.testing.assert_allclose(per_bus.solution.values, [0.0, 0.0]) + + +def test_a_member_with_no_value_is_still_refused_through_a_group() -> None: + with pytest.raises(SpecDataError, match="parameter 'capacity' covers 1 fewer"): + Model.from_spec(GROUPED_SPEC, grouped_sources(pd.Series([3.0], index=GENS[:1]))) + + +# --------------------------------------------------------------------------- +# where predicates +# --------------------------------------------------------------------------- + +WHERE_CASES: dict[str, tuple[str, str, list[Any]]] = { + "dimension-comparison": ("x", "t > 1", [2, 3]), + "lookup-comparison": ("x", "season_of == 'a'", [0, 1]), + "lookup-not-equal-skips-unmapped": ("x", "season_of != 'a'", [2]), + "lookup-pair": ("x", "season_of != other_of", [1]), + "lookup-defined": ("x", "season_of", [0, 1, 2]), + "label-space-lookup": ("x", "tag == 'q'", [1]), + "not": ("x", "NOT (t > 1)", [0, 1]), + "and": ("x", "t > 0 AND t < 3", [1, 2]), + "or": ("x", "t == 0 OR t == 3", [0, 3]), + "position": ("x", "position(t) == -1", [3]), + "position-in-groups": ("x", "position(t, by=season_of) == 0", [0, 2]), + "bool-parameter": ("x", "flag", [0]), + "float-parameter-must-be-finite": ("x", "cost", [0, 2]), + "str-parameter": ("x", "label", [1, 2]), + "parameter-comparison": ("x", "cost > 2", [2, 1]), + "datetime-axis": ("y", "d >= '2030-01-03'", list(WHERE_DATA["d"][2:])), +} + + +@pytest.mark.parametrize("case", WHERE_CASES) +def test_a_where_picks_the_rows_it_names(case: str) -> None: + variable, predicate, labels = WHERE_CASES[case] + spec = with_( + WHERE_SPEC, + variables={variable: {**WHERE_SPEC["variables"][variable], "where": predicate}}, + ) + built = Model.from_spec(spec, WHERE_DATA).variables[variable] + dim = built.dims[0] + present = built.labels[dim][(built.labels != -1).to_numpy()] + assert sorted(present.to_numpy().tolist()) == sorted(labels) + + +@pytest.mark.parametrize( + ("predicate", "match"), + [ + ("position(t) == 7", "names position 7 of 't', which has 4"), + ("position(t, by=season_of) == 1", "shorter than that: \\['b'\\]"), + ], +) +def test_a_position_no_coordinate_holds_is_refused(predicate: str, match: str) -> None: + spec = with_( + WHERE_SPEC, + variables={"x": {**WHERE_SPEC["variables"]["x"], "where": predicate}}, + ) + with pytest.raises(SpecDataError, match=match): + Model.from_spec(spec, WHERE_DATA) + + +# --------------------------------------------------------------------------- +# edges: partial lookups and an empty dimension beside a sum +# --------------------------------------------------------------------------- + +PARTIAL_CASES: dict[str, list[float]] = { + "sum-by": [3.0, 4.0], + "at": [10.0, 20.0, 80.0, np.nan], + "shift-wrap-in-groups": [2.0, 1.0, 4.0, np.nan], + "sum-back-in-groups": [1.0, 3.0, 4.0, np.nan], +} + + +@pytest.mark.parametrize("key", PARTIAL_CASES) +def test_a_member_a_lookup_sends_nowhere_reaches_nothing(key: str) -> None: + data = {**OPERATOR_DATA, "season_of": pd.Series(["a", "a", "b"], index=TT[:3])} + m = solved(operator_spec(), data, retain="all") + _, dims, _ = OPERATORS[key] + name = key.replace("-", "_") + folded = m.spec.expressions[f"probe_{name}"].solution + want = xr.DataArray( + PARTIAL_CASES[key], coords={dims[0]: OPERATOR_DATA[dims[0]]}, dims=dims + ) + xr.testing.assert_allclose(folded, want.rename(folded.name)) + if dims == ["t"]: + assert int(m.constraints[f"link_{name}"].labels.sel(t=3)) == -1 + + +def test_a_sum_beside_an_empty_dimension_is_the_empty_sum() -> None: + spec: dict[str, Any] = { + "dimensions": {"t": {"dtype": "int"}, "s": {"dtype": "str"}}, + "variables": {"x": {"foreach": ["t", "s"], "bounds": {"lower": 0, "upper": 1}}}, + "constraints": {"cap": {"foreach": ["s"], "expression": "sum(x, over=t) <= 1"}}, + "objective": {"sense": "maximize", "expression": "sum(x)"}, + } + m = Model.from_spec(spec, {"t": [0, 1], "s": pd.Index([], name="s", dtype=object)}) + assert "cap" not in m.constraints + + +def test_a_window_width_no_member_carries_is_a_window_of_nothing() -> None: + data = { + **OPERATOR_DATA, + "season_of": pd.Series( + [], index=pd.Index([], name="t", dtype=int), dtype=object + ), + } + m = solved(operator_spec(), data, retain="all") + folded = m.spec.expressions["probe_sum_back_group_width"].solution + assert bool(folded.isnull().all())