From 9b7f3adc333b67eceb56ab74783d121f9b62b136 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 14:33:56 +0200 Subject: [PATCH 1/8] feat(spec): build models from math-spec programs and fold named expressions Port lpspec's linopy lane onto the binder: builder, where, operators, coverage and curves, wired to Bound and SpecDataError. Add Model.add_spec, Model.from_spec and the model.spec accessor with expressions and evaluate. --- linopy/model.py | 81 +++ linopy/spec/__init__.py | 11 +- linopy/spec/accessor.py | 172 ++++++ linopy/spec/binder.py | 2 + linopy/spec/builder.py | 317 ++++++++++++ linopy/spec/context.py | 65 +++ linopy/spec/coverage.py | 125 +++++ linopy/spec/curves.py | 182 +++++++ linopy/spec/operators.py | 337 ++++++++++++ linopy/spec/terms.py | 66 +++ linopy/spec/where.py | 151 ++++++ test/test_spec_builder.py | 1036 +++++++++++++++++++++++++++++++++++++ 12 files changed, 2544 insertions(+), 1 deletion(-) create mode 100644 linopy/spec/accessor.py create mode 100644 linopy/spec/builder.py create mode 100644 linopy/spec/context.py create mode 100644 linopy/spec/coverage.py create mode 100644 linopy/spec/curves.py create mode 100644 linopy/spec/operators.py create mode 100644 linopy/spec/terms.py create mode 100644 linopy/spec/where.py create mode 100644 test/test_spec_builder.py diff --git a/linopy/model.py b/linopy/model.py index 769ec1a5..99247328 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -116,6 +116,7 @@ if TYPE_CHECKING: from linopy.piecewise import PiecewiseFormulation + from linopy.spec import ModelSpec, Retain, SpecLike logger = logging.getLogger(__name__) @@ -193,6 +194,7 @@ class Model: "_piecewise_formulations", "_solver", "_sos_reformulation_state", + "_spec", "__weakref__", ) @@ -296,6 +298,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: @@ -428,6 +431,84 @@ 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.parameters``: those the named + expressions read, all of them, or none. + + 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. + """ + 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: """ diff --git a/linopy/spec/__init__.py b/linopy/spec/__init__.py index c50fd391..fbddd060 100644 --- a/linopy/spec/__init__.py +++ b/linopy/spec/__init__.py @@ -16,7 +16,16 @@ "`pip install math-spec` (Python >= 3.12) and try again." ) +from linopy.spec.accessor import ModelSpec, NamedExpressions, SpecLike from linopy.spec.binder import Bound, Retain, bind from linopy.spec.errors import SpecDataError -__all__ = ["Bound", "Retain", "SpecDataError", "bind"] +__all__ = [ + "Bound", + "ModelSpec", + "NamedExpressions", + "Retain", + "SpecDataError", + "SpecLike", + "bind", +] diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py new file mode 100644 index 00000000..6ee19029 --- /dev/null +++ b/linopy/spec/accessor.py @@ -0,0 +1,172 @@ +""" +``model.spec``: the program a model was built from, and its named expressions as data. + +The model owns the data. The spec text, the retained parameters, the lookups +and the master coordinates all sit on the model, so this accessor holds +nothing a round trip through a file could lose: it re-lowers the text and +reads ``model.parameters``. +""" + +from __future__ import annotations + +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, to_program, to_spec +from math_spec import program as ms + +from linopy.model import Model +from linopy.semantics import is_v1 +from linopy.spec.binder import Bound, Retain, bind +from linopy.spec.builder import build, fold +from linopy.spec.context import Context, Parameters, Resolve +from linopy.spec.errors import SpecDataError + +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. + """ + 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) + bound: Bound = bind(program, sources, retain=retain) + build(model, bound) + model.parameters = bound.retained().assign_coords(dict(bound.coords)) + return ModelSpec(model, program, text) + + +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) + + +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) -> None: + self._model = model + self.program = program + self.text = text + + def __repr__(self) -> str: + names = list(self.program.named_expressions) + return f"ModelSpec(expressions={names})" + + @property + def parameters(self) -> xr.Dataset: + """The parameters and lookups retained on the model, on the master coordinates.""" + return self._model.parameters + + @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 folded over the solution and the retained parameters.""" + return NamedExpressions(self) + + def evaluate( + self, name: str, sources: Mapping[str, Any] | xr.Dataset + ) -> xr.DataArray: + """ + The named expression *name*, with its parameters bound afresh from *sources*. + + For a model built with ``retain="none"``, or an expression reading a + parameter ``retain="report"`` did not keep. *sources* is read the way + ``add_spec`` read it, and must describe the coordinates the model was + built on. + """ + bound = bind(self.program, sources, retain="none") + return fold(name, self._context(bound.parameter)) + + def _retained(self, name: str) -> xr.DataArray: + if name not in self.parameters: + raise SpecDataError( + f"parameter '{name}' is not retained on the model: retain='report' keeps only what " + f"the named expressions read, and retain='none' keeps nothing. Build with " + f"retain='all', or read the expression with evaluate(name, sources)." + ) + return self.parameters[name] + + 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, xr.DataArray]): + """The named expressions of a spec, each folded to data on read.""" + + def __init__(self, spec: ModelSpec) -> None: + self._spec = spec + + def __getitem__(self, name: str) -> xr.DataArray: + return fold(name, self._spec._context(self._spec._retained)) + + 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)})" diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py index f893bf43..e10c3f86 100644 --- a/linopy/spec/binder.py +++ b/linopy/spec/binder.py @@ -180,6 +180,8 @@ def _report_closure(program: ms.Program) -> set[str]: names.add(node.offset) elif isinstance(node, ms.Window) and isinstance(node.width, str): names.add(node.width) + elif isinstance(node, ms.Power): + names |= ms.parameters_of(node.base, node.exponent) elif isinstance(node, ms.Cases): for region in node.regions: names |= region.when.names_read diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py new file mode 100644 index 00000000..da876e7d --- /dev/null +++ b/linopy/spec/builder.py @@ -0,0 +1,317 @@ +""" +Program plus bound data to linopy declarations, and a named expression to its value. + +One evaluator serves both: a build hands every variable to linopy as its +term, a fold hands it in as its solved values, and every other node reads the +same way. Which linopy call each construct becomes is one branch of +:func:`evaluate` or one section below. +""" + +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.model import Model +from linopy.spec import curves, operators, terms +from linopy.spec.binder import Bound +from linopy.spec.context import Context, Parameters +from linopy.spec.coverage import ( + check_bounds_cover, + check_constant_side_covers, + check_divisors_cover, +) +from linopy.spec.errors import SpecDataError +from linopy.spec.terms import Array, Term, Value +from linopy.spec.where import as_linopy_mask, bound_lookup, evaluate_where +from linopy.variables import Variable + +_SIGN = {"==": "=", "<=": "<=", ">=": ">="} +_FLIPPED = {"==": "==", "<=": ">=", ">=": "<="} +_SENSE = {"minimize": "min", "maximize": "max"} + + +def build(model: Model, bound: Bound) -> None: + """ + Add every declaration of the bound program to *model*. + + Variables, special-ordered sets, constraints and the objective, in that + order; then every named expression is checked for divisor coverage, so a + body that cannot be folded is refused at build rather than at read. + """ + ctx = Context( + model, + bound.program, + bound.coords, + bound.lookups, + Parameters(bound.program, bound.parameter), + ) + curves.validate(ctx.program, ctx.parameters) + _variables(ctx) + _sos(ctx) + _constraints(ctx) + _objective(ctx) + for name, body in ctx.program.named_expressions.items(): + check_divisors_cover(f"expression '{name}'", (body,), ctx, None) + + +def fold(name: str, ctx: Context) -> xr.DataArray: + """The named expression *name* as data, folded over the solution and the parameters *ctx* holds.""" + 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] + check_divisors_cover(f"expression '{name}'", (body,), ctx, None) + value = evaluate(body, ctx) + if isinstance(value, xr.DataArray): + stray = [c for c in value.coords if c not in value.dims] + return value.drop_vars(stray).rename(name) + 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" + ) + + +# --------------------------------------------------------------------------- +# declarations +# --------------------------------------------------------------------------- + + +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_divisors_cover(f"constraint '{name}'", (row.lhs, row.rhs), ctx, mask) + check_constant_side_covers(name, row, ctx, mask) + 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_divisors_cover("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 bound, so there is nothing to optimize" + ) + ctx.model.add_objective(expr, overwrite=True, sense=_SENSE[declared.sense]) + + +# --------------------------------------------------------------------------- +# evaluation +# --------------------------------------------------------------------------- + + +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.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(operator.add, 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 _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 bound 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 = bound_lookup(node.partition, node.dimension, ctx.lookups) + 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(bound_lookup(name, over, ctx.lookups) for name in names) diff --git a/linopy/spec/context.py b/linopy/spec/context.py new file mode 100644 index 00000000..c3163dc9 --- /dev/null +++ b/linopy/spec/context.py @@ -0,0 +1,65 @@ +"""The data an evaluation reads: parameters resolved once, and the model, coordinates and lookups beside them.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, field + +import pandas as pd +import xarray as xr +from math_spec import program as ms + +from linopy.model import Model +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) + + +@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: Parameters + solved: bool = field(default=False) diff --git a/linopy/spec/coverage.py b/linopy/spec/coverage.py new file mode 100644 index 00000000..79678cf8 --- /dev/null +++ b/linopy/spec/coverage.py @@ -0,0 +1,125 @@ +""" +Is the data there where a declaration needs it? The positions that ask. + +Everywhere else an absent parameter row is a zero coefficient. Three +positions have no answer for that reading: a bound, where zero is a bound +rather than the absence of one; a constant side, where it binds; and 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 Iterator + +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.where import evaluate_where + +Rows = xr.DataArray | None + + +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_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(ms.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' + ) + + +def check_constant_side_covers( + name: str, row: ms.ConstraintDeclaration, ctx: Context, rows: Rows +) -> None: + """A comparison's constant side must have values wherever the row is built, or the zero is the bound.""" + for side in (row.lhs, row.rhs): + if ms.carries_variable(side): + continue + found = sorted( + ( + (node.name, narrowed) + for node, narrowed in _under_regions(side, ctx, rows) + if isinstance(node, ms.Parameter) + ), + key=lambda pair: pair[0], + ) + for param, narrowed in found: + missing = gaps_under(ctx.parameters[param], narrowed) + if missing: + raise SpecDataError( + f"constraint '{name}': 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_divisors_cover( + subject: str, expressions: tuple[ms.ExpressionNode, ...], ctx: Context, rows: Rows +) -> None: + """ + A divisor must have a value wherever *subject* divides by it. + + The rows that ask are the declaration's own, narrowed by the presence of + every variable in the quotient's numerator and by the region of a + ``cases:`` block. Reached before evaluation, the last moment the gap is + visible: the coefficient fill would turn it into a division by zero. + """ + for expression in expressions: + for quotient, region in _under_regions(expression, ctx, rows): + if not isinstance(quotient, ms.Divide): + continue + params = ms.parameters_of(quotient.divisor) + if not params: + continue + needed = region + 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 + for param in sorted(params): + 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 _under_regions( + node: ms.ExpressionNode, ctx: Context, rows: Rows +) -> Iterator[tuple[ms.ExpressionNode, Rows]]: + """Every node under *node* with the rows it has to cover, narrowed at each ``cases:`` region.""" + yield node, rows + if isinstance(node, ms.Cases): + for region in node.regions: + inside = evaluate_where(region.when, ctx) + yield from _under_regions( + region.value, ctx, inside if rows is None else rows & inside + ) + return + for child in ms.children(node): + yield from _under_regions(child, ctx, rows) diff --git a/linopy/spec/curves.py b/linopy/spec/curves.py new file mode 100644 index 00000000..04c28f01 --- /dev/null +++ b/linopy/spec/curves.py @@ -0,0 +1,182 @@ +""" +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/operators.py b/linopy/spec/operators.py new file mode 100644 index 00000000..1c093a7f --- /dev/null +++ b/linopy/spec/operators.py @@ -0,0 +1,337 @@ +""" +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 import terms +from linopy.spec.terms import Array, Term + +Amount = int | xr.DataArray + + +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.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) + + +@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: + groups = _grouped(over, np.asarray(array.indexes[over]), by) + return _gather_in_groups(array, over, _per_group(offset, by), groups, 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 terms.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 = ( + int(np.nanmax(np.asarray(within))) + if isinstance(within, xr.DataArray) + else int(within) + ) + widest = max(1, min(asked, int(array.sizes[over]))) + probe = _Edge(wrap=wrap, fill=None) + groups = 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 groups is None + else _gather_in_groups(array, over, lag, groups, probe) + ) + live, term = ~lagged.isnull(), terms.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 _merged(values: list[Array]) -> Array: + """The sum of *values* in one step: a running sum would re-concatenate the term axis once per lag.""" + data = [value for value in values if isinstance(value, xr.DataArray)] + if len(data) == len(values): + return reduce(operator.add, data) + from linopy import merge + + held = [value for value in values if not isinstance(value, xr.DataArray)] + return cast(LinearExpression, merge(held)) + + +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 terms.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)) + + +@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) + grouped = np.zeros(len(labels), dtype=bool) + for k, key in enumerate(keys): + if terms.unmapped(key): + continue + grouped[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 held else 1 for key, held in zip(keys, grouped)], dtype=int + ) + + def on_axis(values: np.ndarray) -> xr.DataArray: + return xr.DataArray(values, coords={over: labels}, dims=[over]) + + return _Groups( + labels, + on_axis(grouped), + on_axis(belongs), + on_axis(within), + on_axis(span), + roster, + tuple(peers), + tuple(len(beside) for beside in peers.values()), + ) + + +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 terms.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/terms.py b/linopy/spec/terms.py new file mode 100644 index 00000000..2b7b8b48 --- /dev/null +++ b/linopy/spec/terms.py @@ -0,0 +1,66 @@ +""" +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 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 + + +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) + + +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) diff --git a/linopy/spec/where.py b/linopy/spec/where.py new file mode 100644 index 00000000..ab21fc07 --- /dev/null +++ b/linopy/spec/where.py @@ -0,0 +1,151 @@ +"""A ``where:`` predicate as a boolean array over the coordinates it masks.""" + +from __future__ import annotations + +import operator +from collections.abc import Callable, Mapping +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.operators 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 bound_lookup( + name: str, over: str, lookups: Mapping[str, Mapping[str, xr.DataArray]] +) -> xr.DataArray: + """The lookup *name* as an array over *over*, NaN where a label is unmapped.""" + return lookups[over][name] + + +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 = bound_lookup(node.name, node.over, ctx.lookups) + compared = _PREDICATE_OPS[node.op](arr, node.value) & arr.notnull() + return compared.fillna(False).astype(bool) + if isinstance(node, ms.LookupPairComparisonNode): + left = bound_lookup(node.name, node.over, ctx.lookups) + right = bound_lookup(node.other, node.over, ctx.lookups) + compared = ( + _PREDICATE_OPS[node.op](left, right) & left.notnull() & right.notnull() + ) + return compared.fillna(False).astype(bool) + if isinstance(node, ms.LookupDefinedNode): + return bound_lookup(node.name, node.over, ctx.lookups).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 = bound_lookup(node.by, node.name, ctx.lookups) + 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/test/test_spec_builder.py b/test/test_spec_builder.py new file mode 100644 index 00000000..81e1c897 --- /dev/null +++ b/test/test_spec_builder.py @@ -0,0 +1,1036 @@ +""" +Building linopy models from math-spec programs, and reading named expressions back. + +``EXAMPLE_DISPATCH`` is math-spec's ``examples/dispatch.yaml`` with two named +expressions added, so the end-to-end check runs on a spec the language ships. +Setting ``MATH_SPEC_EXAMPLES`` to a math-spec ``examples`` directory builds +and solves every example in it with synthetic data. +""" + +from __future__ import annotations + +import glob +import os +from collections.abc import Callable, Mapping +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 linopy import Model # noqa: E402 +from linopy.spec import ModelSpec, SpecDataError # noqa: E402 + +pytestmark = [ + pytest.mark.v1, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +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: Mapping[str, Any], **kwargs: Any) -> Model: + m = Model.from_spec(spec, sources, **kwargs) + m.solve(solver_name="highs", output_flag=False, reformulate_sos=True) + return m + + +# --------------------------------------------------------------------------- +# 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 yaml_dict() -> dict[str, Any]: + return math_spec.to_spec(yaml.safe_load(EXAMPLE_DISPATCH)).to_dict() + + +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"} + + +# --------------------------------------------------------------------------- +# end to end +# --------------------------------------------------------------------------- + + +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"] + xr.testing.assert_allclose( + spend, (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") + ) + usage = m.spec.expressions["usage"] + 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) + + +def synthetic_sources(program: Any, n: int = 3) -> dict[str, Any]: + """Dense data for every declaration: labels per dimension, a linear ramp per parameter, cyclic lookups.""" + sources: dict[str, Any] = {} + for dim, decl in program.dimensions.items(): + if decl.dtype == "int": + sources[dim] = pd.Index(range(n), name=dim) + elif decl.dtype == "datetime": + sources[dim] = pd.date_range("2030-01-01", periods=n, freq="h", name=dim) + else: + sources[dim] = pd.Index([f"{dim}{i}" for i in range(n)], name=dim) + for over, lk in program.lookups: + if lk.target is not None: + values = [sources[lk.target][i % n] for i in range(n)] + else: + values = ( + list(range(n)) + if lk.dtype == "int" + else [f"{lk.name}{i}" for i in range(n)] + ) + sources[lk.name] = pd.Series(values, index=sources[over]) + ramp = 1.0 + np.arange(n) + for name, p in program.parameters.items(): + if p.derivation is not None: + continue + shape = [n] * len(p.dims) + if p.dtype == "float": + data = np.broadcast_to(ramp, shape).copy() if p.dims else np.array(1.0) + elif p.dtype == "int": + data = np.ones(shape, dtype=int) + elif p.dtype == "bool": + data = np.ones(shape, dtype=bool) + else: + data = np.full(shape, "a", dtype=object) + if not p.dims: + sources[name] = data.item() + else: + sources[name] = xr.DataArray( + data, coords={d: sources[d] for d in p.dims}, dims=p.dims + ) + return sources + + +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") + + +# --------------------------------------------------------------------------- +# retain and evaluate +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("retain", "kept"), + [ + ("report", {"cost", "p_max"}), + ("all", {"cost", "load", "p_max"}), + ("none", set()), + ], +) +def test_retain_decides_what_the_fold_can_read(retain: str, kept: set[str]) -> None: + m = solved(yaml_dict(), DISPATCH_DATA, retain=retain) + assert set(m.parameters.data_vars) == kept + want = (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") + xr.testing.assert_allclose(m.spec.evaluate("spend", DISPATCH_DATA), want) + if "cost" in kept: + xr.testing.assert_allclose(m.spec.expressions["spend"], want) + else: + with pytest.raises(SpecDataError, match="not retained"): + m.spec.expressions["spend"] + + +def test_evaluate_refuses_data_on_other_labels_than_the_model() -> None: + m = solved(yaml_dict(), DISPATCH_DATA) + reordered = {**DISPATCH_DATA, "generator": GENERATOR[::-1]} + with pytest.raises(SpecDataError, match="not aligned on 'generator'"): + m.spec.evaluate("spend", reordered) + + +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"]) == pytest.approx(1.05**3) + with pytest.raises(RuntimeError, match="no solution yet"): + m.spec.expressions["spend"] + + +# --------------------------------------------------------------------------- +# 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:]) + + +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 + + +@pytest.mark.parametrize( + ("spec", "data", "objective"), + [ + pytest.param( + SPARSE_SPEC, + {"w": W_HOLE_AT_0, "c": FULL_C}, + 19.0, + id="coefficient-reads-as-zero", + ), + pytest.param( + with_( + SPARSE_SPEC, + constraints={ + "cap": {**SPARSE_SPEC["constraints"]["cap"], "where": "c"} + }, + ), + {"w": FULL_W, "c": HOLE_AT_0}, + 19.0, + id="constant-side-behind-a-where-is-no-row", + ), + ], +) +def test_a_missing_row_is_a_zero_coefficient_or_no_row( + spec: dict[str, Any], data: dict[str, Any], objective: float +) -> None: + m = solved(spec, {"t": T, **data}) + assert m.objective.value == pytest.approx(objective) + + +@pytest.mark.parametrize( + ("spec", "data", "match"), + [ + 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, + 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, 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_as_bound_constant_side_or_divisor( + 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 + + +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) + + +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_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]))) + + +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"] + 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()) + + +# --------------------------------------------------------------------------- +# operators, built as a constraint and folded as a named expression +# --------------------------------------------------------------------------- + +TT = pd.Index([0, 1, 2, 3], name="t") +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}"] + xr.testing.assert_allclose(built, want.rename(f"y_{name}")) + xr.testing.assert_allclose(folded, want.rename(f"probe_{name}")) + + +# --------------------------------------------------------------------------- +# piecewise curves +# --------------------------------------------------------------------------- + +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")), +} + + +@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"].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_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 + + +# --------------------------------------------------------------------------- +# where predicates +# --------------------------------------------------------------------------- + +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)"}, +} +DAYS = pd.date_range("2030-01-01", periods=4, freq="D", name="d") +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, np.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), +} +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(DAYS[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, swapped sides, constants, empty dimensions +# --------------------------------------------------------------------------- + +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}"] + 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_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) + + +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"] + assert got.ndim == 0 and float(got) == 42.0 + + +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_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}) From c121d978328f078bd1f03c6af3871f44fd6659bc Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 14:46:21 +0200 Subject: [PATCH 2/8] fix(spec): walk into powers, refuse relabelled evaluate sources, harden windows Coverage and the retain closure now descend into a Power's operands; evaluate() refuses sources labelled unlike the model; an all-null window width is a window of nothing; cases fold through the aligned combine. --- linopy/spec/accessor.py | 12 ++++++ linopy/spec/binder.py | 7 ++-- linopy/spec/builder.py | 2 +- linopy/spec/coverage.py | 7 ++-- linopy/spec/nodes.py | 26 ++++++++++++ linopy/spec/operators.py | 22 +++++----- test/test_spec_builder.py | 85 +++++++++++++++++++++++++++++++++++---- 7 files changed, 136 insertions(+), 25 deletions(-) create mode 100644 linopy/spec/nodes.py diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index 6ee19029..e74e7189 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -129,8 +129,20 @@ def evaluate( parameter ``retain="report"`` did not keep. *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. """ bound = bind(self.program, sources, retain="none") + coords = self.coords + for dim, index in bound.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 bound on the same labels in the same order." + ) return fold(name, self._context(bound.parameter)) def _retained(self, name: str) -> xr.DataArray: diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py index e10c3f86..cbf7341f 100644 --- a/linopy/spec/binder.py +++ b/linopy/spec/binder.py @@ -24,6 +24,7 @@ from math_spec import program as ms from linopy.spec.errors import SpecDataError +from linopy.spec.nodes import parameters_of, walk Retain = Literal["report", "all", "none"] _RETAIN: tuple[str, ...] = get_args(Retain) @@ -174,14 +175,12 @@ def _retained_names(self) -> list[str]: def _report_closure(program: ms.Program) -> set[str]: """Every parameter a named expression reads, by node or by name.""" bodies = tuple(program.named_expressions.values()) - names = set(ms.parameters_of(*bodies)) - for node in ms.walk(*bodies): + names = set(parameters_of(*bodies)) + for node in walk(*bodies): if isinstance(node, ms.Translate) and isinstance(node.offset, str): names.add(node.offset) elif isinstance(node, ms.Window) and isinstance(node.width, str): names.add(node.width) - elif isinstance(node, ms.Power): - names |= ms.parameters_of(node.base, node.exponent) elif isinstance(node, ms.Cases): for region in node.regions: names |= region.when.names_read diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py index da876e7d..c6d39614 100644 --- a/linopy/spec/builder.py +++ b/linopy/spec/builder.py @@ -238,7 +238,7 @@ def evaluate(node: ms.ExpressionNode, ctx: Context) -> Value: _in_region(evaluate(region.value, ctx), evaluate_where(region.when, ctx)) for region in node.regions ) - return functools.reduce(operator.add, regions) + return functools.reduce(lambda a, b: _combine(operator.add, a, b), regions) assert_never(node) diff --git a/linopy/spec/coverage.py b/linopy/spec/coverage.py index 79678cf8..687836e1 100644 --- a/linopy/spec/coverage.py +++ b/linopy/spec/coverage.py @@ -19,6 +19,7 @@ from linopy.spec import terms from linopy.spec.context import Context from linopy.spec.errors import SpecDataError +from linopy.spec.nodes import children, parameters_of from linopy.spec.where import evaluate_where Rows = xr.DataArray | None @@ -36,7 +37,7 @@ 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(ms.parameters_of(declared.lower, declared.upper)) + names = sorted(parameters_of(declared.lower, declared.upper)) missing = sum(gaps_under(ctx.parameters[p], rows) for p in names) if missing: raise SpecDataError( @@ -90,7 +91,7 @@ def check_divisors_cover( for quotient, region in _under_regions(expression, ctx, rows): if not isinstance(quotient, ms.Divide): continue - params = ms.parameters_of(quotient.divisor) + params = parameters_of(quotient.divisor) if not params: continue needed = region @@ -121,5 +122,5 @@ def _under_regions( region.value, ctx, inside if rows is None else rows & inside ) return - for child in ms.children(node): + for child in children(node): yield from _under_regions(child, ctx, rows) diff --git a/linopy/spec/nodes.py b/linopy/spec/nodes.py new file mode 100644 index 00000000..a12c3ecb --- /dev/null +++ b/linopy/spec/nodes.py @@ -0,0 +1,26 @@ +"""Walks over expression nodes that descend into every operand, a ``Power``'s included.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from math_spec import program as ms + + +def children(node: ms.ExpressionNode) -> tuple[ms.ExpressionNode, ...]: + """The operands of *node*: ``math_spec.program.children`` plus a power's base and exponent.""" + if isinstance(node, ms.Power): + return (node.base, node.exponent) + return ms.children(node) + + +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(*children(node)) + + +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)) diff --git a/linopy/spec/operators.py b/linopy/spec/operators.py index 1c093a7f..06a5b0e3 100644 --- a/linopy/spec/operators.py +++ b/linopy/spec/operators.py @@ -160,11 +160,7 @@ def sum_back( """ if by is not None: within = _per_group(within, by) - asked = ( - int(np.nanmax(np.asarray(within))) - if isinstance(within, xr.DataArray) - else int(within) - ) + asked = _widest(within) widest = max(1, min(asked, int(array.sizes[over]))) probe = _Edge(wrap=wrap, fill=None) groups = None if by is None else _grouped(over, np.asarray(array.indexes[over]), by) @@ -184,15 +180,21 @@ def sum_back( 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.""" - data = [value for value in values if isinstance(value, xr.DataArray)] - if len(data) == len(values): - return reduce(operator.add, data) + if isinstance(values[0], xr.DataArray): + return reduce(operator.add, values) from linopy import merge - held = [value for value in values if not isinstance(value, xr.DataArray)] - return cast(LinearExpression, merge(held)) + return cast(LinearExpression, merge(cast(list[Term], values))) def _renamed( diff --git a/test/test_spec_builder.py b/test/test_spec_builder.py index 81e1c897..ad3b5bea 100644 --- a/test/test_spec_builder.py +++ b/test/test_spec_builder.py @@ -263,13 +263,6 @@ def test_retain_decides_what_the_fold_can_read(retain: str, kept: set[str]) -> N m.spec.expressions["spend"] -def test_evaluate_refuses_data_on_other_labels_than_the_model() -> None: - m = solved(yaml_dict(), DISPATCH_DATA) - reordered = {**DISPATCH_DATA, "generator": GENERATOR[::-1]} - with pytest.raises(SpecDataError, match="not aligned on 'generator'"): - m.spec.evaluate("spend", reordered) - - 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"): @@ -1034,3 +1027,81 @@ def test_a_convex_hull_curve_may_bend_either_way_but_not_both() -> None: ) with pytest.raises(SpecDataError, match="exact only for a single bend"): Model.from_spec(spec, {**CURVE_DATA, "bp_x": FULL_X, "bp_y": mixed}) + + +# --------------------------------------------------------------------------- +# a power hides nothing +# --------------------------------------------------------------------------- + + +@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.parameters.data_vars) + xr.testing.assert_allclose( + m.spec.expressions["e"], + xr.DataArray([0.0, 0.0, 4.0], coords={"t": T}, name="e"), + ) + + +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_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"] + assert bool(folded.isnull().all()) From 399908ad24611e9d578262209c547c6b03d2f4d6 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 14:56:22 +0200 Subject: [PATCH 3/8] feat(spec): round trip a spec-built model through netcdf Persist the spec text, the master coordinates and the lookups alongside the model, re-lowering the program from the text on read; math-spec is imported only for a file that carries a spec. Lookups and arrays of labels are stored as codes into a category table, so partial maps keep their holes and dtypes. --- benchmarks/models/__init__.py | 1 + benchmarks/models/spec_pypsa.py | 95 ++++++++++++++ linopy/io.py | 28 ++++- linopy/spec/accessor.py | 9 ++ linopy/spec/netcdf.py | 165 +++++++++++++++++++++++++ linopy/testing.py | 4 + test/test_spec_io.py | 211 ++++++++++++++++++++++++++++++++ 7 files changed, 511 insertions(+), 2 deletions(-) create mode 100644 benchmarks/models/spec_pypsa.py create mode 100644 linopy/spec/netcdf.py create mode 100644 test/test_spec_io.py diff --git a/benchmarks/models/__init__.py b/benchmarks/models/__init__.py index 66c9a7c7..2b9f7eca 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 00000000..5d3fac63 --- /dev/null +++ b/benchmarks/models/spec_pypsa.py @@ -0,0 +1,95 @@ +""" +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. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np +import pandas as pd +import xarray as xr + +from benchmarks.registry import BUILD, FROM_NETCDF, TO_NETCDF, BenchSpec, register + +if TYPE_CHECKING: + import linopy + +SIZES = (5, 40) # labels per dimension; 40 is ~20k variables + +EXAMPLES = os.environ.get("MATH_SPEC_EXAMPLES") +EXAMPLE = Path(EXAMPLES, "pypsa.yaml") if EXAMPLES else None + + +def synthetic_sources(program: Any, n: int) -> dict[str, Any]: + """``n`` labels per dimension, a linear ramp per parameter, cyclic lookups.""" + sources: dict[str, Any] = {} + for dim, decl in program.dimensions.items(): + if decl.dtype == "int": + sources[dim] = pd.Index(range(n), name=dim) + elif decl.dtype == "datetime": + sources[dim] = pd.date_range("2030-01-01", periods=n, freq="h", name=dim) + else: + sources[dim] = pd.Index([f"{dim}{i}" for i in range(n)], name=dim) + for over, lookup in program.lookups: + into = sources[lookup.target] if lookup.target else range(n) + sources[lookup.name] = pd.Series( + [into[i % n] for i in range(n)], index=sources[over] + ) + ramp = 1.0 + np.arange(n) + for name, parameter in program.parameters.items(): + if parameter.derivation is not None: + continue + dims = parameter.dims + shape = [n] * len(dims) + if parameter.dtype == "bool": + values: Any = np.ones(shape, dtype=bool) + elif parameter.dtype == "int": + values = np.ones(shape, dtype=int) + elif parameter.dtype == "str": + values = np.full(shape, "a", dtype=object) + else: + values = np.broadcast_to(ramp, shape).copy() if dims else np.array(1.0) + sources[name] = ( + values.item() + if not dims + else xr.DataArray( + values, coords={d: sources[d] for d in dims}, dims=list(dims) + ) + ) + return sources + + +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 + + path = str(EXAMPLE) + sources = synthetic_sources(math_spec.to_program(path), n) + with linopy.options as options: + options["semantics"] = "v1" # a spec-built model is v1 only + 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/linopy/io.py b/linopy/io.py index 6ee8b36f..f64b90b5 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -46,6 +46,7 @@ NETCDF_VERSION_ATTR = "_linopy_version" EXPR_TYPE_ATTR = "_linopy_expr_type" +SPEC_ATTR = "_linopy_spec" ufunc_kwargs = dict(vectorize=True) @@ -1038,6 +1039,11 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: type) are all persisted and fully restored by :func:`linopy.io.read_netcdf`. + A model built with :meth:`Model.add_spec` also persists its spec: the + YAML text, the master coordinates and the lookups. ``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 @@ -1098,12 +1104,23 @@ 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")] + parameters = m.parameters + specs: list[xr.Dataset] = [] + if m._spec is not None: + from linopy.spec.netcdf import encode + + parameters, spec_ds = encode(m._spec) + specs = [spec_ds] + params = [with_prefix(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") + if m._spec is not None: + ds.attrs[SPEC_ATTR] = m._spec.text if m._relaxed_registry: ds.attrs["_relaxed_registry"] = json.dumps(m._relaxed_registry) if m._piecewise_formulations: @@ -1260,6 +1277,11 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: m.parameters = 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: setattr(m, k, ds.attrs[k]) @@ -1392,6 +1414,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._rebound(new_model) 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/spec/accessor.py b/linopy/spec/accessor.py index e74e7189..525e1862 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -61,6 +61,11 @@ def attach( return ModelSpec(model, program, text) +def restore(model: Model, text: str) -> ModelSpec: + """The accessor for *model*, with the program lowered afresh from *text*.""" + return ModelSpec(model, to_program(yaml.safe_load(text)), text) + + 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): @@ -96,6 +101,10 @@ def __repr__(self) -> str: names = list(self.program.named_expressions) return f"ModelSpec(expressions={names})" + def _rebound(self, model: Model) -> ModelSpec: + """The same spec, read off *model*.""" + return ModelSpec(model, self.program, self.text) + @property def parameters(self) -> xr.Dataset: """The parameters and lookups retained on the model, on the master coordinates.""" diff --git a/linopy/spec/netcdf.py b/linopy/spec/netcdf.py new file mode 100644 index 00000000..de255929 --- /dev/null +++ b/linopy/spec/netcdf.py @@ -0,0 +1,165 @@ +""" +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. + +Labels are the delicate part. A partial lookup holds NaN in an array of +labels, and neither the labels nor their holes survive a netcdf type: the +engines hand back `` tuple[xr.Dataset, xr.Dataset]: + """ + The model's parameters without the coded arrays, and the spec's own dataset. + + The spec dataset holds one array of labels per master coordinate and, per + coded array, its codes and its categories. It 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. + """ + parameters = spec.parameters + arrays: dict[str, xr.DataArray] = { + COORD + dim: _array(index.to_numpy(), (dim,)) + for dim, index in spec.coords.items() + } + for name in _coded(spec): + arrays.update(_encode(name, parameters[name])) + parameters = parameters.drop_vars(name) + return parameters, _prefixed(xr.Dataset(arrays)) + + +def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: + """ + Re-lower *text* onto *model* and put its coded arrays and coordinates back. + + The parameters read from the file are the retained ones minus what + :func:`encode` took out; together with the master coordinates and the + decoded arrays they are the dataset :func:`linopy.spec.accessor.attach` + left on the model when it was built. + """ + sub = _unprefixed(ds) + coords = { + _stripped(name, COORD): _index(sub[name]) + for name in sub.data_vars + if str(name).startswith(COORD) + } + coded = { + _stripped(name, CODES): _decode(sub, _stripped(name, CODES), coords) + for name in sub.data_vars + if str(name).startswith(CODES) + } + model.parameters = model.parameters.assign_coords(coords).assign(coded) + return restore(model, text) + + +def _coded(spec: ModelSpec) -> list[str]: + """The parameters written as codes: every lookup and every array of labels.""" + 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.kind in LABEL_KINDS + ] + + +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 _decode(sub: xr.Dataset, name: str, coords: dict[str, pd.Index]) -> xr.DataArray: + codes = sub[CODES + name] + dtype = np.dtype(codes.attrs[DTYPE]) + 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: dtype or str(values.dtype)}) + + +def _prefixed(ds: xr.Dataset) -> xr.Dataset: + return ds.rename({k: PREFIX + str(k) for k in (*ds.dims, *ds.data_vars)}) + + +def _unprefixed(ds: xr.Dataset) -> xr.Dataset: + sub = ds[[k for k in ds.data_vars if str(k).startswith(PREFIX)]] + return sub.rename({k: str(k)[len(PREFIX) :] for k in (*sub.dims, *sub.data_vars)}) + + +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])) + + +def _index(arr: xr.DataArray) -> pd.Index: + return pd.Index(_values(arr), name=_stripped(arr.name, COORD)) diff --git a/linopy/testing.py b/linopy/testing.py index 5fd16778..c51d3d34 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -134,6 +134,10 @@ 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 a.status == b.status assert a.termination_condition == b.termination_condition diff --git a/test/test_spec_io.py b/test/test_spec_io.py new file mode 100644 index 00000000..16007306 --- /dev/null +++ b/test/test_spec_io.py @@ -0,0 +1,211 @@ +""" +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, + synthetic_sources, +) + +import linopy # noqa: E402 +from linopy import Model, read_netcdf # noqa: E402 +from linopy.io import SPEC_ATTR # 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} + + +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], p.spec.expressions[name]) + for dim, index in m.spec.coords.items(): + pd.testing.assert_index_equal(index, p.spec.coords[dim]) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_a_retain_none_model_evaluates_after_a_round_trip( + tmp_path: Path, engine: str +) -> None: + 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 + assert_arrayequal( + m.spec.evaluate("spend", DISPATCH_DATA), p.spec.evaluate("spend", DISPATCH_DATA) + ) + + +@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_arrayequal(m.spec.lookups[over][name], p.spec.lookups[over][name]) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_labelled_parameters_and_unreached_coordinates_round_trip( + tmp_path: Path, engine: str +) -> None: + m = Model.from_spec(WHERE_SPEC, WHERE_DATA, retain="all") + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert_arrayequal(m.spec.parameters["label"], p.spec.parameters["label"]) + for dim, index in m.spec.coords.items(): + pd.testing.assert_index_equal(index, p.spec.coords[dim]) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_a_solved_model_reports_the_same_expression( + tmp_path: Path, engine: str +) -> None: + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA) + p = roundtrip(m, tmp_path, engine) + + assert p.objective.value == m.objective.value + assert_arrayequal(m.spec.expressions["spend"], p.spec.expressions["spend"]) + assert float(p.spec.expressions["spend"].sum()) == pytest.approx(2500.0) + + +@pytest.mark.parametrize("deep", [True, False]) +def test_a_copy_carries_the_spec(deep: bool) -> None: + m = Model.from_spec(WHERE_SPEC, WHERE_DATA, retain="all") + p = m.copy(deep=deep) + m.parameters = m.parameters.drop_vars("cost") + + assert p.spec.text == m.spec.text + assert "cost" in p.spec.parameters, "the copy's spec reads the copy, not the source" + for over, by_name in m.spec.lookups.items(): + for name, lookup in by_name.items(): + assert_arrayequal(lookup, p.spec.lookups[over][name]) + + +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) + for dim, index in m.spec.coords.items(): + pd.testing.assert_index_equal(index, p.spec.coords[dim]) + for over, by_name in m.spec.lookups.items(): + for name, lookup in by_name.items(): + assert_arrayequal(lookup, p.spec.lookups[over][name]) From e0a29c97ad8880973f1296508ae3ce525994a3a2 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 15:35:19 +0200 Subject: [PATCH 4/8] fix(spec): keep parameter dtypes and one coordinate dtype per dimension Write the in-memory dtype of every parameter and cast it back on read, and stamp the master coordinates onto every container, so no engine leaves a model disagreeing with itself. assert_model_equal now compares dataset dtypes, and synthetic_sources moves to linopy/spec/testing.py for both users. --- benchmarks/models/spec_pypsa.py | 54 +++----------------- linopy/io.py | 2 - linopy/spec/netcdf.py | 83 +++++++++++++++++++++++++------ linopy/spec/testing.py | 72 +++++++++++++++++++++++++++ linopy/testing.py | 19 ++++++- test/test_spec_builder.py | 43 +--------------- test/test_spec_io.py | 87 +++++++++++++++++++++++---------- 7 files changed, 226 insertions(+), 134 deletions(-) create mode 100644 linopy/spec/testing.py diff --git a/benchmarks/models/spec_pypsa.py b/benchmarks/models/spec_pypsa.py index 5d3fac63..fce1719a 100644 --- a/benchmarks/models/spec_pypsa.py +++ b/benchmarks/models/spec_pypsa.py @@ -4,69 +4,28 @@ 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. +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, Any - -import numpy as np -import pandas as pd -import xarray as xr +from typing import TYPE_CHECKING from benchmarks.registry import BUILD, FROM_NETCDF, TO_NETCDF, BenchSpec, register if TYPE_CHECKING: import linopy -SIZES = (5, 40) # labels per dimension; 40 is ~20k variables +SIZES = (5, 40) EXAMPLES = os.environ.get("MATH_SPEC_EXAMPLES") EXAMPLE = Path(EXAMPLES, "pypsa.yaml") if EXAMPLES else None -def synthetic_sources(program: Any, n: int) -> dict[str, Any]: - """``n`` labels per dimension, a linear ramp per parameter, cyclic lookups.""" - sources: dict[str, Any] = {} - for dim, decl in program.dimensions.items(): - if decl.dtype == "int": - sources[dim] = pd.Index(range(n), name=dim) - elif decl.dtype == "datetime": - sources[dim] = pd.date_range("2030-01-01", periods=n, freq="h", name=dim) - else: - sources[dim] = pd.Index([f"{dim}{i}" for i in range(n)], name=dim) - for over, lookup in program.lookups: - into = sources[lookup.target] if lookup.target else range(n) - sources[lookup.name] = pd.Series( - [into[i % n] for i in range(n)], index=sources[over] - ) - ramp = 1.0 + np.arange(n) - for name, parameter in program.parameters.items(): - if parameter.derivation is not None: - continue - dims = parameter.dims - shape = [n] * len(dims) - if parameter.dtype == "bool": - values: Any = np.ones(shape, dtype=bool) - elif parameter.dtype == "int": - values = np.ones(shape, dtype=int) - elif parameter.dtype == "str": - values = np.full(shape, "a", dtype=object) - else: - values = np.broadcast_to(ramp, shape).copy() if dims else np.array(1.0) - sources[name] = ( - values.item() - if not dims - else xr.DataArray( - values, coords={d: sources[d] for d in dims}, dims=list(dims) - ) - ) - return sources - - def build_spec_pypsa(n: int) -> linopy.Model: """Lower ``pypsa.yaml`` and build it with ``n`` labels per dimension.""" import pytest @@ -76,11 +35,12 @@ def build_spec_pypsa(n: int) -> linopy.Model: 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" # a spec-built model is v1 only + options["semantics"] = "v1" return linopy.Model.from_spec(path, sources) diff --git a/linopy/io.py b/linopy/io.py index f64b90b5..6c59d0e4 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1119,8 +1119,6 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ) ds = ds.assign_attrs(scalars) ds.attrs[NETCDF_VERSION_ATTR] = version("linopy") - if m._spec is not None: - ds.attrs[SPEC_ATTR] = m._spec.text if m._relaxed_registry: ds.attrs["_relaxed_registry"] = json.dumps(m._relaxed_registry) if m._piecewise_formulations: diff --git a/linopy/spec/netcdf.py b/linopy/spec/netcdf.py index de255929..e2214b0e 100644 --- a/linopy/spec/netcdf.py +++ b/linopy/spec/netcdf.py @@ -6,24 +6,32 @@ coordinates and the lookups; the program is re-lowered from the text on read, so no lowered ``Program`` ever reaches the file. -Labels are the delicate part. A partial lookup holds NaN in an array of -labels, and neither the labels nor their holes survive a netcdf type: the -engines hand back `` tuple[xr.Dataset, xr.Dataset]: """ The model's parameters without the coded arrays, and the spec's own dataset. - The spec dataset holds one array of labels per master coordinate and, per - coded array, its codes and its categories. It carries no coordinates of + The spec dataset carries the spec text as its one attribute, which the + merge lifts to the file's, and holds one array of labels per master + coordinate and, per coded array, its codes and its categories. It 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. @@ -56,7 +64,12 @@ def encode(spec: ModelSpec) -> tuple[xr.Dataset, xr.Dataset]: for name in _coded(spec): arrays.update(_encode(name, parameters[name])) parameters = parameters.drop_vars(name) - return parameters, _prefixed(xr.Dataset(arrays)) + typed = { + str(name): arr.assign_attrs({DTYPE: str(arr.dtype)}) + for name, arr in parameters.items() + } + written = _prefixed(xr.Dataset(arrays)).assign_attrs({SPEC_ATTR: spec.text}) + return parameters.assign(typed), written def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: @@ -79,17 +92,50 @@ def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: for name in sub.data_vars if str(name).startswith(CODES) } - model.parameters = model.parameters.assign_coords(coords).assign(coded) + typed = {str(name): _cast(arr) for name, arr in model.parameters.items()} + model.parameters = ( + model.parameters.assign(typed).assign_coords(coords).assign(coded) + ) + _restamp(model, coords) return restore(model, text) +def _restamp(model: Model, coords: Mapping[str, pd.Index]) -> None: + """Put the master coordinates on every container that carries a dimension.""" + from linopy.constraints import Constraint, CSRConstraint + + for _, variable in model.variables.items(): + variable._data = _stamped(variable.data, coords) + for _, expression in model.expressions.items(): + expression._data = _stamped(expression.data, coords) + model.objective.expression._data = _stamped(model.objective.expression.data, coords) + for _, constraint in model.constraints.items(): + if isinstance(constraint, Constraint): + constraint._data = _stamped(constraint.data, coords) + elif isinstance(constraint, CSRConstraint): + constraint._coords = [ + coords.get(str(index.name), index) for index in constraint._coords + ] + + +def _stamped(data: xr.Dataset, coords: Mapping[str, pd.Index]) -> xr.Dataset: + """*data* with the master coordinates in place of the ones a dtype narrowed.""" + indexes = data.indexes + stale = { + dim: index + for dim, index in coords.items() + if dim in indexes and indexes[dim].dtype != index.dtype + } + return data.assign_coords(stale) if stale else data + + def _coded(spec: ModelSpec) -> list[str]: - """The parameters written as codes: every lookup and every array of labels.""" + """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.kind in LABEL_KINDS + if name in lookups or arr.dtype == object ] @@ -137,6 +183,11 @@ def _categories(sub: xr.Dataset, name: str, dtype: np.dtype) -> np.ndarray: return np.empty(0, dtype=dtype) +def _cast(arr: xr.DataArray) -> xr.DataArray: + """A parameter at the dtype it had in memory, whatever the engine returned.""" + return arr.astype(np.dtype(arr.attrs.pop(DTYPE))) + + def _array( values: np.ndarray, dims: tuple[Any, ...], dtype: str | None = None ) -> xr.DataArray: diff --git a/linopy/spec/testing.py b/linopy/spec/testing.py new file mode 100644 index 00000000..5bec8301 --- /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/testing.py b/linopy/testing.py index c51d3d34..59a9bb61 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 set(a.variables) == set(b.variables) assert set(a.constraints) == set(b.constraints) diff --git a/test/test_spec_builder.py b/test/test_spec_builder.py index ad3b5bea..9c6fd15e 100644 --- a/test/test_spec_builder.py +++ b/test/test_spec_builder.py @@ -26,6 +26,7 @@ import linopy # noqa: E402 from linopy import Model # noqa: E402 from linopy.spec import ModelSpec, SpecDataError # noqa: E402 +from linopy.spec.testing import synthetic_sources # noqa: E402 pytestmark = [ pytest.mark.v1, @@ -173,48 +174,6 @@ def test_the_dispatch_example_solves_and_its_expressions_fold() -> None: assert m.spec.coords["generator"].equals(GENERATOR) -def synthetic_sources(program: Any, n: int = 3) -> dict[str, Any]: - """Dense data for every declaration: labels per dimension, a linear ramp per parameter, cyclic lookups.""" - sources: dict[str, Any] = {} - for dim, decl in program.dimensions.items(): - if decl.dtype == "int": - sources[dim] = pd.Index(range(n), name=dim) - elif decl.dtype == "datetime": - sources[dim] = pd.date_range("2030-01-01", periods=n, freq="h", name=dim) - else: - sources[dim] = pd.Index([f"{dim}{i}" for i in range(n)], name=dim) - for over, lk in program.lookups: - if lk.target is not None: - values = [sources[lk.target][i % n] for i in range(n)] - else: - values = ( - list(range(n)) - if lk.dtype == "int" - else [f"{lk.name}{i}" for i in range(n)] - ) - sources[lk.name] = pd.Series(values, index=sources[over]) - ramp = 1.0 + np.arange(n) - for name, p in program.parameters.items(): - if p.derivation is not None: - continue - shape = [n] * len(p.dims) - if p.dtype == "float": - data = np.broadcast_to(ramp, shape).copy() if p.dims else np.array(1.0) - elif p.dtype == "int": - data = np.ones(shape, dtype=int) - elif p.dtype == "bool": - data = np.ones(shape, dtype=bool) - else: - data = np.full(shape, "a", dtype=object) - if not p.dims: - sources[name] = data.item() - else: - sources[name] = xr.DataArray( - data, coords={d: sources[d] for d in p.dims}, dims=p.dims - ) - return sources - - EXAMPLES_DIR = os.environ.get("MATH_SPEC_EXAMPLES") EXAMPLES = ( sorted(glob.glob(f"{EXAMPLES_DIR}/*.yaml") + glob.glob(f"{EXAMPLES_DIR}/*/*.yaml")) diff --git a/test/test_spec_io.py b/test/test_spec_io.py index 16007306..4f791186 100644 --- a/test/test_spec_io.py +++ b/test/test_spec_io.py @@ -26,12 +26,12 @@ WHERE_DATA, WHERE_SPEC, solved, - synthetic_sources, ) import linopy # noqa: E402 from linopy import Model, read_netcdf # noqa: E402 from linopy.io import SPEC_ATTR # noqa: E402 +from linopy.spec.testing import synthetic_sources # noqa: E402 from linopy.testing import assert_model_equal # noqa: E402 pytestmark = [ @@ -66,6 +66,25 @@ 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.""" @@ -110,8 +129,6 @@ def test_a_spec_built_model_round_trips( assert set(p.spec.expressions) == set(m.spec.expressions) for name in m.spec.expressions: assert_arrayequal(m.spec.expressions[name], p.spec.expressions[name]) - for dim, index in m.spec.coords.items(): - pd.testing.assert_index_equal(index, p.spec.coords[dim]) @pytest.mark.parametrize("engine", ENGINES) @@ -139,45 +156,67 @@ def test_a_lookup_round_trips_exactly( p = roundtrip(m, tmp_path, engine) assert_model_equal(m, p) - assert_arrayequal(m.spec.lookups[over][name], p.spec.lookups[over][name]) + assert name in p.spec.lookups[over] @pytest.mark.parametrize("engine", ENGINES) -def test_labelled_parameters_and_unreached_coordinates_round_trip( - tmp_path: Path, engine: str -) -> None: - m = Model.from_spec(WHERE_SPEC, WHERE_DATA, retain="all") +@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_arrayequal(m.spec.parameters["label"], p.spec.parameters["label"]) - for dim, index in m.spec.coords.items(): - pd.testing.assert_index_equal(index, p.spec.coords[dim]) + assert p.spec.parameters[name].dtype == m.spec.parameters[name].dtype @pytest.mark.parametrize("engine", ENGINES) -def test_a_solved_model_reports_the_same_expression( +@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: - m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA) + """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 p.objective.value == m.objective.value - assert_arrayequal(m.spec.expressions["spend"], p.spec.expressions["spend"]) - assert float(p.spec.expressions["spend"].sum()) == pytest.approx(2500.0) + 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) - m.parameters = m.parameters.drop_vars("cost") + p.parameters["label"].values[1] = "changed" assert p.spec.text == m.spec.text - assert "cost" in p.spec.parameters, "the copy's spec reads the copy, not the source" - for over, by_name in m.spec.lookups.items(): - for name, lookup in by_name.items(): - assert_arrayequal(lookup, p.spec.lookups[over][name]) + assert p.spec.parameters["label"].values[1] == "changed" + assert m.spec.parameters["label"].values[1] == ("u" if deep else "changed") def test_a_model_without_a_spec_carries_none(tmp_path: Path) -> None: @@ -204,8 +243,4 @@ def test_the_pypsa_example_round_trips(tmp_path: Path, engine: str) -> None: p = roundtrip(m, tmp_path, engine) assert_model_equal(m, p) - for dim, index in m.spec.coords.items(): - pd.testing.assert_index_equal(index, p.spec.coords[dim]) - for over, by_name in m.spec.lookups.items(): - for name, lookup in by_name.items(): - assert_arrayequal(lookup, p.spec.lookups[over][name]) + assert set(p.spec.coords) == set(m.spec.coords) From b94740907baacd9c2775e8f10dcce6a60946b2e8 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 4 Sep 2026 11:40:08 +0200 Subject: [PATCH 5/8] feat(spec): refuse a missing parameter row wherever it is used A missing parameter row was read as a silent zero when it stood as a coefficient, while a bound, constant side or divisor already refused it. Refuse it as a coefficient too, so every position behaves alike and a hole is never filled without the modeller saying so: mask the coordinate out with a where, or fill the value into the data. --- linopy/spec/binder.py | 11 ++-- linopy/spec/builder.py | 9 ++- linopy/spec/coverage.py | 56 +++++++++++++++--- test/test_spec_builder.py | 120 ++++++++++++++++++++++++++++++-------- 4 files changed, 157 insertions(+), 39 deletions(-) diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py index cbf7341f..56081f47 100644 --- a/linopy/spec/binder.py +++ b/linopy/spec/binder.py @@ -80,9 +80,9 @@ def bind( 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. + 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. """ if retain not in _RETAIN: raise SpecDataError( @@ -568,9 +568,8 @@ def _refuse_strangers(name: str, dim: str, labels: pd.Index, known: pd.Index) -> 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 missing row is a zero coefficient, but a label that is not a coordinate is a typo: its " - f"row joins nothing, so the coordinate it was meant for silently reads as absent. Fix the " - f"label, or add it to sources['{dim}']." + f"A label that is not a coordinate is a typo: its row joins nothing, so the coordinate it " + f"was meant for is left uncovered. Fix the label, or add it to sources['{dim}']." ) diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py index c6d39614..3aaba35d 100644 --- a/linopy/spec/builder.py +++ b/linopy/spec/builder.py @@ -25,6 +25,7 @@ from linopy.spec.context import Context, Parameters from linopy.spec.coverage import ( check_bounds_cover, + check_coefficients_cover, check_constant_side_covers, check_divisors_cover, ) @@ -43,8 +44,9 @@ def build(model: Model, bound: Bound) -> None: Add every declaration of the bound program to *model*. Variables, special-ordered sets, constraints and the objective, in that - order; then every named expression is checked for divisor coverage, so a - body that cannot be folded is refused at build rather than at read. + 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, @@ -60,6 +62,7 @@ def build(model: Model, bound: Bound) -> None: _objective(ctx) for name, body in ctx.program.named_expressions.items(): check_divisors_cover(f"expression '{name}'", (body,), ctx, None) + check_coefficients_cover(f"expression '{name}'", (body,), ctx, None) def fold(name: str, ctx: Context) -> xr.DataArray: @@ -127,6 +130,7 @@ def _constraints(ctx: Context) -> None: mask = as_linopy_mask(rows) check_divisors_cover(f"constraint '{name}'", (row.lhs, row.rhs), ctx, mask) check_constant_side_covers(name, row, ctx, mask) + check_coefficients_cover(f"constraint '{name}'", (row.lhs, row.rhs), ctx, mask) lhs, rhs = evaluate(row.lhs, ctx), evaluate(row.rhs, ctx) if _term_free(lhs) and _term_free(rhs): continue @@ -159,6 +163,7 @@ def _objective(ctx: Context) -> None: if declared is None: return check_divisors_cover("the objective", (declared.expression,), ctx, None) + check_coefficients_cover("the objective", (declared.expression,), ctx, None) expr = evaluate(declared.expression, ctx) if not isinstance(expr, Variable | LinearExpression | QuadraticExpression): raise SpecDataError( diff --git a/linopy/spec/coverage.py b/linopy/spec/coverage.py index 687836e1..06b47da2 100644 --- a/linopy/spec/coverage.py +++ b/linopy/spec/coverage.py @@ -1,12 +1,13 @@ """ -Is the data there where a declaration needs it? The positions that ask. - -Everywhere else an absent parameter row is a zero coefficient. Three -positions have no answer for that reading: a bound, where zero is a bound -rather than the absence of one; a constant side, where it binds; and 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. +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 @@ -110,6 +111,45 @@ def check_divisors_cover( ) +def check_coefficients_cover( + subject: str, expressions: tuple[ms.ExpressionNode, ...], ctx: Context, rows: Rows +) -> 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. Decided against the rows the declaration builds, + narrowed at each ``cases:`` region exactly as the other checks are, so a + ``where`` that removed the coordinate has already answered. A shift offset + or window width given by name is a coefficient too, and stands or falls + over its own coordinates. + """ + for expression in expressions: + for node, region in _under_regions(expression, ctx, rows): + for param, needed in _coefficient_uses(node, region): + 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 _coefficient_uses( + node: ms.ExpressionNode, region: Rows +) -> Iterator[tuple[str, Rows]]: + """Each parameter *node* uses as a coefficient, with the rows it has to cover.""" + if isinstance(node, ms.Parameter): + yield node.name, region + elif isinstance(node, ms.Translate) and isinstance(node.offset, str): + yield node.offset, None + elif isinstance(node, ms.Window) and isinstance(node.width, str): + yield node.width, None + + def _under_regions( node: ms.ExpressionNode, ctx: Context, rows: Rows ) -> Iterator[tuple[ms.ExpressionNode, Rows]]: diff --git a/test/test_spec_builder.py b/test/test_spec_builder.py index 9c6fd15e..dcdec4ac 100644 --- a/test/test_spec_builder.py +++ b/test/test_spec_builder.py @@ -272,38 +272,34 @@ def with_(spec: dict[str, Any], **sections: dict[str, Any]) -> dict[str, Any]: return out +NO_W_CONSTRAINT = {"cap": {"foreach": ["t"], "expression": "x <= c"}} + + @pytest.mark.parametrize( - ("spec", "data", "objective"), + ("spec", "data", "match"), [ pytest.param( SPARSE_SPEC, {"w": W_HOLE_AT_0, "c": FULL_C}, - 19.0, - id="coefficient-reads-as-zero", + "constraint 'cap'.*parameter 'w' is used as a coefficient", + id="coefficient-in-a-constraint", ), pytest.param( with_( SPARSE_SPEC, - constraints={ - "cap": {**SPARSE_SPEC["constraints"]["cap"], "where": "c"} - }, + constraints=NO_W_CONSTRAINT, + objective={"sense": "maximize", "expression": "sum(w * x, over=t)"}, ), - {"w": FULL_W, "c": HOLE_AT_0}, - 19.0, - id="constant-side-behind-a-where-is-no-row", + {"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", ), - ], -) -def test_a_missing_row_is_a_zero_coefficient_or_no_row( - spec: dict[str, Any], data: dict[str, Any], objective: float -) -> None: - m = solved(spec, {"t": T, **data}) - assert m.objective.value == pytest.approx(objective) - - -@pytest.mark.parametrize( - ("spec", "data", "match"), - [ pytest.param( SPARSE_SPEC, {"w": FULL_W, "c": HOLE_AT_0}, @@ -333,6 +329,7 @@ def test_a_missing_row_is_a_zero_coefficient_or_no_row( 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}, @@ -340,14 +337,18 @@ def test_a_missing_row_is_a_zero_coefficient_or_no_row( id="divisor-in-the-objective", ), pytest.param( - with_(SPARSE_SPEC, expressions={"ratio": "x / w"}), + 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_as_bound_constant_side_or_divisor( +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): @@ -378,6 +379,79 @@ def test_a_masked_variable_bound_needs_no_row_where_it_is_masked() -> None: assert int((m.variables["x"].labels != -1).sum()) == 2 +# --------------------------------------------------------------------------- +# a shift amount is a coefficient, and a where removes the row that would ask +# --------------------------------------------------------------------------- + + +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: + g = pd.Index([0, 1], name="g") + data = { + "t": T, + "g": g, + "grp": pd.Series([0, 0, 1], index=T), + "v": FULL_C, + "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) + + +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"}}, From 2ae9df9decd7aa4eea37e5bc18e3cbdfd546ece8 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 4 Sep 2026 11:40:39 +0200 Subject: [PATCH 6/8] doc(spec): add a notebook building models from specs A runnable, nbconvert-clean walkthrough of the spec feature: the dispatch program, binding data, folding named expressions, retain and evaluate, the uniform absence rule, lookups and grouped sums, temporal shift, and the netCDF round trip. --- examples/building-models-from-specs.ipynb | 938 ++++++++++++++++++++++ 1 file changed, 938 insertions(+) create mode 100644 examples/building-models-from-specs.ipynb diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb new file mode 100644 index 00000000..be2a0d58 --- /dev/null +++ b/examples/building-models-from-specs.ipynb @@ -0,0 +1,938 @@ +{ + "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` binds\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. Binding 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 binding 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, binds 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\n", + "\n", + "`m.spec` is the accessor onto the program the model was built from. Its\n", + "`expressions` mapping evaluates each named expression **numerically** against\n", + "the solution: every variable is replaced by its solved value, every parameter by\n", + "the data it was bound to, and the arithmetic runs on xarray. This is called\n", + "**folding**.\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", + "usage = m.spec.expressions[\"usage\"]\n", + "print(\"\\nspend per hour:\")\n", + "print(spend)\n", + "print(\"\\nusage (output / capacity):\")\n", + "print(usage)" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "A named expression that reads only data (no variables) folds **before** a solve\n", + "too — it needs a solution only if it actually references a variable. An unknown\n", + "name raises a `KeyError` with a suggestion." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " m.spec.expressions[\"spent\"]\n", + "except KeyError as e:\n", + " print(\"KeyError:\", e)" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "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.parameters` after building:\n", + "\n", + "| `retain` | keeps in `model.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`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "for retain in [\"report\", \"all\", \"none\"]:\n", + " mm = Model.from_spec(DISPATCH, dispatch_data, retain=retain)\n", + " print(f\"retain={retain!r:9} -> parameters kept: {sorted(mm.parameters.data_vars)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "### `evaluate`: fold against fresh data\n", + "\n", + "With `retain=\"none\"` nothing is kept, so `expressions[...]` cannot fold. For\n", + "that case (or any expression whose parameters were not retained) there is\n", + "`spec.evaluate(name, sources)`: it rebinds the parameters from a **fresh** bag\n", + "of data and folds against the model's 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": "21", + "metadata": {}, + "outputs": [], + "source": [ + "lean = Model.from_spec(DISPATCH, dispatch_data, retain=\"none\")\n", + "lean.solve(solver_name=\"highs\", output_flag=False)\n", + "\n", + "# expressions[...] cannot fold: no parameters were retained.\n", + "try:\n", + " lean.spec.expressions[\"spend\"]\n", + "except SpecDataError as e:\n", + " print(\"SpecDataError:\", str(e)[:90], \"...\\n\")\n", + "\n", + "# evaluate rebinds from fresh sources and folds:\n", + "print(lean.spec.evaluate(\"spend\", dispatch_data))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "22", + "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": "23", + "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": "24", + "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": "25", + "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": "26", + "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": "27", + "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": "28", + "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": "29", + "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": "30", + "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": "31", + "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": "32", + "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": "33", + "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": "34", + "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": "35", + "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": "36", + "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": "37", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"folded cost:\", float(bm.spec.expressions[\"cost\"]))" + ] + }, + { + "cell_type": "markdown", + "id": "38", + "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": "39", + "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": "40", + "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": "41", + "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(m2.spec.expressions[name], restored.spec.expressions[name])\n", + " print(f\" {name}: identical\")" + ] + }, + { + "cell_type": "markdown", + "id": "42", + "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": "43", + "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": "44", + "metadata": {}, + "source": [ + "`Model.copy()` carries the spec too, with the accessor rebound 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": "45", + "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(clone.spec.expressions[\"spend\"], m2.spec.expressions[\"spend\"])\n", + "print(\"after solving the copy, folded expressions match the original\")" + ] + }, + { + "cell_type": "markdown", + "id": "46", + "metadata": {}, + "source": [ + "## Where the code lives, and one upstream note\n", + "\n", + "The feature is a small package, `linopy/spec/`, imported only when you call\n", + "`add_spec`/`from_spec` — `import linopy` never pulls in `math_spec`. Roughly:\n", + "\n", + "- `accessor.py` — `model.spec`, folding, `evaluate`.\n", + "- `binder.py` — the three binding 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. One workaround lives here:\n", + " math-spec alpha.73's `program.children()` does not descend into a `Power`\n", + " node, so parameters hidden under `**` would be missed; `nodes.py` walks into\n", + " the base and exponent itself.\n", + "\n", + "### Summary\n", + "\n", + "A spec is the maths over labelled axes; the sources are the numbers. `linopy`\n", + "binds them into an ordinary model, folds named expressions back into arrays\n", + "after solving, refuses a missing parameter row wherever it is used — as a\n", + "coefficient, bound, constant side or divisor alike, with `where:` and filling\n", + "the data as the escape hatches — and round-trips the lot through netCDF by\n", + "keeping the spec as text beside 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 +} From c936682fd7ace33b506be6d84a89d9bbf1058062 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 4 Sep 2026 13:08:18 +0200 Subject: [PATCH 7/8] feat(spec): expose a named expression as three views m.spec.expressions[name] returns a NamedExpression bundling .node (the lowered formula), .expression (the unsolved linopy expression) and .solution (the fold over the model's solution). evaluate() returns the same object. Add ModelSpec.to_latex/to_markdown/to_typst for whole-model typesetting, rendered as Markdown in a notebook. --- doc/release_notes.rst | 9 ++ examples/building-models-from-specs.ipynb | 169 ++++++++++++++-------- linopy/spec/__init__.py | 8 +- linopy/spec/accessor.py | 110 ++++++++++++-- linopy/spec/builder.py | 19 ++- linopy/spec/context.py | 7 +- test/test_spec_builder.py | 94 ++++++++++-- test/test_spec_io.py | 7 +- 8 files changed, 331 insertions(+), 92 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index e3580e05..f51b1b9e 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 bound to data, and ``model.spec`` reads it back. Requires the ``math-spec`` package and v1 semantics. + +* ``model.spec.expressions[name]`` returns a ``NamedExpression`` with three views of a named expression: ``.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 bound afresh. + +* ``model.spec.to_latex`` / ``.to_markdown`` / ``.to_typst`` typeset the whole model; a ``ModelSpec`` and a ``NamedExpression`` 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`. diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb index be2a0d58..58f15100 100644 --- a/examples/building-models-from-specs.ipynb +++ b/examples/building-models-from-specs.ipynb @@ -281,13 +281,19 @@ "id": "14", "metadata": {}, "source": [ - "### Named expressions become data\n", + "### 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 evaluates each named expression **numerically** against\n", - "the solution: every variable is replaced by its solved value, every parameter by\n", - "the data it was bound to, and the arithmetic runs on xarray. This is called\n", - "**folding**.\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 bound. 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 bound 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." @@ -302,11 +308,15 @@ "source": [ "print(repr(m.spec))\n", "spend = m.spec.expressions[\"spend\"]\n", - "usage = m.spec.expressions[\"usage\"]\n", - "print(\"\\nspend per hour:\")\n", - "print(spend)\n", - "print(\"\\nusage (output / capacity):\")\n", - "print(usage)" + "\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)" ] }, { @@ -314,9 +324,11 @@ "id": "16", "metadata": {}, "source": [ - "A named expression that reads only data (no variables) folds **before** a solve\n", - "too — it needs a solution only if it actually references a variable. An unknown\n", - "name raises a `KeyError` with a suggestion." + "### The model as maths\n", + "\n", + "The accessor typesets the whole model, delegating to math-spec:\n", + "`m.spec.to_latex()`, `.to_markdown()` and `.to_typst()`. In a notebook the\n", + "accessor renders as Markdown on its own; here we show it explicitly." ] }, { @@ -325,6 +337,29 @@ "id": "17", "metadata": {}, "outputs": [], + "source": [ + "from IPython.display import Markdown\n", + "\n", + "Markdown(m.spec.to_markdown())" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "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": "19", + "metadata": {}, + "outputs": [], "source": [ "try:\n", " m.spec.expressions[\"spent\"]\n", @@ -334,7 +369,7 @@ }, { "cell_type": "markdown", - "id": "18", + "id": "20", "metadata": {}, "source": [ "## 5. `retain`: what data stays on the model\n", @@ -355,7 +390,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "21", "metadata": {}, "outputs": [], "source": [ @@ -366,15 +401,16 @@ }, { "cell_type": "markdown", - "id": "20", + "id": "22", "metadata": {}, "source": [ "### `evaluate`: fold against fresh data\n", "\n", - "With `retain=\"none\"` nothing is kept, so `expressions[...]` cannot fold. For\n", - "that case (or any expression whose parameters were not retained) there is\n", - "`spec.evaluate(name, sources)`: it rebinds the parameters from a **fresh** bag\n", - "of data and folds against the model's solution.\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 rebound 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", @@ -384,27 +420,27 @@ { "cell_type": "code", "execution_count": null, - "id": "21", + "id": "23", "metadata": {}, "outputs": [], "source": [ "lean = Model.from_spec(DISPATCH, dispatch_data, retain=\"none\")\n", "lean.solve(solver_name=\"highs\", output_flag=False)\n", "\n", - "# expressions[...] cannot fold: no parameters were retained.\n", + "# .solution cannot fold: no parameters were retained.\n", "try:\n", - " lean.spec.expressions[\"spend\"]\n", + " lean.spec.expressions[\"spend\"].solution\n", "except SpecDataError as e:\n", " print(\"SpecDataError:\", str(e)[:90], \"...\\n\")\n", "\n", - "# evaluate rebinds from fresh sources and folds:\n", - "print(lean.spec.evaluate(\"spend\", dispatch_data))" + "# evaluate rebinds from fresh sources; .solution folds:\n", + "print(lean.spec.evaluate(\"spend\", dispatch_data).solution)" ] }, { "cell_type": "code", "execution_count": null, - "id": "22", + "id": "24", "metadata": {}, "outputs": [], "source": [ @@ -418,7 +454,7 @@ }, { "cell_type": "markdown", - "id": "23", + "id": "25", "metadata": {}, "source": [ "## 6. Absence and coverage — one rule, every position\n", @@ -451,7 +487,7 @@ { "cell_type": "code", "execution_count": null, - "id": "24", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -483,7 +519,7 @@ }, { "cell_type": "markdown", - "id": "25", + "id": "27", "metadata": {}, "source": [ "The other three positions refuse the same kind of hole, joining the\n", @@ -494,7 +530,7 @@ { "cell_type": "code", "execution_count": null, - "id": "26", + "id": "28", "metadata": {}, "outputs": [], "source": [ @@ -520,7 +556,7 @@ }, { "cell_type": "markdown", - "id": "27", + "id": "29", "metadata": {}, "source": [ "Two escape hatches fix the coefficient hole above, and both build and solve.\n", @@ -536,7 +572,7 @@ { "cell_type": "code", "execution_count": null, - "id": "28", + "id": "30", "metadata": {}, "outputs": [], "source": [ @@ -557,7 +593,7 @@ }, { "cell_type": "markdown", - "id": "29", + "id": "31", "metadata": {}, "source": [ "And the same masking escape hatch on the variable and constraint together:\n", @@ -568,7 +604,7 @@ { "cell_type": "code", "execution_count": null, - "id": "30", + "id": "32", "metadata": {}, "outputs": [], "source": [ @@ -593,7 +629,7 @@ }, { "cell_type": "markdown", - "id": "31", + "id": "33", "metadata": {}, "source": [ "## 7. Lookups and grouped sums\n", @@ -607,7 +643,7 @@ { "cell_type": "code", "execution_count": null, - "id": "32", + "id": "34", "metadata": {}, "outputs": [], "source": [ @@ -641,7 +677,7 @@ }, { "cell_type": "markdown", - "id": "33", + "id": "35", "metadata": {}, "source": [ "Note `south` has no generators mapped to it. Its group is **empty**, and an\n", @@ -651,7 +687,7 @@ }, { "cell_type": "markdown", - "id": "34", + "id": "36", "metadata": {}, "source": [ "## 8. Temporal operators: `shift`\n", @@ -671,7 +707,7 @@ { "cell_type": "code", "execution_count": null, - "id": "35", + "id": "37", "metadata": {}, "outputs": [], "source": [ @@ -726,7 +762,7 @@ }, { "cell_type": "markdown", - "id": "36", + "id": "38", "metadata": {}, "source": [ "The generator over-produces while power is cheap (hour 2 runs at 30 to fill the\n", @@ -737,16 +773,16 @@ { "cell_type": "code", "execution_count": null, - "id": "37", + "id": "39", "metadata": {}, "outputs": [], "source": [ - "print(\"folded cost:\", float(bm.spec.expressions[\"cost\"]))" + "print(\"folded cost:\", float(bm.spec.expressions[\"cost\"].solution))" ] }, { "cell_type": "markdown", - "id": "38", + "id": "40", "metadata": {}, "source": [ "## 9. Synthetic data for any spec\n", @@ -761,7 +797,7 @@ { "cell_type": "code", "execution_count": null, - "id": "39", + "id": "41", "metadata": {}, "outputs": [], "source": [ @@ -779,7 +815,7 @@ }, { "cell_type": "markdown", - "id": "40", + "id": "42", "metadata": {}, "source": [ "## 10. Persistence: netCDF and copy\n", @@ -799,7 +835,7 @@ { "cell_type": "code", "execution_count": null, - "id": "41", + "id": "43", "metadata": {}, "outputs": [], "source": [ @@ -823,13 +859,15 @@ "\n", "# and the named expressions fold identically after the round-trip:\n", "for name in restored.spec.expressions:\n", - " xr.testing.assert_equal(m2.spec.expressions[name], restored.spec.expressions[name])\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": "42", + "id": "44", "metadata": {}, "source": [ "Even a `retain=\"none\"` model round-trips: the spec text and coordinates survive,\n", @@ -839,7 +877,7 @@ { "cell_type": "code", "execution_count": null, - "id": "43", + "id": "45", "metadata": {}, "outputs": [], "source": [ @@ -854,7 +892,7 @@ }, { "cell_type": "markdown", - "id": "44", + "id": "46", "metadata": {}, "source": [ "`Model.copy()` carries the spec too, with the accessor rebound to the copy. The\n", @@ -866,7 +904,7 @@ { "cell_type": "code", "execution_count": null, - "id": "45", + "id": "47", "metadata": {}, "outputs": [], "source": [ @@ -875,21 +913,24 @@ "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(clone.spec.expressions[\"spend\"], m2.spec.expressions[\"spend\"])\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": "46", + "id": "48", "metadata": {}, "source": [ - "## Where the code lives, and one upstream note\n", + "## Where the code lives, and two upstream notes\n", "\n", "The feature is a small package, `linopy/spec/`, imported only when you call\n", "`add_spec`/`from_spec` — `import linopy` never pulls in `math_spec`. Roughly:\n", "\n", - "- `accessor.py` — `model.spec`, folding, `evaluate`.\n", + "- `accessor.py` — `model.spec`, the `NamedExpression` views, `evaluate`, and\n", + " whole-model typesetting (`to_latex` / `to_markdown` / `to_typst`).\n", "- `binder.py` — the three binding rules; data onto master coordinates.\n", "- `builder.py` — emits variables, constraints, objective; folds expressions.\n", "- `operators.py` — `sum`, `by=`, `shift`, `at`, `sum_back`.\n", @@ -903,14 +944,22 @@ " node, so parameters hidden under `**` would be missed; `nodes.py` walks into\n", " the base and exponent itself.\n", "\n", + "Two upstream requests shape what the typesetting shows:\n", + "[math-spec#384](https://github.com/energy-models/math-spec/issues/384) asks for a\n", + "public hook to typeset a **single** named expression, so a `NamedExpression`\n", + "could render its own formula rather than only the whole model; and the\n", + "whole-model output currently prints the objective, constraints and variable\n", + "domains, not the named expressions themselves.\n", + "\n", "### Summary\n", "\n", "A spec is the maths over labelled axes; the sources are the numbers. `linopy`\n", - "binds them into an ordinary model, folds named expressions back into arrays\n", - "after solving, refuses a missing parameter row wherever it is used — as a\n", - "coefficient, bound, constant side or divisor alike, with `where:` and filling\n", - "the data as the escape hatches — and round-trips the lot through netCDF by\n", - "keeping the spec as text beside factorized labels." + "binds them into an ordinary model, hands each named expression back as three\n", + "views — its formula, its unsolved linopy expression and its solution — refuses a\n", + "missing parameter row wherever it is used (as a coefficient, bound, constant\n", + "side or divisor alike, with `where:` and filling the data as the escape\n", + "hatches), and round-trips the lot through netCDF by keeping the spec as text\n", + "beside factorized labels." ] } ], diff --git a/linopy/spec/__init__.py b/linopy/spec/__init__.py index fbddd060..a6f653ac 100644 --- a/linopy/spec/__init__.py +++ b/linopy/spec/__init__.py @@ -16,13 +16,19 @@ "`pip install math-spec` (Python >= 3.12) and try again." ) -from linopy.spec.accessor import ModelSpec, NamedExpressions, SpecLike +from linopy.spec.accessor import ( + ModelSpec, + NamedExpression, + NamedExpressions, + SpecLike, +) from linopy.spec.binder import Bound, Retain, bind from linopy.spec.errors import SpecDataError __all__ = [ "Bound", "ModelSpec", + "NamedExpression", "NamedExpressions", "Retain", "SpecDataError", diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index 525e1862..cd8d5855 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -9,6 +9,7 @@ from __future__ import annotations +import functools from collections.abc import Iterator, Mapping from pathlib import Path from typing import Any, TypeAlias @@ -16,13 +17,22 @@ import pandas as pd import xarray as xr import yaml -from math_spec import Spec, to_program, to_spec +from math_spec import ( + Spec, + did_you_mean, + to_latex, + to_markdown, + to_program, + to_spec, + to_typst, +) from math_spec import program as ms from linopy.model import Model from linopy.semantics import is_v1 +from linopy.spec import terms from linopy.spec.binder import Bound, Retain, bind -from linopy.spec.builder import build, fold +from linopy.spec.builder import build, evaluate_named, fold from linopy.spec.context import Context, Parameters, Resolve from linopy.spec.errors import SpecDataError @@ -125,12 +135,32 @@ def lookups(self) -> dict[str, dict[str, xr.DataArray]]: @property def expressions(self) -> NamedExpressions: - """Each named expression folded over the solution and the retained parameters.""" + """Each named expression as a :class:`NamedExpression`: its math, its linopy fold and its solution.""" return NamedExpressions(self) + 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 - ) -> xr.DataArray: + ) -> NamedExpression: """ The named expression *name*, with its parameters bound afresh from *sources*. @@ -152,7 +182,7 @@ def evaluate( f"was built on {coords[dim].tolist()[:5]}. evaluate() reads the solution the " f"model holds, so the data must be bound on the same labels in the same order." ) - return fold(name, self._context(bound.parameter)) + return NamedExpression(self, name, self._context(bound.parameter)) def _retained(self, name: str) -> xr.DataArray: if name not in self.parameters: @@ -174,14 +204,21 @@ def _context(self, resolve: Resolve) -> Context: ) -class NamedExpressions(Mapping[str, xr.DataArray]): - """The named expressions of a spec, each folded to data on read.""" +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) -> xr.DataArray: - return fold(name, self._spec._context(self._spec._retained)) + 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._retained) + ) def __iter__(self) -> Iterator[str]: return iter(self._spec.program.named_expressions) @@ -191,3 +228,58 @@ def __len__(self) -> int: def __repr__(self) -> str: return f"NamedExpressions({list(self)})" + + +class NamedExpression: + """ + 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 retained parameters and + the solution the model holds; ``evaluate(name, sources)`` binds fresh data. + + Attributes: + node: The lowered expression body, math-spec's own AST handle. + """ + + def __init__(self, spec: ModelSpec, name: str, ctx: Context) -> None: + self._spec = spec + self._name = 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] + + @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 not retained. + """ + 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}')" + + def _repr_markdown_(self) -> str: + return self._spec.to_markdown() diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py index 3aaba35d..431e4de1 100644 --- a/linopy/spec/builder.py +++ b/linopy/spec/builder.py @@ -65,8 +65,8 @@ def build(model: Model, bound: Bound) -> None: check_coefficients_cover(f"expression '{name}'", (body,), ctx, None) -def fold(name: str, ctx: Context) -> xr.DataArray: - """The named expression *name* as data, folded over the solution and the parameters *ctx* holds.""" +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}'. " @@ -75,9 +75,14 @@ def fold(name: str, ctx: Context) -> xr.DataArray: body = ctx.program.named_expressions[name] check_divisors_cover(f"expression '{name}'", (body,), ctx, None) 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): - stray = [c for c in value.coords if c not in value.dims] - return value.drop_vars(stray).rename(name) + return value if isinstance(value, float | int): return xr.DataArray(float(value), name=name) raise TypeError( @@ -85,6 +90,12 @@ def fold(name: str, ctx: Context) -> xr.DataArray: ) +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) + + # --------------------------------------------------------------------------- # declarations # --------------------------------------------------------------------------- diff --git a/linopy/spec/context.py b/linopy/spec/context.py index c3163dc9..235e8b10 100644 --- a/linopy/spec/context.py +++ b/linopy/spec/context.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Callable, Iterator, Mapping -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace import pandas as pd import xarray as xr @@ -63,3 +63,8 @@ class Context: lookups: Mapping[str, Mapping[str, xr.DataArray]] parameters: Parameters 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) diff --git a/test/test_spec_builder.py b/test/test_spec_builder.py index dcdec4ac..6ee93785 100644 --- a/test/test_spec_builder.py +++ b/test/test_spec_builder.py @@ -25,7 +25,7 @@ import linopy # noqa: E402 from linopy import Model # noqa: E402 -from linopy.spec import ModelSpec, SpecDataError # noqa: E402 +from linopy.spec import ModelSpec, NamedExpression, SpecDataError # noqa: E402 from linopy.spec.testing import synthetic_sources # noqa: E402 pytestmark = [ @@ -161,11 +161,11 @@ 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"] + 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"] + 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 @@ -214,12 +214,12 @@ def test_retain_decides_what_the_fold_can_read(retain: str, kept: set[str]) -> N m = solved(yaml_dict(), DISPATCH_DATA, retain=retain) assert set(m.parameters.data_vars) == kept want = (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") - xr.testing.assert_allclose(m.spec.evaluate("spend", DISPATCH_DATA), want) + xr.testing.assert_allclose(m.spec.evaluate("spend", DISPATCH_DATA).solution, want) if "cost" in kept: - xr.testing.assert_allclose(m.spec.expressions["spend"], want) + xr.testing.assert_allclose(m.spec.expressions["spend"].solution, want) else: with pytest.raises(SpecDataError, match="not retained"): - m.spec.expressions["spend"] + m.spec.expressions["spend"].solution def test_an_unknown_expression_is_a_key_error_with_a_hint() -> None: @@ -242,9 +242,73 @@ def test_a_fold_over_variables_needs_a_solution_and_one_over_data_does_not() -> }, } m = Model.from_spec(spec, {**DISPATCH_DATA, "rate": 1.05, "years": 3.0}) - assert float(m.spec.expressions["growth"]) == pytest.approx(1.05**3) + assert float(m.spec.expressions["growth"].solution) == pytest.approx(1.05**3) with pytest.raises(RuntimeError, match="no solution yet"): - m.spec.expressions["spend"] + 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"] + 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_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() # --------------------------------------------------------------------------- @@ -589,7 +653,7 @@ def test_a_fold_reads_a_masked_slot_the_way_its_absence_says( 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"] + 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: @@ -699,7 +763,7 @@ def test_an_operator_builds_and_folds_alike(operators_model: Model, key: str) -> 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}"] + 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}")) @@ -805,7 +869,7 @@ 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"].item() == pytest.approx(spend) + assert m.spec.expressions["spend"].solution.item() == pytest.approx(spend) assert m.objective.value == pytest.approx(spend) @@ -1001,7 +1065,7 @@ def test_a_member_a_lookup_sends_nowhere_reaches_nothing(key: str) -> None: m = solved(operator_spec(), data, retain="all") _, dims, _ = OPERATORS[key] name = key.replace("-", "_") - folded = m.spec.expressions[f"probe_{name}"] + folded = m.spec.expressions[f"probe_{name}"].solution want = xr.DataArray( PARTIAL_CASES[key], coords={dims[0]: OPERATOR_DATA[dims[0]]}, dims=dims ) @@ -1020,7 +1084,7 @@ def test_a_constant_on_the_left_is_the_same_row() -> None: 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"] + got = Model.from_spec(spec, DISPATCH_DATA).spec.expressions["answer"].solution assert got.ndim == 0 and float(got) == 42.0 @@ -1099,7 +1163,7 @@ def test_an_operator_under_a_power_keeps_its_parameters_retained() -> None: m = Model.from_spec(spec, {"t": T, "w": FULL_W, "c": FULL_C, "lag": 1}) assert {"c", "lag"} <= set(m.parameters.data_vars) xr.testing.assert_allclose( - m.spec.expressions["e"], + m.spec.expressions["e"].solution, xr.DataArray([0.0, 0.0, 4.0], coords={"t": T}, name="e"), ) @@ -1136,5 +1200,5 @@ def test_a_window_width_no_member_carries_is_a_window_of_nothing() -> None: ), } m = solved(operator_spec(), data, retain="all") - folded = m.spec.expressions["probe_sum_back_group_width"] + folded = m.spec.expressions["probe_sum_back_group_width"].solution assert bool(folded.isnull().all()) diff --git a/test/test_spec_io.py b/test/test_spec_io.py index 4f791186..5af661f0 100644 --- a/test/test_spec_io.py +++ b/test/test_spec_io.py @@ -128,7 +128,9 @@ def test_a_spec_built_model_round_trips( 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], p.spec.expressions[name]) + assert_arrayequal( + m.spec.expressions[name].solution, p.spec.expressions[name].solution + ) @pytest.mark.parametrize("engine", ENGINES) @@ -141,7 +143,8 @@ def test_a_retain_none_model_evaluates_after_a_round_trip( assert_model_equal(m, p) assert not p.spec.parameters.data_vars assert_arrayequal( - m.spec.evaluate("spend", DISPATCH_DATA), p.spec.evaluate("spend", DISPATCH_DATA) + m.spec.evaluate("spend", DISPATCH_DATA).solution, + p.spec.evaluate("spend", DISPATCH_DATA).solution, ) From 4875f8f2944bd5749e9874c1fe3ec5c0bdbbb4e8 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 4 Sep 2026 18:49:05 +0200 Subject: [PATCH 8/8] ci: skip spec notebook until math-spec is on PyPI building-models-from-specs.ipynb imports math_spec, which the docs CI environment does not install, so the notebook job failed on import. Skip it like the other special-setup notebooks. --- .github/workflows/test-notebooks.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/test-notebooks.yml b/.github/workflows/test-notebooks.yml index 4050badb..14081651 100644 --- a/.github/workflows/test-notebooks.yml +++ b/.github/workflows/test-notebooks.yml @@ -44,6 +44,10 @@ jobs: echo "Skipping $name (requires credentials or special setup)" continue ;; + building-models-from-specs.ipynb) + echo "Skipping $name (requires math-spec, not yet on PyPI)" + continue + ;; esac echo "::group::Running $name"