From 28ac1ada4951e0deea8f4002fbdc0c482be317f0 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 13:16:25 +0200 Subject: [PATCH 01/35] feat(spec): xarray data binder for math-spec programs Turns a lowered math-spec Program plus user data into master coordinates, padded lookups and on-demand parameter arrays under the three binding rules. Missing rows stay NaN for the builder. Sources are pulled by key, never iterated, and aligned arrays keep their buffer. --- linopy/spec/__init__.py | 22 ++ linopy/spec/binder.py | 614 +++++++++++++++++++++++++++++++ linopy/spec/errors.py | 12 + test/test_spec_binder.py | 772 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 1420 insertions(+) create mode 100644 linopy/spec/__init__.py create mode 100644 linopy/spec/binder.py create mode 100644 linopy/spec/errors.py create mode 100644 test/test_spec_binder.py diff --git a/linopy/spec/__init__.py b/linopy/spec/__init__.py new file mode 100644 index 00000000..c50fd391 --- /dev/null +++ b/linopy/spec/__init__.py @@ -0,0 +1,22 @@ +""" +Build linopy models from math-spec programs. + +The package needs the ``math-spec`` distribution (import name ``math_spec``, +Python >= 3.12). It is imported here and nowhere else in linopy, so +``import linopy`` never pulls it in. +""" + +from __future__ import annotations + +from importlib.util import find_spec + +if find_spec("math_spec") is None: + raise ImportError( + "linopy.spec needs the math-spec package. Install it with " + "`pip install math-spec` (Python >= 3.12) and try again." + ) + +from linopy.spec.binder import Bound, Retain, bind +from linopy.spec.errors import SpecDataError + +__all__ = ["Bound", "Retain", "SpecDataError", "bind"] diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py new file mode 100644 index 00000000..008e915f --- /dev/null +++ b/linopy/spec/binder.py @@ -0,0 +1,614 @@ +""" +Bind user data to a math-spec program. + +The language fixes three binding rules and this module enforces them: a +dimension's members come only from the source keyed by the dimension's +name, their order is the source's order and is never sorted, and a +parameter or lookup source is read for values, never for labels. Parameters +are resolved from ``sources`` on demand and aligned onto the master +coordinates without copying an already aligned array. A coordinate a table +leaves out becomes NaN (``False`` for a ``bool`` parameter); what that means +is the builder's question, not this module's. +""" + +from __future__ import annotations + +from collections.abc import Hashable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any, Literal + +import numpy as np +import pandas as pd +import xarray as xr +from math_spec import did_you_mean +from math_spec import program as ms + +from linopy.spec.errors import SpecDataError + +Retain = Literal["report", "all", "none"] + +_ACCEPTED_KINDS: dict[str, frozenset[str]] = { + "float": frozenset("fiu"), + "int": frozenset("iu"), + "bool": frozenset("b"), + "str": frozenset("OUS"), +} +_KIND_NAMES: dict[str, str] = { + "f": "float", + "i": "int", + "u": "int", + "b": "bool", + "O": "str", + "U": "str", + "S": "str", +} +_SCALARS = (bool, int, float, str, np.number, np.bool_) +_DIMENSION_SHAPES = "a pandas Index, a list, a tuple, a 1-D numpy array, a pandas Series or a 1-D DataArray" +_LOOKUP_SHAPES = "a pandas Series indexed by '{over}', a dict keyed by '{over}' labels, or a 1-D DataArray over '{over}'" +_PARAMETER_SHAPES = ( + "a DataArray over {dims}, a pandas Series whose (Multi)Index levels are {dims}, " + "a DataFrame with columns {columns} or in wide form, a dict keyed by label, or one number" +) + + +def bind( + program: ms.Program, + sources: Mapping[str, Any] | xr.Dataset, + *, + retain: Retain = "report", +) -> Bound: + """ + Bind *sources* to *program*: master coordinates now, parameters on demand. + + Args: + program: The lowered spec. + sources: Data keyed by declared name. Any mapping works; it is read by + key and never iterated beyond ``sources.keys()``. An ``xr.Dataset`` + is accepted too: its indexes are dimension sources, its data + variables parameters and lookups. + retain: Which parameters :meth:`Bound.retained` persists. + + Raises: + SpecDataError: 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. + """ + if isinstance(sources, xr.Dataset): + sources = _dataset_sources(sources) + keys = frozenset(sources.keys()) + _check_keys(program, keys) + coords = _master_coords(program, sources, keys) + lookups = _lookups(program, sources, keys, coords) + return Bound(program, coords, lookups, retain, sources, keys) + + +@dataclass(frozen=True) +class Bound: + """ + A program bound to its data. + + Attributes: + program: The lowered spec the data is bound to. + coords: Master coordinates by dimension, in source order, each index + named after its dimension. A declared dimension nothing reaches + and nothing supplies is absent. + lookups: By dimension, by lookup name, the map as an array over the + dimension's master coordinates, NaN where a label is unmapped. + retain: Which parameters :meth:`retained` persists. + sources: The caller's data, read by key on demand. + keys: The keys ``sources`` carries, read once at bind time. + """ + + program: ms.Program + coords: dict[str, pd.Index] + lookups: dict[str, dict[str, xr.DataArray]] + retain: Retain + sources: Mapping[str, Any] + keys: frozenset[str] + + def parameter(self, name: str) -> xr.DataArray: + """ + The parameter *name* resolved from ``sources`` and aligned to ``coords``. + + Resolved on every call and never cached. An already aligned array is + returned without a copy; a mismatching one is reindexed onto the + master coordinates, leaving NaN (``False`` for ``bool``) where no row + was supplied. + + Raises: + SpecDataError: No data, a shape no reader accepts, a rank other + than declared, a label its dimension lacks, two rows for one + coordinate, a null value in a row, or values of another type + than declared. + """ + declared = self._declaration(name) + if name not in self.keys: + raise SpecDataError(f"no data provided for parameter '{name}'") + arr = _as_array(name, declared, self.sources[name], self.coords) + onto = {d: self.coords[d] for d in declared.dims} + return _aligned(name, arr, onto, _fill(declared)) + + def retained(self) -> xr.Dataset: + """The lookups plus the parameters ``retain`` keeps, as one dataset.""" + arrays = {n: self.parameter(n) for n in self._retained_names()} + for by_name in self.lookups.values(): + arrays.update(by_name) + return xr.Dataset(arrays) + + def _declaration(self, name: str) -> ms.ParameterDeclaration: + if name not in self.program.parameters: + raise SpecDataError( + f"unknown parameter '{name}'. {did_you_mean(name, self.program.parameters)}" + ) + declared = self.program.parameters[name] + if declared.derivation is not None: + raise SpecDataError( + f"parameter '{name}' is emitted by piecewise block '{declared.derivation.block}' " + f"and is filled from the block's own breakpoints, not bound from sources." + ) + return declared + + def _retained_names(self) -> list[str]: + if self.retain == "none": + return [] + declared = [ + n for n, p in self.program.parameters.items() if p.derivation is None + ] + if self.retain == "all": + return declared + closure = _report_closure(self.program) + return [n for n in declared if n in closure] + + +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): + 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.Cases): + for region in node.regions: + names |= region.when.names_read + return names & set(program.parameters) + + +# --------------------------------------------------------------------------- +# sources and keys +# --------------------------------------------------------------------------- + + +def _dataset_sources(ds: xr.Dataset) -> dict[str, Any]: + sources: dict[str, Any] = {str(d): index for d, index in ds.indexes.items()} + sources.update({str(n): ds[n] for n in ds.data_vars}) + return sources + + +def _attachable(program: ms.Program) -> dict[str, str]: + kinds = { + n: "parameter" for n, p in program.parameters.items() if p.derivation is None + } + kinds.update({d: "dimension" for d in program.dimensions}) + kinds.update({lk.name: "lookup" for _, lk in program.lookups}) + return kinds + + +def _check_keys(program: ms.Program, keys: frozenset[str]) -> None: + known = _attachable(program) + unknown = sorted(keys - set(known)) + if not unknown: + return + lead = ( + f"source key {unknown[0]!r} names" + if len(unknown) == 1 + else f"source keys {unknown} name" + ) + raise SpecDataError( + f"{lead} neither a parameter, a dimension nor a lookup this spec declares. " + f"{did_you_mean(unknown[0], known)} Pass only what the spec takes." + ) + + +# --------------------------------------------------------------------------- +# dimensions +# --------------------------------------------------------------------------- + + +def _reached(program: ms.Program) -> set[str]: + dims: set[str] = set() + for p in program.parameters.values(): + dims.update(p.dims) + for v in program.variables.values(): + dims.update(v.dims) + for c in program.constraints.values(): + dims.update(c.dims) + for over, lk in program.lookups: + dims.add(over) + if lk.target is not None: + dims.add(lk.target) + for pw in program.piecewise.values(): + dims.add(pw.over) + return dims + + +def _master_coords( + program: ms.Program, sources: Mapping[str, Any], keys: frozenset[str] +) -> dict[str, pd.Index]: + reached = _reached(program) + coords: dict[str, pd.Index] = {} + for dim in program.dimensions: + if dim in keys: + coords[dim] = _index(dim, sources[dim]) + elif dim in reached: + raise SpecDataError( + f"dimension '{dim}' has no index: pass its labels under key '{dim}' as " + f"{_DIMENSION_SHAPES}. The index is what says which labels exist, and without " + f"one a mistyped label is indistinguishable from a new one." + ) + return coords + + +def _index(dim: str, obj: Any) -> pd.Index: + if isinstance(obj, (pd.Series, xr.DataArray, np.ndarray)): + if obj.ndim != 1: + raise SpecDataError( + f"index for dimension '{dim}' is {obj.ndim}-dimensional; pass {_DIMENSION_SHAPES}." + ) + values: Any = np.asarray(obj) + elif isinstance(obj, (pd.Index, list, tuple)): + values = obj + else: + raise SpecDataError( + f"index for dimension '{dim}': cannot read labels out of {type(obj).__name__}; pass {_DIMENSION_SHAPES}." + ) + index = pd.Index(values, name=dim) + if index.has_duplicates: + twice = index[index.duplicated()].unique().tolist() + raise SpecDataError( + f"dimension '{dim}' lists {_shown(twice)} more than once. A dimension's members are a set: " + f"each label appears once, in the order the source gives it." + ) + return index + + +# --------------------------------------------------------------------------- +# lookups +# --------------------------------------------------------------------------- + + +def _lookups( + program: ms.Program, + sources: Mapping[str, Any], + keys: frozenset[str], + coords: Mapping[str, pd.Index], +) -> dict[str, dict[str, xr.DataArray]]: + out: dict[str, dict[str, xr.DataArray]] = {} + for over, lk in program.lookups: + space = lk.target or lk.name + if lk.name not in keys: + raise SpecDataError( + f"no data provided for lookup '{lk.name}'. Pass it under key '{lk.name}' as " + f"{_LOOKUP_SHAPES.format(over=over)}, holding a '{space}' value for each " + f"'{over}' label it maps and nothing for a label it does not." + ) + series = _lookup_series(lk.name, over, sources[lk.name]) + _check_lookup(series, lk, over, coords) + padded = series.reindex(coords[over]) + array = xr.DataArray( + padded.to_numpy(), dims=[over], coords={over: coords[over]}, name=lk.name + ) + out.setdefault(over, {})[lk.name] = array + return out + + +def _lookup_series(name: str, over: str, obj: Any) -> pd.Series: + if isinstance(obj, xr.DataArray): + if obj.dims != (over,) or over not in obj.indexes: + raise SpecDataError( + f"lookup '{name}' arrived as a DataArray over {list(obj.dims)}, and it is a map " + f"out of '{over}': pass a 1-D DataArray with '{over}' as its labelled dimension." + ) + return obj.to_series() + if isinstance(obj, Mapping): + return pd.Series(dict(obj)).rename_axis(over) + if isinstance(obj, pd.Series): + if obj.index.name not in (None, over): + raise SpecDataError( + f"lookup '{name}' is a Series indexed by '{obj.index.name}', and it is a map out of " + f"'{over}': index it by '{over}' labels." + ) + return obj.rename_axis(over) + raise SpecDataError( + f"lookup '{name}': cannot adapt {type(obj).__name__} to a map; pass {_LOOKUP_SHAPES.format(over=over)}." + ) + + +def _check_lookup( + series: pd.Series, + lk: ms.LookupDeclaration, + over: str, + coords: Mapping[str, pd.Index], +) -> None: + space = lk.target or lk.name + holes = series.isna() + if holes.any(): + at = _coordinates_shown((over,), series.index[holes][:5]) + raise SpecDataError( + f"lookup '{lk.name}' carries {int(holes.sum())} row(s) with a null in '{space}': {at}. A map is " + f"partial by leaving a label out, not by mapping it to nothing: drop the row and the " + f"label is unmapped, which is what every operator reading the lookup already means by it." + ) + if series.index.has_duplicates: + twice = series.index[series.index.duplicated()].unique().tolist() + raise SpecDataError( + f"lookup '{lk.name}' maps {len(twice)} '{over}' label(s) more than once: {_shown(twice)}. " + f"A lookup is single-valued, so each label it maps takes exactly one row." + ) + strays = series.index[~series.index.isin(coords[over])].tolist() + if strays: + raise SpecDataError( + f"lookup '{lk.name}' maps {_shown(strays)}, which are not labels of '{over}'. " + f"'{over}' takes its labels from sources['{over}'], and they are " + f"{_shown(coords[over].tolist(), 8)}. A map maps the labels that exist: a key matching " + f"none of them would place its terms nowhere, so it is a typo on one side or a label " + f"missing from the other." + ) + if lk.target is None: + return + values = pd.Index(series.to_numpy()) + foreign = values[~values.isin(coords[lk.target])].unique().tolist() + if foreign: + raise SpecDataError( + f"dimension '{over}' lookup '{lk.name}' has value(s) that are not '{lk.target}' labels: " + f"{_shown(foreign)}. Every value must be a declared '{lk.target}' label, otherwise " + f"sum(by={lk.name}) drops those terms in the join that places them, and the model " + f"builds and solves without them." + ) + + +# --------------------------------------------------------------------------- +# parameters +# --------------------------------------------------------------------------- + + +def _fill(declared: ms.ParameterDeclaration) -> Any: + return False if declared.dtype == "bool" else np.nan + + +def _as_array( + name: str, + declared: ms.ParameterDeclaration, + obj: Any, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + dims = declared.dims + if isinstance(obj, xr.DataArray): + return _from_dense(name, declared, obj) + if isinstance(obj, pd.DataFrame): + return _from_frame(name, declared, obj, coords) + if isinstance(obj, pd.Series): + return _from_rows(name, declared, obj, coords) + if isinstance(obj, Mapping): + return _from_rows(name, declared, pd.Series(dict(obj)), coords) + if isinstance(obj, _SCALARS): + return _from_scalar(name, declared, obj, coords) + raise SpecDataError( + f"parameter '{name}': cannot adapt {type(obj).__name__} to an array over {list(dims)}; " + f"pass {_parameter_shapes(dims)}." + ) + + +def _parameter_shapes(dims: Sequence[str]) -> str: + return _PARAMETER_SHAPES.format(dims=list(dims), columns=[*dims, "value"]) + + +def _from_scalar( + name: str, + declared: ms.ParameterDeclaration, + obj: Any, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + if pd.isna(obj): + raise SpecDataError( + f"parameter '{name}' is one value and that value is a hole (null or NaN). " + f"A number was meant, or the parameter has no data and should not be passed." + ) + value = ( + np.asarray(obj, dtype=float) if declared.dtype == "float" else np.asarray(obj) + ) + _check_value_dtype(name, declared, value.dtype) + arr = xr.DataArray(value, name=name) + if declared.dims: + arr = arr.expand_dims({d: coords[d] for d in declared.dims}) + return arr + + +def _from_dense( + name: str, declared: ms.ParameterDeclaration, arr: xr.DataArray +) -> xr.DataArray: + dims = declared.dims + _check_value_dtype(name, declared, arr.dtype) + if set(arr.dims) != set(dims) or len(arr.dims) != len(dims): + raise SpecDataError( + f"parameter '{name}' arrived as a DataArray over {list(arr.dims)}, and '{name}' is over " + f"{list(dims)}. The dims must be the declared ones, in any order." + ) + for d in dims: + if d not in arr.indexes: + raise SpecDataError( + f"parameter '{name}' has no coordinate labels along '{d}'. A parameter is read for " + f"values against its labels, so every dimension needs an index coordinate." + ) + _refuse_duplicate_coordinates( + name, dims, arr.indexes[d].duplicated(), arr.indexes[d] + ) + return arr.transpose(*dims) + + +def _from_frame( + name: str, + declared: ms.ParameterDeclaration, + df: pd.DataFrame, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + dims = declared.dims + tidy = df.reset_index() if set(dims) - set(df.columns) else df + if "value" in tidy.columns and set(dims) <= set(tidy.columns): + if not dims: + return _from_rows(name, declared, tidy["value"], coords) + return _from_rows(name, declared, tidy.set_index(list(dims))["value"], coords) + if len(dims) == 2: + return _from_dense(name, declared, xr.DataArray(_wide(name, dims, df))) + raise SpecDataError( + f"parameter '{name}' arrived as a DataFrame with columns {list(df.columns)}; a table for " + f"'{name}' carries columns {[*dims, 'value']}." + ) + + +def _wide(name: str, dims: tuple[str, ...], df: pd.DataFrame) -> pd.DataFrame: + names = (df.index.name, df.columns.name) + if names == (None, None): + return df.rename_axis(index=dims[0], columns=dims[1]) + if set(names) == set(dims): + return df + raise SpecDataError( + f"parameter '{name}' arrived as a wide DataFrame with index '{names[0]}' and columns " + f"'{names[1]}', and '{name}' is over {list(dims)}. Name the index and columns after the " + f"two dims, or pass a table with columns {[*dims, 'value']}." + ) + + +def _from_rows( + name: str, + declared: ms.ParameterDeclaration, + series: pd.Series, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + dims = declared.dims + if not dims: + if len(series) != 1: + raise SpecDataError( + f"parameter '{name}' is declared with no dims, which means one value broadcast " + f"everywhere, but its source has {len(series)} rows. Declare the dims it is indexed " + f"by, or pass one number." + ) + return _from_scalar(name, declared, series.iloc[0], coords) + series = _with_dims(name, dims, series) + holes = series.isna() + if holes.any(): + raise SpecDataError( + f"parameter '{name}' carries {int(holes.sum())} row(s) with no value, null or NaN: " + f"{_coordinates_shown(dims, series.index[holes][:3])}. In a table the absence of a " + f"value is the absence of the row, and such a row says the coordinate exists and denies " + f"it in the same breath. Drop those rows, or supply the values." + ) + _check_value_dtype(name, declared, series.dtype) + _refuse_duplicate_coordinates(name, dims, series.index.duplicated(), series.index) + for d in dims: + labels = series.index.get_level_values(d) + _refuse_strangers(name, d, labels[~labels.isin(coords[d])].unique(), coords[d]) + onto = [coords[d] for d in dims] + full = onto[0] if len(dims) == 1 else pd.MultiIndex.from_product(onto, names=dims) + values = series.reindex(full, fill_value=_fill(declared)).to_numpy() + shape = tuple(len(index) for index in onto) + return xr.DataArray( + values.reshape(shape), dims=dims, coords=dict(zip(dims, onto)), name=name + ) + + +def _with_dims(name: str, dims: tuple[str, ...], series: pd.Series) -> pd.Series: + index = series.index + if index.nlevels != len(dims): + said = "a Series or dict carries one label per level" + raise SpecDataError( + f"parameter '{name}': {said}, and its index has {index.nlevels} level(s) where '{name}' " + f"is over {list(dims)}. Pass {_parameter_shapes(dims)}." + ) + names = list(index.names) + if all(n is None for n in names): + return series.set_axis(index.set_names(list(dims))) + if set(names) != set(dims): + raise SpecDataError( + f"parameter '{name}' is indexed by {names}, and '{name}' is over {list(dims)}. " + f"Name the index levels after the declared dims." + ) + if tuple(names) != dims: + series = series.reorder_levels(list(dims)) + return series + + +def _refuse_duplicate_coordinates( + name: str, dims: tuple[str, ...], duplicated: Any, index: pd.Index +) -> None: + if not duplicated.any(): + return + counts = index[duplicated].value_counts() + shown = "; ".join( + f"{_coordinate(dims, key)} ({n + 1} rows)" for key, n in counts.iloc[:3].items() + ) + raise SpecDataError( + f"parameter '{name}' has more than one row for a coordinate: {shown}. A parameter is a " + f"function of its dims, so which value applies is undefined; aggregate the source to one " + f"row per {list(dims)} before attaching it." + ) + + +def _refuse_strangers( + name: str, dim: str, strangers: pd.Index, labels: pd.Index +) -> None: + if len(strangers) == 0: + return + raise SpecDataError( + f"parameter '{name}' has label(s) in dimension '{dim}' that are not coordinates of it: " + f"{_shown(strangers.tolist())}.\n {dim} has: {_shown(labels.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}']." + ) + + +def _aligned( + name: str, arr: xr.DataArray, onto: Mapping[str, pd.Index], fill: Any +) -> xr.DataArray: + if all(arr.indexes[d].equals(index) for d, index in onto.items()): + return arr + for d, index in onto.items(): + _refuse_strangers(name, d, arr.indexes[d].difference(index), index) + return arr.reindex(onto, fill_value=fill) + + +def _check_value_dtype( + name: str, declared: ms.ParameterDeclaration, dtype: Any +) -> None: + kind = str(dtype.kind) + if kind in _ACCEPTED_KINDS[declared.dtype]: + return + arrived = _KIND_NAMES.get(kind, str(dtype)) + raise SpecDataError( + f"parameter '{name}' is declared '{declared.dtype}' and its values arrived as '{arrived}'. " + f"A declared dtype is a claim about the values, and it is checked here: the file says what " + f"the values are, or the values are not attached.\n" + f" Cast the values to {declared.dtype}, if the declaration is what you meant\n" + f" Or declare what the data has: {{dtype: {arrived}}}" + ) + + +# --------------------------------------------------------------------------- +# wording +# --------------------------------------------------------------------------- + + +def _shown(labels: Sequence[Any], limit: int = 5) -> str: + head = ", ".join(repr(x) for x in labels[:limit]) + return head + (f" (and {len(labels) - limit} more)" if len(labels) > limit else "") + + +def _coordinate(dims: Sequence[str], key: Hashable) -> str: + row = key if isinstance(key, tuple) else (key,) + return ", ".join(f"{d}={v!r}" for d, v in zip(dims, row)) + + +def _coordinates_shown(dims: Sequence[str], rows: Iterable[Hashable]) -> str: + return "; ".join(_coordinate(dims, row) for row in rows) diff --git a/linopy/spec/errors.py b/linopy/spec/errors.py new file mode 100644 index 00000000..7e075add --- /dev/null +++ b/linopy/spec/errors.py @@ -0,0 +1,12 @@ +"""Errors raised while binding data to a math-spec program.""" + +from __future__ import annotations + + +class SpecDataError(ValueError): + """ + Data bound to a valid spec is missing, malformed or the wrong shape. + + Every refusal names the symbol, the dimension(s) and the offending labels, + so the message points back at the ``sources`` entry to fix. + """ diff --git a/test/test_spec_binder.py b/test/test_spec_binder.py new file mode 100644 index 00000000..784d7571 --- /dev/null +++ b/test/test_spec_binder.py @@ -0,0 +1,772 @@ +"""Binding user data to a math-spec program.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +math_spec = pytest.importorskip("math_spec") + +from linopy.spec import SpecDataError, bind # noqa: E402 + +SPEC = { + "dimensions": {"f": {"dtype": "str"}, "t": {"dtype": "int"}, "g": {"dtype": "str"}}, + "lookups": {"grp": {"over": "f", "into": "g"}}, + "parameters": { + "cost": {"dims": ["f"]}, + "cap": {"dims": ["f", "t"]}, + "flag": {"dims": ["f"], "dtype": "bool"}, + "rate": {"dims": []}, + "lead": {"dims": ["f"], "dtype": "int"}, + }, + "variables": { + "x": { + "foreach": ["f", "t"], + "where": "flag", + "bounds": {"lower": 0, "upper": "cap"}, + } + }, + "constraints": { + "k": {"foreach": ["g", "t"], "expression": "sum(x, by=grp) <= 10"}, + "s": { + "foreach": ["f", "t"], + "expression": "shift(x, over=t, offset=lead, edge=0) >= 0", + }, + }, + "objective": {"sense": "maximize", "expression": "sum(x * cost)"}, + "expressions": { + "spend": "sum(x * cost, over=t)", + "total": "sum(spend, over=f) * rate", + }, +} + +F = pd.Index(["b", "a", "c"], name="f") +T = pd.Index([0, 1, 2], name="t") +CAP = xr.DataArray(np.arange(9.0).reshape(3, 3), coords={"f": F, "t": T}, name="cap") +COST = pd.Series([1.0, 2.0, 3.0], index=F) + + +@pytest.fixture(scope="module") +def program() -> Any: + return math_spec.to_program(SPEC) + + +@pytest.fixture +def good() -> dict[str, Any]: + return { + "f": list(F), + "t": list(T), + "g": ["n", "e"], + "cost": COST, + "cap": CAP, + "flag": pd.Series([True, False, True], index=F), + "rate": 0.5, + "lead": pd.Series([1, 0, 1], index=F), + "grp": pd.Series(["n", "e", "n"], index=F), + } + + +def read_all(program: Any, sources: Mapping[str, Any]) -> list[xr.DataArray]: + bound = bind(program, sources) + return [bound.parameter(name) for name in program.parameters] + + +CAP_SHAPES = { + "dataarray": CAP, + "dataarray-transposed": CAP.transpose("t", "f"), + "series": CAP.to_series(), + "series-transposed": CAP.transpose("t", "f").to_series(), + "series-unnamed-levels": CAP.to_series().rename_axis([None, None]), + "tidy-frame": CAP.to_series().reset_index(name="value"), + "tidy-frame-extra-column": CAP.to_series() + .reset_index(name="value") + .assign(note="x"), + "wide-frame": CAP.to_pandas(), + "wide-frame-transposed": CAP.to_pandas().T, + "wide-frame-unnamed": CAP.to_pandas().rename_axis(index=None, columns=None), + "dict": CAP.to_series().to_dict(), +} + + +@pytest.mark.parametrize("cap", CAP_SHAPES.values(), ids=CAP_SHAPES.keys()) +def test_rank_two_shapes_bind_alike( + program: Any, good: dict[str, Any], cap: Any +) -> None: + got = bind(program, {**good, "cap": cap}).parameter("cap") + xr.testing.assert_equal(got, CAP) + assert got.dims == ("f", "t") + + +COST_SHAPES = { + "series": COST, + "series-unnamed": COST.rename_axis(None), + "dataarray": xr.DataArray(COST), + "dict": COST.to_dict(), + "tidy-frame": COST.reset_index(name="value"), +} + + +@pytest.mark.parametrize("cost", COST_SHAPES.values(), ids=COST_SHAPES.keys()) +def test_rank_one_shapes_bind_alike( + program: Any, good: dict[str, Any], cost: Any +) -> None: + got = bind(program, {**good, "cost": cost}).parameter("cost") + xr.testing.assert_equal(got, xr.DataArray(COST, name="cost")) + + +DIMENSION_SHAPES = { + "index": F, + "list": list(F), + "tuple": tuple(F), + "ndarray": F.to_numpy(), + "series": pd.Series(F), + "dataarray": xr.DataArray(list(F), dims=["f"]), +} + + +@pytest.mark.parametrize("f", DIMENSION_SHAPES.values(), ids=DIMENSION_SHAPES.keys()) +def test_dimension_shapes_keep_source_order( + program: Any, good: dict[str, Any], f: Any +) -> None: + coords = bind(program, {**good, "f": f}).coords + assert coords["f"].tolist() == ["b", "a", "c"] + assert coords["f"].name == "f" + assert list(coords) == ["f", "t", "g"] + + +def test_lookup_is_padded_onto_the_dimension( + program: Any, good: dict[str, Any] +) -> None: + bound = bind(program, {**good, "grp": {"a": "n"}}) + grp = bound.lookups["f"]["grp"] + assert grp.dims == ("f",) + assert grp.sel(f="a").item() == "n" + assert pd.isna(grp.sel(f=["b", "c"])).all() + + +LOOKUP_SHAPES = { + "series": pd.Series(["n", "e", "n"], index=F), + "series-unnamed": pd.Series(["n", "e", "n"], index=F.rename(None)), + "dict": {"b": "n", "a": "e", "c": "n"}, + "dataarray": xr.DataArray(["n", "e", "n"], coords={"f": F}), +} + + +@pytest.mark.parametrize("grp", LOOKUP_SHAPES.values(), ids=LOOKUP_SHAPES.keys()) +def test_lookup_shapes_bind_alike(program: Any, good: dict[str, Any], grp: Any) -> None: + got = bind(program, {**good, "grp": grp}).lookups["f"]["grp"] + assert got.values.tolist() == ["n", "e", "n"] + + +def test_missing_rows_become_nan_and_false(program: Any, good: dict[str, Any]) -> None: + sparse = { + **good, + "cost": pd.Series({"a": 1.0}), + "flag": pd.Series({"a": True}), + "cap": CAP.sel(t=[0, 1]), + } + bound = bind(program, sparse) + cost = bound.parameter("cost") + assert cost.sel(f="a").item() == 1.0 + assert cost.sel(f=["b", "c"]).isnull().all() + flag = bound.parameter("flag") + assert flag.dtype == bool + assert flag.values.tolist() == [False, True, False] + cap = bound.parameter("cap") + assert cap.dims == ("f", "t") + assert cap.sel(t=2).isnull().all() + + +@pytest.mark.parametrize( + ("name", "value", "expected_dtype"), + [ + ("cost", 2, float), + ("cap", 1.5, float), + ("flag", True, bool), + ("lead", 3, np.int64), + ], +) +def test_scalar_is_broadcast_over_declared_dims( + program: Any, good: dict[str, Any], name: str, value: Any, expected_dtype: Any +) -> None: + got = bind(program, {**good, name: value}).parameter(name) + assert got.dims == tuple(SPEC["parameters"][name]["dims"]) + assert got.dtype == expected_dtype + assert (got == value).all() + + +def test_scalar_parameter_stays_scalar(program: Any, good: dict[str, Any]) -> None: + got = bind(program, good).parameter("rate") + assert got.dims == () + assert got.item() == 0.5 + + +def test_missing_parameter_is_refused_when_read( + program: Any, good: dict[str, Any] +) -> None: + good.pop("cost") + bound = bind(program, good) + with pytest.raises(SpecDataError, match="no data provided for parameter 'cost'"): + bound.parameter("cost") + + +REFUSALS = [ + pytest.param( + {"f": ["a", "a", "b"]}, + r"dimension 'f' lists 'a' more than once", + id="duplicate-member", + ), + pytest.param( + {"cost": pd.Series({"a": 1.0, "zz": 2.0})}, + r"parameter 'cost'.*'f'.*'zz'", + id="unknown-label", + ), + pytest.param( + {"cap": CAP.assign_coords(t=[0, 1, 9])}, + r"parameter 'cap'.*'t'.*\b9\b", + id="unknown-label-dense", + ), + pytest.param( + {"cost": pd.Series([1.0, 9.0, 2.0], index=pd.Index(["a", "a", "b"], name="f"))}, + r"parameter 'cost' has more than one row for a coordinate: f='a' \(2 rows\)", + id="duplicated-coordinate-row", + ), + pytest.param( + {"cap": xr.DataArray([1.0, 2.0], coords={"f": ["a", "a"]})}, + r"parameter 'cap' arrived as a DataArray over \['f'\]", + id="dense-wrong-dims", + ), + pytest.param( + { + "cap": xr.DataArray( + np.ones((2, 3)), coords={"f": ["a", "a"], "t": [0, 1, 2]} + ) + }, + r"parameter 'cap' has more than one row", + id="dense-duplicate-coordinate", + ), + pytest.param( + {"cap": COST}, + r"parameter 'cap'.*1 level\(s\) where 'cap' is over \['f', 't'\]", + id="wrong-rank", + ), + pytest.param( + { + "cap": pd.Series( + [1.0], index=pd.MultiIndex.from_tuples([("a", 0)], names=["f", "q"]) + ) + }, + r"parameter 'cap' is indexed by \['f', 'q'\]", + id="wrong-level-names", + ), + pytest.param( + {"rate": COST}, + r"parameter 'rate' is declared with no dims.*3 rows", + id="rows-for-scalar", + ), + pytest.param( + {"cost": {"a", "b"}}, + r"parameter 'cost': cannot adapt set", + id="unsupported-shape", + ), + pytest.param( + {"cap": xr.DataArray(np.ones((3, 3)), dims=["f", "t"])}, + r"parameter 'cap' has no coordinate labels along 'f'", + id="dense-without-labels", + ), + pytest.param( + {"cost": pd.Series({"a": 1.0, "b": None})}, + r"parameter 'cost' carries 1 row.*f='b'", + id="null-row", + ), + pytest.param( + {"cost": pd.Series({"a": 1.0, "b": np.nan})}, + r"parameter 'cost' carries 1 row", + id="nan-row", + ), + pytest.param( + {"rate": float("nan")}, + r"parameter 'rate' is one value and that value is a hole", + id="nan-scalar", + ), + pytest.param( + {"lead": pd.Series([1.5, 0.0, 1.0], index=F)}, + r"'lead' is declared 'int'.*'float'", + id="float-for-int", + ), + pytest.param( + {"flag": pd.Series([1, 0, 1], index=F)}, + r"'flag' is declared 'bool'.*'int'", + id="int-for-bool", + ), + pytest.param( + {"flag": 1.0}, r"'flag' is declared 'bool'.*'float'", id="float-scalar-for-bool" + ), + pytest.param( + {"cost": pd.Series(["x", "y", "z"], index=F)}, + r"'cost' is declared 'float'.*'str'", + id="str-for-float", + ), + pytest.param( + {"csot": COST}, r"source key 'csot'.*Did you mean 'cost'", id="unknown-key" + ), + pytest.param({"f": None}, r"dimension 'f' has no index", id="missing-dimension"), + pytest.param( + {"f": {"a": 1}}, + r"index for dimension 'f': cannot read labels out of dict", + id="dimension-shape", + ), + pytest.param( + {"grp": None}, r"no data provided for lookup 'grp'", id="missing-lookup" + ), + pytest.param( + {"grp": {"zz": "n"}}, + r"lookup 'grp' maps 'zz', which are not labels of 'f'", + id="lookup-stray-key", + ), + pytest.param( + {"grp": {"a": "zz"}}, + r"lookup 'grp' has value\(s\) that are not 'g' labels: 'zz'", + id="lookup-stray-value", + ), + pytest.param( + {"grp": pd.Series(["n", "e"], index=pd.Index(["a", "a"], name="f"))}, + r"lookup 'grp' maps 1 'f' label\(s\) more than once: 'a'", + id="lookup-two-values", + ), + pytest.param( + {"grp": {"a": None, "b": "n"}}, + r"lookup 'grp' carries 1 row\(s\) with a null in 'g': f='a'", + id="lookup-null", + ), + pytest.param( + {"grp": pd.DataFrame({"f": ["a"], "g": ["n"]})}, + r"lookup 'grp': cannot adapt DataFrame", + id="lookup-shape", + ), + pytest.param( + {"grp": pd.Series(["n"], index=pd.Index(["a"], name="t"))}, + r"lookup 'grp' is a Series indexed by 't'", + id="lookup-wrong-index", + ), +] + + +@pytest.mark.parametrize(("override", "match"), REFUSALS) +def test_malformed_data_is_refused_naming_the_symbol( + program: Any, good: dict[str, Any], override: dict[str, Any], match: str +) -> None: + sources = {**good, **override} + for key, value in override.items(): + if value is None: + sources.pop(key) + with pytest.raises(SpecDataError, match=match): + read_all(program, sources) + + +def test_int_labels_are_shown_as_written(program: Any, good: dict[str, Any]) -> None: + with pytest.raises(SpecDataError, match=r"\b99\b") as error: + read_all(program, {**good, "cap": CAP.assign_coords(t=[0, 1, 99])}) + assert "int64" not in str(error.value) + + +def test_dataset_is_a_source(program: Any, good: dict[str, Any]) -> None: + ds = xr.Dataset( + { + "cost": xr.DataArray(COST), + "cap": CAP, + "flag": xr.DataArray(good["flag"]), + "rate": 0.5, + "lead": xr.DataArray(good["lead"]), + "grp": xr.DataArray(good["grp"]), + }, + coords={"f": F, "t": T, "g": ["n", "e"]}, + ) + from_dataset = bind(program, ds) + from_mapping = bind(program, good) + assert from_dataset.coords["f"].equals(from_mapping.coords["f"]) + for name in program.parameters: + xr.testing.assert_equal( + from_dataset.parameter(name), from_mapping.parameter(name) + ) + xr.testing.assert_equal( + from_dataset.lookups["f"]["grp"], from_mapping.lookups["f"]["grp"] + ) + + +class Counting(Mapping[str, Any]): + def __init__(self, data: dict[str, Any]) -> None: + self.data = data + self.pulled: list[str] = [] + + def __getitem__(self, key: str) -> Any: + self.pulled.append(key) + return self.data[key] + + def __iter__(self) -> Iterator[str]: + raise AssertionError("sources must not be iterated") + + def __len__(self) -> int: + return len(self.data) + + def keys(self) -> Any: + return self.data.keys() + + +def test_sources_are_pulled_by_key_on_demand( + program: Any, good: dict[str, Any] +) -> None: + sources = Counting(good) + bound = bind(program, sources) + assert set(sources.pulled) == {"f", "t", "g", "grp"} + bound.parameter("cap") + bound.parameter("cap") + assert sources.pulled.count("cap") == 2 + + +@pytest.mark.parametrize( + ("retain", "expected"), + [ + ("report", {"cost", "rate", "grp"}), + ("all", {"cost", "cap", "flag", "rate", "lead", "grp"}), + ("none", {"grp"}), + ], +) +def test_retained_follows_the_named_expressions( + program: Any, good: dict[str, Any], retain: Any, expected: set[str] +) -> None: + retained = bind(program, good, retain=retain).retained() + assert set(retained.data_vars) == expected + assert retained.coords["f"].values.tolist() == ["b", "a", "c"] + + +def test_report_closure_reads_names_and_masks() -> None: + spec = { + "dimensions": {"f": {"dtype": "str"}, "t": {"dtype": "int"}}, + "parameters": { + "cost": {"dims": ["f"]}, + "lag": {"dims": ["f"], "dtype": "int"}, + "on": {"dims": ["f"], "dtype": "bool"}, + "other": {"dims": ["f"]}, + }, + "variables": {"x": {"foreach": ["f", "t"], "bounds": {"lower": 0, "upper": 1}}}, + "objective": {"sense": "maximize", "expression": "sum(x * other)"}, + "expressions": { + "late": { + "foreach": ["f", "t"], + "cases": { + "active": { + "when": "on", + "expression": "shift(x, over=t, offset=lag, edge=0)", + } + }, + "otherwise": "x * cost", + } + }, + } + program = math_spec.to_program(spec) + f = pd.Index(["a"], name="f") + sources = { + "f": f, + "t": [0, 1], + "cost": pd.Series([1.0], index=f), + "lag": pd.Series([1], index=f), + "on": pd.Series([True], index=f), + "other": pd.Series([2.0], index=f), + } + assert set(bind(program, sources).retained().data_vars) == {"cost", "lag", "on"} + + +def test_aligned_array_is_not_copied(program: Any, good: dict[str, Any]) -> None: + bound = bind(program, good) + assert np.shares_memory(CAP.values, bound.parameter("cap").values) + assert np.shares_memory(CAP.values, bound.parameter("cap").values) + permuted = CAP.transpose("t", "f") + assert np.shares_memory( + permuted.values, + bind(program, {**good, "cap": permuted}).parameter("cap").values, + ) + wide = CAP.to_pandas() + assert np.shares_memory( + wide.values, bind(program, {**good, "cap": wide}).parameter("cap").values + ) + + +def test_derived_parameter_is_not_bound_from_sources() -> None: + spec = { + "dimensions": {"bp": {"dtype": "int"}}, + "parameters": {"bp_x": {"dims": ["bp"]}, "bp_y": {"dims": ["bp"]}}, + "variables": { + "x": {"foreach": [], "bounds": {"lower": 0, "upper": 10}}, + "y": {"foreach": []}, + }, + "piecewise": { + "curve": { + "over": "bp", + "method": "lp", + "points": "bp_x", + "links": [["x", "bp_x"], ["y", "bp_y", ">="]], + } + }, + "objective": {"sense": "minimize", "expression": "y"}, + } + program = math_spec.to_program(spec) + derived = [n for n, p in program.parameters.items() if p.derivation is not None] + assert derived + bp = pd.Index([0, 1, 2], name="bp") + sources = { + "bp": bp, + "bp_x": pd.Series([0.0, 5.0, 10.0], index=bp), + "bp_y": pd.Series([0.0, 2.0, 8.0], index=bp), + } + bound = bind(program, sources, retain="all") + assert set(bound.retained().data_vars) == {"bp_x", "bp_y"} + with pytest.raises(SpecDataError, match="emitted by piecewise block 'curve'"): + bound.parameter(derived[0]) + with pytest.raises(SpecDataError, match=derived[0]): + bind(program, {**sources, derived[0]: 1.0}) + + +# --------------------------------------------------------------------------- +# lpspec data-parity cases, eager representation +# --------------------------------------------------------------------------- + +PARITY_SPEC = { + "dimensions": {"f": {"dtype": "str"}}, + "parameters": {"cost": {"dims": ["f"]}, "cap": {"dims": ["f"]}}, + "variables": {"x": {"foreach": ["f"], "bounds": {"lower": 0, "upper": "cap"}}}, + "constraints": {"k": {"foreach": ["f"], "expression": "x <= cap"}}, + "objective": {"sense": "maximize", "expression": "sum(x * cost)"}, +} + +GOOD = { + "f": ["a", "b"], + "cost": pd.Series({"a": 1.0, "b": 2.0}), + "cap": pd.Series({"a": 5.0, "b": 5.0}), +} +ACCEPTED = "accepted" + +PARITY_CASES = [ + pytest.param(GOOD, ACCEPTED, id="valid"), + pytest.param( + {"f": ["a", "b"], "cost": GOOD["cost"]}, + SpecDataError, + id="parameter-missing-entirely", + ), + pytest.param( + {**GOOD, "cost": pd.Series({"a": 1.0})}, ACCEPTED, id="coefficient-sparse" + ), + pytest.param( + { + **GOOD, + "cost": pd.Series( + [1.0, 9.0, 2.0], index=pd.Index(["a", "a", "b"], name="f") + ), + }, + SpecDataError, + id="duplicated-coordinate-row", + ), + pytest.param( + {**GOOD, "cost": pd.Series({"a": 1.0, "zz": 2.0})}, + SpecDataError, + id="label-the-dimension-does-not-have", + ), + pytest.param( + {**GOOD, "cost": pd.Series({"a": 1.0, "b": None})}, + SpecDataError, + id="a-null-value", + ), + pytest.param( + {**GOOD, "cost": pd.Series({"a": 1.0, "b": float("nan")})}, + SpecDataError, + id="a-nan-value", + ), + pytest.param( + {**GOOD, "cap": pd.Series({"a": 5.0, "b": None})}, + SpecDataError, + id="a-hole-in-a-bound", + ), + pytest.param( + {**GOOD, "cost": float("nan")}, SpecDataError, id="a-hole-as-a-scalar" + ), + pytest.param( + {**GOOD, "cost": [1.0, float("nan")]}, SpecDataError, id="a-hole-in-a-sequence" + ), + pytest.param( + {**GOOD, "cost": {"a": 1.0, "b": None}}, SpecDataError, id="a-hole-in-a-dict" + ), + pytest.param( + {**GOOD, "cost": pd.DataFrame({"f": ["a", "b"], "value": [1.0, None]})}, + SpecDataError, + id="a-hole-in-a-tidy-frame", + ), + pytest.param( + {**GOOD, "cost": pd.Series({"a": 1, "b": 2})}, + ACCEPTED, + id="whole-numbers-serve-a-float-declaration", + ), + pytest.param( + {**GOOD, "csot": GOOD["cost"]}, + SpecDataError, + id="a-source-key-the-model-does-not-declare", + ), + pytest.param( + { + **GOOD, + "cost": pd.Series( + [5.0, 5.0], + index=pd.MultiIndex.from_tuples([("a", 0), ("b", 0)], names=["f", "k"]), + ), + }, + SpecDataError, + id="a-series-deeper-than-the-declared-dims", + ), +] + + +@pytest.mark.parametrize(("sources", "verdict"), PARITY_CASES) +def test_parity_with_lpspec_data_verdicts( + sources: dict[str, Any], verdict: Any +) -> None: + program = math_spec.to_program(PARITY_SPEC) + if verdict is ACCEPTED: + read_all(program, sources) + return + with pytest.raises(verdict): + read_all(program, sources) + + +def test_a_hole_is_named_where_it_sits() -> None: + program = math_spec.to_program(PARITY_SPEC) + with pytest.raises(SpecDataError, match="parameter 'cost'") as error: + read_all(program, {**GOOD, "cost": pd.Series({"a": 1.0, "b": None})}) + assert "divisor" not in str(error.value) + assert "f='b'" in str(error.value) + + +def test_a_hole_in_a_scalar_parameter_is_refused() -> None: + spec = { + "dimensions": {"f": {"dtype": "str"}}, + "parameters": {"rate": {"dims": []}}, + "variables": {"x": {"foreach": ["f"], "bounds": {"lower": 0, "upper": 1}}}, + "objective": {"sense": "maximize", "expression": "sum(x * rate)"}, + } + program = math_spec.to_program(spec) + with pytest.raises(SpecDataError, match="hole"): + read_all(program, {"f": ["a", "b"], "rate": pd.DataFrame({"value": [None]})}) + + +@pytest.mark.parametrize( + ("column", "verdict"), + [ + pytest.param( + pd.Series({"a": True, "b": False}), ACCEPTED, id="a-boolean-column" + ), + pytest.param(pd.Series({"a": 1, "b": 0}), SpecDataError, id="a-1-0-int-column"), + pytest.param( + pd.Series({"a": 1.0, "b": 0.0}), SpecDataError, id="a-1-0-float-column" + ), + ], +) +def test_a_flag_binds_by_its_declaration(column: pd.Series, verdict: Any) -> None: + spec = { + "dimensions": {"g": {"dtype": "str"}}, + "parameters": {"active": {"dims": ["g"], "dtype": "bool"}}, + "variables": { + "x": { + "foreach": ["g"], + "where": "active", + "bounds": {"lower": 0, "upper": 1}, + } + }, + "objective": {"sense": "maximize", "expression": "sum(x)"}, + } + program = math_spec.to_program(spec) + sources = {"g": ["a", "b"], "active": column} + if verdict is ACCEPTED: + assert bind(program, sources).parameter("active").dtype == bool + return + with pytest.raises(SpecDataError, match="declared 'bool'"): + read_all(program, sources) + + +LOOKUP_SPEC = { + "dimensions": {"g": {}, "b": {"dtype": "str"}}, + "lookups": {"gen_bus": {"over": "g", "into": "b"}}, + "parameters": {"p_max": {"dims": ["g"]}}, + "variables": {"x": {"foreach": ["g"], "bounds": {"lower": 0, "upper": "p_max"}}}, + "constraints": {"k": {"foreach": ["b"], "expression": "sum(x, by=gen_bus) <= 10"}}, + "objective": {"sense": "maximize", "expression": "sum(x)"}, +} +P_MAX = {"p_max": pd.Series({"w": 5.0, "s": 5.0})} +INDEX = {"g": ["w", "s"], "b": ["n", "e"]} +MAP = {"gen_bus": pd.Series({"w": "n", "s": "e"})} + + +@pytest.mark.parametrize( + ("sources", "match"), + [ + pytest.param( + {**P_MAX, **MAP}, "dimension 'g' has no index", id="a-map-and-no-labels" + ), + pytest.param( + {**P_MAX, **INDEX}, "no data provided for lookup", id="an-index-and-no-map" + ), + pytest.param( + {**P_MAX, **INDEX, "gen_bus": pd.Series({"w": "n", "s": "zz"})}, + "not 'b' labels", + id="a-stray-value", + ), + pytest.param( + { + **P_MAX, + **INDEX, + "gen_bus": pd.Series( + ["n", "e", "e"], index=pd.Index(["w", "w", "s"], name="g") + ), + }, + "more than once", + id="two-values-for-one-label", + ), + pytest.param( + {**P_MAX, **INDEX, "gen_bus": pd.Series({"w": None, "s": "e"})}, + "null in 'b'", + id="mapping-a-label-to-nothing", + ), + pytest.param( + { + **P_MAX, + **INDEX, + "gen_bus": pd.Series( + [None, "n", "n"], index=pd.Index(["w", "w", "s"], name="g") + ), + }, + "null in 'b'", + id="a-label-held-twice-with-a-null", + ), + ], +) +def test_a_lookup_defect_is_refused(sources: dict[str, Any], match: str) -> None: + program = math_spec.to_program(LOOKUP_SPEC) + with pytest.raises(SpecDataError, match=match): + read_all(program, sources) + + +def test_a_stray_lookup_value_over_an_int_target_is_shown_as_written() -> None: + program = math_spec.to_program( + {**LOOKUP_SPEC, "dimensions": {"g": {}, "b": {"dtype": "int"}}} + ) + sources = { + **P_MAX, + "g": ["w", "s"], + "b": [1, 2], + "gen_bus": pd.Series({"w": 1, "s": 99}), + } + with pytest.raises(SpecDataError, match=r"not 'b' labels: 99\b") as error: + read_all(program, sources) + assert "int64" not in str(error.value) From a4e6dc303a74d410ea446ab03f8eea51f26e2d72 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 13:24:49 +0200 Subject: [PATCH 02/35] refactor(spec): compact binder and tests; report duplicate coordinates under the right dimension --- linopy/spec/binder.py | 67 ++++------ test/test_spec_binder.py | 276 +++++++++++++-------------------------- 2 files changed, 114 insertions(+), 229 deletions(-) diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py index 008e915f..0cf91f7d 100644 --- a/linopy/spec/binder.py +++ b/linopy/spec/binder.py @@ -151,13 +151,10 @@ def _declaration(self, name: str) -> ms.ParameterDeclaration: def _retained_names(self) -> list[str]: if self.retain == "none": return [] - declared = [ - n for n, p in self.program.parameters.items() if p.derivation is None - ] - if self.retain == "all": - return declared - closure = _report_closure(self.program) - return [n for n in declared if n in closure] + all_names = set(self.program.parameters) + keep = all_names if self.retain == "all" else _report_closure(self.program) + items = self.program.parameters.items() + return [n for n, p in items if p.derivation is None and n in keep] def _report_closure(program: ms.Program) -> set[str]: @@ -200,11 +197,8 @@ def _check_keys(program: ms.Program, keys: frozenset[str]) -> None: unknown = sorted(keys - set(known)) if not unknown: return - lead = ( - f"source key {unknown[0]!r} names" - if len(unknown) == 1 - else f"source keys {unknown} name" - ) + one = len(unknown) == 1 + lead = f"source key {unknown[0]!r} names" if one else f"source keys {unknown} name" raise SpecDataError( f"{lead} neither a parameter, a dimension nor a lookup this spec declares. " f"{did_you_mean(unknown[0], known)} Pass only what the spec takes." @@ -218,18 +212,13 @@ def _check_keys(program: ms.Program, keys: frozenset[str]) -> None: def _reached(program: ms.Program) -> set[str]: dims: set[str] = set() - for p in program.parameters.values(): - dims.update(p.dims) - for v in program.variables.values(): - dims.update(v.dims) - for c in program.constraints.values(): - dims.update(c.dims) + for declared in (program.parameters, program.variables, program.constraints): + dims.update(d for decl in declared.values() for d in decl.dims) + dims.update(pw.over for pw in program.piecewise.values()) for over, lk in program.lookups: dims.add(over) if lk.target is not None: dims.add(lk.target) - for pw in program.piecewise.values(): - dims.add(pw.over) return dims @@ -296,10 +285,7 @@ def _lookups( series = _lookup_series(lk.name, over, sources[lk.name]) _check_lookup(series, lk, over, coords) padded = series.reindex(coords[over]) - array = xr.DataArray( - padded.to_numpy(), dims=[over], coords={over: coords[over]}, name=lk.name - ) - out.setdefault(over, {})[lk.name] = array + out.setdefault(over, {})[lk.name] = xr.DataArray(padded, name=lk.name) return out @@ -383,7 +369,6 @@ def _as_array( obj: Any, coords: Mapping[str, pd.Index], ) -> xr.DataArray: - dims = declared.dims if isinstance(obj, xr.DataArray): return _from_dense(name, declared, obj) if isinstance(obj, pd.DataFrame): @@ -394,6 +379,7 @@ def _as_array( return _from_rows(name, declared, pd.Series(dict(obj)), coords) if isinstance(obj, _SCALARS): return _from_scalar(name, declared, obj, coords) + dims = declared.dims raise SpecDataError( f"parameter '{name}': cannot adapt {type(obj).__name__} to an array over {list(dims)}; " f"pass {_parameter_shapes(dims)}." @@ -415,9 +401,8 @@ def _from_scalar( f"parameter '{name}' is one value and that value is a hole (null or NaN). " f"A number was meant, or the parameter has no data and should not be passed." ) - value = ( - np.asarray(obj, dtype=float) if declared.dtype == "float" else np.asarray(obj) - ) + dtype = float if declared.dtype == "float" else None + value = np.asarray(obj, dtype=dtype) _check_value_dtype(name, declared, value.dtype) arr = xr.DataArray(value, name=name) if declared.dims: @@ -441,9 +426,7 @@ def _from_dense( f"parameter '{name}' has no coordinate labels along '{d}'. A parameter is read for " f"values against its labels, so every dimension needs an index coordinate." ) - _refuse_duplicate_coordinates( - name, dims, arr.indexes[d].duplicated(), arr.indexes[d] - ) + _refuse_duplicate_coordinates(name, (d,), arr.indexes[d]) return arr.transpose(*dims) @@ -456,9 +439,8 @@ def _from_frame( dims = declared.dims tidy = df.reset_index() if set(dims) - set(df.columns) else df if "value" in tidy.columns and set(dims) <= set(tidy.columns): - if not dims: - return _from_rows(name, declared, tidy["value"], coords) - return _from_rows(name, declared, tidy.set_index(list(dims))["value"], coords) + indexed = tidy.set_index(list(dims)) if dims else tidy + return _from_rows(name, declared, indexed["value"], coords) if len(dims) == 2: return _from_dense(name, declared, xr.DataArray(_wide(name, dims, df))) raise SpecDataError( @@ -505,26 +487,24 @@ def _from_rows( f"it in the same breath. Drop those rows, or supply the values." ) _check_value_dtype(name, declared, series.dtype) - _refuse_duplicate_coordinates(name, dims, series.index.duplicated(), series.index) + _refuse_duplicate_coordinates(name, dims, series.index) for d in dims: labels = series.index.get_level_values(d) _refuse_strangers(name, d, labels[~labels.isin(coords[d])].unique(), coords[d]) onto = [coords[d] for d in dims] full = onto[0] if len(dims) == 1 else pd.MultiIndex.from_product(onto, names=dims) values = series.reindex(full, fill_value=_fill(declared)).to_numpy() - shape = tuple(len(index) for index in onto) - return xr.DataArray( - values.reshape(shape), dims=dims, coords=dict(zip(dims, onto)), name=name - ) + dense = values.reshape(tuple(len(index) for index in onto)) + return xr.DataArray(dense, dims=dims, coords=dict(zip(dims, onto)), name=name) def _with_dims(name: str, dims: tuple[str, ...], series: pd.Series) -> pd.Series: index = series.index if index.nlevels != len(dims): - said = "a Series or dict carries one label per level" raise SpecDataError( - f"parameter '{name}': {said}, and its index has {index.nlevels} level(s) where '{name}' " - f"is over {list(dims)}. Pass {_parameter_shapes(dims)}." + f"parameter '{name}': a Series or dict carries one label per level, and its index has " + f"{index.nlevels} level(s) where '{name}' is over {list(dims)}. " + f"Pass {_parameter_shapes(dims)}." ) names = list(index.names) if all(n is None for n in names): @@ -540,8 +520,9 @@ def _with_dims(name: str, dims: tuple[str, ...], series: pd.Series) -> pd.Series def _refuse_duplicate_coordinates( - name: str, dims: tuple[str, ...], duplicated: Any, index: pd.Index + name: str, dims: tuple[str, ...], index: pd.Index ) -> None: + duplicated = index.duplicated() if not duplicated.any(): return counts = index[duplicated].value_counts() diff --git a/test/test_spec_binder.py b/test/test_spec_binder.py index 784d7571..1afbad49 100644 --- a/test/test_spec_binder.py +++ b/test/test_spec_binder.py @@ -50,6 +50,25 @@ CAP = xr.DataArray(np.arange(9.0).reshape(3, 3), coords={"f": F, "t": T}, name="cap") COST = pd.Series([1.0, 2.0, 3.0], index=F) +DUP_ROWS = pd.Series([1.0, 9.0, 2.0], index=pd.Index(["a", "a", "b"], name="f")) +NULL_ROW = pd.Series({"a": 1.0, "b": None}) +NAN_ROW = pd.Series({"a": 1.0, "b": np.nan}) +NULL_FRAME = pd.DataFrame({"f": ["a", "b"], "value": [1.0, None]}) +STRAY_ROW = pd.Series({"a": 1.0, "zz": 2.0}) +DEEP_INDEX = pd.MultiIndex.from_tuples([("a", 0), ("b", 0)], names=["f", "k"]) +DEEP_ROWS = pd.Series([5.0, 5.0], index=DEEP_INDEX) + + +def sources_from( + base: Mapping[str, Any], override: Mapping[str, Any] +) -> dict[str, Any]: + """*base* with *override* applied; a ``None`` value drops the key instead.""" + merged = {**base, **override} + for key, value in override.items(): + if value is None: + merged.pop(key) + return merged + @pytest.fixture(scope="module") def program() -> Any: @@ -222,9 +241,7 @@ def test_missing_parameter_is_refused_when_read( id="duplicate-member", ), pytest.param( - {"cost": pd.Series({"a": 1.0, "zz": 2.0})}, - r"parameter 'cost'.*'f'.*'zz'", - id="unknown-label", + {"cost": STRAY_ROW}, r"parameter 'cost'.*'f'.*'zz'", id="unknown-label" ), pytest.param( {"cap": CAP.assign_coords(t=[0, 1, 9])}, @@ -232,7 +249,7 @@ def test_missing_parameter_is_refused_when_read( id="unknown-label-dense", ), pytest.param( - {"cost": pd.Series([1.0, 9.0, 2.0], index=pd.Index(["a", "a", "b"], name="f"))}, + {"cost": DUP_ROWS}, r"parameter 'cost' has more than one row for a coordinate: f='a' \(2 rows\)", id="duplicated-coordinate-row", ), @@ -242,12 +259,8 @@ def test_missing_parameter_is_refused_when_read( id="dense-wrong-dims", ), pytest.param( - { - "cap": xr.DataArray( - np.ones((2, 3)), coords={"f": ["a", "a"], "t": [0, 1, 2]} - ) - }, - r"parameter 'cap' has more than one row", + {"cap": xr.DataArray(np.ones((3, 3)), coords={"f": list(F), "t": [0, 0, 1]})}, + r"parameter 'cap' has more than one row for a coordinate: t=0 \(2 rows\)", id="dense-duplicate-coordinate", ), pytest.param( @@ -256,11 +269,7 @@ def test_missing_parameter_is_refused_when_read( id="wrong-rank", ), pytest.param( - { - "cap": pd.Series( - [1.0], index=pd.MultiIndex.from_tuples([("a", 0)], names=["f", "q"]) - ) - }, + {"cap": DEEP_ROWS.rename_axis(["f", "q"])}, r"parameter 'cap' is indexed by \['f', 'q'\]", id="wrong-level-names", ), @@ -280,20 +289,19 @@ def test_missing_parameter_is_refused_when_read( id="dense-without-labels", ), pytest.param( - {"cost": pd.Series({"a": 1.0, "b": None})}, - r"parameter 'cost' carries 1 row.*f='b'", - id="null-row", - ), - pytest.param( - {"cost": pd.Series({"a": 1.0, "b": np.nan})}, - r"parameter 'cost' carries 1 row", - id="nan-row", + {"cost": NULL_ROW}, r"parameter 'cost' carries 1 row.*f='b'", id="null-row" ), + pytest.param({"cost": NAN_ROW}, r"parameter 'cost' carries 1 row", id="nan-row"), pytest.param( {"rate": float("nan")}, r"parameter 'rate' is one value and that value is a hole", id="nan-scalar", ), + pytest.param( + {"rate": pd.DataFrame({"value": [None]})}, + r"parameter 'rate' is one value and that value is a hole", + id="nan-scalar-frame", + ), pytest.param( {"lead": pd.Series([1.5, 0.0, 1.0], index=F)}, r"'lead' is declared 'int'.*'float'", @@ -361,12 +369,8 @@ def test_missing_parameter_is_refused_when_read( def test_malformed_data_is_refused_naming_the_symbol( program: Any, good: dict[str, Any], override: dict[str, Any], match: str ) -> None: - sources = {**good, **override} - for key, value in override.items(): - if value is None: - sources.pop(key) with pytest.raises(SpecDataError, match=match): - read_all(program, sources) + read_all(program, sources_from(good, override)) def test_int_labels_are_shown_as_written(program: Any, good: dict[str, Any]) -> None: @@ -376,18 +380,9 @@ def test_int_labels_are_shown_as_written(program: Any, good: dict[str, Any]) -> def test_dataset_is_a_source(program: Any, good: dict[str, Any]) -> None: - ds = xr.Dataset( - { - "cost": xr.DataArray(COST), - "cap": CAP, - "flag": xr.DataArray(good["flag"]), - "rate": 0.5, - "lead": xr.DataArray(good["lead"]), - "grp": xr.DataArray(good["grp"]), - }, - coords={"f": F, "t": T, "g": ["n", "e"]}, - ) - from_dataset = bind(program, ds) + dims = {"f": F, "t": T, "g": ["n", "e"]} + values = {k: xr.DataArray(v) for k, v in good.items() if k not in dims} + from_dataset = bind(program, xr.Dataset(values, coords=dims)) from_mapping = bind(program, good) assert from_dataset.coords["f"].equals(from_mapping.coords["f"]) for name in program.parameters: @@ -482,19 +477,14 @@ def test_report_closure_reads_names_and_masks() -> None: assert set(bind(program, sources).retained().data_vars) == {"cost", "lag", "on"} -def test_aligned_array_is_not_copied(program: Any, good: dict[str, Any]) -> None: - bound = bind(program, good) - assert np.shares_memory(CAP.values, bound.parameter("cap").values) - assert np.shares_memory(CAP.values, bound.parameter("cap").values) - permuted = CAP.transpose("t", "f") - assert np.shares_memory( - permuted.values, - bind(program, {**good, "cap": permuted}).parameter("cap").values, - ) - wide = CAP.to_pandas() - assert np.shares_memory( - wide.values, bind(program, {**good, "cap": wide}).parameter("cap").values - ) +@pytest.mark.parametrize("shape", ["dataarray", "dataarray-transposed", "wide-frame"]) +def test_aligned_array_is_not_copied( + program: Any, good: dict[str, Any], shape: str +) -> None: + source = CAP_SHAPES[shape] + bound = bind(program, {**good, "cap": source}) + assert np.shares_memory(np.asarray(source), bound.parameter("cap").values) + assert np.shares_memory(np.asarray(source), bound.parameter("cap").values) def test_derived_parameter_is_not_bound_from_sources() -> None: @@ -552,88 +542,30 @@ def test_derived_parameter_is_not_bound_from_sources() -> None: ACCEPTED = "accepted" PARITY_CASES = [ - pytest.param(GOOD, ACCEPTED, id="valid"), - pytest.param( - {"f": ["a", "b"], "cost": GOOD["cost"]}, - SpecDataError, - id="parameter-missing-entirely", - ), - pytest.param( - {**GOOD, "cost": pd.Series({"a": 1.0})}, ACCEPTED, id="coefficient-sparse" - ), - pytest.param( - { - **GOOD, - "cost": pd.Series( - [1.0, 9.0, 2.0], index=pd.Index(["a", "a", "b"], name="f") - ), - }, - SpecDataError, - id="duplicated-coordinate-row", - ), - pytest.param( - {**GOOD, "cost": pd.Series({"a": 1.0, "zz": 2.0})}, - SpecDataError, - id="label-the-dimension-does-not-have", - ), - pytest.param( - {**GOOD, "cost": pd.Series({"a": 1.0, "b": None})}, - SpecDataError, - id="a-null-value", - ), - pytest.param( - {**GOOD, "cost": pd.Series({"a": 1.0, "b": float("nan")})}, - SpecDataError, - id="a-nan-value", - ), - pytest.param( - {**GOOD, "cap": pd.Series({"a": 5.0, "b": None})}, - SpecDataError, - id="a-hole-in-a-bound", - ), - pytest.param( - {**GOOD, "cost": float("nan")}, SpecDataError, id="a-hole-as-a-scalar" - ), - pytest.param( - {**GOOD, "cost": [1.0, float("nan")]}, SpecDataError, id="a-hole-in-a-sequence" - ), - pytest.param( - {**GOOD, "cost": {"a": 1.0, "b": None}}, SpecDataError, id="a-hole-in-a-dict" - ), - pytest.param( - {**GOOD, "cost": pd.DataFrame({"f": ["a", "b"], "value": [1.0, None]})}, - SpecDataError, - id="a-hole-in-a-tidy-frame", - ), - pytest.param( - {**GOOD, "cost": pd.Series({"a": 1, "b": 2})}, - ACCEPTED, - id="whole-numbers-serve-a-float-declaration", - ), - pytest.param( - {**GOOD, "csot": GOOD["cost"]}, - SpecDataError, - id="a-source-key-the-model-does-not-declare", - ), - pytest.param( - { - **GOOD, - "cost": pd.Series( - [5.0, 5.0], - index=pd.MultiIndex.from_tuples([("a", 0), ("b", 0)], names=["f", "k"]), - ), - }, - SpecDataError, - id="a-series-deeper-than-the-declared-dims", - ), + pytest.param({}, ACCEPTED, id="valid"), + pytest.param({"cap": None}, SpecDataError, id="parameter-missing-entirely"), + pytest.param({"cost": pd.Series({"a": 1.0})}, ACCEPTED, id="coefficient-sparse"), + pytest.param({"cost": DUP_ROWS}, SpecDataError, id="duplicated-coordinate-row"), + pytest.param({"cost": STRAY_ROW}, SpecDataError, id="label-not-in-the-dimension"), + pytest.param({"cost": NULL_ROW}, SpecDataError, id="a-null-value"), + pytest.param({"cost": NAN_ROW}, SpecDataError, id="a-nan-value"), + pytest.param({"cap": NULL_ROW}, SpecDataError, id="a-hole-in-a-bound"), + pytest.param({"cost": float("nan")}, SpecDataError, id="a-hole-as-a-scalar"), + pytest.param({"cost": [1.0, np.nan]}, SpecDataError, id="a-hole-in-a-sequence"), + pytest.param({"cost": {"a": 1.0, "b": None}}, SpecDataError, id="a-hole-in-a-dict"), + pytest.param({"cost": NULL_FRAME}, SpecDataError, id="a-hole-in-a-tidy-frame"), + pytest.param({"cost": pd.Series({"a": 1, "b": 2})}, ACCEPTED, id="whole-numbers"), + pytest.param({"csot": COST}, SpecDataError, id="an-undeclared-source-key"), + pytest.param({"cost": DEEP_ROWS}, SpecDataError, id="a-series-too-deep"), ] -@pytest.mark.parametrize(("sources", "verdict"), PARITY_CASES) +@pytest.mark.parametrize(("override", "verdict"), PARITY_CASES) def test_parity_with_lpspec_data_verdicts( - sources: dict[str, Any], verdict: Any + override: dict[str, Any], verdict: Any ) -> None: program = math_spec.to_program(PARITY_SPEC) + sources = sources_from(GOOD, override) if verdict is ACCEPTED: read_all(program, sources) return @@ -644,29 +576,25 @@ def test_parity_with_lpspec_data_verdicts( def test_a_hole_is_named_where_it_sits() -> None: program = math_spec.to_program(PARITY_SPEC) with pytest.raises(SpecDataError, match="parameter 'cost'") as error: - read_all(program, {**GOOD, "cost": pd.Series({"a": 1.0, "b": None})}) + read_all(program, {**GOOD, "cost": NULL_ROW}) assert "divisor" not in str(error.value) assert "f='b'" in str(error.value) -def test_a_hole_in_a_scalar_parameter_is_refused() -> None: - spec = { - "dimensions": {"f": {"dtype": "str"}}, - "parameters": {"rate": {"dims": []}}, - "variables": {"x": {"foreach": ["f"], "bounds": {"lower": 0, "upper": 1}}}, - "objective": {"sense": "maximize", "expression": "sum(x * rate)"}, - } - program = math_spec.to_program(spec) - with pytest.raises(SpecDataError, match="hole"): - read_all(program, {"f": ["a", "b"], "rate": pd.DataFrame({"value": [None]})}) +FLAG_SPEC = { + "dimensions": {"g": {"dtype": "str"}}, + "parameters": {"active": {"dims": ["g"], "dtype": "bool"}}, + "variables": { + "x": {"foreach": ["g"], "where": "active", "bounds": {"lower": 0, "upper": 1}} + }, + "objective": {"sense": "maximize", "expression": "sum(x)"}, +} @pytest.mark.parametrize( ("column", "verdict"), [ - pytest.param( - pd.Series({"a": True, "b": False}), ACCEPTED, id="a-boolean-column" - ), + pytest.param(pd.Series({"a": True, "b": False}), ACCEPTED, id="a-bool-column"), pytest.param(pd.Series({"a": 1, "b": 0}), SpecDataError, id="a-1-0-int-column"), pytest.param( pd.Series({"a": 1.0, "b": 0.0}), SpecDataError, id="a-1-0-float-column" @@ -674,19 +602,7 @@ def test_a_hole_in_a_scalar_parameter_is_refused() -> None: ], ) def test_a_flag_binds_by_its_declaration(column: pd.Series, verdict: Any) -> None: - spec = { - "dimensions": {"g": {"dtype": "str"}}, - "parameters": {"active": {"dims": ["g"], "dtype": "bool"}}, - "variables": { - "x": { - "foreach": ["g"], - "where": "active", - "bounds": {"lower": 0, "upper": 1}, - } - }, - "objective": {"sense": "maximize", "expression": "sum(x)"}, - } - program = math_spec.to_program(spec) + program = math_spec.to_program(FLAG_SPEC) sources = {"g": ["a", "b"], "active": column} if verdict is ACCEPTED: assert bind(program, sources).parameter("active").dtype == bool @@ -703,70 +619,58 @@ def test_a_flag_binds_by_its_declaration(column: pd.Series, verdict: Any) -> Non "constraints": {"k": {"foreach": ["b"], "expression": "sum(x, by=gen_bus) <= 10"}}, "objective": {"sense": "maximize", "expression": "sum(x)"}, } -P_MAX = {"p_max": pd.Series({"w": 5.0, "s": 5.0})} -INDEX = {"g": ["w", "s"], "b": ["n", "e"]} -MAP = {"gen_bus": pd.Series({"w": "n", "s": "e"})} +G_TWICE = pd.Index(["w", "w", "s"], name="g") +LOOKUP_GOOD = { + "p_max": pd.Series({"w": 5.0, "s": 5.0}), + "g": ["w", "s"], + "b": ["n", "e"], + "gen_bus": pd.Series({"w": "n", "s": "e"}), +} @pytest.mark.parametrize( - ("sources", "match"), + ("override", "match"), [ pytest.param( - {**P_MAX, **MAP}, "dimension 'g' has no index", id="a-map-and-no-labels" + {"g": None, "b": None}, "dimension 'g' has no index", id="a-map-no-labels" ), pytest.param( - {**P_MAX, **INDEX}, "no data provided for lookup", id="an-index-and-no-map" + {"gen_bus": None}, "no data provided for lookup", id="an-index-no-map" ), pytest.param( - {**P_MAX, **INDEX, "gen_bus": pd.Series({"w": "n", "s": "zz"})}, + {"gen_bus": pd.Series({"w": "n", "s": "zz"})}, "not 'b' labels", id="a-stray-value", ), pytest.param( - { - **P_MAX, - **INDEX, - "gen_bus": pd.Series( - ["n", "e", "e"], index=pd.Index(["w", "w", "s"], name="g") - ), - }, + {"gen_bus": pd.Series(["n", "e", "e"], index=G_TWICE)}, "more than once", id="two-values-for-one-label", ), pytest.param( - {**P_MAX, **INDEX, "gen_bus": pd.Series({"w": None, "s": "e"})}, + {"gen_bus": pd.Series({"w": None, "s": "e"})}, "null in 'b'", id="mapping-a-label-to-nothing", ), pytest.param( - { - **P_MAX, - **INDEX, - "gen_bus": pd.Series( - [None, "n", "n"], index=pd.Index(["w", "w", "s"], name="g") - ), - }, + {"gen_bus": pd.Series([None, "n", "n"], index=G_TWICE)}, "null in 'b'", id="a-label-held-twice-with-a-null", ), ], ) -def test_a_lookup_defect_is_refused(sources: dict[str, Any], match: str) -> None: +def test_a_lookup_defect_is_refused(override: dict[str, Any], match: str) -> None: program = math_spec.to_program(LOOKUP_SPEC) with pytest.raises(SpecDataError, match=match): - read_all(program, sources) + read_all(program, sources_from(LOOKUP_GOOD, override)) def test_a_stray_lookup_value_over_an_int_target_is_shown_as_written() -> None: program = math_spec.to_program( {**LOOKUP_SPEC, "dimensions": {"g": {}, "b": {"dtype": "int"}}} ) - sources = { - **P_MAX, - "g": ["w", "s"], - "b": [1, 2], - "gen_bus": pd.Series({"w": 1, "s": 99}), - } + numbered = {"b": [1, 2], "gen_bus": pd.Series({"w": 1, "s": 99})} + sources = sources_from(LOOKUP_GOOD, numbered) with pytest.raises(SpecDataError, match=r"not 'b' labels: 99\b") as error: read_all(program, sources) assert "int64" not in str(error.value) From 9bf025a7a509aaac62c55d13e387aa5b044ba766 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 13:43:14 +0200 Subject: [PATCH 03/35] fix(spec): review fixes for the binder Validate retain, check a scalar's dtype before casting, bind empty sources as all-NaN, check label-space lookup dtypes, report unknown labels in source order on every path, re-stamp coordinates onto the master dtype without copying, and pin the remaining binder rules. --- linopy/spec/binder.py | 95 ++++++++++++++++----------- test/test_spec_binder.py | 135 ++++++++++++++++++++++++++++++++++++++- 2 files changed, 192 insertions(+), 38 deletions(-) diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py index 0cf91f7d..f893bf43 100644 --- a/linopy/spec/binder.py +++ b/linopy/spec/binder.py @@ -14,8 +14,8 @@ from __future__ import annotations from collections.abc import Hashable, Iterable, Mapping, Sequence -from dataclasses import dataclass -from typing import Any, Literal +from dataclasses import dataclass, field +from typing import Any, Literal, get_args import numpy as np import pandas as pd @@ -26,12 +26,14 @@ from linopy.spec.errors import SpecDataError Retain = Literal["report", "all", "none"] +_RETAIN: tuple[str, ...] = get_args(Retain) _ACCEPTED_KINDS: dict[str, frozenset[str]] = { "float": frozenset("fiu"), "int": frozenset("iu"), "bool": frozenset("b"), "str": frozenset("OUS"), + "datetime": frozenset("M"), } _KIND_NAMES: dict[str, str] = { "f": "float", @@ -41,6 +43,13 @@ "O": "str", "U": "str", "S": "str", + "M": "datetime", +} +_EMPTY_DTYPES: dict[str, Any] = { + "float": float, + "int": int, + "bool": bool, + "str": object, } _SCALARS = (bool, int, float, str, np.number, np.bool_) _DIMENSION_SHAPES = "a pandas Index, a list, a tuple, a 1-D numpy array, a pandas Series or a 1-D DataArray" @@ -69,10 +78,15 @@ def bind( retain: Which parameters :meth:`Bound.retained` persists. Raises: - SpecDataError: 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. + 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. """ + if retain not in _RETAIN: + raise SpecDataError( + f"retain={retain!r} is not one of {_shown(_RETAIN)}. {did_you_mean(retain, _RETAIN)}" + ) if isinstance(sources, xr.Dataset): sources = _dataset_sources(sources) keys = frozenset(sources.keys()) @@ -82,7 +96,7 @@ def bind( return Bound(program, coords, lookups, retain, sources, keys) -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) class Bound: """ A program bound to its data. @@ -96,15 +110,14 @@ class Bound: dimension's master coordinates, NaN where a label is unmapped. retain: Which parameters :meth:`retained` persists. sources: The caller's data, read by key on demand. - keys: The keys ``sources`` carries, read once at bind time. """ program: ms.Program - coords: dict[str, pd.Index] - lookups: dict[str, dict[str, xr.DataArray]] + coords: Mapping[str, pd.Index] + lookups: Mapping[str, Mapping[str, xr.DataArray]] retain: Retain sources: Mapping[str, Any] - keys: frozenset[str] + _keys: frozenset[str] = field(repr=False) def parameter(self, name: str) -> xr.DataArray: """ @@ -122,7 +135,7 @@ def parameter(self, name: str) -> xr.DataArray: than declared. """ declared = self._declaration(name) - if name not in self.keys: + if name not in self._keys: raise SpecDataError(f"no data provided for parameter '{name}'") arr = _as_array(name, declared, self.sources[name], self.coords) onto = {d: self.coords[d] for d in declared.dims} @@ -151,10 +164,11 @@ def _declaration(self, name: str) -> ms.ParameterDeclaration: def _retained_names(self) -> list[str]: if self.retain == "none": return [] - all_names = set(self.program.parameters) - keep = all_names if self.retain == "all" else _report_closure(self.program) - items = self.program.parameters.items() - return [n for n, p in items if p.derivation is None and n in keep] + parameters = self.program.parameters + keep = ( + set(parameters) if self.retain == "all" else _report_closure(self.program) + ) + return [n for n, p in parameters.items() if p.derivation is None and n in keep] def _report_closure(program: ms.Program) -> set[str]: @@ -341,6 +355,8 @@ def _check_lookup( f"none of them would place its terms nowhere, so it is a typo on one side or a label " f"missing from the other." ) + if lk.dtype is not None: + _check_value_dtype(lk.name, lk.dtype, series.dtype, kind="lookup") if lk.target is None: return values = pd.Index(series.to_numpy()) @@ -401,9 +417,10 @@ def _from_scalar( f"parameter '{name}' is one value and that value is a hole (null or NaN). " f"A number was meant, or the parameter has no data and should not be passed." ) - dtype = float if declared.dtype == "float" else None - value = np.asarray(obj, dtype=dtype) - _check_value_dtype(name, declared, value.dtype) + value = np.asarray(obj) + _check_value_dtype(name, declared.dtype, value.dtype) + if declared.dtype == "float": + value = value.astype(float) arr = xr.DataArray(value, name=name) if declared.dims: arr = arr.expand_dims({d: coords[d] for d in declared.dims}) @@ -414,7 +431,7 @@ def _from_dense( name: str, declared: ms.ParameterDeclaration, arr: xr.DataArray ) -> xr.DataArray: dims = declared.dims - _check_value_dtype(name, declared, arr.dtype) + _check_value_dtype(name, declared.dtype, arr.dtype) if set(arr.dims) != set(dims) or len(arr.dims) != len(dims): raise SpecDataError( f"parameter '{name}' arrived as a DataArray over {list(arr.dims)}, and '{name}' is over " @@ -437,7 +454,9 @@ def _from_frame( coords: Mapping[str, pd.Index], ) -> xr.DataArray: dims = declared.dims - tidy = df.reset_index() if set(dims) - set(df.columns) else df + tidy = df + if not set(dims) <= set(df.columns) and set(dims) <= _headers(df): + tidy = df.reset_index() if "value" in tidy.columns and set(dims) <= set(tidy.columns): indexed = tidy.set_index(list(dims)) if dims else tidy return _from_rows(name, declared, indexed["value"], coords) @@ -449,6 +468,10 @@ def _from_frame( ) +def _headers(df: pd.DataFrame) -> set[Any]: + return set(df.columns) | set(df.index.names) + + def _wide(name: str, dims: tuple[str, ...], df: pd.DataFrame) -> pd.DataFrame: names = (df.index.name, df.columns.name) if names == (None, None): @@ -478,6 +501,8 @@ def _from_rows( ) return _from_scalar(name, declared, series.iloc[0], coords) series = _with_dims(name, dims, series) + if series.empty: + series = series.astype(_EMPTY_DTYPES[declared.dtype]) holes = series.isna() if holes.any(): raise SpecDataError( @@ -486,11 +511,10 @@ def _from_rows( f"value is the absence of the row, and such a row says the coordinate exists and denies " f"it in the same breath. Drop those rows, or supply the values." ) - _check_value_dtype(name, declared, series.dtype) + _check_value_dtype(name, declared.dtype, series.dtype) _refuse_duplicate_coordinates(name, dims, series.index) for d in dims: - labels = series.index.get_level_values(d) - _refuse_strangers(name, d, labels[~labels.isin(coords[d])].unique(), coords[d]) + _refuse_strangers(name, d, series.index.get_level_values(d), coords[d]) onto = [coords[d] for d in dims] full = onto[0] if len(dims) == 1 else pd.MultiIndex.from_product(onto, names=dims) values = series.reindex(full, fill_value=_fill(declared)).to_numpy() @@ -536,14 +560,13 @@ def _refuse_duplicate_coordinates( ) -def _refuse_strangers( - name: str, dim: str, strangers: pd.Index, labels: pd.Index -) -> None: - if len(strangers) == 0: +def _refuse_strangers(name: str, dim: str, labels: pd.Index, known: pd.Index) -> None: + strangers = labels[~labels.isin(known)].unique().tolist() + if not strangers: return raise SpecDataError( f"parameter '{name}' has label(s) in dimension '{dim}' that are not coordinates of it: " - f"{_shown(strangers.tolist())}.\n {dim} has: {_shown(labels.tolist(), 10)}\n" + 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}']." @@ -554,24 +577,24 @@ def _aligned( name: str, arr: xr.DataArray, onto: Mapping[str, pd.Index], fill: Any ) -> xr.DataArray: if all(arr.indexes[d].equals(index) for d, index in onto.items()): - return arr + stale = {d: i for d, i in onto.items() if arr.indexes[d].dtype != i.dtype} + return arr.assign_coords(stale) if stale else arr for d, index in onto.items(): - _refuse_strangers(name, d, arr.indexes[d].difference(index), index) + _refuse_strangers(name, d, arr.indexes[d], index) return arr.reindex(onto, fill_value=fill) def _check_value_dtype( - name: str, declared: ms.ParameterDeclaration, dtype: Any + name: str, declared: str, dtype: Any, kind: str = "parameter" ) -> None: - kind = str(dtype.kind) - if kind in _ACCEPTED_KINDS[declared.dtype]: + if str(dtype.kind) in _ACCEPTED_KINDS[declared]: return - arrived = _KIND_NAMES.get(kind, str(dtype)) + arrived = _KIND_NAMES.get(str(dtype.kind), str(dtype)) raise SpecDataError( - f"parameter '{name}' is declared '{declared.dtype}' and its values arrived as '{arrived}'. " + f"{kind} '{name}' is declared '{declared}' and its values arrived as '{arrived}'. " f"A declared dtype is a claim about the values, and it is checked here: the file says what " f"the values are, or the values are not attached.\n" - f" Cast the values to {declared.dtype}, if the declaration is what you meant\n" + f" Cast the values to {declared}, if the declaration is what you meant\n" f" Or declare what the data has: {{dtype: {arrived}}}" ) diff --git a/test/test_spec_binder.py b/test/test_spec_binder.py index 1afbad49..5e9d82f2 100644 --- a/test/test_spec_binder.py +++ b/test/test_spec_binder.py @@ -127,6 +127,7 @@ def test_rank_two_shapes_bind_alike( "dataarray": xr.DataArray(COST), "dict": COST.to_dict(), "tidy-frame": COST.reset_index(name="value"), + "indexed-frame": COST.to_frame("value"), } @@ -186,6 +187,7 @@ def test_missing_rows_become_nan_and_false(program: Any, good: dict[str, Any]) - sparse = { **good, "cost": pd.Series({"a": 1.0}), + "lead": pd.Series({"a": 1}), "flag": pd.Series({"a": True}), "cap": CAP.sel(t=[0, 1]), } @@ -193,6 +195,10 @@ def test_missing_rows_become_nan_and_false(program: Any, good: dict[str, Any]) - cost = bound.parameter("cost") assert cost.sel(f="a").item() == 1.0 assert cost.sel(f=["b", "c"]).isnull().all() + lead = bound.parameter("lead") + assert lead.dtype == np.float64 + assert lead.sel(f="a").item() == 1.0 + assert lead.sel(f=["b", "c"]).isnull().all() flag = bound.parameter("flag") assert flag.dtype == bool assert flag.values.tolist() == [False, True, False] @@ -225,6 +231,23 @@ def test_scalar_parameter_stays_scalar(program: Any, good: dict[str, Any]) -> No assert got.item() == 0.5 +EMPTY_SOURCES = { + "dict": {}, + "object-series": pd.Series(dtype=object), + "float-series": pd.Series(dtype=float), +} + + +@pytest.mark.parametrize("cost", EMPTY_SOURCES.values(), ids=EMPTY_SOURCES.keys()) +def test_empty_source_binds_as_all_nan( + program: Any, good: dict[str, Any], cost: Any +) -> None: + got = bind(program, {**good, "cost": cost}).parameter("cost") + assert got.dtype == np.float64 + assert got.isnull().all() + assert got.indexes["f"].equals(F) + + def test_missing_parameter_is_refused_when_read( program: Any, good: dict[str, Any] ) -> None: @@ -234,6 +257,19 @@ def test_missing_parameter_is_refused_when_read( bound.parameter("cost") +def test_undeclared_parameter_is_refused_with_a_hint( + program: Any, good: dict[str, Any] +) -> None: + with pytest.raises(SpecDataError, match="unknown parameter 'csot'.*'cost'"): + bind(program, good).parameter("csot") + + +def test_retain_is_validated_before_binding(program: Any, good: dict[str, Any]) -> None: + with pytest.raises(SpecDataError, match=r"'report', 'all', 'none'") as error: + bind(program, good, retain="reports") # type: ignore[arg-type] + assert "Did you mean 'report'?" in str(error.value) + + REFUSALS = [ pytest.param( {"f": ["a", "a", "b"]}, @@ -248,6 +284,16 @@ def test_missing_parameter_is_refused_when_read( r"parameter 'cap'.*'t'.*\b9\b", id="unknown-label-dense", ), + pytest.param( + {"cap": CAP.assign_coords(t=[9, 0, 7])}, + r"not coordinates of it: 9, 7\.", + id="unknown-labels-dense-in-source-order", + ), + pytest.param( + {"cap": CAP.to_series().reset_index(name="value").assign(t=[9, 0, 7] * 3)}, + r"not coordinates of it: 9, 7\.", + id="unknown-labels-rows-in-source-order", + ), pytest.param( {"cost": DUP_ROWS}, r"parameter 'cost' has more than one row for a coordinate: f='a' \(2 rows\)", @@ -283,6 +329,16 @@ def test_missing_parameter_is_refused_when_read( r"parameter 'cost': cannot adapt set", id="unsupported-shape", ), + pytest.param( + {"cost": pd.DataFrame({"f": ["a"], "amount": [1.0]})}, + r"parameter 'cost' arrived as a DataFrame with columns \['f', 'amount'\]", + id="frame-without-value-column", + ), + pytest.param( + {"cap": CAP.to_pandas().rename_axis(index="f", columns="q")}, + r"parameter 'cap' arrived as a wide DataFrame with index 'f' and columns 'q'", + id="wide-frame-wrong-axis-names", + ), pytest.param( {"cap": xr.DataArray(np.ones((3, 3)), dims=["f", "t"])}, r"parameter 'cap' has no coordinate labels along 'f'", @@ -315,6 +371,17 @@ def test_missing_parameter_is_refused_when_read( pytest.param( {"flag": 1.0}, r"'flag' is declared 'bool'.*'float'", id="float-scalar-for-bool" ), + pytest.param( + {"rate": "1.5"}, r"'rate' is declared 'float'.*'str'", id="numeric-str-scalar" + ), + pytest.param( + {"rate": True}, + r"'rate' is declared 'float'.*'bool'", + id="bool-scalar-for-float", + ), + pytest.param( + {"rate": "abc"}, r"'rate' is declared 'float'.*'str'", id="str-scalar-for-float" + ), pytest.param( {"cost": pd.Series(["x", "y", "z"], index=F)}, r"'cost' is declared 'float'.*'str'", @@ -329,6 +396,16 @@ def test_missing_parameter_is_refused_when_read( r"index for dimension 'f': cannot read labels out of dict", id="dimension-shape", ), + pytest.param( + {"f": np.ones((2, 2))}, + r"index for dimension 'f' is 2-dimensional", + id="dimension-rank", + ), + pytest.param( + {"grp": xr.DataArray(["n"], coords={"t": [0]})}, + r"lookup 'grp' arrived as a DataArray over \['t'\]", + id="lookup-wrong-dataarray-dim", + ), pytest.param( {"grp": None}, r"no data provided for lookup 'grp'", id="missing-lookup" ), @@ -446,12 +523,14 @@ def test_report_closure_reads_names_and_masks() -> None: "parameters": { "cost": {"dims": ["f"]}, "lag": {"dims": ["f"], "dtype": "int"}, + "span": {"dims": ["f"], "dtype": "int"}, "on": {"dims": ["f"], "dtype": "bool"}, "other": {"dims": ["f"]}, }, "variables": {"x": {"foreach": ["f", "t"], "bounds": {"lower": 0, "upper": 1}}}, "objective": {"sense": "maximize", "expression": "sum(x * other)"}, "expressions": { + "recent": "sum_back(x, over=t, within=span)", "late": { "foreach": ["f", "t"], "cases": { @@ -461,7 +540,7 @@ def test_report_closure_reads_names_and_masks() -> None: } }, "otherwise": "x * cost", - } + }, }, } program = math_spec.to_program(spec) @@ -471,10 +550,18 @@ def test_report_closure_reads_names_and_masks() -> None: "t": [0, 1], "cost": pd.Series([1.0], index=f), "lag": pd.Series([1], index=f), + "span": pd.Series([2], index=f), "on": pd.Series([True], index=f), "other": pd.Series([2.0], index=f), } - assert set(bind(program, sources).retained().data_vars) == {"cost", "lag", "on"} + retained = bind(program, sources).retained() + assert set(retained.data_vars) == {"cost", "lag", "span", "on"} + + +def test_unreached_dimension_needs_no_source() -> None: + dimensions = {**PARITY_SPEC["dimensions"], "z": {"dtype": "int"}} + program = math_spec.to_program({**PARITY_SPEC, "dimensions": dimensions}) + assert list(bind(program, GOOD).coords) == ["f"] @pytest.mark.parametrize("shape", ["dataarray", "dataarray-transposed", "wide-frame"]) @@ -487,6 +574,15 @@ def test_aligned_array_is_not_copied( assert np.shares_memory(np.asarray(source), bound.parameter("cap").values) +def test_master_coordinate_dtype_wins_without_a_copy( + program: Any, good: dict[str, Any] +) -> None: + source = CAP.assign_coords(t=T.astype("int32")) + got = bind(program, {**good, "cap": source}).parameter("cap") + assert got.indexes["t"].dtype == np.int64 + assert np.shares_memory(np.asarray(source), got.values) + + def test_derived_parameter_is_not_bound_from_sources() -> None: spec = { "dimensions": {"bp": {"dtype": "int"}}, @@ -665,6 +761,41 @@ def test_a_lookup_defect_is_refused(override: dict[str, Any], match: str) -> Non read_all(program, sources_from(LOOKUP_GOOD, override)) +TAG_SPEC = { + **LOOKUP_SPEC, + "lookups": {"tag": {"over": "g", "dtype": "int"}}, + "constraints": {"k": {"foreach": ["g"], "expression": "x <= 10"}}, +} +TAG_GOOD = sources_from(LOOKUP_GOOD, {"gen_bus": None, "b": None}) + + +def test_a_label_space_lookup_is_padded_onto_the_dimension() -> None: + program = math_spec.to_program(TAG_SPEC) + bound = bind(program, {**TAG_GOOD, "tag": {"s": 7}}) + tag = bound.lookups["g"]["tag"] + assert tag.dims == ("g",) + assert tag.indexes["g"].tolist() == ["w", "s"] + assert np.isnan(tag.sel(g="w").item()) + assert tag.sel(g="s").item() == 7 + + +@pytest.mark.parametrize( + ("tag", "match"), + [ + pytest.param({"w": None, "s": 7}, "null in 'tag': g='w'", id="a-null"), + pytest.param( + {"w": "x", "s": "y"}, "lookup 'tag' is declared 'int'.*'str'", id="a-str" + ), + ], +) +def test_a_label_space_lookup_defect_is_refused( + tag: dict[str, Any], match: str +) -> None: + program = math_spec.to_program(TAG_SPEC) + with pytest.raises(SpecDataError, match=match): + bind(program, {**TAG_GOOD, "tag": tag}) + + def test_a_stray_lookup_value_over_an_int_target_is_shown_as_written() -> None: program = math_spec.to_program( {**LOOKUP_SPEC, "dimensions": {"g": {}, "b": {"dtype": "int"}}} From 56109093924257facdfd873fe89129c4bf9f1c72 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 14:02:07 +0200 Subject: [PATCH 04/35] ci(spec): skip doctest collection of linopy/spec without math-spec; type the test spec dicts --- conftest.py | 9 +++++++++ pyproject.toml | 2 +- test/test_spec_binder.py | 4 ++-- 3 files changed, 12 insertions(+), 3 deletions(-) create mode 100644 conftest.py diff --git a/conftest.py b/conftest.py new file mode 100644 index 00000000..3fd48ab4 --- /dev/null +++ b/conftest.py @@ -0,0 +1,9 @@ +"""Root pytest configuration for ``--doctest-modules`` collection of ``linopy/``.""" + +from __future__ import annotations + +from importlib.util import find_spec + +collect_ignore: list[str] = [] +if find_spec("math_spec") is None: + collect_ignore.append("linopy/spec") diff --git a/pyproject.toml b/pyproject.toml index b9c0c051..aa9d4902 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -161,7 +161,7 @@ omit = ["test/*"] exclude_also = ["if TYPE_CHECKING:"] [tool.mypy] -exclude = ['dev/*', 'examples/*', '^benchmark/', 'doc/*'] +exclude = ['dev/*', 'examples/*', '^benchmark/', 'doc/*', '^conftest\.py$'] ignore_missing_imports = true no_implicit_optional = true warn_unused_ignores = true diff --git a/test/test_spec_binder.py b/test/test_spec_binder.py index 5e9d82f2..c791eb6e 100644 --- a/test/test_spec_binder.py +++ b/test/test_spec_binder.py @@ -14,7 +14,7 @@ from linopy.spec import SpecDataError, bind # noqa: E402 -SPEC = { +SPEC: dict[str, Any] = { "dimensions": {"f": {"dtype": "str"}, "t": {"dtype": "int"}, "g": {"dtype": "str"}}, "lookups": {"grp": {"over": "f", "into": "g"}}, "parameters": { @@ -622,7 +622,7 @@ def test_derived_parameter_is_not_bound_from_sources() -> None: # lpspec data-parity cases, eager representation # --------------------------------------------------------------------------- -PARITY_SPEC = { +PARITY_SPEC: dict[str, Any] = { "dimensions": {"f": {"dtype": "str"}}, "parameters": {"cost": {"dims": ["f"]}, "cap": {"dims": ["f"]}}, "variables": {"x": {"foreach": ["f"], "bounds": {"lower": 0, "upper": "cap"}}}, From cc69e5af6ecf1b8b30d373707c731952b7c1e44b Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 20:34:35 +0200 Subject: [PATCH 05/35] ci(spec): install math-spec from a pinned git commit via a dependency group math-spec is not on PyPI and needs Python >= 3.12. A PEP 735 dependency group keeps the git pin out of the wheel metadata; the 3.12 and 3.13 test jobs install it so the binder tests and their coverage run in CI. --- .github/workflows/test.yml | 2 +- pyproject.toml | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 89c303af..3e951ee6 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -82,7 +82,7 @@ jobs: - name: Install package and dependencies run: | python -m pip install uv - uv pip install --system "$(ls dist/*.whl)[dev,solvers,oetc]" + uv pip install --system "$(ls dist/*.whl)[dev,solvers,oetc]" --group spec - name: Test with pytest env: diff --git a/pyproject.toml b/pyproject.toml index aa9d4902..cf676fed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -118,6 +118,14 @@ gpu = [ # "cupdlpx>=0.1.2", pip package currently unstable, install manually ] +[dependency-groups] +# math-spec is not on PyPI yet and needs Python >= 3.12. A dependency group +# keeps the git pin out of the published wheel metadata, which PyPI rejects. +# Install with `uv sync --group spec` or `uv pip install --group spec`. +spec = [ + "math-spec @ git+https://github.com/energy-models/math-spec.git@1377f27b759cfbc42bb205338e79751542dabe84 ; python_version >= '3.12'", +] + [tool.uv] # cuopt-cu12 pulls cudf-cu12, which pins pandas<3.0.4, while benchmarks pins # pandas==3.0.5. Resolve the two extras in separate forks instead of together. From 27c9ab5bd18de7ad4a494aa13fef199741027465 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 20:37:01 +0200 Subject: [PATCH 06/35] doc: mention the spec dependency group in the contributing guide --- doc/contributing.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/doc/contributing.rst b/doc/contributing.rst index e0d71cc3..97e47ed6 100644 --- a/doc/contributing.rst +++ b/doc/contributing.rst @@ -45,6 +45,9 @@ To run the test suite: # Install development dependencies uv sync --extra dev --extra solvers + # Also run the math-spec binder tests (needs Python >= 3.12) + uv sync --extra dev --extra solvers --group spec + # Run all tests pytest From 9b7f3adc333b67eceb56ab74783d121f9b62b136 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 14:33:56 +0200 Subject: [PATCH 07/35] 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 08/35] 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 09/35] 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 10/35] 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 11/35] 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 12/35] 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 13/35] 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 14/35] 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" From aafba1656bb4c7e99fd9295a1150b93864fd89e8 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 11:24:03 +0200 Subject: [PATCH 15/35] wip: datarecord + pypsa dependency groups --- pyproject.toml | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index cf676fed..5d34316d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -125,6 +125,22 @@ gpu = [ spec = [ "math-spec @ git+https://github.com/energy-models/math-spec.git@1377f27b759cfbc42bb205338e79751542dabe84 ; python_version >= '3.12'", ] +# datarecord feeds model data as sources into Model.from_spec (adapter showcased +# in dev-scripts). Pre-1.0 git dep, gated on 3.12 like spec. narwhals 2.21.0 has +# a join regression that breaks datarecord's name-uniqueness check, so it is +# excluded until a fix ships. Install with `uv sync --group datarecord`. +datarecord = [ + { include-group = "spec" }, + "datarecord @ git+https://github.com/energy-models/datarecord.git@3b1dd503ace7c9ae12a9adf38f8222e876f09311 ; python_version >= '3.12'", + "narwhals!=2.21.0 ; python_version >= '3.12'", + "pyarrow ; python_version >= '3.12'", +] +# Runs dev-scripts/spec/pypsa_spec_lowering.py: a PyPSA example network lowered +# through math-spec's examples/pypsa.yaml. +pypsa = [ + { include-group = "spec" }, + "pypsa>=1.3 ; python_version >= '3.12'", +] [tool.uv] # cuopt-cu12 pulls cudf-cu12, which pins pandas<3.0.4, while benchmarks pins From e6623afdb7c6dc9234e3a31e26b22f9277febd68 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 19:45:57 +0200 Subject: [PATCH 16/35] fix(spec): restamp CSRConstraint grid after the csr module refactor (#944) --- linopy/spec/netcdf.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/linopy/spec/netcdf.py b/linopy/spec/netcdf.py index e2214b0e..bfcfc805 100644 --- a/linopy/spec/netcdf.py +++ b/linopy/spec/netcdf.py @@ -103,6 +103,7 @@ def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: 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 + from linopy.csr import Grid for _, variable in model.variables.items(): variable._data = _stamped(variable.data, coords) @@ -113,9 +114,12 @@ def _restamp(model: Model, coords: Mapping[str, pd.Index]) -> None: 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 - ] + constraint._grid = Grid( + { + d: coords.get(d, index) + for d, index in constraint._grid.indexes.items() + } + ) def _stamped(data: xr.Dataset, coords: Mapping[str, pd.Index]) -> xr.Dataset: From 0254d47e2e502efa03bc794d5c6e8b80a0cab688 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 19:55:58 +0200 Subject: [PATCH 17/35] fix(spec): bind pandas extension strings as numpy object arrays pandas 3 hands strings over as StringDtype, Arrow-backed when pyarrow is installed. xarray keeps the extension array, refuses it in positional indexing and reports no np.dtype, so the netcdf dtype round trip broke. --- linopy/spec/binder.py | 18 ++++++++++++++++-- test/test_spec_binder.py | 16 ++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py index 56081f47..ec5eb967 100644 --- a/linopy/spec/binder.py +++ b/linopy/spec/binder.py @@ -138,7 +138,7 @@ def parameter(self, name: str) -> xr.DataArray: declared = self._declaration(name) if name not in self._keys: raise SpecDataError(f"no data provided for parameter '{name}'") - arr = _as_array(name, declared, self.sources[name], self.coords) + arr = _numpy(_as_array(name, declared, self.sources[name], self.coords)) onto = {d: self.coords[d] for d in declared.dims} return _aligned(name, arr, onto, _fill(declared)) @@ -300,10 +300,24 @@ def _lookups( series = _lookup_series(lk.name, over, sources[lk.name]) _check_lookup(series, lk, over, coords) padded = series.reindex(coords[over]) - out.setdefault(over, {})[lk.name] = xr.DataArray(padded, name=lk.name) + out.setdefault(over, {})[lk.name] = _numpy(xr.DataArray(padded, name=lk.name)) return out +def _numpy(arr: xr.DataArray) -> xr.DataArray: + """ + *arr* backed by a numpy array. + + xarray keeps a pandas extension array as it arrives, and pandas 3 hands + strings over as one. Its ``dtype`` is no ``np.dtype``, so nothing + downstream that records or restores a dtype can name it, and xarray's + positional indexing refuses the Arrow-backed variant. + """ + if isinstance(arr.dtype, np.dtype): + return arr + return arr.copy(data=arr.to_numpy()) + + def _lookup_series(name: str, over: str, obj: Any) -> pd.Series: if isinstance(obj, xr.DataArray): if obj.dims != (over,) or over not in obj.indexes: diff --git a/test/test_spec_binder.py b/test/test_spec_binder.py index c791eb6e..d5ed8b1d 100644 --- a/test/test_spec_binder.py +++ b/test/test_spec_binder.py @@ -183,6 +183,22 @@ def test_lookup_shapes_bind_alike(program: Any, good: dict[str, Any], grp: Any) assert got.values.tolist() == ["n", "e", "n"] +@pytest.mark.parametrize("storage", ["python", "pyarrow"]) +@pytest.mark.parametrize("shape", ["series", "dataarray"]) +def test_extension_strings_bind_as_numpy_objects( + program: Any, good: dict[str, Any], storage: str, shape: str +) -> None: + if storage == "pyarrow": + pytest.importorskip("pyarrow") + series = pd.Series(["n", "e"], index=F[:2], dtype=pd.StringDtype(storage)) + grp = xr.DataArray(series) if shape == "dataarray" else series + got = bind(program, {**good, "grp": grp}).lookups["f"]["grp"] + assert got.dtype == np.dtype(object) + assert got.values[:2].tolist() == ["n", "e"] + assert pd.isna(got.values[2]) + assert got.sel(f=["b", "a"]).values.tolist() == ["n", "e"] + + def test_missing_rows_become_nan_and_false(program: Any, good: dict[str, Any]) -> None: sparse = { **good, From 3776e66fa20cb31680a173ca4c66d0cc77dab635 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 20:23:58 +0200 Subject: [PATCH 18/35] ci(spec): type-check with math-spec installed; declare pyyaml and pyarrow in the spec group --- .github/workflows/test.yml | 2 +- pyproject.toml | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3e951ee6..234c0db1 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -120,7 +120,7 @@ jobs: - name: Install package and dependencies run: | python -m pip install uv - uv pip install --system "$(ls dist/*.whl)[dev]" + uv pip install --system "$(ls dist/*.whl)[dev]" --group spec - name: Run type checker (mypy) run: | diff --git a/pyproject.toml b/pyproject.toml index 5d34316d..03d4c097 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -124,6 +124,8 @@ gpu = [ # Install with `uv sync --group spec` or `uv pip install --group spec`. spec = [ "math-spec @ git+https://github.com/energy-models/math-spec.git@1377f27b759cfbc42bb205338e79751542dabe84 ; python_version >= '3.12'", + "pyyaml ; python_version >= '3.12'", + "pyarrow ; python_version >= '3.12'", ] # datarecord feeds model data as sources into Model.from_spec (adapter showcased # in dev-scripts). Pre-1.0 git dep, gated on 3.12 like spec. narwhals 2.21.0 has From 0cfbb3a8def9d0ef9d043bade7f060ee1519664d Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 20:23:58 +0200 Subject: [PATCH 19/35] refac(spec): one owner per seam; generic netcdf dtype and coordinate repair moves to io parameters.py owns resolution and derivation, groups.py the axis partition, nodes.amounts_of the parameter-named amounts, Context.lookup the lookups. io records and restores parameter dtypes for every model and owns restamp_coords and the module-level prefix helpers spec/netcdf reuses. --- linopy/io.py | 178 ++++++++++++++++++++++++++------------ linopy/spec/accessor.py | 3 +- linopy/spec/binder.py | 9 +- linopy/spec/builder.py | 9 +- linopy/spec/context.py | 44 ++-------- linopy/spec/coverage.py | 9 +- linopy/spec/groups.py | 68 +++++++++++++++ linopy/spec/netcdf.py | 93 +++++--------------- linopy/spec/nodes.py | 8 ++ linopy/spec/operators.py | 104 ++++++++-------------- linopy/spec/parameters.py | 50 +++++++++++ linopy/spec/terms.py | 27 ------ linopy/spec/where.py | 23 ++--- test/test_io.py | 19 ++++ 14 files changed, 355 insertions(+), 289 deletions(-) create mode 100644 linopy/spec/groups.py create mode 100644 linopy/spec/parameters.py diff --git a/linopy/io.py b/linopy/io.py index 1ddc2e7e..5cc9d001 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -11,7 +11,7 @@ import shutil import time import warnings -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from importlib.metadata import version from io import BufferedWriter from pathlib import Path @@ -45,6 +45,7 @@ logger = logging.getLogger(__name__) NETCDF_VERSION_ATTR = "_linopy_version" +DTYPE_ATTR = "_linopy_dtype" EXPR_TYPE_ATTR = "_linopy_expr_type" SPEC_ATTR = "_linopy_spec" CONTAINER_ORDER_ATTR = "_linopy_{}_order" @@ -1020,6 +1021,124 @@ def non_bool_dict( return {k: int(v) if isinstance(v, bool) else v for k, v in d.items()} +def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: + """*ds* with every dim, coordinate, variable and attribute of it namespaced under *prefix*.""" + to_rename = set([*ds.dims, *ds.coords, *ds]) + ds = ds.rename({d: f"{prefix}-{d}" for d in to_rename}) + ds.attrs = {f"{prefix}-{k}": v for k, v in ds.attrs.items()} + + # Flatten multiindexes + for dim in ds.dims: + if isinstance(ds[dim].to_index(), pd.MultiIndex): + prefix_len = len(prefix) + 1 # leave original index level name + names = [n[prefix_len:] for n in ds[dim].to_index().names] + ds = ds.reset_index(dim) + # scipy netCDF3 backend cannot write unicode-array attrs. + ds.attrs[f"{dim}_multiindex"] = json.dumps(list(names)) + + return ds + + +def has_prefix(k: str, prefix: str) -> bool: + return k.rsplit("-", 1)[0] == prefix + + +def remove_prefix(k: str, prefix: str) -> str: + return k[len(prefix) + 1 :] + + +def parse_multiindex_attr(value: str | Iterable[str]) -> list[str]: + # str = JSON (new); iterable = legacy list from older linopy. + if isinstance(value, str): + return [str(n) for n in json.loads(value)] + return [str(n) for n in value] + + +def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: + """The part of *ds* :func:`with_prefix` wrote under *prefix*, its names given back.""" + ds = ds[[k for k in ds if has_prefix(str(k), prefix)]] + multiindexes = [] + for dim in ds.dims: + attr = ds.attrs.get(f"{dim}_multiindex") + if attr is None: + continue + for name in parse_multiindex_attr(attr): + multiindexes.append(prefix + "-" + name) + ds = ds.drop_vars(set(ds.coords) - set(ds.dims) - set(multiindexes)) + to_rename = set([*ds.dims, *ds.coords, *ds]) + ds = ds.rename({d: remove_prefix(d, prefix) for d in to_rename}) + ds.attrs = { + remove_prefix(k, prefix): v + for k, v in ds.attrs.items() + if has_prefix(k, prefix) + } + + for dim in ds.dims: + if f"{dim}_multiindex" in ds.attrs: + names = parse_multiindex_attr(ds.attrs.pop(f"{dim}_multiindex")) + ds = ds.set_index({dim: names}) # type: ignore[dict-item] + + return ds + + +def record_dtypes(ds: xr.Dataset) -> xr.Dataset: + """ + *ds* with each array's in-memory dtype written as an attribute. + + No netcdf type holds a dtype as written: an engine narrows an int64 to + int32 and hands a bool back as int8, so the dtype travels beside the + values and :func:`restore_dtypes` puts it back. + """ + typed = { + str(name): arr.assign_attrs({DTYPE_ATTR: str(arr.dtype)}) + for name, arr in ds.items() + } + return ds.assign(typed) + + +def restore_dtypes(ds: xr.Dataset) -> xr.Dataset: + """*ds* with each array back at the dtype :func:`record_dtypes` recorded; one written without is left as it is.""" + cast = { + str(name): arr.astype(np.dtype(arr.attrs.pop(DTYPE_ATTR))) + for name, arr in ds.items() + if DTYPE_ATTR in arr.attrs + } + return ds.assign(cast) + + +def restamp_coords(m: Model, coords: Mapping[str, pd.Index]) -> None: + """Put *coords* on every container of *m* that carries one of those dimensions.""" + from linopy.constraints import Constraint, CSRConstraint + from linopy.csr import Grid + + for _, variable in m.variables.items(): + variable._data = _stamped(variable.data, coords) + for _, expression in m.expressions.items(): + expression._data = _stamped(expression.data, coords) + m.objective.expression._data = _stamped(m.objective.expression.data, coords) + for _, constraint in m.constraints.items(): + if isinstance(constraint, Constraint): + constraint._data = _stamped(constraint.data, coords) + elif isinstance(constraint, CSRConstraint): + constraint._grid = Grid( + { + d: coords.get(d, index) + for d, index in constraint._grid.indexes.items() + } + ) + + +def _stamped(data: xr.Dataset, coords: Mapping[str, pd.Index]) -> xr.Dataset: + """*data* with *coords* in place of the ones a dtype narrowed.""" + 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 to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: """ Write out the model to a netcdf file. @@ -1067,22 +1186,6 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: stacklevel=2, ) - def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: - to_rename = set([*ds.dims, *ds.coords, *ds]) - ds = ds.rename({d: f"{prefix}-{d}" for d in to_rename}) - ds.attrs = {f"{prefix}-{k}": v for k, v in ds.attrs.items()} - - # Flatten multiindexes - for dim in ds.dims: - if isinstance(ds[dim].to_index(), pd.MultiIndex): - prefix_len = len(prefix) + 1 # leave original index level name - names = [n[prefix_len:] for n in ds[dim].to_index().names] - ds = ds.reset_index(dim) - # scipy netCDF3 backend cannot write unicode-array attrs. - ds.attrs[f"{dim}_multiindex"] = json.dumps(list(names)) - - return ds - vars = [ with_prefix(var.data, f"variables-{name}") for name, var in m.variables.items() ] @@ -1113,7 +1216,7 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: parameters, spec_ds = encode(m._spec) specs = [spec_ds] - params = [with_prefix(parameters, "parameters")] + params = [with_prefix(record_dtypes(parameters), "parameters")] scalars = {k: getattr(m, k) for k in m.scalar_attrs} ds = xr.merge( @@ -1191,43 +1294,6 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: m = Model() ds = xr.load_dataset(path, **kwargs) - def has_prefix(k: str, prefix: str) -> bool: - return k.rsplit("-", 1)[0] == prefix - - def remove_prefix(k: str, prefix: str) -> str: - return k[len(prefix) + 1 :] - - def parse_multiindex_attr(value: str | Iterable[str]) -> list[str]: - # str = JSON (new); iterable = legacy list from older linopy. - if isinstance(value, str): - return [str(n) for n in json.loads(value)] - return [str(n) for n in value] - - def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: - ds = ds[[k for k in ds if has_prefix(str(k), prefix)]] - multiindexes = [] - for dim in ds.dims: - attr = ds.attrs.get(f"{dim}_multiindex") - if attr is None: - continue - for name in parse_multiindex_attr(attr): - multiindexes.append(prefix + "-" + name) - ds = ds.drop_vars(set(ds.coords) - set(ds.dims) - set(multiindexes)) - to_rename = set([*ds.dims, *ds.coords, *ds]) - ds = ds.rename({d: remove_prefix(d, prefix) for d in to_rename}) - ds.attrs = { - remove_prefix(k, prefix): v - for k, v in ds.attrs.items() - if has_prefix(k, prefix) - } - - for dim in ds.dims: - if f"{dim}_multiindex" in ds.attrs: - names = parse_multiindex_attr(ds.attrs.pop(f"{dim}_multiindex")) - ds = ds.set_index({dim: names}) # type: ignore[dict-item] - - return ds - def container_names(kind: str) -> list[str]: found = {str(k).rsplit("-", 1)[0] for k in ds if str(k).startswith(kind)} order_attr = ds.attrs.get(CONTAINER_ORDER_ATTR.format(kind)) @@ -1293,7 +1359,7 @@ def container_names(kind: str) -> list[str]: ) m.objective._value = objective.attrs.pop("value", None) - m.parameters = get_prefix(ds, "parameters") + m.parameters = restore_dtypes(get_prefix(ds, "parameters")) if SPEC_ATTR in ds.attrs: from linopy.spec.netcdf import decode diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index cd8d5855..621a7a23 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -33,8 +33,9 @@ from linopy.spec import terms from linopy.spec.binder import Bound, Retain, bind from linopy.spec.builder import build, evaluate_named, fold -from linopy.spec.context import Context, Parameters, Resolve +from linopy.spec.context import Context from linopy.spec.errors import SpecDataError +from linopy.spec.parameters import Parameters, Resolve SpecLike: TypeAlias = str | Path | Mapping[str, Any] | Spec diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py index ec5eb967..b1b0e5f0 100644 --- a/linopy/spec/binder.py +++ b/linopy/spec/binder.py @@ -24,7 +24,7 @@ from math_spec import program as ms from linopy.spec.errors import SpecDataError -from linopy.spec.nodes import parameters_of, walk +from linopy.spec.nodes import amounts_of, parameters_of, walk Retain = Literal["report", "all", "none"] _RETAIN: tuple[str, ...] = get_args(Retain) @@ -177,11 +177,8 @@ def _report_closure(program: ms.Program) -> set[str]: bodies = tuple(program.named_expressions.values()) 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.Cases): + names.update(amounts_of(node)) + if isinstance(node, ms.Cases): for region in node.regions: names |= region.when.names_read return names & set(program.parameters) diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py index 431e4de1..a4ae4a48 100644 --- a/linopy/spec/builder.py +++ b/linopy/spec/builder.py @@ -22,7 +22,7 @@ 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.context import Context from linopy.spec.coverage import ( check_bounds_cover, check_coefficients_cover, @@ -30,8 +30,9 @@ check_divisors_cover, ) from linopy.spec.errors import SpecDataError +from linopy.spec.parameters import Parameters from linopy.spec.terms import Array, Term, Value -from linopy.spec.where import as_linopy_mask, bound_lookup, evaluate_where +from linopy.spec.where import as_linopy_mask, evaluate_where from linopy.variables import Variable _SIGN = {"==": "=", "<=": "<=", ">=": ">="} @@ -323,11 +324,11 @@ def _partition(node: ms.Translate | ms.Window, ctx: Context) -> xr.DataArray | N """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) + array = ctx.lookup(node.partition, node.dimension) return array.rename(ctx.program.dimension(node.dimension).targets[node.partition]) def _lookup_arrays( over: str, names: tuple[str, ...], ctx: Context ) -> tuple[xr.DataArray, ...]: - return tuple(bound_lookup(name, over, ctx.lookups) for name in names) + return tuple(ctx.lookup(name, over) for name in names) diff --git a/linopy/spec/context.py b/linopy/spec/context.py index 235e8b10..ea9d53c6 100644 --- a/linopy/spec/context.py +++ b/linopy/spec/context.py @@ -1,8 +1,8 @@ -"""The data an evaluation reads: parameters resolved once, and the model, coordinates and lookups beside them.""" +"""The data an evaluation reads: the parameters, and the model, coordinates and lookups beside them.""" from __future__ import annotations -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Mapping from dataclasses import dataclass, field, replace import pandas as pd @@ -10,40 +10,6 @@ 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) @@ -61,10 +27,14 @@ class Context: program: ms.Program coords: Mapping[str, pd.Index] lookups: Mapping[str, Mapping[str, xr.DataArray]] - parameters: Parameters + parameters: Mapping[str, xr.DataArray] solved: bool = field(default=False) @property def unsolved(self) -> Context: """The same context with the fold's switch off, so a variable enters as its linopy term.""" return replace(self, solved=False) + + def lookup(self, name: str, over: str) -> xr.DataArray: + """The lookup *name* as an array over *over*, NaN where a label is unmapped.""" + return self.lookups[over][name] diff --git a/linopy/spec/coverage.py b/linopy/spec/coverage.py index 06b47da2..b74540b7 100644 --- a/linopy/spec/coverage.py +++ b/linopy/spec/coverage.py @@ -20,7 +20,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.nodes import amounts_of, children, parameters_of from linopy.spec.where import evaluate_where Rows = xr.DataArray | None @@ -144,10 +144,9 @@ def _coefficient_uses( """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 + return + for name in amounts_of(node): + yield name, None def _under_regions( diff --git a/linopy/spec/groups.py b/linopy/spec/groups.py new file mode 100644 index 00000000..c39a4c6b --- /dev/null +++ b/linopy/spec/groups.py @@ -0,0 +1,68 @@ +"""How a lookup partitions an axis: the shape every group-wise operator reads.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import xarray as xr + + +def unmapped(key: object) -> bool: + """Whether a lookup left this member in no group: ``None``, or the NaN that never equals itself.""" + return key is None or key != key + + +@dataclass(frozen=True) +class Groups: + labels: np.ndarray + grouped: xr.DataArray + belongs: xr.DataArray + within: xr.DataArray + size: xr.DataArray + roster: np.ndarray + names: tuple[object, ...] + counts: tuple[int, ...] + + +def grouped(over: str, labels: np.ndarray, groups: xr.DataArray) -> Groups: + """ + How the lookup *groups* partitions the axis *over*. + + A coordinate the lookup sends nowhere belongs to no group: its ``within`` + is 0, its ``size`` 1 and its ``grouped`` False. + """ + keys = np.asarray(groups.sel({over: labels}).values, dtype=object) + peers: dict[object, list[int]] = {} + within = np.zeros(len(labels), dtype=int) + held = np.zeros(len(labels), dtype=bool) + for k, key in enumerate(keys): + if unmapped(key): + continue + held[k] = True + beside = peers.setdefault(key, []) + within[k] = len(beside) + beside.append(k) + order = {key: g for g, key in enumerate(peers)} + widest = max((len(beside) for beside in peers.values()), default=1) + roster = np.zeros((max(len(peers), 1), widest), dtype=int) + for key, beside in peers.items(): + roster[order[key], : len(beside)] = beside + belongs = np.array([order.get(key, 0) for key in keys], dtype=int) + span = np.array( + [len(peers[key]) if inside else 1 for key, inside in zip(keys, held)], dtype=int + ) + + def on_axis(values: np.ndarray) -> xr.DataArray: + return xr.DataArray(values, coords={over: labels}, dims=[over]) + + return Groups( + labels, + on_axis(held), + on_axis(belongs), + on_axis(within), + on_axis(span), + roster, + tuple(peers), + tuple(len(beside) for beside in peers.values()), + ) diff --git a/linopy/spec/netcdf.py b/linopy/spec/netcdf.py index bfcfc805..9cfcfc22 100644 --- a/linopy/spec/netcdf.py +++ b/linopy/spec/netcdf.py @@ -6,10 +6,9 @@ coordinates and the lookups; the program is re-lowered from the text on read, so no lowered ``Program`` ever reaches the file. -No netcdf type holds a dtype as written. An engine narrows an int64 to -int32, hands a bool back as int8 and a string array back as `` tuple[xr.Dataset, xr.Dataset]: for name in _coded(spec): arrays.update(_encode(name, parameters[name])) parameters = parameters.drop_vars(name) - 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 + written = with_prefix(xr.Dataset(arrays), PREFIX).assign_attrs( + {SPEC_ATTR: spec.text} + ) + return parameters, written def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: @@ -81,7 +82,7 @@ def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: decoded arrays they are the dataset :func:`linopy.spec.accessor.attach` left on the model when it was built. """ - sub = _unprefixed(ds) + sub = get_prefix(ds, PREFIX) coords = { _stripped(name, COORD): _index(sub[name]) for name in sub.data_vars @@ -92,47 +93,11 @@ def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: for name in sub.data_vars if str(name).startswith(CODES) } - 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) + model.parameters = model.parameters.assign_coords(coords).assign(coded) + restamp_coords(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 - from linopy.csr import Grid - - 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._grid = Grid( - { - d: coords.get(d, index) - for d, index in constraint._grid.indexes.items() - } - ) - - -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 objects.""" lookups = {name for by_name in spec.lookups.values() for name in by_name} @@ -159,7 +124,7 @@ def _encode(name: str, arr: xr.DataArray) -> dict[str, xr.DataArray]: def _decode(sub: xr.Dataset, name: str, coords: dict[str, pd.Index]) -> xr.DataArray: codes = sub[CODES + name] - dtype = np.dtype(codes.attrs[DTYPE]) + dtype = np.dtype(codes.attrs[DTYPE_ATTR]) categories = _categories(sub, name, dtype) positions = codes.to_numpy().astype(int) mapped = positions >= 0 @@ -187,24 +152,12 @@ 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: - 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)}) + return xr.DataArray( + values, dims=dims, attrs={DTYPE_ATTR: dtype or str(values.dtype)} + ) def _stripped(name: Any, prefix: str) -> str: @@ -213,7 +166,7 @@ def _stripped(name: Any, prefix: str) -> str: 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])) + return arr.to_numpy().astype(np.dtype(arr.attrs[DTYPE_ATTR])) def _index(arr: xr.DataArray) -> pd.Index: diff --git a/linopy/spec/nodes.py b/linopy/spec/nodes.py index a12c3ecb..68a93f16 100644 --- a/linopy/spec/nodes.py +++ b/linopy/spec/nodes.py @@ -21,6 +21,14 @@ def walk(*nodes: ms.ExpressionNode) -> Iterator[ms.ExpressionNode]: yield from walk(*children(node)) +def amounts_of(node: ms.ExpressionNode) -> Iterator[str]: + """The parameters *node* names as an amount: a translation's offset or a window's width.""" + if isinstance(node, ms.Translate) and isinstance(node.offset, str): + yield node.offset + elif isinstance(node, ms.Window) and isinstance(node.width, str): + yield node.width + + def parameters_of(*nodes: ms.ExpressionNode) -> frozenset[str]: """Every parameter named anywhere under *nodes*.""" return frozenset(n.name for n in walk(*nodes) if isinstance(n, ms.Parameter)) diff --git a/linopy/spec/operators.py b/linopy/spec/operators.py index 06a5b0e3..17fd2d46 100644 --- a/linopy/spec/operators.py +++ b/linopy/spec/operators.py @@ -20,12 +20,35 @@ import xarray as xr from linopy.expressions import LinearExpression -from linopy.spec import terms +from linopy.spec.groups import Groups, grouped from linopy.spec.terms import Array, Term +from linopy.variables import Variable Amount = int | xr.DataArray +def filled(expression: Array, fill: float) -> Array: + """*expression* with every absence in it standing as *fill*.""" + if isinstance(expression, Variable): + expression = expression.to_linexpr() + return expression.fillna(fill) + + +def vacated( + shifted: Array, operand: Array, over: str, vacated: xr.DataArray, fill: float +) -> Array: + """ + *shifted*, with the positions the shift vacated filled, and only those. + + The fill lands where the shift vacated and the operand carries the + coordinate; every other slot keeps the absence it arrived with, so no row + is invented at a coordinate the operand never had. + """ + carried = (~operand.isnull()).any(over) + keep = carried & (~shifted.isnull() | vacated) + return filled(shifted, fill).where(keep) + + def sum_over(array: Array, over: str) -> Array: """Sum *array* over *over*; a term beside an empty dimension is built as the constant zero.""" if not isinstance(array, xr.DataArray) and any( @@ -125,8 +148,8 @@ def shift( """ 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) + partition = grouped(over, np.asarray(array.indexes[over]), by) + return _gather_in_groups(array, over, _per_group(offset, by), partition, edge) if isinstance(offset, xr.DataArray) and offset.ndim: return _gather_by_offset(array, over, offset, edge) amount: dict[Hashable, int] = {over: int(offset)} @@ -139,9 +162,7 @@ def shift( shifted = array.shift(amount) if fill is None: return shifted - return terms.vacated( - shifted, array, over, _off_the_axis(array, over, amount[over]), fill - ) + return vacated(shifted, array, over, _off_the_axis(array, over, amount[over]), fill) def sum_back( @@ -163,16 +184,18 @@ def sum_back( 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) + partition = ( + None if by is None else grouped(over, np.asarray(array.indexes[over]), by) + ) lagged_terms: list[Array] = [] reached: list[xr.DataArray] = [] for lag in range(widest): lagged = ( _gather_by_offset(array, over, lag, probe) - if groups is None - else _gather_in_groups(array, over, lag, groups, probe) + if partition is None + else _gather_in_groups(array, over, lag, partition, probe) ) - live, term = ~lagged.isnull(), terms.filled(lagged, 0.0) + live, term = ~lagged.isnull(), filled(lagged, 0.0) if isinstance(within, xr.DataArray): live, term = live & (within > lag), term * (within > lag).astype(float) lagged_terms.append(term) @@ -230,7 +253,7 @@ def gathered(ordinals: xr.DataArray) -> Array: 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) + return vacated(moved, array, over, ~inside, edge.fill) def _per_group(offset: Amount, groups: xr.DataArray) -> Amount: @@ -241,63 +264,8 @@ def _per_group(offset: Amount, groups: xr.DataArray) -> Amount: 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: Array, over: str, offset: Amount, groups: Groups, edge: _Edge ) -> Array: """ Translate *array* inside each group rather than along the axis. @@ -322,7 +290,7 @@ def peer(group: np.ndarray, position: np.ndarray) -> np.ndarray: ) if edge.fill is None: return gathered - return terms.vacated(gathered, array, over, groups.grouped & ~inside, edge.fill) + return vacated(gathered, array, over, groups.grouped & ~inside, edge.fill) def _off_the_axis(array: Array, over: str, offset: int) -> xr.DataArray: diff --git a/linopy/spec/parameters.py b/linopy/spec/parameters.py new file mode 100644 index 00000000..afc7fd1b --- /dev/null +++ b/linopy/spec/parameters.py @@ -0,0 +1,50 @@ +""" +Every parameter of a program by name, resolved once. + +A declared parameter is resolved from the caller's data and aligned by the +binder; one a ``piecewise:`` expansion emitted is derived from the block's own +breakpoints. Which of the two a name is, is the declaration's answer, and this +module is where it is asked. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping + +import xarray as xr +from math_spec import program as ms + +from linopy.spec import curves + +Resolve = Callable[[str], xr.DataArray] + + +class Parameters(Mapping[str, xr.DataArray]): + """ + Every parameter of a program by name, each resolved on first read and then held. + + A declared parameter comes from *resolve*; one a ``piecewise:`` expansion + emitted is derived from the block's own breakpoints the way its + derivation says, so a caller never supplies it. + """ + + def __init__(self, program: ms.Program, resolve: Resolve) -> None: + self._program = program + self._resolve = resolve + self._arrays: dict[str, xr.DataArray] = {} + + def __getitem__(self, name: str) -> xr.DataArray: + if name not in self._arrays: + derivation = self._program.parameter(name).derivation + self._arrays[name] = ( + self._resolve(name) + if derivation is None + else curves.derive(derivation, self, self._program) + ) + return self._arrays[name] + + def __iter__(self) -> Iterator[str]: + return iter(self._program.parameters) + + def __len__(self) -> int: + return len(self._program.parameters) diff --git a/linopy/spec/terms.py b/linopy/spec/terms.py index 2b7b8b48..2432122e 100644 --- a/linopy/spec/terms.py +++ b/linopy/spec/terms.py @@ -24,11 +24,6 @@ def present(variable: Variable) -> xr.DataArray: 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 @@ -42,25 +37,3 @@ def solution(variable: Variable, absence: str) -> xr.DataArray: 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 index ab21fc07..fc301d0c 100644 --- a/linopy/spec/where.py +++ b/linopy/spec/where.py @@ -3,7 +3,7 @@ from __future__ import annotations import operator -from collections.abc import Callable, Mapping +from collections.abc import Callable from typing import assert_never import numpy as np @@ -13,7 +13,7 @@ from linopy.spec import terms from linopy.spec.context import Context from linopy.spec.errors import SpecDataError -from linopy.spec.operators import _grouped +from linopy.spec.groups import grouped _PREDICATE_OPS: dict[str, Callable[..., xr.DataArray]] = { "==": operator.eq, @@ -39,13 +39,6 @@ def as_linopy_mask(mask: xr.DataArray) -> xr.DataArray | 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. @@ -75,18 +68,18 @@ def _node(node: ms.WhereNode, ctx: Context) -> xr.DataArray: if isinstance(node, ms.DimensionPositionNode): return _position(node, ctx) if isinstance(node, ms.LookupComparisonNode): - arr = bound_lookup(node.name, node.over, ctx.lookups) + arr = ctx.lookup(node.name, node.over) compared = _PREDICATE_OPS[node.op](arr, node.value) & arr.notnull() return compared.fillna(False).astype(bool) if isinstance(node, ms.LookupPairComparisonNode): - left = bound_lookup(node.name, node.over, ctx.lookups) - right = bound_lookup(node.other, node.over, ctx.lookups) + left = ctx.lookup(node.name, node.over) + right = ctx.lookup(node.other, node.over) compared = ( _PREDICATE_OPS[node.op](left, right) & left.notnull() & right.notnull() ) return compared.fillna(False).astype(bool) if isinstance(node, ms.LookupDefinedNode): - return bound_lookup(node.name, node.over, ctx.lookups).notnull() + return ctx.lookup(node.name, node.over).notnull() if isinstance(node, ms.NotNode): return ~_node(node.operand, ctx) if isinstance(node, ms.AndNode): @@ -108,7 +101,7 @@ def _defined(arr: xr.DataArray, dtype: str) -> xr.DataArray: 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) + groups = ctx.lookup(node.by, node.name) arr = _group_offsets(node, groups, np.asarray(labels)) compared = _PREDICATE_OPS[node.op](arr, 0) & arr.notnull() return compared.fillna(False).astype(bool) @@ -129,7 +122,7 @@ 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) + 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 diff --git a/test/test_io.py b/test/test_io.py index 825ca16a..0317cba1 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -120,6 +120,25 @@ def test_model_to_netcdf(model: Model, tmp_path: Path) -> None: assert_model_equal(m, p) +@pytest.mark.parametrize("engine", ["netcdf4", "scipy"]) +def test_model_to_netcdf_keeps_parameter_dtypes( + model: Model, tmp_path: Path, engine: str +) -> None: + if engine == "netcdf4" and not HAS_NETCDF4: + pytest.skip("needs the netCDF4 backend") + model.parameters["count"] = xr.DataArray( + np.array([1, 2, 3, 4], dtype=np.int64), dims=["x"] + ) + model.parameters["flag"] = xr.DataArray(np.array([True, False]), dims=["y"]) + fn = tmp_path / f"dtypes-{engine}.nc" + model.to_netcdf(fn, engine=engine) + p = read_netcdf(fn) + + for name in ("count", "flag"): + assert p.parameters[name].dtype == model.parameters[name].dtype + assert p.parameters[name].equals(model.parameters[name]) + + @pytest.fixture def unsorted_model() -> Model: m = Model() From 966e0b55087f711552550045f3d8c08c559da907 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 20:53:51 +0200 Subject: [PATCH 20/35] refac(spec): evaluator in its own module; coverage checked in one tree walk evaluate.py holds the recursive evaluator, builder.py the declarations. check_coverage collects divisor, constant-side and coefficient obligations in one walk, so cases: masks are evaluated once per declaration. Public docstrings in numpy style so the API pages render. --- linopy/spec/accessor.py | 48 +++++--- linopy/spec/binder.py | 68 ++++++----- linopy/spec/builder.py | 227 +++--------------------------------- linopy/spec/coverage.py | 251 ++++++++++++++++++++++------------------ linopy/spec/curves.py | 14 ++- linopy/spec/evaluate.py | 198 +++++++++++++++++++++++++++++++ 6 files changed, 429 insertions(+), 377 deletions(-) create mode 100644 linopy/spec/evaluate.py diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index 621a7a23..2c6aefb0 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -32,9 +32,10 @@ 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, evaluate_named, fold +from linopy.spec.builder import build from linopy.spec.context import Context from linopy.spec.errors import SpecDataError +from linopy.spec.evaluate import evaluate_named, fold from linopy.spec.parameters import Parameters, Resolve SpecLike: TypeAlias = str | Path | Mapping[str, Any] | Spec @@ -49,11 +50,14 @@ def attach( """ 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. + 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( @@ -98,9 +102,12 @@ 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. + 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: @@ -170,9 +177,11 @@ def evaluate( ``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. + Raises + ------ + SpecDataError + *sources* label a dimension differently than the + model was built on. """ bound = bind(self.program, sources, retain="none") coords = self.coords @@ -239,8 +248,10 @@ class NamedExpression: 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. + Attributes + ---------- + node + The lowered expression body, math-spec's own AST handle. """ def __init__(self, spec: ModelSpec, name: str, ctx: Context) -> None: @@ -270,9 +281,12 @@ 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. + 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) diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py index b1b0e5f0..0f580e4e 100644 --- a/linopy/spec/binder.py +++ b/linopy/spec/binder.py @@ -70,19 +70,25 @@ def bind( """ Bind *sources* to *program*: master coordinates now, parameters on demand. - Args: - program: The lowered spec. - sources: Data keyed by declared name. Any mapping works; it is read by - key and never iterated beyond ``sources.keys()``. An ``xr.Dataset`` - is accepted too: its indexes are dimension sources, its data - variables parameters and lookups. - retain: Which parameters :meth:`Bound.retained` persists. - - Raises: - SpecDataError: A ``retain`` outside its three values, a key naming - nothing the spec declares, a reached dimension or a lookup with no - source, a duplicated dimension member, or a lookup breaking the - rules a map has. + Parameters + ---------- + program + The lowered spec. + sources + Data keyed by declared name. Any mapping works; it is read by + key and never iterated beyond ``sources.keys()``. An ``xr.Dataset`` + is accepted too: its indexes are dimension sources, its data + variables parameters and lookups. + retain + Which parameters :meth:`Bound.retained` persists. + + Raises + ------ + SpecDataError + A ``retain`` outside its three values, a key naming + nothing the spec declares, a reached dimension or a lookup with no + source, a duplicated dimension member, or a lookup breaking the + rules a map has. """ if retain not in _RETAIN: raise SpecDataError( @@ -102,15 +108,21 @@ class Bound: """ A program bound to its data. - Attributes: - program: The lowered spec the data is bound to. - coords: Master coordinates by dimension, in source order, each index - named after its dimension. A declared dimension nothing reaches - and nothing supplies is absent. - lookups: By dimension, by lookup name, the map as an array over the - dimension's master coordinates, NaN where a label is unmapped. - retain: Which parameters :meth:`retained` persists. - sources: The caller's data, read by key on demand. + Attributes + ---------- + program + The lowered spec the data is bound to. + coords + Master coordinates by dimension, in source order, each index + named after its dimension. A declared dimension nothing reaches + and nothing supplies is absent. + lookups + By dimension, by lookup name, the map as an array over the + dimension's master coordinates, NaN where a label is unmapped. + retain + Which parameters :meth:`retained` persists. + sources + The caller's data, read by key on demand. """ program: ms.Program @@ -129,11 +141,13 @@ def parameter(self, name: str) -> xr.DataArray: master coordinates, leaving NaN (``False`` for ``bool``) where no row was supplied. - Raises: - SpecDataError: No data, a shape no reader accepts, a rank other - than declared, a label its dimension lacks, two rows for one - coordinate, a null value in a row, or values of another type - than declared. + Raises + ------ + SpecDataError + No data, a shape no reader accepts, a rank other + than declared, a label its dimension lacks, two rows for one + coordinate, a null value in a row, or values of another type + than declared. """ declared = self._declaration(name) if name not in self._keys: diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py index a4ae4a48..ac7a54c1 100644 --- a/linopy/spec/builder.py +++ b/linopy/spec/builder.py @@ -1,37 +1,26 @@ """ -Program plus bound data to linopy declarations, and a named expression to its value. +Program plus bound data to linopy declarations. -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. +A build hands every variable to linopy as its term, then adds special-ordered +sets, constraints and the objective; which linopy call each construct becomes +is one branch of :func:`linopy.spec.evaluate.evaluate`. """ from __future__ import annotations -import 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 import curves from linopy.spec.binder import Bound from linopy.spec.context import Context -from linopy.spec.coverage import ( - check_bounds_cover, - check_coefficients_cover, - check_constant_side_covers, - check_divisors_cover, -) +from linopy.spec.coverage import check_bounds_cover, check_coverage from linopy.spec.errors import SpecDataError +from linopy.spec.evaluate import carried, evaluate from linopy.spec.parameters import Parameters -from linopy.spec.terms import Array, Term, Value +from linopy.spec.terms import Term, Value from linopy.spec.where import as_linopy_mask, evaluate_where from linopy.variables import Variable @@ -62,44 +51,7 @@ def build(model: Model, bound: Bound) -> None: _constraints(ctx) _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 evaluate_named(name: str, ctx: Context) -> Value: - """The named expression *name* as its linopy term, array or number over *ctx*, its divisors checked first.""" - if name not in ctx.program.named_expressions: - raise KeyError( - f"unknown named expression '{name}'. " - + did_you_mean(name, ctx.program.named_expressions) - ) - body = ctx.program.named_expressions[name] - 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): - return value - if isinstance(value, float | int): - return xr.DataArray(float(value), name=name) - raise TypeError( - f"expression '{name}' folded to a {type(value).__name__}, not to data" - ) - - -def _named(value: xr.DataArray, name: str) -> xr.DataArray: - """*value* with its stray non-dimension coordinates dropped and renamed to *name*.""" - stray = [c for c in value.coords if c not in value.dims] - return value.drop_vars(stray).rename(name) - - -# --------------------------------------------------------------------------- -# declarations -# --------------------------------------------------------------------------- + check_coverage(f"expression '{name}'", (body,), ctx, None) def _variables(ctx: Context) -> None: @@ -140,15 +92,15 @@ 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) - check_coefficients_cover(f"constraint '{name}'", (row.lhs, row.rhs), ctx, mask) + check_coverage( + f"constraint '{name}'", (row.lhs, row.rhs), ctx, mask, comparison=True + ) lhs, rhs = evaluate(row.lhs, ctx), evaluate(row.rhs, ctx) if _term_free(lhs) and _term_free(rhs): continue term, other, sense = _sides(lhs, rhs, row.sense) if isinstance(other, xr.DataArray): - term, other = _carried(term, other) + term, other = carried(term, other) ctx.model.add_constraints(term, _SIGN[sense], other, name=name, mask=mask) @@ -174,161 +126,10 @@ def _objective(ctx: Context) -> None: declared = ctx.program.objective if declared is None: return - check_divisors_cover("the objective", (declared.expression,), ctx, None) - check_coefficients_cover("the objective", (declared.expression,), ctx, None) + check_coverage("the objective", (declared.expression,), ctx, None) expr = evaluate(declared.expression, ctx) if not isinstance(expr, Variable | LinearExpression | QuadraticExpression): raise SpecDataError( "the objective carries no variable term once the data is 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(lambda a, b: _combine(operator.add, a, b), regions) - assert_never(node) - - -def _variable(name: str, ctx: Context) -> Value: - variable = ctx.model.variables[name] - absence = ctx.program.variable(name).absence - if not ctx.solved: - return terms.variable_term(variable, absence) - if "solution" not in variable.data: - raise RuntimeError( - f"variable '{name}' has no solution yet: solve the model before reading a named expression" - ) - return terms.solution(variable, absence) - - -def _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 = ctx.lookup(node.partition, node.dimension) - return array.rename(ctx.program.dimension(node.dimension).targets[node.partition]) - - -def _lookup_arrays( - over: str, names: tuple[str, ...], ctx: Context -) -> tuple[xr.DataArray, ...]: - return tuple(ctx.lookup(name, over) for name in names) diff --git a/linopy/spec/coverage.py b/linopy/spec/coverage.py index b74540b7..401c6898 100644 --- a/linopy/spec/coverage.py +++ b/linopy/spec/coverage.py @@ -12,7 +12,8 @@ from __future__ import annotations -from collections.abc import Iterator +from collections.abc import Sequence +from dataclasses import dataclass, field import xarray as xr from math_spec import program as ms @@ -24,6 +25,16 @@ from linopy.spec.where import evaluate_where Rows = xr.DataArray | None +Obligation = tuple[str, Rows] + + +@dataclass +class Obligations: + """Every parameter use under a declaration, each with the rows it has to cover, gathered in one walk.""" + + divisors: list[Obligation] = field(default_factory=list) + constants: list[Obligation] = field(default_factory=list) + coefficients: list[Obligation] = field(default_factory=list) def gaps_under(array: xr.DataArray, rows: Rows) -> int: @@ -34,132 +45,144 @@ def gaps_under(array: xr.DataArray, rows: Rows) -> int: return int(missing.sum()) -def check_bounds_cover( - name: str, declared: ms.VariableDeclaration, ctx: Context, rows: Rows +def check_coverage( + subject: str, + expressions: Sequence[ms.ExpressionNode], + ctx: Context, + rows: Rows, + *, + comparison: bool = False, ) -> None: - """A bound parameter must have a value at every coordinate the variable occupies.""" - names = sorted(parameters_of(declared.lower, declared.upper)) - missing = sum(gaps_under(ctx.parameters[p], rows) for p in names) - if missing: - raise SpecDataError( - f"variable '{name}': {missing} rows have NULL bounds, a bound parameter is missing " - f"values for some coordinates. The two ways out build different models, so neither " - f"is picked:\n" - f" supply the value the variable exists there, bounded (`inf` is a value)\n" - f' where: "" the variable does not exist there at all' - ) + """ + Refuse *subject* if a parameter it reads leaves a row it builds uncovered. + One walk over *expressions* collects what every parameter has to cover, + narrowed at each ``cases:`` region; divisors are judged first, then, for a + *comparison*, the side without a variable term, then every coefficient. + """ + found = obligations_of(expressions, ctx, rows, comparison=comparison) + check_divisors(subject, found.divisors, ctx) + check_constant_sides(subject, found.constants, ctx) + check_coefficients(subject, found.coefficients, ctx) + + +def obligations_of( + expressions: Sequence[ms.ExpressionNode], + ctx: Context, + rows: Rows, + *, + comparison: bool = False, +) -> Obligations: + """What the parameters under *expressions* have to cover, a side of a *comparison* without a variable being its constant side.""" + found = Obligations() + for expression in expressions: + constant = comparison and not ms.carries_variable(expression) + _collect(expression, ctx, rows, constant, found) + return found -def 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 + +def _collect( + node: ms.ExpressionNode, + ctx: Context, + rows: Rows, + constant: bool, + into: Obligations, ) -> None: + if isinstance(node, ms.Divide): + into.divisors.extend(_divisor_uses(node, ctx, rows)) + if isinstance(node, ms.Parameter): + if constant: + into.constants.append((node.name, rows)) + into.coefficients.append((node.name, rows)) + into.coefficients.extend((name, None) for name in amounts_of(node)) + if isinstance(node, ms.Cases): + for region in node.regions: + inside = evaluate_where(region.when, ctx) + narrowed = inside if rows is None else rows & inside + _collect(region.value, ctx, narrowed, constant, into) + return + for child in children(node): + _collect(child, ctx, rows, constant, into) + + +def _divisor_uses(quotient: ms.Divide, ctx: Context, rows: Rows) -> list[Obligation]: + """Each parameter in the divisor, with the rows the quotient is divided over: the region, narrowed by the presence of every numerator variable.""" + params = parameters_of(quotient.divisor) + if not params: + return [] + needed = rows + for variable in sorted(ms.variables_of(quotient.numerator)): + present = terms.present(ctx.model.variables[variable]) + needed = present if needed is None else needed & present + return [(param, needed) for param in sorted(params)] + + +def check_divisors(subject: str, found: Sequence[Obligation], ctx: Context) -> None: """ A divisor must have a value wherever *subject* divides by it. - 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. + 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 = 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 check_coefficients_cover( - subject: str, expressions: tuple[ms.ExpressionNode, ...], ctx: Context, rows: Rows + for param, needed in found: + missing = gaps_under(ctx.parameters[param], needed) + if missing: + raise SpecDataError( + f"{subject}: parameter '{param}' is used as a divisor but covers {missing} " + f"fewer coordinates than it is divided over. A missing row means a zero " + f"coefficient everywhere else, and zero is not a divisor: the term would drop " + f"and the row would silently stop constraining.\n" + f" Supply the missing rows, or mask the coordinates out with a where." + ) + + +def check_constant_sides( + subject: str, found: Sequence[Obligation], ctx: Context ) -> None: + """A comparison's constant side must have values wherever the row is built, or the zero is the bound.""" + for param, needed in sorted(found, key=lambda pair: pair[0]): + missing = gaps_under(ctx.parameters[param], needed) + if missing: + raise SpecDataError( + f"{subject}: parameter '{param}' covers {missing} fewer coordinates " + f"than the rows built here. A missing row is read as 0, and on the constant side " + f"that zero is a bound rather than an absence: the row still exists, and it binds.\n" + f" Supply the missing rows, if the value is what was meant.\n" + f" Mask them out with a where, if the row should not exist there." + ) + + +def check_coefficients(subject: str, found: Sequence[Obligation], ctx: Context) -> None: """ A coefficient parameter must reach every row it is built over. A missing coefficient row would otherwise read as a zero, dropping its term - while the row stays. 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. + while the row stays. A shift offset or window width given by name is a + coefficient too, and stands or falls over its own coordinates. """ - for 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 - return - for name in amounts_of(node): - yield name, None + for param, needed in found: + missing = gaps_under(ctx.parameters[param], needed) + if missing: + raise SpecDataError( + f"{subject}: parameter '{param}' is used as a coefficient but leaves " + f"{missing} of the rows built here uncovered. A missing row reads as a zero " + f"coefficient, dropping the term while the row stays.\n" + f" Supply the missing rows, if a value other than 0 was meant.\n" + f" Mask them out with a where, if the row should not exist there." + ) -def _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 children(node): - yield from _under_regions(child, ctx, rows) +def check_bounds_cover( + name: str, declared: ms.VariableDeclaration, ctx: Context, rows: Rows +) -> None: + """A bound parameter must have a value at every coordinate the variable occupies.""" + names = sorted(parameters_of(declared.lower, declared.upper)) + missing = sum(gaps_under(ctx.parameters[p], rows) for p in names) + if missing: + raise SpecDataError( + f"variable '{name}': {missing} rows have NULL bounds, a bound parameter is missing " + f"values for some coordinates. The two ways out build different models, so neither " + f"is picked:\n" + f" supply the value the variable exists there, bounded (`inf` is a value)\n" + f' where: "" the variable does not exist there at all' + ) diff --git a/linopy/spec/curves.py b/linopy/spec/curves.py index 04c28f01..ae9c8428 100644 --- a/linopy/spec/curves.py +++ b/linopy/spec/curves.py @@ -51,12 +51,14 @@ def validate(program: ms.Program, parameters: Mapping[str, xr.DataArray]) -> Non """ 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. + 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) diff --git a/linopy/spec/evaluate.py b/linopy/spec/evaluate.py new file mode 100644 index 00000000..99d712ed --- /dev/null +++ b/linopy/spec/evaluate.py @@ -0,0 +1,198 @@ +"""The recursive evaluator: one expression node to its linopy term, array or number.""" + +from __future__ import annotations + +import functools +import operator +from collections.abc import Callable +from typing import assert_never + +import xarray as xr +from math_spec import did_you_mean +from math_spec import program as ms + +from linopy.expressions import LinearExpression, QuadraticExpression +from linopy.spec import operators, terms +from linopy.spec.context import Context +from linopy.spec.coverage import check_divisors, obligations_of +from linopy.spec.errors import SpecDataError +from linopy.spec.terms import Array, Term, Value +from linopy.spec.where import evaluate_where +from linopy.variables import Variable + + +def evaluate_named(name: str, ctx: Context) -> Value: + """The named expression *name* as its linopy term, array or number over *ctx*, its divisors checked first.""" + if name not in ctx.program.named_expressions: + raise KeyError( + f"unknown named expression '{name}'. " + + did_you_mean(name, ctx.program.named_expressions) + ) + body = ctx.program.named_expressions[name] + found = obligations_of((body,), ctx, None) + check_divisors(f"expression '{name}'", found.divisors, ctx) + value = evaluate(body, ctx) + return _named(value, name) if isinstance(value, xr.DataArray) else value + + +def fold(name: str, ctx: Context) -> xr.DataArray: + """The named expression *name* as data, folded over the solution and the parameters *ctx* holds.""" + value = evaluate_named(name, ctx) + if isinstance(value, xr.DataArray): + return value + if isinstance(value, float | int): + return xr.DataArray(float(value), name=name) + raise TypeError( + f"expression '{name}' folded to a {type(value).__name__}, not to data" + ) + + +def _named(value: xr.DataArray, name: str) -> xr.DataArray: + """*value* with its stray non-dimension coordinates dropped and renamed to *name*.""" + stray = [c for c in value.coords if c not in value.dims] + return value.drop_vars(stray).rename(name) + + +def evaluate(node: ms.ExpressionNode, ctx: Context) -> Value: + """One node as a linopy term, an array or a number.""" + if isinstance(node, ms.Constant): + return node.value + if isinstance(node, ms.Variable): + return _variable(node.name, ctx) + if isinstance(node, ms.Parameter): + return terms.coefficient(ctx.parameters[node.name]) + if isinstance(node, ms.Negate): + return -evaluate(node.operand, ctx) + if isinstance(node, ms.Add): + return _combine( + operator.add, evaluate(node.left, ctx), evaluate(node.right, ctx) + ) + if isinstance(node, ms.Multiply): + return _combine( + operator.mul, evaluate(node.left, ctx), evaluate(node.right, ctx) + ) + if isinstance(node, ms.Divide): + return _combine( + operator.truediv, evaluate(node.numerator, ctx), evaluate(node.divisor, ctx) + ) + if isinstance(node, ms.Power): + return _combine( + operator.pow, evaluate(node.base, ctx), evaluate(node.exponent, ctx) + ) + if isinstance(node, ms.Sum): + summed = _array(evaluate(node.operand, ctx)) + for dimension in node.over: + summed = operators.sum_over(summed, dimension) + return summed + if isinstance(node, ms.GroupSum): + return operators.grouped_sum( + _array(evaluate(node.operand, ctx)), + _lookup_arrays(node.over, node.coordinate, ctx), + into=node.into, + labels=ctx.coords, + ) + if isinstance(node, ms.At): + return operators.at( + _array(evaluate(node.operand, ctx)), + _lookup_arrays(node.over, node.coordinate, ctx), + into=node.into, + ) + if isinstance(node, ms.Translate): + return operators.shift( + _array(evaluate(node.operand, ctx)), + over=node.dimension, + offset=_amount(node.offset, ctx), + wrap=node.wrap, + fill=node.fill, + by=_partition(node, ctx), + ) + if isinstance(node, ms.Window): + return operators.sum_back( + _array(evaluate(node.operand, ctx)), + over=node.dimension, + within=_amount(node.width, ctx), + wrap=node.wrap, + by=_partition(node, ctx), + ) + if isinstance(node, ms.Cases): + regions = ( + _in_region(evaluate(region.value, ctx), evaluate_where(region.when, ctx)) + for region in node.regions + ) + return functools.reduce(lambda a, b: _combine(operator.add, a, b), regions) + assert_never(node) + + +def _variable(name: str, ctx: Context) -> Value: + variable = ctx.model.variables[name] + absence = ctx.program.variable(name).absence + if not ctx.solved: + return terms.variable_term(variable, absence) + if "solution" not in variable.data: + raise RuntimeError( + f"variable '{name}' has no solution yet: solve the model before reading a named expression" + ) + return terms.solution(variable, absence) + + +def _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 = ctx.lookup(node.partition, node.dimension) + return array.rename(ctx.program.dimension(node.dimension).targets[node.partition]) + + +def _lookup_arrays( + over: str, names: tuple[str, ...], ctx: Context +) -> tuple[xr.DataArray, ...]: + return tuple(ctx.lookup(name, over) for name in names) From 88fe4111a7cc85ba4576d4c2d656ace4f0f4e7e1 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 20:53:52 +0200 Subject: [PATCH 21/35] test(spec): split builder tests into accessor, operators and curves; shared material in conftest --- test/conftest.py | 188 +++++++- test/test_spec_accessor.py | 238 ++++++++++ test/test_spec_builder.py | 850 ++---------------------------------- test/test_spec_curves.py | 160 +++++++ test/test_spec_operators.py | 310 +++++++++++++ 5 files changed, 928 insertions(+), 818 deletions(-) create mode 100644 test/test_spec_accessor.py create mode 100644 test/test_spec_curves.py create mode 100644 test/test_spec_operators.py diff --git a/test/conftest.py b/test/conftest.py index d636778d..cbc57673 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -5,7 +5,8 @@ import os import warnings from collections.abc import Generator -from typing import TYPE_CHECKING +from importlib.util import find_spec +from typing import TYPE_CHECKING, Any import pandas as pd import pytest @@ -157,3 +158,188 @@ def u(m: Model) -> Variable: idx.name = "dim_3" m.add_variables(coords=[idx], name="u") return m.variables["u"] + + +if find_spec("math_spec") is not None: + import math_spec + import xarray as xr + import yaml + + from linopy import Model + + EXAMPLE_DISPATCH = """ +description: Least-cost dispatch of a generator fleet against an hourly load. + +dimensions: + snapshot: { dtype: int, description: dispatch periods } + generator: { description: generating units } + +parameters: + p_max: { dims: [generator], description: installed capacity } + load: { dims: [snapshot], description: demand to be met } + cost: { dims: [generator], description: marginal cost } + +variables: + p: + description: output of a generator in a snapshot + foreach: [snapshot, generator] + where: "p_max > 0" + bounds: { lower: 0, upper: p_max } + +constraints: + power_balance: + foreach: [snapshot] + expression: sum(p, over=generator) == load + +objective: + sense: minimize + expression: sum(p * cost) + +expressions: + spend: sum(p * cost, over=generator) + usage: p / p_max +""" + + GENERATOR = pd.Index(["wind", "gas"], name="generator") + SNAPSHOT = pd.Index([0, 1, 2], name="snapshot") + DISPATCH_DATA: dict[str, Any] = { + "snapshot": SNAPSHOT, + "generator": GENERATOR, + "p_max": pd.Series([100.0, 200.0], index=GENERATOR), + "load": pd.Series([80.0, 150.0, 50.0], index=SNAPSHOT), + "cost": pd.Series([0.0, 50.0], index=GENERATOR), + } + DISPATCH_P = xr.DataArray( + [[80.0, 0.0], [100.0, 50.0], [50.0, 0.0]], + coords={"snapshot": SNAPSHOT, "generator": GENERATOR}, + ) + + def solved(spec: Any, sources: Any, **kwargs: Any) -> Model: + m = Model.from_spec(spec, sources, **kwargs) + m.solve(solver_name="highs", output_flag=False, reformulate_sos=True) + return m + + def yaml_dict() -> dict[str, Any]: + return math_spec.to_spec(yaml.safe_load(EXAMPLE_DISPATCH)).to_dict() + + def with_(spec: dict[str, Any], **sections: dict[str, Any]) -> dict[str, Any]: + out = dict(spec) + for section, entries in sections.items(): + out[section] = {**spec.get(section, {}), **entries} + return out + + TT = pd.Index([0, 1, 2, 3], name="t") + S = pd.Index(["a", "b"], name="s") + DAYS = pd.date_range("2030-01-01", periods=4, freq="D", name="d") + + WHERE_SPEC: dict[str, Any] = { + "dimensions": { + "t": {"dtype": "int"}, + "s": {"dtype": "str"}, + "d": {"dtype": "datetime"}, + }, + "lookups": { + "season_of": {"over": "t", "into": "s"}, + "other_of": {"over": "t", "into": "s"}, + "tag": {"over": "t", "dtype": "str"}, + }, + "parameters": { + "flag": {"dims": ["t"], "dtype": "bool"}, + "cost": {"dims": ["t"]}, + "label": {"dims": ["t"], "dtype": "str"}, + "day_cost": {"dims": ["d"]}, + }, + "variables": { + "x": {"foreach": ["t"], "bounds": {"lower": 0, "upper": 1}}, + "y": {"foreach": ["d"], "bounds": {"lower": 0, "upper": 1}}, + }, + "objective": {"sense": "minimize", "expression": "sum(x) + sum(y)"}, + } + WHERE_DATA: dict[str, Any] = { + "t": TT, + "s": S, + "d": DAYS, + "season_of": pd.Series(["a", "a", "b"], index=TT[:3]), + "other_of": pd.Series(["a", "b", "b", "a"], index=TT), + "tag": pd.Series(["p", "q"], index=TT[:2]), + "flag": pd.Series([True, False], index=TT[:2]), + "cost": pd.Series([1.0, float("inf"), 3.0], index=TT[:3]), + "label": pd.Series(["u", "v"], index=TT[1:3]), + "day_cost": pd.Series([1.0, 2.0, 3.0, 4.0], index=DAYS), + } + + BP = pd.Index([0, 1, 2, 3], name="bp") + UNITS = pd.Index(["hydro", "gas"], name="generator") + CURVE_SPEC: dict[str, Any] = { + "dimensions": { + "snapshot": {"dtype": "int"}, + "generator": {"dtype": "str"}, + "bp": {"dtype": "int"}, + }, + "parameters": { + "p_max": {"dims": ["generator"]}, + "load": {"dims": ["snapshot"]}, + "bp_x": {"dims": ["generator", "bp"]}, + "bp_y": {"dims": ["generator", "bp"]}, + }, + "variables": { + "p": { + "foreach": ["snapshot", "generator"], + "bounds": {"lower": 0, "upper": "p_max"}, + }, + "op_cost": {"foreach": ["snapshot", "generator"], "bounds": {"lower": 0}}, + }, + "piecewise": { + "cost_curve": { + "over": "bp", + "links": [["p", "bp_x"], ["op_cost", "bp_y", ">="]], + "method": "lp", + } + }, + "expressions": {"spend": "sum(op_cost, over=generator)"}, + "constraints": { + "balance": { + "foreach": ["snapshot"], + "expression": "sum(p, over=generator) == load", + } + }, + "objective": {"sense": "minimize", "expression": "sum(op_cost)"}, + } + MASKED_CURVE_SPEC = with_( + CURVE_SPEC, + piecewise={ + "cost_curve": {**CURVE_SPEC["piecewise"]["cost_curve"], "points": "bp_x"} + }, + ) + + def curve(points: dict[tuple[str, int], float]) -> pd.Series: + index = pd.MultiIndex.from_tuples(list(points), names=["generator", "bp"]) + return pd.Series(list(points.values()), index=index) + + FULL_X = curve( + {(g, k): x for g in UNITS for k, x in enumerate([0.0, 20.0, 50.0, 80.0])} + ) + FULL_Y = curve( + {(g, k): y for g in UNITS for k, y in enumerate([0.0, 150.0, 450.0, 900.0])} + ) + RAGGED_X = curve( + { + ("hydro", 0): 0.0, + ("hydro", 1): 40.0, + **{("gas", k): x for k, x in enumerate([0.0, 20.0, 50.0, 80.0])}, + } + ) + RAGGED_Y = curve( + { + ("hydro", 0): 0.0, + ("hydro", 1): 200.0, + **{("gas", k): y for k, y in enumerate([0.0, 150.0, 450.0, 900.0])}, + } + ) + CURVE_DATA: dict[str, Any] = { + "snapshot": [0], + "generator": UNITS, + "bp": BP, + "p_max": pd.Series([40.0, 80.0], index=UNITS), + "load": pd.Series([50.0], index=pd.Index([0], name="snapshot")), + } diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py new file mode 100644 index 00000000..de593cb2 --- /dev/null +++ b/test/test_spec_accessor.py @@ -0,0 +1,238 @@ +""" +``model.spec``, ``ModelSpec``, ``NamedExpression``, ``evaluate``, typesetting, +and the ``add_spec``/``from_spec`` argument handling that builds them. +""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import pandas as pd +import pytest +import xarray as xr + +math_spec = pytest.importorskip("math_spec") +yaml = pytest.importorskip("yaml") + +import linopy # noqa: E402 +from conftest import ( # noqa: E402 + DISPATCH_DATA, + DISPATCH_P, + EXAMPLE_DISPATCH, + GENERATOR, + solved, + yaml_dict, +) +from linopy import Model # noqa: E402 +from linopy.spec import ModelSpec, NamedExpression, SpecDataError # noqa: E402 + +pytestmark = [ + pytest.mark.v1, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +# --------------------------------------------------------------------------- +# inputs and model integration +# --------------------------------------------------------------------------- + + +SPEC_FORMS: dict[str, Callable[[Path], Any]] = { + "path": lambda path: path, + "path-string": str, + "yaml-text": lambda path: path.read_text(), + "dict": lambda path: math_spec.to_spec(path).to_dict(), + "spec": lambda path: math_spec.to_spec(path), +} + + +@pytest.mark.parametrize("form", SPEC_FORMS.values(), ids=SPEC_FORMS.keys()) +def test_spec_forms_build_the_same_model( + tmp_path: Path, form: Callable[[Path], Any] +) -> None: + path = tmp_path / "dispatch.yaml" + path.write_text(EXAMPLE_DISPATCH) + m = Model.from_spec(form(path), DISPATCH_DATA) + assert list(m.variables) == ["p"] + assert list(m.constraints) == ["power_balance"] + reread = math_spec.to_program(yaml.safe_load(m.spec.text)) + assert reread.constraints == m.spec.program.constraints + assert isinstance(m.spec, ModelSpec) + + +def test_a_lowered_program_is_refused() -> None: + program = math_spec.to_program(yaml_dict()) + with pytest.raises(TypeError, match="not a lowered Program"): + Model().add_spec(program, DISPATCH_DATA) + + +def test_add_spec_needs_an_empty_model() -> None: + m = Model() + m.add_variables(name="x") + with pytest.raises(ValueError, match="empty model"): + m.add_spec(yaml_dict(), DISPATCH_DATA) + + +def test_legacy_semantics_is_refused() -> None: + with linopy.options as options: + options["semantics"] = "legacy" + with pytest.raises(ValueError, match="v1"): + Model.from_spec(yaml_dict(), DISPATCH_DATA) + + +def test_a_model_without_a_spec_has_no_accessor() -> None: + with pytest.raises(AttributeError, match="not built from a spec"): + _ = Model().spec + + +def test_from_spec_passes_model_kwargs_and_chains() -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA, force_dim_names=True) + assert m.force_dim_names + assert Model().add_spec( + yaml_dict(), DISPATCH_DATA + ).spec.program.variables.keys() == {"p"} + + +# --------------------------------------------------------------------------- +# retain and evaluate +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("retain", "kept"), + [ + ("report", {"cost", "p_max"}), + ("all", {"cost", "load", "p_max"}), + ("none", set()), + ], +) +def test_retain_decides_what_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).solution, want) + if "cost" in kept: + xr.testing.assert_allclose(m.spec.expressions["spend"].solution, want) + else: + with pytest.raises(SpecDataError, match="not retained"): + m.spec.expressions["spend"].solution + + +def test_an_unknown_expression_is_a_key_error_with_a_hint() -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + with pytest.raises(KeyError, match="unknown named expression 'spent'.*spend"): + m.spec.expressions["spent"] + + +def test_a_fold_over_variables_needs_a_solution_and_one_over_data_does_not() -> None: + spec = { + **yaml_dict(), + "parameters": { + **yaml_dict()["parameters"], + "rate": {"dims": []}, + "years": {"dims": []}, + }, + "expressions": { + "spend": "sum(p * cost, over=generator)", + "growth": "rate ** years", + }, + } + m = Model.from_spec(spec, {**DISPATCH_DATA, "rate": 1.05, "years": 3.0}) + assert float(m.spec.expressions["growth"].solution) == pytest.approx(1.05**3) + with pytest.raises(RuntimeError, match="no solution yet"): + m.spec.expressions["spend"].solution + + +# --------------------------------------------------------------------------- +# three views: math, the linopy expression and the solution +# --------------------------------------------------------------------------- + +VIEWS_SPEC: dict[str, Any] = { + **math_spec.to_spec(yaml.safe_load(EXAMPLE_DISPATCH)).to_dict(), + "expressions": { + "spend": "sum(p * cost, over=generator)", + "bare": "p", + "levels": "cost * 2", + "answer": "6 * 7", + }, +} + + +@pytest.mark.parametrize( + ("name", "kind"), + [ + ("spend", linopy.LinearExpression), + ("bare", linopy.Variable), + ("levels", xr.DataArray), + ("answer", float), + ], +) +def test_expression_is_the_unsolved_linopy_term(name: str, kind: type) -> None: + m = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA) + assert isinstance(m.spec.expressions[name].expression, kind) + + +def test_expression_reads_unsolved_but_solution_waits_for_a_solve() -> None: + e = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.expressions["spend"] + assert isinstance(e.expression, linopy.LinearExpression) + with pytest.raises(RuntimeError, match="no solution yet"): + e.solution + + +def test_the_named_expression_bundles_the_three_views() -> None: + m = solved(VIEWS_SPEC, DISPATCH_DATA) + e = m.spec.expressions["spend"] + assert e.node is m.spec.program.named_expressions["spend"] + 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() + + +def test_a_constant_expression_folds_to_a_scalar() -> None: + spec = {**yaml_dict(), "expressions": {"answer": "6 * 7"}} + got = Model.from_spec(spec, DISPATCH_DATA).spec.expressions["answer"].solution + assert got.ndim == 0 and float(got) == 42.0 + + +OTHER = pd.Index(["x", "y"], name="generator") + + +@pytest.mark.parametrize( + ("generator", "match"), + [ + pytest.param(GENERATOR[::-1], "as \\['gas', 'wind'\\]", id="reordered"), + pytest.param(OTHER, "as \\['x', 'y'\\]", id="relabelled"), + ], +) +def test_evaluate_refuses_sources_on_other_labels_than_the_model( + generator: pd.Index, match: str +) -> None: + m = solved({**yaml_dict(), "expressions": {"twice": "cost * 2"}}, DISPATCH_DATA) + sources = { + **DISPATCH_DATA, + "generator": generator, + "p_max": pd.Series([100.0, 200.0], index=generator), + "cost": pd.Series([0.0, 50.0], index=generator), + } + with pytest.raises(SpecDataError, match=f"dimension 'generator' {match}"): + m.spec.evaluate("twice", sources) diff --git a/test/test_spec_builder.py b/test/test_spec_builder.py index 6ee93785..e5a3d879 100644 --- a/test/test_spec_builder.py +++ b/test/test_spec_builder.py @@ -1,21 +1,16 @@ """ -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. +Building declarations from math-spec programs: variables, constraints, the +objective, SOS-constrained curves, coverage refusals and side-swapped +expressions. """ from __future__ import annotations import glob import os -from 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 @@ -24,8 +19,23 @@ yaml = pytest.importorskip("yaml") import linopy # noqa: E402 +from conftest import ( # noqa: E402, F401 + CURVE_DATA, + CURVE_SPEC, + DISPATCH_DATA, + DISPATCH_P, + EXAMPLE_DISPATCH, + FULL_X, + FULL_Y, + GENERATOR, + WHERE_DATA, + WHERE_SPEC, + solved, + with_, + yaml_dict, +) from linopy import Model # noqa: E402 -from linopy.spec import ModelSpec, NamedExpression, SpecDataError # noqa: E402 +from linopy.spec import SpecDataError # noqa: E402 from linopy.spec.testing import synthetic_sources # noqa: E402 pytestmark = [ @@ -33,147 +43,6 @@ 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"].solution - xr.testing.assert_allclose( - spend, (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") - ) - usage = m.spec.expressions["usage"].solution - xr.testing.assert_allclose(usage, (DISPATCH_P / [100.0, 200.0]).rename("usage")) - assert ( - set(m.spec.expressions) == {"spend", "usage"} and len(m.spec.expressions) == 2 - ) - assert set(m.spec.parameters.data_vars) == {"cost", "p_max"} - assert m.spec.coords["generator"].equals(GENERATOR) - - EXAMPLES_DIR = os.environ.get("MATH_SPEC_EXAMPLES") EXAMPLES = ( sorted(glob.glob(f"{EXAMPLES_DIR}/*.yaml") + glob.glob(f"{EXAMPLES_DIR}/*/*.yaml")) @@ -197,118 +66,21 @@ def test_every_math_spec_example_builds_and_solves(path: str) -> None: 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).solution, want) - if "cost" in kept: - xr.testing.assert_allclose(m.spec.expressions["spend"].solution, want) - else: - with pytest.raises(SpecDataError, match="not retained"): - m.spec.expressions["spend"].solution - - -def test_an_unknown_expression_is_a_key_error_with_a_hint() -> None: - m = Model.from_spec(yaml_dict(), DISPATCH_DATA) - with pytest.raises(KeyError, match="unknown named expression 'spent'.*spend"): - m.spec.expressions["spent"] - - -def test_a_fold_over_variables_needs_a_solution_and_one_over_data_does_not() -> None: - spec = { - **yaml_dict(), - "parameters": { - **yaml_dict()["parameters"], - "rate": {"dims": []}, - "years": {"dims": []}, - }, - "expressions": { - "spend": "sum(p * cost, over=generator)", - "growth": "rate ** years", - }, - } - m = Model.from_spec(spec, {**DISPATCH_DATA, "rate": 1.05, "years": 3.0}) - assert float(m.spec.expressions["growth"].solution) == pytest.approx(1.05**3) - with pytest.raises(RuntimeError, match="no solution yet"): - m.spec.expressions["spend"].solution - - -# --------------------------------------------------------------------------- -# three views: math, the linopy expression and the solution -# --------------------------------------------------------------------------- - -VIEWS_SPEC: dict[str, Any] = { - **math_spec.to_spec(yaml.safe_load(EXAMPLE_DISPATCH)).to_dict(), - "expressions": { - "spend": "sum(p * cost, over=generator)", - "bare": "p", - "levels": "cost * 2", - "answer": "6 * 7", - }, -} - - -@pytest.mark.parametrize( - ("name", "kind"), - [ - ("spend", linopy.LinearExpression), - ("bare", linopy.Variable), - ("levels", xr.DataArray), - ("answer", float), - ], -) -def test_expression_is_the_unsolved_linopy_term(name: str, kind: type) -> None: - m = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA) - assert isinstance(m.spec.expressions[name].expression, kind) - - -def test_expression_reads_unsolved_but_solution_waits_for_a_solve() -> None: - e = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.expressions["spend"] - assert isinstance(e.expression, linopy.LinearExpression) - with pytest.raises(RuntimeError, match="no solution yet"): - e.solution - - -def test_the_named_expression_bundles_the_three_views() -> None: - m = solved(VIEWS_SPEC, DISPATCH_DATA) - e = m.spec.expressions["spend"] - assert e.node is m.spec.program.named_expressions["spend"] - assert isinstance(e.expression, linopy.LinearExpression) +def test_the_dispatch_example_solves_and_its_expressions_fold() -> None: + m = solved(yaml_dict(), DISPATCH_DATA) + assert m.objective.value == pytest.approx(2500.0) + xr.testing.assert_allclose(m.solution["p"], DISPATCH_P) + spend = m.spec.expressions["spend"].solution xr.testing.assert_allclose( - e.solution, (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") + spend, (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") + 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 ) - - -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() + assert set(m.spec.parameters.data_vars) == {"cost", "p_max"} + assert m.spec.coords["generator"].equals(GENERATOR) # --------------------------------------------------------------------------- @@ -328,14 +100,6 @@ def test_the_whole_model_typesets() -> None: 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 - - NO_W_CONSTRAINT = {"cap": {"foreach": ["t"], "expression": "x <= c"}} @@ -443,43 +207,6 @@ 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"}}, @@ -586,44 +313,6 @@ def test_a_scalar_where_gates_a_whole_variable(on: bool, objective: float) -> No 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, @@ -662,286 +351,9 @@ def test_a_fold_reads_a_masked_slot_the_way_its_absence_says( # --------------------------------------------------------------------------- -# operators, built as a constraint and folded as a named expression +# piecewise curves as SOS constraints # --------------------------------------------------------------------------- -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}"].solution - 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"].solution.item() == pytest.approx(spend) - assert m.objective.value == pytest.approx(spend) - - -def without(series: pd.Series, *keys: tuple[str, int]) -> pd.Series: - return series.drop(index=list(keys)) - - -@pytest.mark.parametrize( - ("spec", "data", "match"), - [ - pytest.param( - CURVE_SPEC, - {"bp_x": without(FULL_X, ("gas", 3)), "bp_y": FULL_Y}, - "parameter 'bp_x' has no value at \\(generator='gas', bp=3\\)", - id="a-hole-in-a-whole-curve", - ), - pytest.param( - MASKED_CURVE_SPEC, - {"bp_x": RAGGED_X, "bp_y": without(RAGGED_Y, ("gas", 3))}, - "Shorten it 'bp_x' claims this breakpoint", - id="a-hole-inside-the-mask", - ), - pytest.param( - MASKED_CURVE_SPEC, - {"bp_x": without(FULL_X, ("gas", 1)), "bp_y": FULL_Y}, - "Not so at generator='gas'", - id="a-mask-with-a-gap", - ), - pytest.param( - CURVE_SPEC, - { - "bp_x": curve( - { - (g, k): x - for g in UNITS - for k, x in enumerate([0.0, 20.0, 20.0, 80.0]) - } - ), - "bp_y": FULL_Y, - }, - "strictly increasing", - id="breakpoints-that-do-not-increase", - ), - pytest.param( - CURVE_SPEC, - { - "bp_x": FULL_X, - "bp_y": curve( - { - (g, k): y - for g in UNITS - for k, y in enumerate([0.0, 300.0, 500.0, 600.0]) - } - ), - }, - "exact only for a convex curve", - id="a-concave-curve-under-lp", - ), - pytest.param( - MASKED_CURVE_SPEC, - {"bp_x": without(RAGGED_X, ("hydro", 1)), "bp_y": RAGGED_Y}, - "This curve carries 1", - id="a-one-point-curve-under-lp", - ), - ], -) -def test_a_curve_the_method_cannot_build_is_refused( - spec: dict[str, Any], data: dict[str, Any], match: str -) -> None: - with pytest.raises(SpecDataError, match=match): - Model.from_spec(spec, {**CURVE_DATA, **data}) - def test_a_sos2_curve_is_built_as_a_special_ordered_set() -> None: spec = with_( @@ -959,120 +371,9 @@ def test_a_sos2_curve_is_built_as_a_special_ordered_set() -> None: # --------------------------------------------------------------------------- -# where predicates +# a constant on the left is the same row, a power hides nothing # --------------------------------------------------------------------------- -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}"].solution - want = xr.DataArray( - PARTIAL_CASES[key], coords={dims[0]: OPERATOR_DATA[dims[0]]}, dims=dims - ) - xr.testing.assert_allclose(folded, want.rename(folded.name)) - if dims == ["t"]: - assert int(m.constraints[f"link_{name}"].labels.sel(t=3)) == -1 - def test_a_constant_on_the_left_is_the_same_row() -> None: flipped = with_( @@ -1082,55 +383,6 @@ def test_a_constant_on_the_left_is_the_same_row() -> None: 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"].solution - 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}) - - -# --------------------------------------------------------------------------- -# a power hides nothing -# --------------------------------------------------------------------------- - - @pytest.mark.parametrize( ("expression", "match"), [ @@ -1166,39 +418,3 @@ def test_an_operator_under_a_power_keeps_its_parameters_retained() -> None: m.spec.expressions["e"].solution, 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"].solution - assert bool(folded.isnull().all()) diff --git a/test/test_spec_curves.py b/test/test_spec_curves.py new file mode 100644 index 00000000..d0c3feb0 --- /dev/null +++ b/test/test_spec_curves.py @@ -0,0 +1,160 @@ +""" +Piecewise curve derivation and validation: whole and ragged breakpoint +tables, the checks that refuse a curve a method cannot build, and the +convex-hull method's single-bend requirement. +""" + +from __future__ import annotations + +from typing import Any + +import pandas as pd +import pytest + +math_spec = pytest.importorskip("math_spec") +yaml = pytest.importorskip("yaml") + +import linopy # noqa: E402 +from conftest import ( # noqa: E402 + CURVE_DATA, + CURVE_SPEC, + FULL_X, + FULL_Y, + MASKED_CURVE_SPEC, + RAGGED_X, + RAGGED_Y, + UNITS, + curve, + solved, + with_, +) +from linopy import Model # noqa: E402 +from linopy.spec import SpecDataError # noqa: E402 + +pytestmark = [ + pytest.mark.v1, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +# --------------------------------------------------------------------------- +# piecewise curves +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("spec", "data", "spend"), + [ + pytest.param( + CURVE_SPEC, {"bp_x": FULL_X, "bp_y": FULL_Y}, 400.0, id="whole-curves" + ), + pytest.param( + MASKED_CURVE_SPEC, + {"bp_x": RAGGED_X, "bp_y": RAGGED_Y}, + 275.0, + id="ragged-curves-under-points", + ), + ], +) +def test_a_piecewise_cost_lands_on_the_curve( + spec: dict[str, Any], data: dict[str, Any], spend: float +) -> None: + m = solved(spec, {**CURVE_DATA, **data}, retain="all") + assert m.spec.expressions["spend"].solution.item() == pytest.approx(spend) + assert m.objective.value == pytest.approx(spend) + + +def without(series: pd.Series, *keys: tuple[str, int]) -> pd.Series: + return series.drop(index=list(keys)) + + +@pytest.mark.parametrize( + ("spec", "data", "match"), + [ + pytest.param( + CURVE_SPEC, + {"bp_x": without(FULL_X, ("gas", 3)), "bp_y": FULL_Y}, + "parameter 'bp_x' has no value at \\(generator='gas', bp=3\\)", + id="a-hole-in-a-whole-curve", + ), + pytest.param( + MASKED_CURVE_SPEC, + {"bp_x": RAGGED_X, "bp_y": without(RAGGED_Y, ("gas", 3))}, + "Shorten it 'bp_x' claims this breakpoint", + id="a-hole-inside-the-mask", + ), + pytest.param( + MASKED_CURVE_SPEC, + {"bp_x": without(FULL_X, ("gas", 1)), "bp_y": FULL_Y}, + "Not so at generator='gas'", + id="a-mask-with-a-gap", + ), + pytest.param( + CURVE_SPEC, + { + "bp_x": curve( + { + (g, k): x + for g in UNITS + for k, x in enumerate([0.0, 20.0, 20.0, 80.0]) + } + ), + "bp_y": FULL_Y, + }, + "strictly increasing", + id="breakpoints-that-do-not-increase", + ), + pytest.param( + CURVE_SPEC, + { + "bp_x": FULL_X, + "bp_y": curve( + { + (g, k): y + for g in UNITS + for k, y in enumerate([0.0, 300.0, 500.0, 600.0]) + } + ), + }, + "exact only for a convex curve", + id="a-concave-curve-under-lp", + ), + pytest.param( + MASKED_CURVE_SPEC, + {"bp_x": without(RAGGED_X, ("hydro", 1)), "bp_y": RAGGED_Y}, + "This curve carries 1", + id="a-one-point-curve-under-lp", + ), + ], +) +def test_a_curve_the_method_cannot_build_is_refused( + spec: dict[str, Any], data: dict[str, Any], match: str +) -> None: + with pytest.raises(SpecDataError, match=match): + Model.from_spec(spec, {**CURVE_DATA, **data}) + + +def test_a_convex_hull_curve_may_bend_either_way_but_not_both() -> None: + spec = with_( + CURVE_SPEC, + piecewise={ + "cost_curve": { + "over": "bp", + "links": [["p", "bp_x"], ["op_cost", "bp_y"]], + "method": "convex", + } + }, + ) + concave = curve( + {(g, k): y for g in UNITS for k, y in enumerate([0.0, 300.0, 500.0, 600.0])} + ) + mixed = curve( + {(g, k): y for g in UNITS for k, y in enumerate([0.0, 300.0, 350.0, 600.0])} + ) + assert ( + "cost_curve_lam" + in Model.from_spec( + spec, {**CURVE_DATA, "bp_x": FULL_X, "bp_y": concave} + ).variables + ) + with pytest.raises(SpecDataError, match="exact only for a single bend"): + Model.from_spec(spec, {**CURVE_DATA, "bp_x": FULL_X, "bp_y": mixed}) diff --git a/test/test_spec_operators.py b/test/test_spec_operators.py new file mode 100644 index 00000000..4e470eed --- /dev/null +++ b/test/test_spec_operators.py @@ -0,0 +1,310 @@ +""" +Operators built as a constraint and folded as a named expression: sum, +grouped sum, ``at``, shift/translate, sum_back windows, and the where/mask +predicates that gate which rows exist. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +math_spec = pytest.importorskip("math_spec") +yaml = pytest.importorskip("yaml") + +import linopy # noqa: E402 +from conftest import TT, WHERE_DATA, WHERE_SPEC, solved, with_ # noqa: E402 +from linopy import Model # noqa: E402 +from linopy.spec import SpecDataError # noqa: E402 + +pytestmark = [ + pytest.mark.v1, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +# --------------------------------------------------------------------------- +# operators, built as a constraint and folded as a named expression +# --------------------------------------------------------------------------- + +S = pd.Index(["a", "b"], name="s") +V = np.array([1.0, 2.0, 4.0, 8.0]) +OPERATORS: dict[str, tuple[str, list[str], list[float]]] = { + "shift-edge-0": ("shift(x, over=t, offset=1, edge=0)", ["t"], [0, 1, 2, 4]), + "shift-ahead-edge-0": ("shift(x, over=t, offset=-1, edge=0)", ["t"], [2, 4, 8, 0]), + "shift-wrap": ("shift(x, over=t, offset=1, edge='wrap')", ["t"], [8, 1, 2, 4]), + "shift-wrap-in-groups": ( + "shift(x, over=t, offset=1, edge='wrap', by=season_of)", + ["t"], + [2, 1, 8, 4], + ), + "shift-by-group-offset": ( + "shift(x, over=t, offset=lag, edge=0, by=season_of)", + ["t"], + [0, 1, 0, 0], + ), + "sum-back": ("sum_back(x, over=t, within=2)", ["t"], [1, 3, 6, 12]), + "sum-back-wrap": ( + "sum_back(x, over=t, within=2, edge='wrap')", + ["t"], + [9, 3, 6, 12], + ), + "sum-back-in-groups": ( + "sum_back(x, over=t, within=2, by=season_of)", + ["t"], + [1, 3, 4, 12], + ), + "sum-back-group-width": ( + "sum_back(x, over=t, within=width, by=season_of)", + ["t"], + [1, 2, 4, 12], + ), + "sum-by": ("sum(x, by=season_of)", ["s"], [3, 12]), + "at": ("x * at(z, by=season_of)", ["t"], [10, 20, 80, 160]), + "cases": ("x_state", ["t"], [100, 1, 2, 4]), +} + + +def operator_spec() -> dict[str, Any]: + spec: dict[str, Any] = { + "dimensions": {"t": {"dtype": "int"}, "s": {"dtype": "str"}}, + "lookups": {"season_of": {"over": "t", "into": "s"}}, + "parameters": { + "v": {"dims": ["t"]}, + "z": {"dims": ["s"]}, + "lag": {"dims": ["s"], "dtype": "int"}, + "width": {"dims": ["s"], "dtype": "int"}, + }, + "variables": {"x": {"foreach": ["t"], "bounds": {"lower": 0, "upper": 100}}}, + "constraints": {"fix": {"foreach": ["t"], "expression": "x == v"}}, + "expressions": { + "x_state": { + "foreach": ["t"], + "cases": {"first": {"when": "position(t) == 0", "expression": 100}}, + "otherwise": "shift(x, over=t, offset=1)", + } + }, + "objective": {"sense": "minimize", "expression": "sum(x)"}, + } + for key, (expression, dims, _) in OPERATORS.items(): + name = key.replace("-", "_") + spec["variables"][f"y_{name}"] = { + "foreach": dims, + "bounds": {"lower": -1000, "upper": 1000}, + } + spec["constraints"][f"link_{name}"] = { + "foreach": dims, + "expression": f"y_{name} == {expression}", + } + spec["expressions"][f"probe_{name}"] = expression + return spec + + +OPERATOR_DATA: dict[str, Any] = { + "t": TT, + "s": S, + "season_of": pd.Series(["a", "a", "b", "b"], index=TT), + "v": pd.Series(V, index=TT), + "z": pd.Series([10.0, 20.0], index=S), + "lag": pd.Series([1, 2], index=S), + "width": pd.Series([1, 2], index=S), +} + + +@pytest.fixture(scope="module") +def operators_model() -> Model: + with linopy.options as options: + options["semantics"] = "v1" + return solved(operator_spec(), OPERATOR_DATA, retain="all") + + +@pytest.mark.parametrize("key", OPERATORS) +def test_an_operator_builds_and_folds_alike(operators_model: Model, key: str) -> None: + _, dims, expected = OPERATORS[key] + name = key.replace("-", "_") + want = xr.DataArray(expected, coords={dims[0]: OPERATOR_DATA[dims[0]]}, dims=dims) + built = operators_model.solution[f"y_{name}"] + folded = operators_model.spec.expressions[f"probe_{name}"].solution + xr.testing.assert_allclose(built, want.rename(f"y_{name}")) + xr.testing.assert_allclose(folded, want.rename(f"probe_{name}")) + + +AMOUNT_SPEC: dict[str, Any] = { + "dimensions": {"t": {"dtype": "int"}, "g": {"dtype": "int"}}, + "lookups": {"grp": {"over": "t", "into": "g"}}, + "parameters": {"v": {"dims": ["t"]}, "lag": {"dims": ["g"], "dtype": "int"}}, + "variables": { + "x": {"foreach": ["t"], "bounds": {"lower": 0, "upper": 100}}, + "y": {"foreach": ["t"], "bounds": {"lower": -100, "upper": 100}}, + }, + "constraints": { + "fix": {"foreach": ["t"], "expression": "x == v"}, + "link": { + "foreach": ["t"], + "expression": "y == shift(x, over=t, offset=lag, edge=0, by=grp)", + }, + }, + "objective": {"sense": "minimize", "expression": "sum(x)"}, +} + + +def test_a_missing_shift_amount_is_refused() -> None: + t = pd.Index([0, 1, 2], name="t") + g = pd.Index([0, 1], name="g") + data = { + "t": t, + "g": g, + "grp": pd.Series([0, 0, 1], index=t), + "v": pd.Series([0.0, 4.0, 5.0], index=t), + "lag": pd.Series([1], index=g[:1]), + } + with pytest.raises(SpecDataError, match="parameter 'lag' is used as a coefficient"): + Model.from_spec(AMOUNT_SPEC, data) + + +# --------------------------------------------------------------------------- +# grouped sum +# --------------------------------------------------------------------------- + +GROUPED_SPEC: dict[str, Any] = { + "dimensions": {"generator": {}, "bus": {"dtype": "str"}}, + "lookups": {"gen_bus": {"over": "generator", "into": "bus"}}, + "parameters": {"capacity": {"dims": ["generator"]}}, + "variables": { + "imports": {"foreach": ["bus"], "bounds": {"lower": 0, "upper": 100}} + }, + "constraints": { + "import_limit": { + "foreach": ["bus"], + "expression": "imports <= sum(capacity, by=gen_bus)", + } + }, + "objective": {"sense": "maximize", "expression": "sum(imports, over=bus)"}, +} +GENS = pd.Index(["g1", "g2"], name="generator") + + +def grouped_sources(capacity: pd.Series) -> dict[str, Any]: + return { + "bus": ["north", "south"], + "generator": GENS, + "gen_bus": pd.Series(["north", "north"], index=GENS), + "capacity": capacity, + } + + +def test_an_empty_group_on_the_constant_side_is_a_zero_and_not_a_gap() -> None: + m = solved(GROUPED_SPEC, grouped_sources(pd.Series([3.0, 4.0], index=GENS))) + assert m.objective.value == pytest.approx(7.0) + assert float(m.solution["imports"].sel(bus="south")) == pytest.approx(0.0) + + +def test_a_member_with_no_value_is_still_refused_through_a_group() -> None: + with pytest.raises(SpecDataError, match="parameter 'capacity' covers 1 fewer"): + Model.from_spec(GROUPED_SPEC, grouped_sources(pd.Series([3.0], index=GENS[:1]))) + + +# --------------------------------------------------------------------------- +# where predicates +# --------------------------------------------------------------------------- + +WHERE_CASES: dict[str, tuple[str, str, list[Any]]] = { + "dimension-comparison": ("x", "t > 1", [2, 3]), + "lookup-comparison": ("x", "season_of == 'a'", [0, 1]), + "lookup-not-equal-skips-unmapped": ("x", "season_of != 'a'", [2]), + "lookup-pair": ("x", "season_of != other_of", [1]), + "lookup-defined": ("x", "season_of", [0, 1, 2]), + "label-space-lookup": ("x", "tag == 'q'", [1]), + "not": ("x", "NOT (t > 1)", [0, 1]), + "and": ("x", "t > 0 AND t < 3", [1, 2]), + "or": ("x", "t == 0 OR t == 3", [0, 3]), + "position": ("x", "position(t) == -1", [3]), + "position-in-groups": ("x", "position(t, by=season_of) == 0", [0, 2]), + "bool-parameter": ("x", "flag", [0]), + "float-parameter-must-be-finite": ("x", "cost", [0, 2]), + "str-parameter": ("x", "label", [1, 2]), + "parameter-comparison": ("x", "cost > 2", [2, 1]), + "datetime-axis": ("y", "d >= '2030-01-03'", list(WHERE_DATA["d"][2:])), +} + + +@pytest.mark.parametrize("case", WHERE_CASES) +def test_a_where_picks_the_rows_it_names(case: str) -> None: + variable, predicate, labels = WHERE_CASES[case] + spec = with_( + WHERE_SPEC, + variables={variable: {**WHERE_SPEC["variables"][variable], "where": predicate}}, + ) + built = Model.from_spec(spec, WHERE_DATA).variables[variable] + dim = built.dims[0] + present = built.labels[dim][(built.labels != -1).to_numpy()] + assert sorted(present.to_numpy().tolist()) == sorted(labels) + + +@pytest.mark.parametrize( + ("predicate", "match"), + [ + ("position(t) == 7", "names position 7 of 't', which has 4"), + ("position(t, by=season_of) == 1", "shorter than that: \\['b'\\]"), + ], +) +def test_a_position_no_coordinate_holds_is_refused(predicate: str, match: str) -> None: + spec = with_( + WHERE_SPEC, + variables={"x": {**WHERE_SPEC["variables"]["x"], "where": predicate}}, + ) + with pytest.raises(SpecDataError, match=match): + Model.from_spec(spec, WHERE_DATA) + + +# --------------------------------------------------------------------------- +# edges: partial lookups and an empty dimension beside a sum +# --------------------------------------------------------------------------- + +PARTIAL_CASES: dict[str, list[float]] = { + "sum-by": [3.0, 4.0], + "at": [10.0, 20.0, 80.0, np.nan], + "shift-wrap-in-groups": [2.0, 1.0, 4.0, np.nan], + "sum-back-in-groups": [1.0, 3.0, 4.0, np.nan], +} + + +@pytest.mark.parametrize("key", PARTIAL_CASES) +def test_a_member_a_lookup_sends_nowhere_reaches_nothing(key: str) -> None: + data = {**OPERATOR_DATA, "season_of": pd.Series(["a", "a", "b"], index=TT[:3])} + m = solved(operator_spec(), data, retain="all") + _, dims, _ = OPERATORS[key] + name = key.replace("-", "_") + folded = m.spec.expressions[f"probe_{name}"].solution + want = xr.DataArray( + PARTIAL_CASES[key], coords={dims[0]: OPERATOR_DATA[dims[0]]}, dims=dims + ) + xr.testing.assert_allclose(folded, want.rename(folded.name)) + if dims == ["t"]: + assert int(m.constraints[f"link_{name}"].labels.sel(t=3)) == -1 + + +def test_a_sum_beside_an_empty_dimension_is_the_empty_sum() -> None: + spec: dict[str, Any] = { + "dimensions": {"t": {"dtype": "int"}, "s": {"dtype": "str"}}, + "variables": {"x": {"foreach": ["t", "s"], "bounds": {"lower": 0, "upper": 1}}}, + "constraints": {"cap": {"foreach": ["s"], "expression": "sum(x, over=t) <= 1"}}, + "objective": {"sense": "maximize", "expression": "sum(x)"}, + } + m = Model.from_spec(spec, {"t": [0, 1], "s": pd.Index([], name="s", dtype=object)}) + assert "cap" not in m.constraints + + +def test_a_window_width_no_member_carries_is_a_window_of_nothing() -> None: + data = { + **OPERATOR_DATA, + "season_of": pd.Series( + [], index=pd.Index([], name="t", dtype=int), dtype=object + ), + } + m = solved(operator_spec(), data, retain="all") + folded = m.spec.expressions["probe_sum_back_group_width"].solution + assert bool(folded.isnull().all()) From 2cb5c9f3b9799c1680771e388e5b7abe5bd2496c Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 20:53:52 +0200 Subject: [PATCH 22/35] doc(spec): wire the notebook and API into the docs; run the notebook in CI with the spec group --- .github/workflows/test-notebooks.yml | 6 +----- doc/api.rst | 20 ++++++++++++++++++++ doc/building-models-from-specs.nblink | 3 +++ doc/index.rst | 1 + doc/release_notes.rst | 4 ++-- doc/user-guide.rst | 14 ++++++++++++++ examples/building-models-from-specs.ipynb | 3 +++ 7 files changed, 44 insertions(+), 7 deletions(-) create mode 100644 doc/building-models-from-specs.nblink diff --git a/.github/workflows/test-notebooks.yml b/.github/workflows/test-notebooks.yml index 14081651..cfed5914 100644 --- a/.github/workflows/test-notebooks.yml +++ b/.github/workflows/test-notebooks.yml @@ -30,7 +30,7 @@ jobs: - name: Install package and dependencies run: | python -m pip install uv - uv pip install --system -e ".[docs]" + uv pip install --system -e ".[docs]" --group spec - name: Execute notebooks run: | @@ -44,10 +44,6 @@ 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" diff --git a/doc/api.rst b/doc/api.rst index 0656a99a..eb9d0d7a 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -115,6 +115,26 @@ IO model.Model.to_netcdf io.read_netcdf +Building from specs +------------------- + +Build a model from a `math-spec +`__ YAML program bound to +data. Requires the ``spec`` dependency group. + +.. autosummary:: + :toctree: generated/ + + model.Model.add_spec + model.Model.from_spec + model.Model.spec + spec.ModelSpec + spec.NamedExpressions + spec.NamedExpression + spec.bind + spec.Bound + spec.SpecDataError + Variable ======== diff --git a/doc/building-models-from-specs.nblink b/doc/building-models-from-specs.nblink new file mode 100644 index 00000000..f9918a8d --- /dev/null +++ b/doc/building-models-from-specs.nblink @@ -0,0 +1,3 @@ +{ + "path": "../examples/building-models-from-specs.ipynb" +} diff --git a/doc/index.rst b/doc/index.rst index b3c75447..e07d3199 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -116,6 +116,7 @@ This package is published under MIT license. coordinate-alignment migrating-to-v1 manipulating-models + building-models-from-specs .. toctree:: :hidden: diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 7d9208ab..b1baab6e 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -23,9 +23,9 @@ Upcoming Version *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.from_spec`` / ``model.add_spec`` build a model from a `math-spec `__ YAML program bound to data, and ``model.spec`` (a ``linopy.spec.ModelSpec``) reads it back. Requires the ``spec`` dependency group (``uv sync --group spec`` / ``uv pip install --group spec``, Python >= 3.12) and v1 semantics. Data is bound onto the spec's dimensions and parameters with ``linopy.spec.bind``, raising a ``linopy.spec.SpecDataError`` on mismatched or missing data; ``linopy.spec.Bound`` carries the bound result. See :doc:`building-models-from-specs` for a worked example. -* ``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.expressions`` (a ``linopy.spec.NamedExpressions`` mapping) returns a ``linopy.spec.NamedExpression`` for each declared name, with three views: ``.node`` (the lowered formula), ``.expression`` (the unsolved linopy expression — a ``LinearExpression``, bare ``Variable``, array or scalar) and ``.solution`` (the expression folded over the solved model). ``model.spec.evaluate(name, sources)`` returns the same object with its parameters 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. diff --git a/doc/user-guide.rst b/doc/user-guide.rst index 92995e3f..33589f38 100644 --- a/doc/user-guide.rst +++ b/doc/user-guide.rst @@ -83,6 +83,20 @@ bound, swap a constraint, or copy it for what-if analysis. variables. +Building a model from a spec +----------------------------- + +Instead of calling ``add_variables`` / ``add_constraints`` directly, +you can declare a model as a `math-spec +`__ YAML program bound to +data, and let linopy build it. + +- :doc:`building-models-from-specs` — ``Model.from_spec`` and + ``model.add_spec``, binding data to a spec, and reading named + expressions back through ``model.spec`` after solving. Requires the + ``spec`` dependency group and v1 semantics. + + Where to go next ---------------- diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb index 58f15100..3e7cf7fd 100644 --- a/examples/building-models-from-specs.ipynb +++ b/examples/building-models-from-specs.ipynb @@ -980,6 +980,9 @@ "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" + }, + "nbsphinx": { + "execute": "never" } }, "nbformat": 4, From d5bd61ac61a67b03f5c82003c54aad80deb0a86f Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 7 Sep 2026 21:08:53 +0200 Subject: [PATCH 23/35] feat(spec): emit EvolvingAPIWarning once per session from add_spec, from_spec and bind warn_evolving_api moves to linopy.constants so piecewise and spec share the once-per-key dedup; the pytest filter silences the spec prefix. --- doc/release_notes.rst | 2 +- linopy/constants.py | 17 +++++++++++++++ linopy/model.py | 7 +++++++ linopy/piecewise.py | 33 ++++-------------------------- linopy/spec/accessor.py | 4 +++- linopy/spec/binder.py | 8 ++++++++ pyproject.toml | 1 + test/test_piecewise_constraints.py | 2 +- test/test_spec_accessor.py | 13 ++++++++++++ 9 files changed, 55 insertions(+), 32 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index b1baab6e..7e6805e6 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -23,7 +23,7 @@ Upcoming Version *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`` (a ``linopy.spec.ModelSpec``) reads it back. Requires the ``spec`` dependency group (``uv sync --group spec`` / ``uv pip install --group spec``, Python >= 3.12) and v1 semantics. Data is bound onto the spec's dimensions and parameters with ``linopy.spec.bind``, raising a ``linopy.spec.SpecDataError`` on mismatched or missing data; ``linopy.spec.Bound`` carries the bound result. See :doc:`building-models-from-specs` for a worked example. +* ``Model.from_spec`` / ``model.add_spec`` build a model from a `math-spec `__ YAML program bound to data, and ``model.spec`` (a ``linopy.spec.ModelSpec``) reads it back. Requires the ``spec`` dependency group (``uv sync --group spec`` / ``uv pip install --group spec``, Python >= 3.12) and v1 semantics. Data is bound onto the spec's dimensions and parameters with ``linopy.spec.bind``, raising a ``linopy.spec.SpecDataError`` on mismatched or missing data; ``linopy.spec.Bound`` carries the bound result. The spec API emits an :class:`linopy.EvolvingAPIWarning` once per session while it stabilises. See :doc:`building-models-from-specs` for a worked example. * ``model.spec.expressions`` (a ``linopy.spec.NamedExpressions`` mapping) returns a ``linopy.spec.NamedExpression`` for each declared name, with three views: ``.node`` (the lowered formula), ``.expression`` (the unsolved linopy expression — a ``LinearExpression``, bare ``Variable``, array or scalar) and ``.solution`` (the expression folded over the solved model). ``model.spec.evaluate(name, sources)`` returns the same object with its parameters bound afresh. diff --git a/linopy/constants.py b/linopy/constants.py index 7936ef1c..91211e61 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -4,6 +4,7 @@ """ import logging +import warnings from dataclasses import dataclass, field from enum import StrEnum from typing import Any, Literal, Self, TypeAlias, get_args @@ -124,6 +125,22 @@ class EvolvingAPIWarning(FutureWarning): """ +_emitted_evolving_warnings: set[str] = set() + + +def warn_evolving_api(key: str, message: str, stacklevel: int = 3) -> None: + """ + Emit an :class:`EvolvingAPIWarning` at most once per session per ``key``. + + ``stacklevel`` counts from the ``warnings.warn`` call: 3 points at the + caller of the function that calls this helper. + """ + if key in _emitted_evolving_warnings: + return + _emitted_evolving_warnings.add(key) + warnings.warn(message, category=EvolvingAPIWarning, stacklevel=stacklevel) + + class ModelStatus(StrEnum): """ Model status. diff --git a/linopy/model.py b/linopy/model.py index 7d63971f..416128cd 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -497,6 +497,13 @@ def add_spec( under legacy semantics. linopy.spec.SpecDataError If the data does not fit the spec. + + Warns + ----- + EvolvingAPIWarning + Once per session: the spec API is newly added and may change in + minor releases. Silence with ``warnings.filterwarnings("ignore", + category=linopy.EvolvingAPIWarning)``. """ from linopy.spec.accessor import attach diff --git a/linopy/piecewise.py b/linopy/piecewise.py index 5a07dba4..5349e778 100644 --- a/linopy/piecewise.py +++ b/linopy/piecewise.py @@ -8,7 +8,6 @@ from __future__ import annotations import logging -import warnings from collections.abc import Sequence from dataclasses import dataclass from numbers import Real @@ -47,8 +46,8 @@ PWL_SELECT_SUFFIX, SEGMENT_DIM, SIGNS, - EvolvingAPIWarning, sign_replace_dict, + warn_evolving_api, ) from linopy.semantics import check_user_nan_breakpoints @@ -61,30 +60,6 @@ logger = logging.getLogger(__name__) -# Each user-facing piecewise entry point fires its EvolvingAPIWarning at -# most once per process. Without dedup, a single model build emits the -# verbose warning hundreds of times and drowns out other output. -_EvolvingApiKey: TypeAlias = Literal[ - "tangent_lines", "add_piecewise_formulation", "Slopes" -] -_emitted_evolving_warnings: set[_EvolvingApiKey] = set() - - -def _warn_evolving_api(key: _EvolvingApiKey, message: str, stacklevel: int = 3) -> None: - """ - Emit an :class:`EvolvingAPIWarning` at most once per session per ``key``. - - ``stacklevel`` defaults to 3 (helper → entry-point function → user - code). Pass a larger value when called from one frame deeper than - a function — e.g. from a dataclass ``__post_init__``, which is - itself invoked by an auto-generated ``__init__``. - """ - if key in _emitted_evolving_warnings: - return - _emitted_evolving_warnings.add(key) - warnings.warn(message, category=EvolvingAPIWarning, stacklevel=stacklevel) - - # Accepted input types for breakpoint-like data BreaksLike: TypeAlias = ( Sequence[float] @@ -172,7 +147,7 @@ class Slopes: def __post_init__(self) -> None: # ``stacklevel=4``: warn → _warn_evolving_api → __post_init__ → # dataclass-generated ``__init__`` → user code. - _warn_evolving_api( + warn_evolving_api( "Slopes", "piecewise: Slopes is a new API; the constructor signature and " "the dispatch rules for inheriting an x grid from sibling tuples " @@ -826,7 +801,7 @@ def tangent_lines( Silence with ``warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning)``. """ - _warn_evolving_api( + warn_evolving_api( "tangent_lines", "piecewise: tangent_lines is a new API; the returned expression " "shape and the piece-dim name may be refined in minor releases. " @@ -1272,7 +1247,7 @@ def add_piecewise_formulation( with ``warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning)``. """ - _warn_evolving_api( + warn_evolving_api( "add_piecewise_formulation", "piecewise: add_piecewise_formulation is a new API; some details " "(e.g. the per-tuple sign convention, active+sign semantics) " diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index 2c6aefb0..c2a6d485 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -28,10 +28,11 @@ ) from math_spec import program as ms +from linopy.constants import warn_evolving_api from linopy.model import Model from linopy.semantics import is_v1 from linopy.spec import terms -from linopy.spec.binder import Bound, Retain, bind +from linopy.spec.binder import EVOLVING_MESSAGE, Bound, Retain, bind from linopy.spec.builder import build from linopy.spec.context import Context from linopy.spec.errors import SpecDataError @@ -59,6 +60,7 @@ def attach( *spec* is a lowered ``Program``, which has no YAML form to keep on the model. """ + warn_evolving_api("spec", EVOLVING_MESSAGE, stacklevel=4) if not is_v1(): raise ValueError( "a spec-built model uses linopy's v1 semantics, and the current setting is " diff --git a/linopy/spec/binder.py b/linopy/spec/binder.py index 0f580e4e..0d2bf723 100644 --- a/linopy/spec/binder.py +++ b/linopy/spec/binder.py @@ -23,6 +23,7 @@ from math_spec import did_you_mean from math_spec import program as ms +from linopy.constants import warn_evolving_api from linopy.spec.errors import SpecDataError from linopy.spec.nodes import amounts_of, parameters_of, walk @@ -60,6 +61,12 @@ "a DataFrame with columns {columns} or in wide form, a dict keyed by label, or one number" ) +EVOLVING_MESSAGE = ( + "spec: Model.add_spec, Model.from_spec, model.spec and linopy.spec.bind are " + "newly added and their details may change in minor releases. Silence with " + '`warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning)`.' +) + def bind( program: ms.Program, @@ -90,6 +97,7 @@ def bind( source, a duplicated dimension member, or a lookup breaking the rules a map has. """ + warn_evolving_api("spec", EVOLVING_MESSAGE) if retain not in _RETAIN: raise SpecDataError( f"retain={retain!r} is not one of {_shown(_RETAIN)}. {did_you_mean(retain, _RETAIN)}" diff --git a/pyproject.toml b/pyproject.toml index 03d4c097..9be61ff7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -177,6 +177,7 @@ filterwarnings = [ # collection of ``linopy/variables.py`` in the source tree on # Windows CI. "ignore:piecewise:FutureWarning", + "ignore:spec:FutureWarning", ] [tool.coverage.run] diff --git a/test/test_piecewise_constraints.py b/test/test_piecewise_constraints.py index 788a0674..6c4745e0 100644 --- a/test/test_piecewise_constraints.py +++ b/test/test_piecewise_constraints.py @@ -3266,7 +3266,7 @@ def _reset_dedup(self) -> Generator[None, None, None]: Warnings dedup is module-global so order between tests would otherwise matter. Clear before each test. """ - from linopy.piecewise import _emitted_evolving_warnings + from linopy.constants import _emitted_evolving_warnings _emitted_evolving_warnings.clear() yield diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index de593cb2..c5d02e10 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -5,6 +5,7 @@ from __future__ import annotations +import warnings from collections.abc import Callable from pathlib import Path from typing import Any @@ -236,3 +237,15 @@ def test_evaluate_refuses_sources_on_other_labels_than_the_model( } with pytest.raises(SpecDataError, match=f"dimension 'generator' {match}"): m.spec.evaluate("twice", sources) + + +def test_spec_api_warns_once_per_session() -> None: + from linopy import EvolvingAPIWarning + from linopy.constants import _emitted_evolving_warnings + + _emitted_evolving_warnings.discard("spec") + with pytest.warns(EvolvingAPIWarning, match="spec: Model.add_spec"): + Model.from_spec(EXAMPLE_DISPATCH, DISPATCH_DATA) + with warnings.catch_warnings(): + warnings.simplefilter("error", EvolvingAPIWarning) + Model.from_spec(EXAMPLE_DISPATCH, DISPATCH_DATA) From d2f82a6ff10e845e2594e96dbbf506debfdf42a0 Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 8 Sep 2026 11:33:36 +0200 Subject: [PATCH 24/35] ci(docs): install spec group on Read the Docs autodoc imports linopy.spec, which needs math-spec. RTD installed only the docs extra, so the Sphinx build failed with ImportError. Add a post_install job installing the spec dependency group (pip >= 25.1). --- .readthedocs.yaml | 6 ++++++ linopy/spec/{binder.py => attach.py} | 0 test/{test_spec_binder.py => test_spec_attach.py} | 0 3 files changed, 6 insertions(+) rename linopy/spec/{binder.py => attach.py} (100%) rename test/{test_spec_binder.py => test_spec_attach.py} (100%) diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 5eac0cca..0249aaa7 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -8,6 +8,12 @@ build: jobs: pre_system_dependencies: - git fetch --unshallow # Needed to get version tags + post_install: + # The spec API is documented via autodoc, which imports linopy.spec and so + # needs math-spec. It lives in the `spec` dependency group, not the docs + # extra; --group needs pip >= 25.1, hence the upgrade. + - python -m pip install --upgrade pip + - python -m pip install --group spec python: install: - method: pip diff --git a/linopy/spec/binder.py b/linopy/spec/attach.py similarity index 100% rename from linopy/spec/binder.py rename to linopy/spec/attach.py diff --git a/test/test_spec_binder.py b/test/test_spec_attach.py similarity index 100% rename from test/test_spec_binder.py rename to test/test_spec_attach.py From b774e7776e64ad46d0e4d14170dfbce3983d6dad Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 8 Sep 2026 11:33:36 +0200 Subject: [PATCH 25/35] refac(spec): rename data attachment to attach/Attached lpspec reserves "bound" for a variable/constraint limit and calls the data operation "attach". Rename bind()->attach(), Bound->Attached and binder.py->attach.py so "bound" names one thing. Variable-limit names (_bound, check_bounds_cover) stay untouched. --- doc/api.rst | 6 +- doc/release_notes.rst | 4 +- doc/user-guide.rst | 4 +- examples/building-models-from-specs.ipynb | 22 ++--- linopy/io.py | 2 +- linopy/spec/__init__.py | 6 +- linopy/spec/accessor.py | 23 ++--- linopy/spec/attach.py | 24 ++--- linopy/spec/builder.py | 18 ++-- linopy/spec/errors.py | 4 +- linopy/spec/evaluate.py | 2 +- linopy/spec/netcdf.py | 2 +- test/test_spec_attach.py | 102 +++++++++++----------- 13 files changed, 112 insertions(+), 107 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index eb9d0d7a..f74c7c5c 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -119,7 +119,7 @@ Building from specs ------------------- Build a model from a `math-spec -`__ YAML program bound to +`__ YAML program attached to data. Requires the ``spec`` dependency group. .. autosummary:: @@ -131,8 +131,8 @@ data. Requires the ``spec`` dependency group. spec.ModelSpec spec.NamedExpressions spec.NamedExpression - spec.bind - spec.Bound + spec.attach + spec.Attached spec.SpecDataError diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 7e6805e6..e3929d9d 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -23,9 +23,9 @@ Upcoming Version *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`` (a ``linopy.spec.ModelSpec``) reads it back. Requires the ``spec`` dependency group (``uv sync --group spec`` / ``uv pip install --group spec``, Python >= 3.12) and v1 semantics. Data is bound onto the spec's dimensions and parameters with ``linopy.spec.bind``, raising a ``linopy.spec.SpecDataError`` on mismatched or missing data; ``linopy.spec.Bound`` carries the bound result. The spec API emits an :class:`linopy.EvolvingAPIWarning` once per session while it stabilises. See :doc:`building-models-from-specs` for a worked example. +* ``Model.from_spec`` / ``model.add_spec`` build a model from a `math-spec `__ YAML program attached to data, and ``model.spec`` (a ``linopy.spec.ModelSpec``) reads it back. Requires the ``spec`` dependency group (``uv sync --group spec`` / ``uv pip install --group spec``, Python >= 3.12) and v1 semantics. Data is attached onto the spec's dimensions and parameters with ``linopy.spec.attach``, raising a ``linopy.spec.SpecDataError`` on mismatched or missing data; ``linopy.spec.Attached`` carries the attached result. The spec API emits an :class:`linopy.EvolvingAPIWarning` once per session while it stabilises. See :doc:`building-models-from-specs` for a worked example. -* ``model.spec.expressions`` (a ``linopy.spec.NamedExpressions`` mapping) returns a ``linopy.spec.NamedExpression`` for each declared name, with three views: ``.node`` (the lowered formula), ``.expression`` (the unsolved linopy expression — a ``LinearExpression``, bare ``Variable``, array or scalar) and ``.solution`` (the expression folded over the solved model). ``model.spec.evaluate(name, sources)`` returns the same object with its parameters bound afresh. +* ``model.spec.expressions`` (a ``linopy.spec.NamedExpressions`` mapping) returns a ``linopy.spec.NamedExpression`` for each declared name, with three views: ``.node`` (the lowered formula), ``.expression`` (the unsolved linopy expression — a ``LinearExpression``, bare ``Variable``, array or scalar) and ``.solution`` (the expression folded over the solved model). ``model.spec.evaluate(name, sources)`` returns the same object with its parameters attached afresh. * ``model.spec.to_latex`` / ``.to_markdown`` / ``.to_typst`` typeset the whole model; a ``ModelSpec`` and a ``NamedExpression`` render as Markdown in a notebook. diff --git a/doc/user-guide.rst b/doc/user-guide.rst index 33589f38..fcf43a07 100644 --- a/doc/user-guide.rst +++ b/doc/user-guide.rst @@ -88,11 +88,11 @@ Building a model from a spec Instead of calling ``add_variables`` / ``add_constraints`` directly, you can declare a model as a `math-spec -`__ YAML program bound to +`__ YAML program attached to data, and let linopy build it. - :doc:`building-models-from-specs` — ``Model.from_spec`` and - ``model.add_spec``, binding data to a spec, and reading named + ``model.add_spec``, attaching data to a spec, and reading named expressions back through ``model.spec`` after solving. Requires the ``spec`` dependency group and v1 semantics. diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb index 3e7cf7fd..d25781b5 100644 --- a/examples/building-models-from-specs.ipynb +++ b/examples/building-models-from-specs.ipynb @@ -14,7 +14,7 @@ "\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", + "labelled axes; you supply the labels and the values separately. `linopy` attaches\n", "the two together and emits variables, constraints and an objective that align\n", "and broadcast by dimension, exactly as if you had written them by hand.\n", "\n", @@ -22,7 +22,7 @@ "\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", + "3. Attaching data and building a model with `Model.from_spec`.\n", "4. Solving, and folding **named expressions** back into arrays.\n", "5. `retain` modes and `evaluate` — what data stays on the model.\n", "6. **Absence and coverage** — the rule that decides when a missing row is\n", @@ -161,7 +161,7 @@ "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", + "Three attachment rules are worth knowing, because they make the result\n", "predictable:\n", "\n", "1. A dimension's members come **only** from the source keyed by that\n", @@ -198,7 +198,7 @@ "source": [ "## 3. Building the model\n", "\n", - "`Model.from_spec(spec, sources)` lowers the spec, binds the data and emits a\n", + "`Model.from_spec(spec, sources)` lowers the spec, attaches the data and emits a\n", "normal linopy `Model`. The `spec` argument is flexible: a path, YAML text, a\n", "`dict`, or a `math_spec.Spec`. (A pre-lowered `Program` is refused — it has no\n", "YAML form to keep on the model.)\n", @@ -289,10 +289,10 @@ "\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", + " and parameters already attached. A `LinearExpression`, a bare `Variable`, an\n", " array, or a scalar (a named expression is affine, so never quadratic).\n", "- `.solution` — the expression **folded** over the solution: every variable\n", - " replaced by its solved value, every parameter by the data it was bound to, the\n", + " replaced by its solved value, every parameter by the data it was attached to, the\n", " arithmetic run on xarray.\n", "\n", "`spend` = `sum(p * cost, over=generator)` folds to the cost incurred each hour;\n", @@ -409,7 +409,7 @@ "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", + "parameters are reattached from a **fresh** bag of data, folding against the model's\n", "solution.\n", "\n", "The catch: `evaluate` reads the solution the model already holds, so the fresh\n", @@ -433,7 +433,7 @@ "except SpecDataError as e:\n", " print(\"SpecDataError:\", str(e)[:90], \"...\\n\")\n", "\n", - "# evaluate rebinds from fresh sources; .solution folds:\n", + "# evaluate reattaches from fresh sources; .solution folds:\n", "print(lean.spec.evaluate(\"spend\", dispatch_data).solution)" ] }, @@ -895,7 +895,7 @@ "id": "46", "metadata": {}, "source": [ - "`Model.copy()` carries the spec too, with the accessor rebound to the copy. The\n", + "`Model.copy()` carries the spec too, with the accessor reattached to the copy. The\n", "copy is a fresh, unsolved model (like any linopy copy), so solve it before\n", "folding an expression that reads a variable — the folded result then matches the\n", "original." @@ -931,7 +931,7 @@ "\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", + "- `attach.py` — the three attachment rules; data onto master coordinates.\n", "- `builder.py` — emits variables, constraints, objective; folds expressions.\n", "- `operators.py` — `sum`, `by=`, `shift`, `at`, `sum_back`.\n", "- `where.py` — `where:` predicates as boolean masks.\n", @@ -954,7 +954,7 @@ "### Summary\n", "\n", "A spec is the maths over labelled axes; the sources are the numbers. `linopy`\n", - "binds them into an ordinary model, hands each named expression back as three\n", + "attaches them into an ordinary model, hands each named expression back as three\n", "views — its formula, its unsolved linopy expression and its solution — refuses a\n", "missing parameter row wherever it is used (as a coefficient, bound, constant\n", "side or divisor alike, with `where:` and filling the data as the escape\n", diff --git a/linopy/io.py b/linopy/io.py index 5cc9d001..5e1b09d2 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1499,7 +1499,7 @@ 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._spec = m._spec._reattach(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/__init__.py b/linopy/spec/__init__.py index a6f653ac..e9bd758e 100644 --- a/linopy/spec/__init__.py +++ b/linopy/spec/__init__.py @@ -22,16 +22,16 @@ NamedExpressions, SpecLike, ) -from linopy.spec.binder import Bound, Retain, bind +from linopy.spec.attach import Attached, Retain, attach from linopy.spec.errors import SpecDataError __all__ = [ - "Bound", + "Attached", "ModelSpec", "NamedExpression", "NamedExpressions", "Retain", "SpecDataError", "SpecLike", - "bind", + "attach", ] diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index c2a6d485..4a790c55 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -32,7 +32,8 @@ from linopy.model import Model from linopy.semantics import is_v1 from linopy.spec import terms -from linopy.spec.binder import EVOLVING_MESSAGE, Bound, Retain, bind +from linopy.spec.attach import EVOLVING_MESSAGE, Attached, Retain +from linopy.spec.attach import attach as attach_data from linopy.spec.builder import build from linopy.spec.context import Context from linopy.spec.errors import SpecDataError @@ -72,9 +73,9 @@ def attach( 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)) + attached: Attached = attach_data(program, sources, retain=retain) + build(model, attached) + model.parameters = attached.retained().assign_coords(dict(attached.coords)) return ModelSpec(model, program, text) @@ -121,7 +122,7 @@ def __repr__(self) -> str: names = list(self.program.named_expressions) return f"ModelSpec(expressions={names})" - def _rebound(self, model: Model) -> ModelSpec: + def _reattach(self, model: Model) -> ModelSpec: """The same spec, read off *model*.""" return ModelSpec(model, self.program, self.text) @@ -172,7 +173,7 @@ def evaluate( self, name: str, sources: Mapping[str, Any] | xr.Dataset ) -> NamedExpression: """ - The named expression *name*, with its parameters bound afresh from *sources*. + The named expression *name*, with its parameters attached 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 @@ -185,16 +186,16 @@ def evaluate( *sources* label a dimension differently than the model was built on. """ - bound = bind(self.program, sources, retain="none") + attached = attach_data(self.program, sources, retain="none") coords = self.coords - for dim, index in bound.coords.items(): + for dim, index in attached.coords.items(): if dim in coords and not index.equals(coords[dim]): raise SpecDataError( f"sources describe dimension '{dim}' as {index.tolist()[:5]}, and the model " f"was built on {coords[dim].tolist()[:5]}. evaluate() reads the solution the " - f"model holds, so the data must be bound on the same labels in the same order." + f"model holds, so the data must be attached on the same labels in the same order." ) - return NamedExpression(self, name, self._context(bound.parameter)) + return NamedExpression(self, name, self._context(attached.parameter)) def _retained(self, name: str) -> xr.DataArray: if name not in self.parameters: @@ -248,7 +249,7 @@ class NamedExpression: 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. + the solution the model holds; ``evaluate(name, sources)`` attaches fresh data. Attributes ---------- diff --git a/linopy/spec/attach.py b/linopy/spec/attach.py index 0d2bf723..be0fe001 100644 --- a/linopy/spec/attach.py +++ b/linopy/spec/attach.py @@ -1,7 +1,7 @@ """ -Bind user data to a math-spec program. +Attach user data to a math-spec program. -The language fixes three binding rules and this module enforces them: a +The language fixes three attachment rules and this module enforces them: a dimension's members come only from the source keyed by the dimension's name, their order is the source's order and is never sorted, and a parameter or lookup source is read for values, never for labels. Parameters @@ -62,20 +62,20 @@ ) EVOLVING_MESSAGE = ( - "spec: Model.add_spec, Model.from_spec, model.spec and linopy.spec.bind are " + "spec: Model.add_spec, Model.from_spec, model.spec and linopy.spec.attach are " "newly added and their details may change in minor releases. Silence with " '`warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning)`.' ) -def bind( +def attach( program: ms.Program, sources: Mapping[str, Any] | xr.Dataset, *, retain: Retain = "report", -) -> Bound: +) -> Attached: """ - Bind *sources* to *program*: master coordinates now, parameters on demand. + Attach *sources* to *program*: master coordinates now, parameters on demand. Parameters ---------- @@ -87,7 +87,7 @@ def bind( is accepted too: its indexes are dimension sources, its data variables parameters and lookups. retain - Which parameters :meth:`Bound.retained` persists. + Which parameters :meth:`Attached.retained` persists. Raises ------ @@ -108,18 +108,18 @@ def bind( _check_keys(program, keys) coords = _master_coords(program, sources, keys) lookups = _lookups(program, sources, keys, coords) - return Bound(program, coords, lookups, retain, sources, keys) + return Attached(program, coords, lookups, retain, sources, keys) @dataclass(frozen=True, eq=False) -class Bound: +class Attached: """ - A program bound to its data. + A program attached to its data. Attributes ---------- program - The lowered spec the data is bound to. + The lowered spec the data is attached to. coords Master coordinates by dimension, in source order, each index named after its dimension. A declared dimension nothing reaches @@ -180,7 +180,7 @@ def _declaration(self, name: str) -> ms.ParameterDeclaration: if declared.derivation is not None: raise SpecDataError( f"parameter '{name}' is emitted by piecewise block '{declared.derivation.block}' " - f"and is filled from the block's own breakpoints, not bound from sources." + f"and is filled from the block's own breakpoints, not attached from sources." ) return declared diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py index ac7a54c1..3d825a84 100644 --- a/linopy/spec/builder.py +++ b/linopy/spec/builder.py @@ -1,5 +1,5 @@ """ -Program plus bound data to linopy declarations. +Program plus attached data to linopy declarations. A build hands every variable to linopy as its term, then adds special-ordered sets, constraints and the objective; which linopy call each construct becomes @@ -14,7 +14,7 @@ from linopy.expressions import LinearExpression, QuadraticExpression from linopy.model import Model from linopy.spec import curves -from linopy.spec.binder import Bound +from linopy.spec.attach import Attached from linopy.spec.context import Context from linopy.spec.coverage import check_bounds_cover, check_coverage from linopy.spec.errors import SpecDataError @@ -29,9 +29,9 @@ _SENSE = {"minimize": "min", "maximize": "max"} -def build(model: Model, bound: Bound) -> None: +def build(model: Model, attached: Attached) -> None: """ - Add every declaration of the bound program to *model*. + Add every declaration of the attached program to *model*. Variables, special-ordered sets, constraints and the objective, in that order; then every named expression is checked for divisor and coefficient @@ -40,10 +40,10 @@ def build(model: Model, bound: Bound) -> None: """ ctx = Context( model, - bound.program, - bound.coords, - bound.lookups, - Parameters(bound.program, bound.parameter), + attached.program, + attached.coords, + attached.lookups, + Parameters(attached.program, attached.parameter), ) curves.validate(ctx.program, ctx.parameters) _variables(ctx) @@ -130,6 +130,6 @@ def _objective(ctx: Context) -> 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" + "the objective carries no variable term once the data is attached, so there is nothing to optimize" ) ctx.model.add_objective(expr, overwrite=True, sense=_SENSE[declared.sense]) diff --git a/linopy/spec/errors.py b/linopy/spec/errors.py index 7e075add..88f9fa57 100644 --- a/linopy/spec/errors.py +++ b/linopy/spec/errors.py @@ -1,11 +1,11 @@ -"""Errors raised while binding data to a math-spec program.""" +"""Errors raised while attaching data to a math-spec program.""" from __future__ import annotations class SpecDataError(ValueError): """ - Data bound to a valid spec is missing, malformed or the wrong shape. + Data attached to a valid spec is missing, malformed or the wrong shape. Every refusal names the symbol, the dimension(s) and the offending labels, so the message points back at the ``sources`` entry to fix. diff --git a/linopy/spec/evaluate.py b/linopy/spec/evaluate.py index 99d712ed..a2645e11 100644 --- a/linopy/spec/evaluate.py +++ b/linopy/spec/evaluate.py @@ -143,7 +143,7 @@ def _combine(op: Callable[[Value, Value], Value], left: Value, right: Value) -> 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." + f"coordinates, so the data was attached against other labels than the model was built on." ) elif isinstance(left, xr.DataArray) and isinstance( right, Variable | LinearExpression | QuadraticExpression diff --git a/linopy/spec/netcdf.py b/linopy/spec/netcdf.py index 9cfcfc22..2a62e1ea 100644 --- a/linopy/spec/netcdf.py +++ b/linopy/spec/netcdf.py @@ -13,7 +13,7 @@ empty string, indistinguishable from a label. So a lookup, and any array of objects, is written instead as integer codes into its own table of categories, ``-1`` where a label is missing. Decoding indexes the table and -fills the holes back in, which reproduces what the binder built, values and +fills the holes back in, which reproduces what attach built, values and dtype alike. The master coordinates are canonical: a container's coordinates for a diff --git a/test/test_spec_attach.py b/test/test_spec_attach.py index d5ed8b1d..56b0175d 100644 --- a/test/test_spec_attach.py +++ b/test/test_spec_attach.py @@ -12,7 +12,7 @@ math_spec = pytest.importorskip("math_spec") -from linopy.spec import SpecDataError, bind # noqa: E402 +from linopy.spec import SpecDataError, attach # noqa: E402 SPEC: dict[str, Any] = { "dimensions": {"f": {"dtype": "str"}, "t": {"dtype": "int"}, "g": {"dtype": "str"}}, @@ -91,8 +91,8 @@ def good() -> dict[str, Any]: def read_all(program: Any, sources: Mapping[str, Any]) -> list[xr.DataArray]: - bound = bind(program, sources) - return [bound.parameter(name) for name in program.parameters] + attached = attach(program, sources) + return [attached.parameter(name) for name in program.parameters] CAP_SHAPES = { @@ -113,10 +113,10 @@ def read_all(program: Any, sources: Mapping[str, Any]) -> list[xr.DataArray]: @pytest.mark.parametrize("cap", CAP_SHAPES.values(), ids=CAP_SHAPES.keys()) -def test_rank_two_shapes_bind_alike( +def test_rank_two_shapes_attach_alike( program: Any, good: dict[str, Any], cap: Any ) -> None: - got = bind(program, {**good, "cap": cap}).parameter("cap") + got = attach(program, {**good, "cap": cap}).parameter("cap") xr.testing.assert_equal(got, CAP) assert got.dims == ("f", "t") @@ -132,10 +132,10 @@ def test_rank_two_shapes_bind_alike( @pytest.mark.parametrize("cost", COST_SHAPES.values(), ids=COST_SHAPES.keys()) -def test_rank_one_shapes_bind_alike( +def test_rank_one_shapes_attach_alike( program: Any, good: dict[str, Any], cost: Any ) -> None: - got = bind(program, {**good, "cost": cost}).parameter("cost") + got = attach(program, {**good, "cost": cost}).parameter("cost") xr.testing.assert_equal(got, xr.DataArray(COST, name="cost")) @@ -153,7 +153,7 @@ def test_rank_one_shapes_bind_alike( def test_dimension_shapes_keep_source_order( program: Any, good: dict[str, Any], f: Any ) -> None: - coords = bind(program, {**good, "f": f}).coords + coords = attach(program, {**good, "f": f}).coords assert coords["f"].tolist() == ["b", "a", "c"] assert coords["f"].name == "f" assert list(coords) == ["f", "t", "g"] @@ -162,8 +162,8 @@ def test_dimension_shapes_keep_source_order( def test_lookup_is_padded_onto_the_dimension( program: Any, good: dict[str, Any] ) -> None: - bound = bind(program, {**good, "grp": {"a": "n"}}) - grp = bound.lookups["f"]["grp"] + attached = attach(program, {**good, "grp": {"a": "n"}}) + grp = attached.lookups["f"]["grp"] assert grp.dims == ("f",) assert grp.sel(f="a").item() == "n" assert pd.isna(grp.sel(f=["b", "c"])).all() @@ -178,21 +178,23 @@ def test_lookup_is_padded_onto_the_dimension( @pytest.mark.parametrize("grp", LOOKUP_SHAPES.values(), ids=LOOKUP_SHAPES.keys()) -def test_lookup_shapes_bind_alike(program: Any, good: dict[str, Any], grp: Any) -> None: - got = bind(program, {**good, "grp": grp}).lookups["f"]["grp"] +def test_lookup_shapes_attach_alike( + program: Any, good: dict[str, Any], grp: Any +) -> None: + got = attach(program, {**good, "grp": grp}).lookups["f"]["grp"] assert got.values.tolist() == ["n", "e", "n"] @pytest.mark.parametrize("storage", ["python", "pyarrow"]) @pytest.mark.parametrize("shape", ["series", "dataarray"]) -def test_extension_strings_bind_as_numpy_objects( +def test_extension_strings_attach_as_numpy_objects( program: Any, good: dict[str, Any], storage: str, shape: str ) -> None: if storage == "pyarrow": pytest.importorskip("pyarrow") series = pd.Series(["n", "e"], index=F[:2], dtype=pd.StringDtype(storage)) grp = xr.DataArray(series) if shape == "dataarray" else series - got = bind(program, {**good, "grp": grp}).lookups["f"]["grp"] + got = attach(program, {**good, "grp": grp}).lookups["f"]["grp"] assert got.dtype == np.dtype(object) assert got.values[:2].tolist() == ["n", "e"] assert pd.isna(got.values[2]) @@ -207,18 +209,18 @@ def test_missing_rows_become_nan_and_false(program: Any, good: dict[str, Any]) - "flag": pd.Series({"a": True}), "cap": CAP.sel(t=[0, 1]), } - bound = bind(program, sparse) - cost = bound.parameter("cost") + attached = attach(program, sparse) + cost = attached.parameter("cost") assert cost.sel(f="a").item() == 1.0 assert cost.sel(f=["b", "c"]).isnull().all() - lead = bound.parameter("lead") + lead = attached.parameter("lead") assert lead.dtype == np.float64 assert lead.sel(f="a").item() == 1.0 assert lead.sel(f=["b", "c"]).isnull().all() - flag = bound.parameter("flag") + flag = attached.parameter("flag") assert flag.dtype == bool assert flag.values.tolist() == [False, True, False] - cap = bound.parameter("cap") + cap = attached.parameter("cap") assert cap.dims == ("f", "t") assert cap.sel(t=2).isnull().all() @@ -235,14 +237,14 @@ def test_missing_rows_become_nan_and_false(program: Any, good: dict[str, Any]) - def test_scalar_is_broadcast_over_declared_dims( program: Any, good: dict[str, Any], name: str, value: Any, expected_dtype: Any ) -> None: - got = bind(program, {**good, name: value}).parameter(name) + got = attach(program, {**good, name: value}).parameter(name) assert got.dims == tuple(SPEC["parameters"][name]["dims"]) assert got.dtype == expected_dtype assert (got == value).all() def test_scalar_parameter_stays_scalar(program: Any, good: dict[str, Any]) -> None: - got = bind(program, good).parameter("rate") + got = attach(program, good).parameter("rate") assert got.dims == () assert got.item() == 0.5 @@ -255,10 +257,10 @@ def test_scalar_parameter_stays_scalar(program: Any, good: dict[str, Any]) -> No @pytest.mark.parametrize("cost", EMPTY_SOURCES.values(), ids=EMPTY_SOURCES.keys()) -def test_empty_source_binds_as_all_nan( +def test_empty_source_attaches_as_all_nan( program: Any, good: dict[str, Any], cost: Any ) -> None: - got = bind(program, {**good, "cost": cost}).parameter("cost") + got = attach(program, {**good, "cost": cost}).parameter("cost") assert got.dtype == np.float64 assert got.isnull().all() assert got.indexes["f"].equals(F) @@ -268,21 +270,23 @@ def test_missing_parameter_is_refused_when_read( program: Any, good: dict[str, Any] ) -> None: good.pop("cost") - bound = bind(program, good) + attached = attach(program, good) with pytest.raises(SpecDataError, match="no data provided for parameter 'cost'"): - bound.parameter("cost") + attached.parameter("cost") def test_undeclared_parameter_is_refused_with_a_hint( program: Any, good: dict[str, Any] ) -> None: with pytest.raises(SpecDataError, match="unknown parameter 'csot'.*'cost'"): - bind(program, good).parameter("csot") + attach(program, good).parameter("csot") -def test_retain_is_validated_before_binding(program: Any, good: dict[str, Any]) -> None: +def test_retain_is_validated_before_attaching( + program: Any, good: dict[str, Any] +) -> None: with pytest.raises(SpecDataError, match=r"'report', 'all', 'none'") as error: - bind(program, good, retain="reports") # type: ignore[arg-type] + attach(program, good, retain="reports") # type: ignore[arg-type] assert "Did you mean 'report'?" in str(error.value) @@ -475,8 +479,8 @@ def test_int_labels_are_shown_as_written(program: Any, good: dict[str, Any]) -> def test_dataset_is_a_source(program: Any, good: dict[str, Any]) -> None: dims = {"f": F, "t": T, "g": ["n", "e"]} values = {k: xr.DataArray(v) for k, v in good.items() if k not in dims} - from_dataset = bind(program, xr.Dataset(values, coords=dims)) - from_mapping = bind(program, good) + from_dataset = attach(program, xr.Dataset(values, coords=dims)) + from_mapping = attach(program, good) assert from_dataset.coords["f"].equals(from_mapping.coords["f"]) for name in program.parameters: xr.testing.assert_equal( @@ -510,10 +514,10 @@ def test_sources_are_pulled_by_key_on_demand( program: Any, good: dict[str, Any] ) -> None: sources = Counting(good) - bound = bind(program, sources) + attached = attach(program, sources) assert set(sources.pulled) == {"f", "t", "g", "grp"} - bound.parameter("cap") - bound.parameter("cap") + attached.parameter("cap") + attached.parameter("cap") assert sources.pulled.count("cap") == 2 @@ -528,7 +532,7 @@ def test_sources_are_pulled_by_key_on_demand( def test_retained_follows_the_named_expressions( program: Any, good: dict[str, Any], retain: Any, expected: set[str] ) -> None: - retained = bind(program, good, retain=retain).retained() + retained = attach(program, good, retain=retain).retained() assert set(retained.data_vars) == expected assert retained.coords["f"].values.tolist() == ["b", "a", "c"] @@ -570,14 +574,14 @@ def test_report_closure_reads_names_and_masks() -> None: "on": pd.Series([True], index=f), "other": pd.Series([2.0], index=f), } - retained = bind(program, sources).retained() + retained = attach(program, sources).retained() assert set(retained.data_vars) == {"cost", "lag", "span", "on"} def test_unreached_dimension_needs_no_source() -> None: dimensions = {**PARITY_SPEC["dimensions"], "z": {"dtype": "int"}} program = math_spec.to_program({**PARITY_SPEC, "dimensions": dimensions}) - assert list(bind(program, GOOD).coords) == ["f"] + assert list(attach(program, GOOD).coords) == ["f"] @pytest.mark.parametrize("shape", ["dataarray", "dataarray-transposed", "wide-frame"]) @@ -585,16 +589,16 @@ def test_aligned_array_is_not_copied( program: Any, good: dict[str, Any], shape: str ) -> None: source = CAP_SHAPES[shape] - bound = bind(program, {**good, "cap": source}) - assert np.shares_memory(np.asarray(source), bound.parameter("cap").values) - assert np.shares_memory(np.asarray(source), bound.parameter("cap").values) + attached = attach(program, {**good, "cap": source}) + assert np.shares_memory(np.asarray(source), attached.parameter("cap").values) + assert np.shares_memory(np.asarray(source), attached.parameter("cap").values) def test_master_coordinate_dtype_wins_without_a_copy( program: Any, good: dict[str, Any] ) -> None: source = CAP.assign_coords(t=T.astype("int32")) - got = bind(program, {**good, "cap": source}).parameter("cap") + got = attach(program, {**good, "cap": source}).parameter("cap") assert got.indexes["t"].dtype == np.int64 assert np.shares_memory(np.asarray(source), got.values) @@ -626,12 +630,12 @@ def test_derived_parameter_is_not_bound_from_sources() -> None: "bp_x": pd.Series([0.0, 5.0, 10.0], index=bp), "bp_y": pd.Series([0.0, 2.0, 8.0], index=bp), } - bound = bind(program, sources, retain="all") - assert set(bound.retained().data_vars) == {"bp_x", "bp_y"} + attached = attach(program, sources, retain="all") + assert set(attached.retained().data_vars) == {"bp_x", "bp_y"} with pytest.raises(SpecDataError, match="emitted by piecewise block 'curve'"): - bound.parameter(derived[0]) + attached.parameter(derived[0]) with pytest.raises(SpecDataError, match=derived[0]): - bind(program, {**sources, derived[0]: 1.0}) + attach(program, {**sources, derived[0]: 1.0}) # --------------------------------------------------------------------------- @@ -713,11 +717,11 @@ def test_a_hole_is_named_where_it_sits() -> None: ), ], ) -def test_a_flag_binds_by_its_declaration(column: pd.Series, verdict: Any) -> None: +def test_a_flag_attaches_by_its_declaration(column: pd.Series, verdict: Any) -> None: program = math_spec.to_program(FLAG_SPEC) sources = {"g": ["a", "b"], "active": column} if verdict is ACCEPTED: - assert bind(program, sources).parameter("active").dtype == bool + assert attach(program, sources).parameter("active").dtype == bool return with pytest.raises(SpecDataError, match="declared 'bool'"): read_all(program, sources) @@ -787,8 +791,8 @@ def test_a_lookup_defect_is_refused(override: dict[str, Any], match: str) -> Non def test_a_label_space_lookup_is_padded_onto_the_dimension() -> None: program = math_spec.to_program(TAG_SPEC) - bound = bind(program, {**TAG_GOOD, "tag": {"s": 7}}) - tag = bound.lookups["g"]["tag"] + attached = attach(program, {**TAG_GOOD, "tag": {"s": 7}}) + tag = attached.lookups["g"]["tag"] assert tag.dims == ("g",) assert tag.indexes["g"].tolist() == ["w", "s"] assert np.isnan(tag.sel(g="w").item()) @@ -809,7 +813,7 @@ def test_a_label_space_lookup_defect_is_refused( ) -> None: program = math_spec.to_program(TAG_SPEC) with pytest.raises(SpecDataError, match=match): - bind(program, {**TAG_GOOD, "tag": tag}) + attach(program, {**TAG_GOOD, "tag": tag}) def test_a_stray_lookup_value_over_an_int_target_is_shown_as_written() -> None: From 46c6f2789adeeaf3f0aedaa07130d1feb7e4097e Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 8 Sep 2026 11:53:52 +0200 Subject: [PATCH 26/35] doc(spec): execute the spec notebook on Read the Docs RTD now installs the spec group, so the notebook can import math-spec and solve. Drop `nbsphinx.execute: never` so it runs at build time and its outputs render (nbstripout keeps them out of git). --- examples/building-models-from-specs.ipynb | 3 --- 1 file changed, 3 deletions(-) diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb index d25781b5..40590ef3 100644 --- a/examples/building-models-from-specs.ipynb +++ b/examples/building-models-from-specs.ipynb @@ -980,9 +980,6 @@ "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" - }, - "nbsphinx": { - "execute": "never" } }, "nbformat": 4, From b7a8115e19ac2dca609ab9e56ba4bd3c081ebbb9 Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 8 Sep 2026 12:31:48 +0200 Subject: [PATCH 27/35] refac(spec): track math-spec a75, unwrapping ExpressionDeclaration named_expressions now maps each name to an ExpressionDeclaration that carries the body and whether the math reads it, so read .expression at the five sites that take a body: the evaluator, the builder's coverage check, the report closure and the accessor's node view. --- linopy/spec/accessor.py | 2 +- linopy/spec/attach.py | 2 +- linopy/spec/builder.py | 4 ++-- linopy/spec/evaluate.py | 2 +- pyproject.toml | 2 +- test/test_spec_accessor.py | 2 +- 6 files changed, 7 insertions(+), 7 deletions(-) diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index 4a790c55..e9e74b39 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -265,7 +265,7 @@ def __init__(self, spec: ModelSpec, name: str, ctx: Context) -> None: @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] + return self._spec.program.named_expressions[self._name].expression @functools.cached_property def expression(self) -> terms.Value: diff --git a/linopy/spec/attach.py b/linopy/spec/attach.py index be0fe001..e98ac540 100644 --- a/linopy/spec/attach.py +++ b/linopy/spec/attach.py @@ -196,7 +196,7 @@ 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()) + bodies = tuple(d.expression for d in program.named_expressions.values()) names = set(parameters_of(*bodies)) for node in walk(*bodies): names.update(amounts_of(node)) diff --git a/linopy/spec/builder.py b/linopy/spec/builder.py index 3d825a84..6170c55c 100644 --- a/linopy/spec/builder.py +++ b/linopy/spec/builder.py @@ -50,8 +50,8 @@ def build(model: Model, attached: Attached) -> None: _sos(ctx) _constraints(ctx) _objective(ctx) - for name, body in ctx.program.named_expressions.items(): - check_coverage(f"expression '{name}'", (body,), ctx, None) + for name, declared in ctx.program.named_expressions.items(): + check_coverage(f"expression '{name}'", (declared.expression,), ctx, None) def _variables(ctx: Context) -> None: diff --git a/linopy/spec/evaluate.py b/linopy/spec/evaluate.py index a2645e11..c7bb05bd 100644 --- a/linopy/spec/evaluate.py +++ b/linopy/spec/evaluate.py @@ -28,7 +28,7 @@ def evaluate_named(name: str, ctx: Context) -> Value: f"unknown named expression '{name}'. " + did_you_mean(name, ctx.program.named_expressions) ) - body = ctx.program.named_expressions[name] + body = ctx.program.named_expressions[name].expression found = obligations_of((body,), ctx, None) check_divisors(f"expression '{name}'", found.divisors, ctx) value = evaluate(body, ctx) diff --git a/pyproject.toml b/pyproject.toml index 9be61ff7..482334e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,7 +123,7 @@ gpu = [ # keeps the git pin out of the published wheel metadata, which PyPI rejects. # Install with `uv sync --group spec` or `uv pip install --group spec`. spec = [ - "math-spec @ git+https://github.com/energy-models/math-spec.git@1377f27b759cfbc42bb205338e79751542dabe84 ; python_version >= '3.12'", + "math-spec @ git+https://github.com/energy-models/math-spec.git@67aeedb988ee95d196456b4cbe50821e3799edaa ; python_version >= '3.12'", "pyyaml ; python_version >= '3.12'", "pyarrow ; python_version >= '3.12'", ] diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index c5d02e10..a0381ddb 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -184,7 +184,7 @@ def test_expression_reads_unsolved_but_solution_waits_for_a_solve() -> None: 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 e.node is m.spec.program.named_expressions["spend"].expression assert isinstance(e.expression, linopy.LinearExpression) xr.testing.assert_allclose( e.solution, (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") From f243de29384b292c1ce2a6f7774a52f0e8b5008a Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 8 Sep 2026 12:31:56 +0200 Subject: [PATCH 28/35] feat(spec): read a constraint's dual in a reported expression math-spec a75 lets a reported expression hold a dual(constraint) node. Evaluate it as the solved constraint's dual over its own frame, and refuse it before a solve the way an unsolved variable is refused. --- linopy/spec/evaluate.py | 10 ++++++++++ test/test_spec_accessor.py | 16 ++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/linopy/spec/evaluate.py b/linopy/spec/evaluate.py index c7bb05bd..6571e90d 100644 --- a/linopy/spec/evaluate.py +++ b/linopy/spec/evaluate.py @@ -59,6 +59,8 @@ def evaluate(node: ms.ExpressionNode, ctx: Context) -> Value: return node.value if isinstance(node, ms.Variable): return _variable(node.name, ctx) + if isinstance(node, ms.Dual): + return _dual(node.constraint, ctx) if isinstance(node, ms.Parameter): return terms.coefficient(ctx.parameters[node.name]) if isinstance(node, ms.Negate): @@ -135,6 +137,14 @@ def _variable(name: str, ctx: Context) -> Value: return terms.solution(variable, absence) +def _dual(constraint: str, ctx: Context) -> xr.DataArray: + if not ctx.solved: + raise RuntimeError( + f"constraint '{constraint}' has no dual yet: solve the model before reading a dual" + ) + return ctx.model.constraints[constraint].dual + + def _combine(op: Callable[[Value, Value], Value], left: Value, right: Value) -> Value: """*left* and *right* combined by *op*, once two arrays agree on their shared coordinates and a hole beside a term has become its absence.""" if isinstance(left, xr.DataArray) and isinstance(right, xr.DataArray): diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index a0381ddb..3693a342 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -239,6 +239,22 @@ def test_evaluate_refuses_sources_on_other_labels_than_the_model( m.spec.evaluate("twice", sources) +def test_a_reported_dual_folds_to_the_constraint_dual() -> None: + spec = {**yaml_dict(), "expressions": {"price": "dual(power_balance)"}} + m = solved(spec, DISPATCH_DATA) + xr.testing.assert_allclose( + m.spec.expressions["price"].solution, + m.constraints["power_balance"].dual.rename("price"), + ) + + +def test_a_dual_needs_a_solution() -> None: + spec = {**yaml_dict(), "expressions": {"price": "dual(power_balance)"}} + m = Model.from_spec(spec, DISPATCH_DATA) + with pytest.raises(RuntimeError, match="no dual yet"): + m.spec.expressions["price"].expression + + def test_spec_api_warns_once_per_session() -> None: from linopy import EvolvingAPIWarning from linopy.constants import _emitted_evolving_warnings From 2b12d1072bfc11f1f1a757ed6c19e203a42de4e1 Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 8 Sep 2026 13:11:27 +0200 Subject: [PATCH 29/35] feat(spec): typeset a single named expression Add to_latex/to_markdown/to_typst on NamedExpression via math-spec's typeset_declaration, rendering one expression as a bare line, and fix its _repr_markdown_ to show only itself. Tests, release note and the notebook updated. --- doc/release_notes.rst | 2 +- examples/building-models-from-specs.ipynb | 96 +++++++++++++---------- linopy/spec/accessor.py | 17 +++- test/test_spec_accessor.py | 21 +++++ 4 files changed, 93 insertions(+), 43 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index e3929d9d..e420159c 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -27,7 +27,7 @@ Upcoming Version * ``model.spec.expressions`` (a ``linopy.spec.NamedExpressions`` mapping) returns a ``linopy.spec.NamedExpression`` for each declared name, with three views: ``.node`` (the lowered formula), ``.expression`` (the unsolved linopy expression — a ``LinearExpression``, bare ``Variable``, array or scalar) and ``.solution`` (the expression folded over the solved model). ``model.spec.evaluate(name, sources)`` returns the same object with its parameters attached afresh. -* ``model.spec.to_latex`` / ``.to_markdown`` / ``.to_typst`` typeset the whole model; a ``ModelSpec`` and a ``NamedExpression`` render as Markdown in a notebook. +* ``model.spec.to_latex`` / ``.to_markdown`` / ``.to_typst`` typeset the whole model, and the same three methods on a ``NamedExpression`` typeset that one expression as a single line (math only, no document); both a ``ModelSpec`` and a ``NamedExpression`` render as Markdown in a notebook. *Numerical scaling* diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb index 40590ef3..62b7957e 100644 --- a/examples/building-models-from-specs.ipynb +++ b/examples/building-models-from-specs.ipynb @@ -328,7 +328,12 @@ "\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." + "accessor renders as Markdown on its own; here we show it explicitly.\n", + "\n", + "Each `NamedExpression` typesets on its own too, with the same three methods.\n", + "These render **one** expression as a single line — math only, no surrounding\n", + "document — so the string drops straight into a docstring or a table cell, and a\n", + "`NamedExpression` renders as its own formula in a notebook." ] }, { @@ -344,9 +349,23 @@ ] }, { - "cell_type": "markdown", + "cell_type": "code", + "execution_count": null, "id": "18", "metadata": {}, + "outputs": [], + "source": [ + "print(\"spend.to_latex():\")\n", + "print(spend.to_latex())\n", + "\n", + "# renders as its own formula in a notebook:\n", + "Markdown(f\"$$\\n{spend.to_markdown()}\\n$$\")" + ] + }, + { + "cell_type": "markdown", + "id": "19", + "metadata": {}, "source": [ "A named expression that reads only data (no variables) has a `.solution`\n", "**before** a solve too — it needs a solution only if it actually references a\n", @@ -357,7 +376,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "20", "metadata": {}, "outputs": [], "source": [ @@ -369,7 +388,7 @@ }, { "cell_type": "markdown", - "id": "20", + "id": "21", "metadata": {}, "source": [ "## 5. `retain`: what data stays on the model\n", @@ -390,7 +409,7 @@ { "cell_type": "code", "execution_count": null, - "id": "21", + "id": "22", "metadata": {}, "outputs": [], "source": [ @@ -401,7 +420,7 @@ }, { "cell_type": "markdown", - "id": "22", + "id": "23", "metadata": {}, "source": [ "### `evaluate`: fold against fresh data\n", @@ -420,7 +439,7 @@ { "cell_type": "code", "execution_count": null, - "id": "23", + "id": "24", "metadata": {}, "outputs": [], "source": [ @@ -440,7 +459,7 @@ { "cell_type": "code", "execution_count": null, - "id": "24", + "id": "25", "metadata": {}, "outputs": [], "source": [ @@ -454,7 +473,7 @@ }, { "cell_type": "markdown", - "id": "25", + "id": "26", "metadata": {}, "source": [ "## 6. Absence and coverage — one rule, every position\n", @@ -487,7 +506,7 @@ { "cell_type": "code", "execution_count": null, - "id": "26", + "id": "27", "metadata": {}, "outputs": [], "source": [ @@ -519,7 +538,7 @@ }, { "cell_type": "markdown", - "id": "27", + "id": "28", "metadata": {}, "source": [ "The other three positions refuse the same kind of hole, joining the\n", @@ -530,7 +549,7 @@ { "cell_type": "code", "execution_count": null, - "id": "28", + "id": "29", "metadata": {}, "outputs": [], "source": [ @@ -556,7 +575,7 @@ }, { "cell_type": "markdown", - "id": "29", + "id": "30", "metadata": {}, "source": [ "Two escape hatches fix the coefficient hole above, and both build and solve.\n", @@ -572,7 +591,7 @@ { "cell_type": "code", "execution_count": null, - "id": "30", + "id": "31", "metadata": {}, "outputs": [], "source": [ @@ -593,7 +612,7 @@ }, { "cell_type": "markdown", - "id": "31", + "id": "32", "metadata": {}, "source": [ "And the same masking escape hatch on the variable and constraint together:\n", @@ -604,7 +623,7 @@ { "cell_type": "code", "execution_count": null, - "id": "32", + "id": "33", "metadata": {}, "outputs": [], "source": [ @@ -629,7 +648,7 @@ }, { "cell_type": "markdown", - "id": "33", + "id": "34", "metadata": {}, "source": [ "## 7. Lookups and grouped sums\n", @@ -643,7 +662,7 @@ { "cell_type": "code", "execution_count": null, - "id": "34", + "id": "35", "metadata": {}, "outputs": [], "source": [ @@ -677,7 +696,7 @@ }, { "cell_type": "markdown", - "id": "35", + "id": "36", "metadata": {}, "source": [ "Note `south` has no generators mapped to it. Its group is **empty**, and an\n", @@ -687,7 +706,7 @@ }, { "cell_type": "markdown", - "id": "36", + "id": "37", "metadata": {}, "source": [ "## 8. Temporal operators: `shift`\n", @@ -707,7 +726,7 @@ { "cell_type": "code", "execution_count": null, - "id": "37", + "id": "38", "metadata": {}, "outputs": [], "source": [ @@ -762,7 +781,7 @@ }, { "cell_type": "markdown", - "id": "38", + "id": "39", "metadata": {}, "source": [ "The generator over-produces while power is cheap (hour 2 runs at 30 to fill the\n", @@ -773,7 +792,7 @@ { "cell_type": "code", "execution_count": null, - "id": "39", + "id": "40", "metadata": {}, "outputs": [], "source": [ @@ -782,7 +801,7 @@ }, { "cell_type": "markdown", - "id": "40", + "id": "41", "metadata": {}, "source": [ "## 9. Synthetic data for any spec\n", @@ -797,7 +816,7 @@ { "cell_type": "code", "execution_count": null, - "id": "41", + "id": "42", "metadata": {}, "outputs": [], "source": [ @@ -815,7 +834,7 @@ }, { "cell_type": "markdown", - "id": "42", + "id": "43", "metadata": {}, "source": [ "## 10. Persistence: netCDF and copy\n", @@ -835,7 +854,7 @@ { "cell_type": "code", "execution_count": null, - "id": "43", + "id": "44", "metadata": {}, "outputs": [], "source": [ @@ -867,7 +886,7 @@ }, { "cell_type": "markdown", - "id": "44", + "id": "45", "metadata": {}, "source": [ "Even a `retain=\"none\"` model round-trips: the spec text and coordinates survive,\n", @@ -877,7 +896,7 @@ { "cell_type": "code", "execution_count": null, - "id": "45", + "id": "46", "metadata": {}, "outputs": [], "source": [ @@ -892,7 +911,7 @@ }, { "cell_type": "markdown", - "id": "46", + "id": "47", "metadata": {}, "source": [ "`Model.copy()` carries the spec too, with the accessor reattached to the copy. The\n", @@ -904,7 +923,7 @@ { "cell_type": "code", "execution_count": null, - "id": "47", + "id": "48", "metadata": {}, "outputs": [], "source": [ @@ -921,16 +940,18 @@ }, { "cell_type": "markdown", - "id": "48", + "id": "49", "metadata": {}, "source": [ - "## Where the code lives, and two upstream notes\n", + "## Where the code lives, and an 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`, the `NamedExpression` views, `evaluate`, and\n", - " whole-model typesetting (`to_latex` / `to_markdown` / `to_typst`).\n", + " typesetting: the whole model (`m.spec.to_latex` / `.to_markdown` /\n", + " `.to_typst`) and a single named expression (the same three methods on a\n", + " `NamedExpression`, via math-spec's `typeset_declaration`).\n", "- `attach.py` — the three attachment rules; data onto master coordinates.\n", "- `builder.py` — emits variables, constraints, objective; folds expressions.\n", "- `operators.py` — `sum`, `by=`, `shift`, `at`, `sum_back`.\n", @@ -944,13 +965,6 @@ " 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", diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index e9e74b39..80a2c3c4 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -25,6 +25,7 @@ to_program, to_spec, to_typst, + typeset_declaration, ) from math_spec import program as ms @@ -293,6 +294,20 @@ def solution(self) -> xr.DataArray: """ return fold(self._name, self._ctx) + def to_latex(self, **options: Any) -> str: + """This expression typeset as a single LaTeX line, no document around it.""" + return typeset_declaration(self._spec._schema, self._name, "latex", **options) + + def to_markdown(self, **options: Any) -> str: + """This expression typeset as a single Markdown math line, no ``$$`` around it.""" + return typeset_declaration( + self._spec._schema, self._name, "markdown", **options + ) + + def to_typst(self, **options: Any) -> str: + """This expression typeset as a single Typst line, no document around it.""" + return typeset_declaration(self._spec._schema, self._name, "typst", **options) + def __repr__(self) -> str: value = self.__dict__.get("solution", self.__dict__.get("expression")) if isinstance(value, xr.DataArray): @@ -300,4 +315,4 @@ def __repr__(self) -> str: return f"NamedExpression('{self._name}')" def _repr_markdown_(self) -> str: - return self._spec.to_markdown() + return f"$$\n{self.to_markdown()}\n$$" diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index 3693a342..b6f570b4 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -209,6 +209,27 @@ def test_the_whole_model_typesets() -> None: assert spec._repr_markdown_() == spec.to_markdown() +@pytest.mark.parametrize("fmt", ["to_latex", "to_markdown", "to_typst"]) +def test_a_named_expression_typesets_to_one_line(fmt: str) -> None: + e = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.expressions["spend"] + line = getattr(e, fmt)() + assert "spend" in line + assert "\n" not in line + assert "align" not in line and "$$" not in line + + +def test_a_named_expression_repr_markdown_wraps_only_itself() -> None: + e = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.expressions["spend"] + assert e._repr_markdown_() == f"$$\n{e.to_markdown()}\n$$" + + +def test_a_named_expression_typeset_passes_options() -> None: + e = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.expressions["spend"] + assert e.to_latex( + symbols={"notation": "latex", "names": {"spend": "S"}} + ).startswith("S") + + def test_a_constant_expression_folds_to_a_scalar() -> None: spec = {**yaml_dict(), "expressions": {"answer": "6 * 7"}} got = Model.from_spec(spec, DISPATCH_DATA).spec.expressions["answer"].solution From 74de65470d6163fe0ce42c756af50d6b70cef33f Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 8 Sep 2026 13:18:14 +0200 Subject: [PATCH 30/35] feat(spec): typeset any declaration, not just named expressions Extract a Declaration base carrying to_latex/to_markdown/to_typst (reused by NamedExpression) and add ModelSpec.declaration(name) to typeset a named expression, constraint or variable as one bare line. Tests, release note and the notebook updated. --- doc/release_notes.rst | 2 +- examples/building-models-from-specs.ipynb | 55 +---------------- linopy/spec/__init__.py | 2 + linopy/spec/accessor.py | 73 ++++++++++++++++------- test/test_spec_accessor.py | 17 ++++++ 5 files changed, 76 insertions(+), 73 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index e420159c..0c4d72e5 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -27,7 +27,7 @@ Upcoming Version * ``model.spec.expressions`` (a ``linopy.spec.NamedExpressions`` mapping) returns a ``linopy.spec.NamedExpression`` for each declared name, with three views: ``.node`` (the lowered formula), ``.expression`` (the unsolved linopy expression — a ``LinearExpression``, bare ``Variable``, array or scalar) and ``.solution`` (the expression folded over the solved model). ``model.spec.evaluate(name, sources)`` returns the same object with its parameters attached afresh. -* ``model.spec.to_latex`` / ``.to_markdown`` / ``.to_typst`` typeset the whole model, and the same three methods on a ``NamedExpression`` typeset that one expression as a single line (math only, no document); both a ``ModelSpec`` and a ``NamedExpression`` render as Markdown in a notebook. +* ``model.spec.to_latex`` / ``.to_markdown`` / ``.to_typst`` typeset the whole model, and ``model.spec.declaration(name)`` returns a ``linopy.spec.Declaration`` whose same three methods typeset one named expression, constraint or variable as a single line (math only, no document); a ``NamedExpression`` carries those methods too. A ``ModelSpec``, a ``Declaration`` and a ``NamedExpression`` all render as Markdown in a notebook. *Numerical scaling* diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb index 62b7957e..0b6bdfcb 100644 --- a/examples/building-models-from-specs.ipynb +++ b/examples/building-models-from-specs.ipynb @@ -323,18 +323,7 @@ "cell_type": "markdown", "id": "16", "metadata": {}, - "source": [ - "### 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.\n", - "\n", - "Each `NamedExpression` typesets on its own too, with the same three methods.\n", - "These render **one** expression as a single line — math only, no surrounding\n", - "document — so the string drops straight into a docstring or a table cell, and a\n", - "`NamedExpression` renders as its own formula in a notebook." - ] + "source": "### The model as maths\n\nThe accessor typesets the whole model, delegating to math-spec:\n`m.spec.to_latex()`, `.to_markdown()` and `.to_typst()`. In a notebook the\naccessor renders as Markdown on its own; here we show it explicitly.\n\nAny single declaration typesets on its own too. `m.spec.declaration(name)`\ntakes a named expression, a constraint or a variable and hands back a\n`Declaration` with the same three methods; a `NamedExpression` carries them\ndirectly. These render **one** line — math only, no surrounding document — so\nthe string drops straight into a docstring or a table cell, and both a\n`Declaration` and a `NamedExpression` render as their own formula in a notebook." }, { "cell_type": "code", @@ -354,13 +343,7 @@ "id": "18", "metadata": {}, "outputs": [], - "source": [ - "print(\"spend.to_latex():\")\n", - "print(spend.to_latex())\n", - "\n", - "# renders as its own formula in a notebook:\n", - "Markdown(f\"$$\\n{spend.to_markdown()}\\n$$\")" - ] + "source": "print(\"spend.to_latex(): \", spend.to_latex())\nprint(\"power_balance.to_latex():\", m.spec.declaration(\"power_balance\").to_latex())\nprint(\"p.to_latex(): \", m.spec.declaration(\"p\").to_latex())\n\n# each renders as its own formula in a notebook:\nMarkdown(f\"$$\\n{m.spec.declaration('power_balance').to_markdown()}\\n$$\")" }, { "cell_type": "markdown", @@ -942,39 +925,7 @@ "cell_type": "markdown", "id": "49", "metadata": {}, - "source": [ - "## Where the code lives, and an 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`, the `NamedExpression` views, `evaluate`, and\n", - " typesetting: the whole model (`m.spec.to_latex` / `.to_markdown` /\n", - " `.to_typst`) and a single named expression (the same three methods on a\n", - " `NamedExpression`, via math-spec's `typeset_declaration`).\n", - "- `attach.py` — the three attachment rules; data onto master coordinates.\n", - "- `builder.py` — emits variables, constraints, objective; folds expressions.\n", - "- `operators.py` — `sum`, `by=`, `shift`, `at`, `sum_back`.\n", - "- `where.py` — `where:` predicates as boolean masks.\n", - "- `coverage.py` / `terms.py` — the absence rule from section 6: a missing row\n", - " is refused wherever it is used.\n", - "- `curves.py` — the data side of `piecewise:` blocks.\n", - "- `netcdf.py` — the factorize-based persistence from section 10.\n", - "- `nodes.py` — walks over expression nodes. 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", - "attaches them into an ordinary model, hands each named expression back as three\n", - "views — its formula, its unsolved linopy expression and its solution — refuses a\n", - "missing parameter row wherever it is used (as a coefficient, bound, constant\n", - "side or divisor alike, with `where:` and filling the data as the escape\n", - "hatches), and round-trips the lot through netCDF by keeping the spec as text\n", - "beside factorized labels." - ] + "source": "## Where the code lives, and an upstream note\n\nThe feature is a small package, `linopy/spec/`, imported only when you call\n`add_spec`/`from_spec` — `import linopy` never pulls in `math_spec`. Roughly:\n\n- `accessor.py` — `model.spec`, the `NamedExpression` views, `evaluate`, and\n typesetting: the whole model (`m.spec.to_latex` / `.to_markdown` /\n `.to_typst`) and any single declaration — a named expression, constraint or\n variable — via `m.spec.declaration(name)` and math-spec's\n `typeset_declaration`.\n- `attach.py` — the three attachment rules; data onto master coordinates.\n- `builder.py` — emits variables, constraints, objective; folds expressions.\n- `operators.py` — `sum`, `by=`, `shift`, `at`, `sum_back`.\n- `where.py` — `where:` predicates as boolean masks.\n- `coverage.py` / `terms.py` — the absence rule from section 6: a missing row\n is refused wherever it is used.\n- `curves.py` — the data side of `piecewise:` blocks.\n- `netcdf.py` — the factorize-based persistence from section 10.\n- `nodes.py` — walks over expression nodes. 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\nA spec is the maths over labelled axes; the sources are the numbers. `linopy`\nattaches them into an ordinary model, hands each named expression back as three\nviews — its formula, its unsolved linopy expression and its solution — refuses a\nmissing parameter row wherever it is used (as a coefficient, bound, constant\nside or divisor alike, with `where:` and filling the data as the escape\nhatches), and round-trips the lot through netCDF by keeping the spec as text\nbeside factorized labels." } ], "metadata": { diff --git a/linopy/spec/__init__.py b/linopy/spec/__init__.py index e9bd758e..d3d76bbb 100644 --- a/linopy/spec/__init__.py +++ b/linopy/spec/__init__.py @@ -17,6 +17,7 @@ ) from linopy.spec.accessor import ( + Declaration, ModelSpec, NamedExpression, NamedExpressions, @@ -27,6 +28,7 @@ __all__ = [ "Attached", + "Declaration", "ModelSpec", "NamedExpression", "NamedExpressions", diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index 80a2c3c4..b7796bc9 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -150,6 +150,26 @@ def expressions(self) -> NamedExpressions: """Each named expression as a :class:`NamedExpression`: its math, its linopy fold and its solution.""" return NamedExpressions(self) + def declaration(self, name: str) -> Declaration: + """ + One declaration typeset on its own: a named expression, constraint or variable. + + Its math as a single line, no document around it. A named expression + also carries its linopy fold and solution through :attr:`expressions`; + this handle is the typesetting one every declaration shares. + """ + if name not in self._declarations: + raise KeyError( + f"unknown declaration '{name}'. " + + did_you_mean(name, self._declarations) + ) + return Declaration(self, name) + + @property + def _declarations(self) -> list[str]: + p = self.program + return [*p.named_expressions, *p.constraints, *p.variables] + def to_latex(self, **options: Any) -> str: """The whole model typeset as a LaTeX document.""" return to_latex(self._schema, **options) @@ -244,7 +264,38 @@ def __repr__(self) -> str: return f"NamedExpressions({list(self)})" -class NamedExpression: +class Declaration: + """ + One declaration of a spec, typeset on its own: math only, no document. + + A named expression, a constraint or a variable, reached by name through + :meth:`ModelSpec.declaration`. :class:`NamedExpression` adds the linopy + fold and the solution on top of this. + """ + + def __init__(self, spec: ModelSpec, name: str) -> None: + self._spec = spec + self._name = name + + def to_latex(self, **options: Any) -> str: + """This declaration typeset as a single LaTeX line, no document around it.""" + return typeset_declaration(self._spec._schema, self._name, "latex", **options) + + def to_markdown(self, **options: Any) -> str: + """This declaration typeset as a single Markdown math line, no ``$$`` around it.""" + return typeset_declaration( + self._spec._schema, self._name, "markdown", **options + ) + + def to_typst(self, **options: Any) -> str: + """This declaration typeset as a single Typst line, no document around it.""" + return typeset_declaration(self._spec._schema, self._name, "typst", **options) + + def _repr_markdown_(self) -> str: + return f"$$\n{self.to_markdown()}\n$$" + + +class NamedExpression(Declaration): """ One named expression, in three views: its math, its linopy fold and its solution. @@ -259,8 +310,7 @@ class NamedExpression: """ def __init__(self, spec: ModelSpec, name: str, ctx: Context) -> None: - self._spec = spec - self._name = name + super().__init__(spec, name) self._ctx = ctx @property @@ -294,25 +344,8 @@ def solution(self) -> xr.DataArray: """ return fold(self._name, self._ctx) - def to_latex(self, **options: Any) -> str: - """This expression typeset as a single LaTeX line, no document around it.""" - return typeset_declaration(self._spec._schema, self._name, "latex", **options) - - def to_markdown(self, **options: Any) -> str: - """This expression typeset as a single Markdown math line, no ``$$`` around it.""" - return typeset_declaration( - self._spec._schema, self._name, "markdown", **options - ) - - def to_typst(self, **options: Any) -> str: - """This expression typeset as a single Typst line, no document around it.""" - return typeset_declaration(self._spec._schema, self._name, "typst", **options) - 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 f"$$\n{self.to_markdown()}\n$$" diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index b6f570b4..fbf74a48 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -230,6 +230,23 @@ def test_a_named_expression_typeset_passes_options() -> None: ).startswith("S") +@pytest.mark.parametrize("name", ["power_balance", "p"]) +@pytest.mark.parametrize("fmt", ["to_latex", "to_markdown", "to_typst"]) +def test_a_constraint_or_variable_typesets_to_one_line(name: str, fmt: str) -> None: + d = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec.declaration(name) + line = getattr(d, fmt)() + assert line + assert "\n" not in line + assert "align" not in line and "$$" not in line + + +def test_declaration_reaches_every_kind_and_an_unknown_name_is_a_key_error() -> None: + spec = Model.from_spec(VIEWS_SPEC, DISPATCH_DATA).spec + assert spec.declaration("spend").to_latex() == spec.expressions["spend"].to_latex() + with pytest.raises(KeyError, match="unknown declaration 'spent'.*spend"): + spec.declaration("spent") + + def test_a_constant_expression_folds_to_a_scalar() -> None: spec = {**yaml_dict(), "expressions": {"answer": "6 * 7"}} got = Model.from_spec(spec, DISPATCH_DATA).spec.expressions["answer"].solution From 84e35ab8c5ace29dc78c1db61d666b1aca903b45 Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 8 Sep 2026 14:09:33 +0200 Subject: [PATCH 31/35] feat(spec): summarise the whole model in the ModelSpec repr Show dimensions, variables, constraints, objective and named expressions, one capped line each, instead of only the expression names. --- linopy/spec/accessor.py | 24 ++++++++++++++++++++++-- test/test_spec_accessor.py | 18 ++++++++++++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index b7796bc9..42166034 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -102,6 +102,14 @@ def _source(spec: SpecLike) -> tuple[str, ms.Program]: return loaded.to_yaml(), to_program(loaded) +def _row(label: str, items: list[str], cap: int = 8) -> str: + """One aligned summary line, capped with a ``(+N more)`` tail.""" + shown = items[:cap] + if len(items) > cap: + shown = shown + [f"(+{len(items) - cap} more)"] + return f" {label + ':':<13}{', '.join(shown) if shown else '—'}" + + class ModelSpec: """ The spec a model was built from. @@ -120,8 +128,20 @@ def __init__(self, model: Model, program: ms.Program, text: str) -> None: self.text = text def __repr__(self) -> str: - names = list(self.program.named_expressions) - return f"ModelSpec(expressions={names})" + p = self.program + coords = self.coords + desc = str(self._schema.get("description", "")).strip().splitlines() + head = f"ModelSpec: {desc[0]}" if desc else "ModelSpec" + rows = [ + head, + _row("Dimensions", [f"{d} ({len(coords[d])})" for d in p.dimensions]), + _row("Variables", list(p.variables)), + _row("Constraints", list(p.constraints)), + ] + if p.objective is not None: + rows.append(_row("Objective", [p.objective.sense])) + rows.append(_row("Expressions", list(p.named_expressions))) + return "\n".join(rows) def _reattach(self, model: Model) -> ModelSpec: """The same spec, read off *model*.""" diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index fbf74a48..e3c8df45 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -24,6 +24,7 @@ EXAMPLE_DISPATCH, GENERATOR, solved, + with_, yaml_dict, ) from linopy import Model # noqa: E402 @@ -201,6 +202,23 @@ def test_evaluate_returns_a_named_expression() -> None: ) +def test_repr_summarises_every_section() -> None: + text = repr(Model.from_spec(yaml_dict(), DISPATCH_DATA).spec) + assert text.startswith("ModelSpec: Least-cost dispatch") + assert "Dimensions: snapshot (3), generator (2)" in text + assert "Variables: p" in text + assert "Constraints: power_balance" in text + assert "Objective: minimize" in text + assert "Expressions: spend, usage" in text + + +def test_repr_caps_long_sections() -> None: + spec = with_(yaml_dict(), expressions={f"e{i}": "p / p_max" for i in range(12)}) + text = repr(Model.from_spec(spec, DISPATCH_DATA).spec) + assert "(+6 more)" in text + assert "e11" not in text + + def test_the_whole_model_typesets() -> None: spec = Model.from_spec(yaml_dict(), DISPATCH_DATA).spec assert "align" in spec.to_latex() From 8b88ac721f4188f4463c0194707c716bd8f69266 Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 8 Sep 2026 14:47:14 +0200 Subject: [PATCH 32/35] feat(spec): show the spec in the Model repr Header says the model is built from a math-spec and prints the spec's description. Spec named expressions are listed with their static dims; spec-owned items are tagged [spec] once hand-added ones exist. --- linopy/constraints.py | 7 +++++-- linopy/model.py | 28 +++++++++++++++++++++---- linopy/spec/accessor.py | 15 +++++++++++-- linopy/spec/nodes.py | 24 +++++++++++++++++++++ linopy/variables.py | 7 +++++-- test/test_spec_accessor.py | 43 ++++++++++++++++++++++++++++++++++++++ 6 files changed, 114 insertions(+), 10 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index f3b301cc..6cc02cdd 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -2118,7 +2118,9 @@ def _formatted_names(self) -> dict[str, str]: """ return {format_string_as_variable_name(n): n for n in self} - def _format_items(self, exclude: set[str] | None = None) -> str: + def _format_items( + self, exclude: set[str] | None = None, tag: set[str] | None = None + ) -> str: """Format constraint items, optionally excluding names in a group.""" r = "" count = 0 @@ -2131,7 +2133,8 @@ def _format_items(self, exclude: set[str] | None = None) -> str: if ds.coords else "" ) - r += f" * {name}{coords}\n" + suffix = " [spec]" if tag and name in tag else "" + r += f" * {name}{coords}{suffix}\n" if count == 0: r += "\n" return r diff --git a/linopy/model.py b/linopy/model.py index 416128cd..2e127973 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -706,13 +706,33 @@ def __repr__(self) -> str: from linopy.piecewise import _repr_summary as pwl_repr_summary var_names, con_names = _get_piecewise_groups(self) - var_string = self.variables._format_items(exclude=var_names) - con_string = self.constraints._format_items(exclude=con_names) - expr_string = self.expressions._format_items() model_string = f"Linopy {self.type} model" + var_tag: set[str] | None = None + con_tag: set[str] | None = None + expr_string = self.expressions._format_items() + if self._spec is not None: + model_string += ", built from a math-spec" + program = self._spec.program + spec_vars = set(program.variables) + spec_cons = set(program.constraints) + if any(v not in spec_vars for v in self.variables): + var_tag = spec_vars + if any(c not in spec_cons for c in self.constraints): + con_tag = spec_cons + eager = expr_string if len(self.expressions) else "" + spec = "".join( + f" * {name} ({', '.join(e.dims)}) [spec]\n" + for name, e in self._spec.expressions.items() + ) + expr_string = eager + spec or "\n" + var_string = self.variables._format_items(exclude=var_names, tag=var_tag) + con_string = self.constraints._format_items(exclude=con_names, tag=con_tag) + header = f"{model_string}\n{'=' * len(model_string)}\n" + if self._spec is not None and self._spec.description: + header += f"{self._spec.description}\n" return ( - f"{model_string}\n{'=' * len(model_string)}\n\n" + f"{header}\n" f"Variables:\n----------\n{var_string}\n" f"Expressions:\n------------\n{expr_string}\n" f"Constraints:\n------------\n{con_string}" diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index 42166034..fff7e30f 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -39,6 +39,7 @@ from linopy.spec.context import Context from linopy.spec.errors import SpecDataError from linopy.spec.evaluate import evaluate_named, fold +from linopy.spec.nodes import dims_of from linopy.spec.parameters import Parameters, Resolve SpecLike: TypeAlias = str | Path | Mapping[str, Any] | Spec @@ -130,8 +131,7 @@ def __init__(self, model: Model, program: ms.Program, text: str) -> None: def __repr__(self) -> str: p = self.program coords = self.coords - desc = str(self._schema.get("description", "")).strip().splitlines() - head = f"ModelSpec: {desc[0]}" if desc else "ModelSpec" + head = f"ModelSpec: {self.description}" if self.description else "ModelSpec" rows = [ head, _row("Dimensions", [f"{d} ({len(coords[d])})" for d in p.dimensions]), @@ -152,6 +152,12 @@ def parameters(self) -> xr.Dataset: """The parameters and lookups retained on the model, on the master coordinates.""" return self._model.parameters + @property + def description(self) -> str: + """The spec's own description, its first line, or an empty string.""" + lines = str(self._schema.get("description", "")).strip().splitlines() + return lines[0] if lines else "" + @property def coords(self) -> dict[str, pd.Index]: """Master coordinates by dimension, as the model was built on them.""" @@ -338,6 +344,11 @@ def node(self) -> ms.ExpressionNode: """The expression body as lowered, math-spec's own AST handle.""" return self._spec.program.named_expressions[self._name].expression + @property + def dims(self) -> tuple[str, ...]: + """The dimensions the expression spans, read off the spec without binding data.""" + return dims_of(self.node, self._spec.program) + @functools.cached_property def expression(self) -> terms.Value: """ diff --git a/linopy/spec/nodes.py b/linopy/spec/nodes.py index 68a93f16..7c158a4e 100644 --- a/linopy/spec/nodes.py +++ b/linopy/spec/nodes.py @@ -32,3 +32,27 @@ def amounts_of(node: ms.ExpressionNode) -> Iterator[str]: def parameters_of(*nodes: ms.ExpressionNode) -> frozenset[str]: """Every parameter named anywhere under *nodes*.""" return frozenset(n.name for n in walk(*nodes) if isinstance(n, ms.Parameter)) + + +def dims_of(node: ms.ExpressionNode, program: ms.Program) -> tuple[str, ...]: + """The dimensions *node* spans, in the program's dimension order, before any data is bound.""" + spanned = _dims(node, program) + return tuple(d for d in program.dimensions if d in spanned) + + +def _dims(node: ms.ExpressionNode, program: ms.Program) -> frozenset[str]: + if isinstance(node, ms.Constant): + return frozenset() + if isinstance(node, ms.Variable): + return frozenset(program.variables[node.name].dims) + if isinstance(node, ms.Parameter): + return frozenset(program.parameters[node.name].dims) + if isinstance(node, ms.Dual): + return frozenset(program.constraints[node.constraint].dims) + if isinstance(node, ms.Sum): + return _dims(node.operand, program) - set(node.over) + if isinstance(node, ms.GroupSum | ms.At): + return (_dims(node.operand, program) - {node.over}) | set(node.into) + if isinstance(node, ms.Cases): + return frozenset().union(*(_dims(r.value, program) for r in node.regions)) + return frozenset().union(*(_dims(c, program) for c in children(node))) diff --git a/linopy/variables.py b/linopy/variables.py index e0ad70bc..ea4ada53 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -1809,7 +1809,9 @@ def __dir__(self) -> list[str]: ] return base_attributes + formatted_names - def _format_items(self, exclude: set[str] | None = None) -> str: + def _format_items( + self, exclude: set[str] | None = None, tag: set[str] | None = None + ) -> str: """Format variable items, optionally excluding names in a group.""" r = "" count = 0 @@ -1828,7 +1830,8 @@ def _format_items(self, exclude: set[str] | None = None) -> str: coords += f" - sos{sos_type} on {sos_dim}" if ds.attrs.get("semi_continuous", False): coords += " - semi-continuous" - r += f" * {name}{coords}\n" + suffix = " [spec]" if tag and name in tag else "" + r += f" * {name}{coords}{suffix}\n" if count == 0: r += "\n" return r diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index e3c8df45..3716b190 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -219,6 +219,39 @@ def test_repr_caps_long_sections() -> None: assert "e11" not in text +def test_model_repr_shows_the_spec_and_tags_only_expressions() -> None: + text = repr(Model.from_spec(yaml_dict(), DISPATCH_DATA)) + assert "Linopy LP model, built from a math-spec" in text + assert "Least-cost dispatch of a generator fleet against an hourly load." in text + assert " * spend (snapshot) [spec]" in text + assert " * usage (snapshot, generator) [spec]" in text + assert " * p (snapshot, generator)\n" in text + assert " * power_balance (snapshot)\n" in text + assert "" not in text + + +def test_model_repr_of_a_spec_without_a_description() -> None: + spec = {k: v for k, v in yaml_dict().items() if k != "description"} + m = Model.from_spec(spec, DISPATCH_DATA) + assert m.spec.description == "" + assert repr(m).startswith("Linopy LP model, built from a math-spec\n=") + + +def test_hybrid_model_tags_spec_variables_constraints_and_expressions() -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + v = m.add_variables(lower=0, coords=[GENERATOR], name="reserve") + m.add_expressions(v * 2.0, name="reserve_cost") + m.add_constraints(v <= 10.0, name="reserve_cap") + text = repr(m) + assert " * p (snapshot, generator) [spec]" in text + assert " * reserve (generator)\n" in text + assert " * power_balance (snapshot) [spec]" in text + assert " * reserve_cap (generator)\n" in text + assert " * reserve_cost (generator)\n" in text + assert " * spend (snapshot) [spec]" in text + assert "" not in text + + def test_the_whole_model_typesets() -> None: spec = Model.from_spec(yaml_dict(), DISPATCH_DATA).spec assert "align" in spec.to_latex() @@ -321,3 +354,13 @@ def test_spec_api_warns_once_per_session() -> None: with warnings.catch_warnings(): warnings.simplefilter("error", EvolvingAPIWarning) Model.from_spec(EXAMPLE_DISPATCH, DISPATCH_DATA) + + +@pytest.mark.parametrize( + ("name", "dims"), [("spend", ("snapshot",)), ("usage", ("snapshot", "generator"))] +) +def test_named_expression_dims_are_static(name: str, dims: tuple[str, ...]) -> None: + m = Model.from_spec(yaml_dict(), DISPATCH_DATA) + expr = m.spec.expressions[name] + assert expr.dims == dims + assert set(expr.expression.coord_dims) == set(dims) From 52edc8add734629c0f377df765aa4822acba2002 Mon Sep 17 00:00:00 2001 From: Felix <117816358+FBumann@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:23:39 +0200 Subject: [PATCH 33/35] chore(spec): pin math-spec alpha.76 and drop the Power walk workaround (#946) math-spec's `program.children()` had no branch for `Power`, so every walk stopped there and a parameter written `d ** 2` was invisible to it. linopy worked around it with its own `children()` in `linopy/spec/nodes.py`. energy-models/math-spec#404 fixed it upstream, so the `spec` group pins the release that carries it, `v0.0.0-alpha.76`, and the walks call `math_spec.program.children` directly. The coverage tests for a constant side and a divisor hidden under a power stay as they are and now exercise the upstream walk. Co-Authored-By: Claude Opus 5 (1M context) --- examples/building-models-from-specs.ipynb | 2 +- linopy/spec/coverage.py | 4 ++-- linopy/spec/nodes.py | 13 +++---------- pyproject.toml | 2 +- 4 files changed, 7 insertions(+), 14 deletions(-) diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb index 0b6bdfcb..09495cd4 100644 --- a/examples/building-models-from-specs.ipynb +++ b/examples/building-models-from-specs.ipynb @@ -925,7 +925,7 @@ "cell_type": "markdown", "id": "49", "metadata": {}, - "source": "## Where the code lives, and an upstream note\n\nThe feature is a small package, `linopy/spec/`, imported only when you call\n`add_spec`/`from_spec` — `import linopy` never pulls in `math_spec`. Roughly:\n\n- `accessor.py` — `model.spec`, the `NamedExpression` views, `evaluate`, and\n typesetting: the whole model (`m.spec.to_latex` / `.to_markdown` /\n `.to_typst`) and any single declaration — a named expression, constraint or\n variable — via `m.spec.declaration(name)` and math-spec's\n `typeset_declaration`.\n- `attach.py` — the three attachment rules; data onto master coordinates.\n- `builder.py` — emits variables, constraints, objective; folds expressions.\n- `operators.py` — `sum`, `by=`, `shift`, `at`, `sum_back`.\n- `where.py` — `where:` predicates as boolean masks.\n- `coverage.py` / `terms.py` — the absence rule from section 6: a missing row\n is refused wherever it is used.\n- `curves.py` — the data side of `piecewise:` blocks.\n- `netcdf.py` — the factorize-based persistence from section 10.\n- `nodes.py` — walks over expression nodes. 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\nA spec is the maths over labelled axes; the sources are the numbers. `linopy`\nattaches them into an ordinary model, hands each named expression back as three\nviews — its formula, its unsolved linopy expression and its solution — refuses a\nmissing parameter row wherever it is used (as a coefficient, bound, constant\nside or divisor alike, with `where:` and filling the data as the escape\nhatches), and round-trips the lot through netCDF by keeping the spec as text\nbeside factorized labels." + "source": "## Where the code lives\n\nThe feature is a small package, `linopy/spec/`, imported only when you call\n`add_spec`/`from_spec` — `import linopy` never pulls in `math_spec`. Roughly:\n\n- `accessor.py` — `model.spec`, the `NamedExpression` views, `evaluate`, and\n typesetting: the whole model (`m.spec.to_latex` / `.to_markdown` /\n `.to_typst`) and any single declaration — a named expression, constraint or\n variable — via `m.spec.declaration(name)` and math-spec's\n `typeset_declaration`.\n- `attach.py` — the three attachment rules; data onto master coordinates.\n- `builder.py` — emits variables, constraints, objective; folds expressions.\n- `operators.py` — `sum`, `by=`, `shift`, `at`, `sum_back`.\n- `where.py` — `where:` predicates as boolean masks.\n- `coverage.py` / `terms.py` — the absence rule from section 6: a missing row\n is refused wherever it is used.\n- `curves.py` — the data side of `piecewise:` blocks.\n- `netcdf.py` — the factorize-based persistence from section 10.\n- `nodes.py` — walks over expression nodes, and the dimensions a node\n spans before any data is bound.\n\n### Summary\n\nA spec is the maths over labelled axes; the sources are the numbers. `linopy`\nattaches them into an ordinary model, hands each named expression back as three\nviews — its formula, its unsolved linopy expression and its solution — refuses a\nmissing parameter row wherever it is used (as a coefficient, bound, constant\nside or divisor alike, with `where:` and filling the data as the escape\nhatches), and round-trips the lot through netCDF by keeping the spec as text\nbeside factorized labels." } ], "metadata": { diff --git a/linopy/spec/coverage.py b/linopy/spec/coverage.py index 401c6898..11da5e04 100644 --- a/linopy/spec/coverage.py +++ b/linopy/spec/coverage.py @@ -21,7 +21,7 @@ from linopy.spec import terms from linopy.spec.context import Context from linopy.spec.errors import SpecDataError -from linopy.spec.nodes import amounts_of, children, parameters_of +from linopy.spec.nodes import amounts_of, parameters_of from linopy.spec.where import evaluate_where Rows = xr.DataArray | None @@ -101,7 +101,7 @@ def _collect( narrowed = inside if rows is None else rows & inside _collect(region.value, ctx, narrowed, constant, into) return - for child in children(node): + for child in ms.children(node): _collect(child, ctx, rows, constant, into) diff --git a/linopy/spec/nodes.py b/linopy/spec/nodes.py index 7c158a4e..9c95c968 100644 --- a/linopy/spec/nodes.py +++ b/linopy/spec/nodes.py @@ -1,4 +1,4 @@ -"""Walks over expression nodes that descend into every operand, a ``Power``'s included.""" +"""Walks over a program's expression nodes, and the dimensions a node spans.""" from __future__ import annotations @@ -7,18 +7,11 @@ 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)) + yield from walk(*ms.children(node)) def amounts_of(node: ms.ExpressionNode) -> Iterator[str]: @@ -55,4 +48,4 @@ def _dims(node: ms.ExpressionNode, program: ms.Program) -> frozenset[str]: return (_dims(node.operand, program) - {node.over}) | set(node.into) if isinstance(node, ms.Cases): return frozenset().union(*(_dims(r.value, program) for r in node.regions)) - return frozenset().union(*(_dims(c, program) for c in children(node))) + return frozenset().union(*(_dims(c, program) for c in ms.children(node))) diff --git a/pyproject.toml b/pyproject.toml index 482334e6..f8b3ec14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,7 +123,7 @@ gpu = [ # keeps the git pin out of the published wheel metadata, which PyPI rejects. # Install with `uv sync --group spec` or `uv pip install --group spec`. spec = [ - "math-spec @ git+https://github.com/energy-models/math-spec.git@67aeedb988ee95d196456b4cbe50821e3799edaa ; python_version >= '3.12'", + "math-spec @ git+https://github.com/energy-models/math-spec.git@v0.0.0-alpha.76 ; python_version >= '3.12'", "pyyaml ; python_version >= '3.12'", "pyarrow ; python_version >= '3.12'", ] From 48b267d420907929bc0a4ab1604b18d047875e02 Mon Sep 17 00:00:00 2001 From: Felix <117816358+FBumann@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:24:34 +0200 Subject: [PATCH 34/35] refac(spec): the spec owns its parameters, and one read path for a named expression (#949) Two review findings from #922. **The spec owns its parameters, not `model.parameters`.** A build wrote the retained parameters and lookups straight into `model.parameters`, by assignment. That clobbered whatever the caller had put there, and left nothing able to say which of the model's parameters the spec owned -- `netcdf._coded` compensated by scanning every parameter on the model for an object dtype, so a caller's own labelled array was factorized into the spec's sub-dataset and came back as spec data. `ModelSpec` now holds its own dataset, written under the `spec-` prefix that already carried its coordinates. **One read path for a named expression.** Reading `model.spec.expressions[name]` went through the retained parameters and raised if `retain` had dropped one; `evaluate(name, sources)` went through the data instead. The two differed by one argument, the `Resolve` callable, so `ModelSpec` keeps the `Attached` a build made and falls back to it. `retain` now decides what a netcdf file holds, not what a session can read, and the refusal is left for a model read back from a file, whose sources are gone. Co-Authored-By: Claude Opus 5 (1M context) --- doc/release_notes.rst | 2 +- examples/building-models-from-specs.ipynb | 28 +++++-- linopy/io.py | 13 ++-- linopy/model.py | 7 +- linopy/spec/accessor.py | 93 ++++++++++++++++------- linopy/spec/netcdf.py | 76 ++++++++++-------- linopy/testing.py | 1 + test/test_spec_accessor.py | 26 +++++-- test/test_spec_builder.py | 2 +- test/test_spec_io.py | 32 +++++++- 10 files changed, 196 insertions(+), 84 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 0c4d72e5..fc0d88f8 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -23,7 +23,7 @@ Upcoming Version *Build a model from a math-spec program* -* ``Model.from_spec`` / ``model.add_spec`` build a model from a `math-spec `__ YAML program attached to data, and ``model.spec`` (a ``linopy.spec.ModelSpec``) reads it back. Requires the ``spec`` dependency group (``uv sync --group spec`` / ``uv pip install --group spec``, Python >= 3.12) and v1 semantics. Data is attached onto the spec's dimensions and parameters with ``linopy.spec.attach``, raising a ``linopy.spec.SpecDataError`` on mismatched or missing data; ``linopy.spec.Attached`` carries the attached result. The spec API emits an :class:`linopy.EvolvingAPIWarning` once per session while it stabilises. See :doc:`building-models-from-specs` for a worked example. +* ``Model.from_spec`` / ``model.add_spec`` build a model from a `math-spec `__ YAML program attached to data, and ``model.spec`` (a ``linopy.spec.ModelSpec``) reads it back. Requires the ``spec`` dependency group (``uv sync --group spec`` / ``uv pip install --group spec``, Python >= 3.12) and v1 semantics. Data is attached onto the spec's dimensions and parameters with ``linopy.spec.attach``, raising a ``linopy.spec.SpecDataError`` on mismatched or missing data; ``linopy.spec.Attached`` carries the attached result. The parameters the spec retains live in ``model.spec.parameters``, its own dataset; ``model.parameters`` stays the caller's and a build never writes to it. ``retain`` decides what a netcdf file holds, not what a session can read: a parameter it dropped is resolved from the sources the model was built with, and only a model read back from a file can run out of data. The spec API emits an :class:`linopy.EvolvingAPIWarning` once per session while it stabilises. See :doc:`building-models-from-specs` for a worked example. * ``model.spec.expressions`` (a ``linopy.spec.NamedExpressions`` mapping) returns a ``linopy.spec.NamedExpression`` for each declared name, with three views: ``.node`` (the lowered formula), ``.expression`` (the unsolved linopy expression — a ``LinearExpression``, bare ``Variable``, array or scalar) and ``.solution`` (the expression folded over the solved model). ``model.spec.evaluate(name, sources)`` returns the same object with its parameters attached afresh. diff --git a/examples/building-models-from-specs.ipynb b/examples/building-models-from-specs.ipynb index 09495cd4..cfb2ec04 100644 --- a/examples/building-models-from-specs.ipynb +++ b/examples/building-models-from-specs.ipynb @@ -343,7 +343,14 @@ "id": "18", "metadata": {}, "outputs": [], - "source": "print(\"spend.to_latex(): \", spend.to_latex())\nprint(\"power_balance.to_latex():\", m.spec.declaration(\"power_balance\").to_latex())\nprint(\"p.to_latex(): \", m.spec.declaration(\"p\").to_latex())\n\n# each renders as its own formula in a notebook:\nMarkdown(f\"$$\\n{m.spec.declaration('power_balance').to_markdown()}\\n$$\")" + "source": [ + "print(\"spend.to_latex(): \", spend.to_latex())\n", + "print(\"power_balance.to_latex():\", m.spec.declaration(\"power_balance\").to_latex())\n", + "print(\"p.to_latex(): \", m.spec.declaration(\"p\").to_latex())\n", + "\n", + "# each renders as its own formula in a notebook:\n", + "Markdown(f\"$$\\n{m.spec.declaration('power_balance').to_markdown()}\\n$$\")" + ] }, { "cell_type": "markdown", @@ -377,16 +384,25 @@ "## 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", + "parameters linopy keeps in `model.spec.parameters` after building.\n", + "That is the spec's own dataset — `model.parameters` stays yours, and a\n", + "build never writes to it:\n", "\n", - "| `retain` | keeps in `model.parameters` |\n", + "| `retain` | keeps in `model.spec.parameters` |\n", "|------------|-------------------------------------------------|\n", "| `\"report\"` | only parameters the named expressions read (default) |\n", "| `\"all\"` | every parameter |\n", "| `\"none\"` | nothing |\n", "\n", "`spend` reads `cost`, `usage` reads `p_max`, neither reads `load` — so\n", - "`\"report\"` keeps `cost` and `p_max` but drops `load`." + "`\"report\"` keeps `cost` and `p_max` but drops `load`.\n", + "\n", + "Dropping is about *storage*, not about what you can read. A parameter\n", + "`retain` left out is resolved from the `sources` you built with, which the\n", + "model keeps hold of — so every `retain` folds the same in this session.\n", + "It is writing the model to netCDF that leaves the sources behind: read that\n", + "file back and only what `retain` kept is still there, with\n", + "`m.spec.evaluate(name, sources)` as the way in for the rest." ] }, { @@ -398,7 +414,9 @@ "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)}\")" + " print(\n", + " f\"retain={retain!r:9} -> parameters kept: {sorted(mm.spec.parameters.data_vars)}\"\n", + " )" ] }, { diff --git a/linopy/io.py b/linopy/io.py index 5e1b09d2..275cb2b7 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1160,8 +1160,9 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: :func:`linopy.io.read_netcdf`. The insertion order of each container is stored as a JSON list in the ``_linopy__order`` attribute. - A model built with :meth:`Model.add_spec` also persists its spec: the - YAML text, the master coordinates and the lookups. ``read_netcdf`` + A model built with :meth:`Model.add_spec` also persists its spec under a + ``spec-`` prefix of its own: the YAML text, the master coordinates and the + parameters the spec retained, apart from ``m.parameters``. ``read_netcdf`` lowers the program from the text again, so reading such a file needs the ``math-spec`` package; a file without a spec does not. @@ -1209,14 +1210,12 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: if m.objective.value is not None: objective = objective.assign_attrs(value=m.objective.value) obj = [with_prefix(objective, "objective")] - 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(record_dtypes(parameters), "parameters")] + specs = [encode(m._spec)] + params = [with_prefix(record_dtypes(m.parameters), "parameters")] scalars = {k: getattr(m, k) for k in m.scalar_attrs} ds = xr.merge( @@ -1499,7 +1498,7 @@ def _copy_con_data(con: ConstraintBase) -> xr.Dataset: new_model._parameters = m._parameters.copy(deep=deep) if m._spec is not None: - new_model._spec = m._spec._reattach(new_model) + new_model._spec = m._spec._reattach(new_model, deep=deep) new_model._blocks = m._blocks.copy(deep=deep) if m._blocks is not None else None for attr in m.scalar_attrs: diff --git a/linopy/model.py b/linopy/model.py index 2e127973..41c1e2bd 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -482,8 +482,11 @@ def add_spec( 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. + Which parameters to keep in ``model.spec.parameters``: those the + named expressions read, all of them, or none. ``model.parameters`` + stays the caller's and is never written to. This decides what a + netcdf file holds, not what this session can read: ``model.spec`` + falls back to ``sources`` for a parameter it did not keep. Returns ------- diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index fff7e30f..bb588a2a 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -1,10 +1,18 @@ """ ``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``. +The spec owns its data. The spec text, the retained parameters, the lookups +and the master coordinates sit on the accessor rather than in +``model.parameters``, which stays the caller's: a spec never overwrites what +was put there, and nothing reading a spec-built model has to guess which of +its parameters the spec owns. All of it round trips through a file, written +under the ``spec-`` prefix. + +A parameter is resolved the same way however much of it was retained: from +the retained dataset, else from the sources the model was built with, which +the accessor keeps for as long as the model lives. So ``retain`` decides what +a *file* holds, not what a session can read, and it is only after a round trip +that a parameter can be out of reach. """ from __future__ import annotations @@ -77,13 +85,18 @@ def attach( text, program = _source(spec) attached: Attached = attach_data(program, sources, retain=retain) build(model, attached) - model.parameters = attached.retained().assign_coords(dict(attached.coords)) - return ModelSpec(model, program, text) + parameters = attached.retained().assign_coords(dict(attached.coords)) + return ModelSpec(model, program, text, parameters, attached) -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 restore(model: Model, text: str, parameters: xr.Dataset) -> ModelSpec: + """ + The accessor for *model*, with the program lowered afresh from *text*. + + Read from a file, so the sources the model was built with are gone and + only what ``retain`` kept can be read back. + """ + return ModelSpec(model, to_program(yaml.safe_load(text)), text, parameters, None) def _source(spec: SpecLike) -> tuple[str, ms.Program]: @@ -123,10 +136,19 @@ class ModelSpec: The spec as YAML, verbatim where a file or text was passed. """ - def __init__(self, model: Model, program: ms.Program, text: str) -> None: + def __init__( + self, + model: Model, + program: ms.Program, + text: str, + parameters: xr.Dataset, + attached: Attached | None, + ) -> None: self._model = model self.program = program self.text = text + self._parameters = parameters + self._attached = attached def __repr__(self) -> str: p = self.program @@ -143,14 +165,20 @@ def __repr__(self) -> str: rows.append(_row("Expressions", list(p.named_expressions))) return "\n".join(rows) - def _reattach(self, model: Model) -> ModelSpec: - """The same spec, read off *model*.""" - return ModelSpec(model, self.program, self.text) + def _reattach(self, model: Model, deep: bool = True) -> ModelSpec: + """The same spec, read off *model*, holding its own copy of the parameters.""" + return ModelSpec( + model, + self.program, + self.text, + self._parameters.copy(deep=deep), + self._attached, + ) @property def parameters(self) -> xr.Dataset: - """The parameters and lookups retained on the model, on the master coordinates.""" - return self._model.parameters + """The parameters and lookups the spec retained, on the master coordinates.""" + return self._parameters @property def description(self) -> str: @@ -222,8 +250,9 @@ def evaluate( """ The named expression *name*, with its parameters attached 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 + For reading the spec against other data than the model was built with, + and for a model read from a file, whose own sources are gone. + ``model.spec.expressions`` needs neither. *sources* is read the way ``add_spec`` read it, and must describe the coordinates the model was built on. @@ -244,14 +273,18 @@ def evaluate( ) return NamedExpression(self, name, self._context(attached.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 _resolve(self, name: str) -> xr.DataArray: + """The parameter *name*: retained if it was kept, else read from the sources again.""" + if name in self.parameters: + return self.parameters[name] + if self._attached is not None: + return self._attached.parameter(name) + raise SpecDataError( + f"parameter '{name}' was not retained and this model no longer holds the sources " + f"it was built with, which is what a model read from a file looks like. Build with " + f"retain='all' before writing it out, or read the expression with " + f"evaluate(name, sources)." + ) def _context(self, resolve: Resolve) -> Context: return Context( @@ -277,7 +310,7 @@ def __getitem__(self, name: str) -> NamedExpression: + did_you_mean(name, self._spec.program.named_expressions) ) return NamedExpression( - self._spec, name, self._spec._context(self._spec._retained) + self._spec, name, self._spec._context(self._spec._resolve) ) def __iter__(self) -> Iterator[str]: @@ -326,8 +359,9 @@ class NamedExpression(Declaration): One named expression, in three views: its math, its linopy fold and its solution. The object pins the data sources it was made with for its lifetime, so the - three views agree. ``expressions[name]`` reads the retained parameters and - the solution the model holds; ``evaluate(name, sources)`` attaches fresh data. + three views agree. ``expressions[name]`` reads the model's own data -- + what ``retain`` kept, and the sources behind it for the rest; + ``evaluate(name, sources)`` attaches fresh data instead. Attributes ---------- @@ -371,7 +405,8 @@ def solution(self) -> xr.DataArray: RuntimeError The model reads a variable but holds no solution yet. SpecDataError - A parameter the body reads was not retained. + A parameter the body reads was neither retained nor + still reachable through the model's sources. """ return fold(self._name, self._ctx) diff --git a/linopy/spec/netcdf.py b/linopy/spec/netcdf.py index 2a62e1ea..7f1a2813 100644 --- a/linopy/spec/netcdf.py +++ b/linopy/spec/netcdf.py @@ -41,6 +41,7 @@ PREFIX = "spec" COORD = "coords__" +PARAM = "param__" CODES = "codes__" CATEGORIES = "cats__" CATEGORY_DIM = "category__" @@ -48,39 +49,40 @@ HOLES: dict[str, Any] = {"f": np.nan, "O": np.nan, "M": np.datetime64("NaT")} -def encode(spec: ModelSpec) -> tuple[xr.Dataset, xr.Dataset]: +def encode(spec: ModelSpec) -> xr.Dataset: """ - The model's parameters without the coded arrays, and the spec's own dataset. - - 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. + The spec's own dataset: its text, its master coordinates and its parameters. + + The spec text is the dataset's one attribute, which the merge lifts to the + file's. Beside it sits one array of labels per master coordinate and, per + parameter, either its values or -- where it is coded -- its codes and its + categories. The dataset carries no coordinates of its own: an index + coordinate is dropped on read together with the dimension it indexes once + no data variable is left over that dimension, and a master coordinate + nothing else reaches has exactly that shape. So a parameter is written + over bare dimensions and put back on the master coordinates on read. """ - 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) - written = with_prefix(xr.Dataset(arrays), PREFIX).assign_attrs( - {SPEC_ATTR: spec.text} - ) - return parameters, written + coded = _coded(spec) + for name, arr in spec.parameters.items(): + if str(name) in coded: + arrays.update(_encode(str(name), arr)) + else: + arrays[PARAM + str(name)] = _array(arr.to_numpy(), arr.dims, str(arr.dtype)) + return with_prefix(xr.Dataset(arrays), PREFIX).assign_attrs({SPEC_ATTR: spec.text}) def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: """ - Re-lower *text* onto *model* and put its coded arrays and coordinates back. + Re-lower *text* onto *model* and read back the dataset :func:`encode` wrote. - 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. + The master coordinates, the plainly written parameters and the coded ones + together are the dataset :func:`linopy.spec.accessor.attach` gave the spec + when the model was built. ``model.parameters`` is not touched: it holds + what the caller put there and nothing of the spec. """ sub = get_prefix(ds, PREFIX) coords = { @@ -88,24 +90,30 @@ def decode(model: Model, ds: xr.Dataset, text: str) -> ModelSpec: for name in sub.data_vars if str(name).startswith(COORD) } - coded = { - _stripped(name, CODES): _decode(sub, _stripped(name, CODES), coords) + arrays = { + _stripped(name, PARAM): _plain(sub[name], _stripped(name, PARAM), coords) for name in sub.data_vars - if str(name).startswith(CODES) + if str(name).startswith(PARAM) } - model.parameters = model.parameters.assign_coords(coords).assign(coded) + arrays.update( + { + _stripped(name, CODES): _decode(sub, _stripped(name, CODES), coords) + for name in sub.data_vars + if str(name).startswith(CODES) + } + ) restamp_coords(model, coords) - return restore(model, text) + return restore(model, text, xr.Dataset(arrays).assign_coords(coords)) -def _coded(spec: ModelSpec) -> list[str]: +def _coded(spec: ModelSpec) -> set[str]: """The parameters written as codes: every lookup and every array of objects.""" lookups = {name for by_name in spec.lookups.values() for name in by_name} - return [ + return { str(name) for name, arr in spec.parameters.items() if name in lookups or arr.dtype == object - ] + } def _encode(name: str, arr: xr.DataArray) -> dict[str, xr.DataArray]: @@ -122,6 +130,14 @@ def _encode(name: str, arr: xr.DataArray) -> dict[str, xr.DataArray]: return written +def _plain(arr: xr.DataArray, name: str, coords: dict[str, pd.Index]) -> xr.DataArray: + """A parameter written as its own values, back on the master coordinates at its own dtype.""" + dims = tuple(str(d) for d in arr.dims) + return xr.DataArray( + _values(arr), coords={d: coords[d] for d in dims}, dims=dims, name=name + ) + + def _decode(sub: xr.Dataset, name: str, coords: dict[str, pd.Index]) -> xr.DataArray: codes = sub[CODES + name] dtype = np.dtype(codes.attrs[DTYPE_ATTR]) diff --git a/linopy/testing.py b/linopy/testing.py index 0333d1cb..2cf5d28c 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -154,6 +154,7 @@ def assert_model_equal(a: Model, b: Model) -> None: assert (a._spec is None) == (b._spec is None) if a._spec is not None and b._spec is not None: assert a._spec.text == b._spec.text + assert_datasetequal(a._spec.parameters, b._spec.parameters) assert a.status == b.status assert a.termination_condition == b.termination_condition diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index 3716b190..ed8457d5 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -10,6 +10,7 @@ from pathlib import Path from typing import Any +import numpy as np import pandas as pd import pytest import xarray as xr @@ -109,16 +110,27 @@ def test_from_spec_passes_model_kwargs_and_chains() -> None: ("none", set()), ], ) -def test_retain_decides_what_the_fold_can_read(retain: str, kept: set[str]) -> None: +def test_retain_decides_what_is_kept_and_not_what_can_be_read( + retain: str, kept: set[str] +) -> None: + """A parameter retain dropped is read from the sources the model still holds.""" m = solved(yaml_dict(), DISPATCH_DATA, retain=retain) - assert set(m.parameters.data_vars) == kept + assert set(m.spec.parameters.data_vars) == kept + assert not m.parameters.data_vars want = (DISPATCH_P * [0.0, 50.0]).sum("generator").rename("spend") + xr.testing.assert_allclose(m.spec.expressions["spend"].solution, want) xr.testing.assert_allclose(m.spec.evaluate("spend", DISPATCH_DATA).solution, want) - if "cost" in kept: - xr.testing.assert_allclose(m.spec.expressions["spend"].solution, want) - else: - with pytest.raises(SpecDataError, match="not retained"): - m.spec.expressions["spend"].solution + + +def test_the_spec_keeps_its_parameters_off_the_model() -> None: + """``model.parameters`` is the caller's: a build neither reads nor writes it.""" + own = xr.DataArray(np.array(["a", "b", "c"], dtype=object), dims=["own"]) + m = Model() + m.parameters["cost"] = own + m.add_spec(yaml_dict(), DISPATCH_DATA, retain="all") + + assert m.parameters["cost"].equals(own) + assert m.spec.parameters["cost"].dims == ("generator",) def test_an_unknown_expression_is_a_key_error_with_a_hint() -> None: diff --git a/test/test_spec_builder.py b/test/test_spec_builder.py index e5a3d879..b7ee54f5 100644 --- a/test/test_spec_builder.py +++ b/test/test_spec_builder.py @@ -413,7 +413,7 @@ def test_an_operator_under_a_power_keeps_its_parameters_retained() -> None: 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) + assert {"c", "lag"} <= set(m.spec.parameters.data_vars) xr.testing.assert_allclose( m.spec.expressions["e"].solution, xr.DataArray([0.0, 0.0, 4.0], coords={"t": T}, name="e"), diff --git a/test/test_spec_io.py b/test/test_spec_io.py index 5af661f0..30139a14 100644 --- a/test/test_spec_io.py +++ b/test/test_spec_io.py @@ -31,6 +31,7 @@ import linopy # noqa: E402 from linopy import Model, read_netcdf # noqa: E402 from linopy.io import SPEC_ATTR # noqa: E402 +from linopy.spec import SpecDataError # noqa: E402 from linopy.spec.testing import synthetic_sources # noqa: E402 from linopy.testing import assert_model_equal # noqa: E402 @@ -137,17 +138,34 @@ def test_a_spec_built_model_round_trips( def test_a_retain_none_model_evaluates_after_a_round_trip( tmp_path: Path, engine: str ) -> None: + """A file is where retain bites: the sources the built model still read are gone.""" m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="none") p = roundtrip(m, tmp_path, engine) assert_model_equal(m, p) assert not p.spec.parameters.data_vars + with pytest.raises(SpecDataError, match="no longer holds the sources"): + p.spec.expressions["spend"].solution assert_arrayequal( - m.spec.evaluate("spend", DISPATCH_DATA).solution, + m.spec.expressions["spend"].solution, p.spec.evaluate("spend", DISPATCH_DATA).solution, ) +@pytest.mark.parametrize("engine", ENGINES) +def test_the_caller_parameters_and_the_spec_ones_stay_apart( + tmp_path: Path, engine: str +) -> None: + """A model parameter of the caller's is written beside the spec's, not into them.""" + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="all") + m.parameters["cost"] = xr.DataArray([1, 2, 3], dims=["own"]) + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert_arrayequal(p.parameters["cost"], m.parameters["cost"]) + assert_arrayequal(p.spec.parameters["cost"], m.spec.parameters["cost"]) + + @pytest.mark.parametrize("engine", ENGINES) @pytest.mark.parametrize("mapped", [3, 2, 0], ids=["full", "partial", "empty"]) @pytest.mark.parametrize("name", LOOKUP_OVER) @@ -215,13 +233,23 @@ def test_a_copy_carries_the_spec(deep: bool) -> None: """The copy's spec reads the copy, and only a deep copy owns its buffers.""" m = Model.from_spec(WHERE_SPEC, WHERE_DATA, retain="all") p = m.copy(deep=deep) - p.parameters["label"].values[1] = "changed" + p.spec.parameters["label"].values[1] = "changed" assert p.spec.text == m.spec.text assert p.spec.parameters["label"].values[1] == "changed" assert m.spec.parameters["label"].values[1] == ("u" if deep else "changed") +def test_a_copy_can_still_read_what_retain_dropped() -> None: + """A copy keeps the sources, so it folds an unretained parameter like its original.""" + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="none") + + assert_arrayequal( + m.copy(include_solution=True).spec.expressions["spend"].solution, + m.spec.expressions["spend"].solution, + ) + + def test_a_model_without_a_spec_carries_none(tmp_path: Path) -> None: m = Model() x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") From 87b2ff11137f02250ff176c03d28a9c4ea8c876d Mon Sep 17 00:00:00 2001 From: Felix <117816358+FBumann@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:25:24 +0200 Subject: [PATCH 35/35] fix(spec): four fixes from reviewing the spec builder (#950) Four bugs found reviewing #922, each reproduced by a test that fails without the fix. **`read_netcdf` relabelled a hand-added variable.** `restamp_coords` put a spec's master coordinates on every container carrying a dimension of that name, asking only whether the dtypes differed and never whether the labels were the same. A variable on `snapshot = [10, 11, 12]` beside a spec built on `[0, 1, 2]` came back silently relabelled; one of another length failed the read with `conflicting sizes`. An index is now restamped only where it holds the master's own labels, which is what `Index.equals` asks. **A lookup that maps nothing.** `grouped_sum` promises the empty sum, 0, for a group no member reaches, but delivered it by filtering to the mapped members and grouping what was left -- nothing, which xarray refuses to group. Answered directly now, as `sum_over` already answers a term beside an empty dimension. **`repr(model.spec)` on an unreached dimension.** `attach` deliberately lets a declared dimension nothing reaches go without a source, so it never lands in `coords`; the repr indexed it anyway and raised `KeyError`. Shown as `unreached`. **A half-built model on a retain failure.** `retain="all"` resolved parameters after `build()`, so a missing one raised with the model half-built and `_spec` unset -- and the retry was then refused by the empty-model guard. The retained set needs nothing the build produces, so it is resolved first. Co-Authored-By: Claude Opus 5 (1M context) --- doc/release_notes.rst | 8 ++++++++ linopy/io.py | 33 +++++++++++++++++++++++++++------ linopy/spec/accessor.py | 11 +++++++++-- linopy/spec/operators.py | 36 ++++++++++++++++++++++++++++++++++++ test/test_spec_accessor.py | 24 ++++++++++++++++++++++++ test/test_spec_io.py | 18 ++++++++++++++++++ test/test_spec_operators.py | 23 +++++++++++++++++++++++ 7 files changed, 145 insertions(+), 8 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index fc0d88f8..9210ff80 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -40,6 +40,14 @@ Upcoming Version *Other* +* ``Model.add_spec`` resolves the parameters it retains before it builds. A ``retain="all"`` build that could not read a parameter no declaration uses raised after the variables and constraints were already added, leaving a model that the "builds into an empty model" guard then refused to build into again. + +* ``repr(model.spec)`` no longer raises ``KeyError`` for a dimension the spec declares but nothing reaches, which needs no source and so has no coordinates; it is shown as ``unreached``. + +* A grouped sum through a lookup that maps no member at all now holds the empty sum, ``0``, on every declared group, as its documented rule says. It raised xarray's ``ValueError: must not be empty`` instead. + +* ``read_netcdf`` no longer rewrites the coordinates of a container that merely shares a dimension's *name* with a spec-built model's master coordinates. A hand-added variable on its own labels kept them; before, it was silently relabelled onto the master ones, or the read failed outright when the two lengths differed. + * ``add_piecewise_formulation`` gained a ``mask`` parameter declaring which breakpoint slots hold a real breakpoint. It is needed for **ragged** curves — entities with different numbers of breakpoints — which are stored densely with the surplus slots left absent. Under v1 that absence must be declared (``mask=x_pts.notnull()``) rather than read off the NaN padding. (https://github.com/PyPSA/linopy/issues/884) *Internal* diff --git a/linopy/io.py b/linopy/io.py index 275cb2b7..22781750 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1107,7 +1107,14 @@ def restore_dtypes(ds: xr.Dataset) -> xr.Dataset: def restamp_coords(m: Model, coords: Mapping[str, pd.Index]) -> None: - """Put *coords* on every container of *m* that carries one of those dimensions.""" + """ + Put *coords* on every container of *m* that was built on them. + + Only on those: a container may carry a dimension of that name and its own + labels -- a hand-added variable beside a spec-built one -- and restamping + it would rewrite labels it never had, or fail outright over a length the + master coordinate does not share. + """ from linopy.constraints import Constraint, CSRConstraint from linopy.csr import Grid @@ -1122,7 +1129,7 @@ def restamp_coords(m: Model, coords: Mapping[str, pd.Index]) -> None: elif isinstance(constraint, CSRConstraint): constraint._grid = Grid( { - d: coords.get(d, index) + d: _restamped(index, coords.get(str(d))) for d, index in constraint._grid.indexes.items() } ) @@ -1130,15 +1137,29 @@ def restamp_coords(m: Model, coords: Mapping[str, pd.Index]) -> None: def _stamped(data: xr.Dataset, coords: Mapping[str, pd.Index]) -> xr.Dataset: """*data* with *coords* 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 + str(dim): restamped + for dim, index in data.indexes.items() + if (restamped := _restamped(index, coords.get(str(dim)))) is not index } return data.assign_coords(stale) if stale else data +def _restamped(found: pd.Index, master: pd.Index | None) -> pd.Index: + """ + *master* where *found* is it as a netcdf type gave it back, else *found* itself. + + A narrowed int or a widened bool holds the same labels at another dtype + and is the one to replace -- which is what ``Index.equals`` asks, since it + compares labels and not dtypes. An index of another length, or of other + labels entirely, belongs to a container that was never built on *master* + and is left alone. + """ + if master is None or found.dtype == master.dtype: + return found + return master if found.equals(master) else found + + def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: """ Write out the model to a netcdf file. diff --git a/linopy/spec/accessor.py b/linopy/spec/accessor.py index bb588a2a..b8cec579 100644 --- a/linopy/spec/accessor.py +++ b/linopy/spec/accessor.py @@ -84,8 +84,10 @@ def attach( ) text, program = _source(spec) attached: Attached = attach_data(program, sources, retain=retain) - build(model, attached) + # Resolved before the build, so a parameter no declaration reads cannot fail + # halfway through one and leave a model too full to build into again. parameters = attached.retained().assign_coords(dict(attached.coords)) + build(model, attached) return ModelSpec(model, program, text, parameters, attached) @@ -116,6 +118,11 @@ def _source(spec: SpecLike) -> tuple[str, ms.Program]: return loaded.to_yaml(), to_program(loaded) +def _dimension(dim: str, coords: Mapping[str, pd.Index]) -> str: + """A dimension and how many labels it holds; a declared one nothing reaches holds none.""" + return f"{dim} ({len(coords[dim])})" if dim in coords else f"{dim} (unreached)" + + def _row(label: str, items: list[str], cap: int = 8) -> str: """One aligned summary line, capped with a ``(+N more)`` tail.""" shown = items[:cap] @@ -156,7 +163,7 @@ def __repr__(self) -> str: head = f"ModelSpec: {self.description}" if self.description else "ModelSpec" rows = [ head, - _row("Dimensions", [f"{d} ({len(coords[d])})" for d in p.dimensions]), + _row("Dimensions", [_dimension(d, coords) for d in p.dimensions]), _row("Variables", list(p.variables)), _row("Constraints", list(p.constraints)), ] diff --git a/linopy/spec/operators.py b/linopy/spec/operators.py index 17fd2d46..7b86565a 100644 --- a/linopy/spec/operators.py +++ b/linopy/spec/operators.py @@ -81,6 +81,8 @@ def grouped_sum( mappings = _renamed(mappings, into) present = _present(mappings) dim = str(mappings[0].dims[0]) + if not bool(present.any()): + return _empty_groups(array, dim, into=into, labels=labels) if not bool(present.all()): keep = present.to_numpy() mappings = tuple(m.isel({dim: keep}) for m in mappings) @@ -92,6 +94,40 @@ def grouped_sum( return summed.reindex({d: labels[d] for d in into}).fillna(0.0) +def _empty_groups( + array: Array, + dim: str, + *, + into: tuple[str, ...], + labels: Mapping[str, pd.Index], +) -> Array: + """ + The grouped sum of an operand no member of *dim* is mapped out of. + + Every declared group holds the empty sum, which is 0. Grouping cannot say + so itself: filtering the operand down to its mapped members leaves nothing, + and an empty dimension is one xarray refuses to group over. + """ + kept = [d for d in _coord_dims(array) if d != dim] + zeros = xr.DataArray( + np.zeros([array.sizes[d] for d in kept] + [len(labels[d]) for d in into]), + coords={ + **{d: array.indexes[d] for d in kept}, + **{d: labels[d] for d in into}, + }, + dims=kept + list(into), + ) + if isinstance(array, xr.DataArray): + return zeros + return LinearExpression.from_constant(array.model, zeros) + + +def _coord_dims(array: Array) -> list[str]: + """The dimensions the operand is labelled over, without a term's own ``_term``.""" + dims = array.dims if isinstance(array, xr.DataArray) else array.coord_dims + return [str(d) for d in dims] + + @overload def at( array: xr.DataArray, mappings: tuple[xr.DataArray, ...], *, into: tuple[str, ...] diff --git a/test/test_spec_accessor.py b/test/test_spec_accessor.py index ed8457d5..91605cec 100644 --- a/test/test_spec_accessor.py +++ b/test/test_spec_accessor.py @@ -133,6 +133,30 @@ def test_the_spec_keeps_its_parameters_off_the_model() -> None: assert m.spec.parameters["cost"].dims == ("generator",) +def test_a_build_that_cannot_retain_leaves_the_model_buildable() -> None: + """retain='all' reaches parameters no declaration does, and must not half-build on one.""" + spec = with_(yaml_dict(), parameters={"spare": {"dims": ["generator"]}}) + m = Model() + with pytest.raises(SpecDataError, match="no data provided for parameter 'spare'"): + m.add_spec(spec, DISPATCH_DATA, retain="all") + assert not len(m.variables) and not len(m.constraints) + + spare = pd.Series([1.0, 2.0], index=GENERATOR) + m.add_spec(spec, {**DISPATCH_DATA, "spare": spare}, retain="all") + assert "spare" in m.spec.parameters + + +def test_a_declared_dimension_with_no_source_still_reprs() -> None: + """A dimension nothing reaches needs no source, so the repr must do without its labels.""" + spec = with_( + yaml_dict(), dimensions={"spare": {"dtype": "int", "description": "unreached"}} + ) + m = Model.from_spec(spec, DISPATCH_DATA) + + assert "spare (unreached)" in repr(m.spec) + assert "snapshot (3)" in repr(m.spec) + + def test_an_unknown_expression_is_a_key_error_with_a_hint() -> None: m = Model.from_spec(yaml_dict(), DISPATCH_DATA) with pytest.raises(KeyError, match="unknown named expression 'spent'.*spend"): diff --git a/test/test_spec_io.py b/test/test_spec_io.py index 30139a14..4dddf94e 100644 --- a/test/test_spec_io.py +++ b/test/test_spec_io.py @@ -190,6 +190,24 @@ def test_a_parameter_keeps_its_dtype(tmp_path: Path, engine: str, name: str) -> assert p.spec.parameters[name].dtype == m.spec.parameters[name].dtype +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize( + "labels", [[10, 11, 12], [0, 1]], ids=["relabelled", "shorter"] +) +def test_a_hand_added_variable_keeps_its_own_labels( + tmp_path: Path, engine: str, labels: list[int] +) -> None: + """A container sharing a master dimension's name but not its labels is left alone.""" + own = pd.Index(labels, name="snapshot") + m = Model.from_spec(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="all") + m.add_variables(coords=[own], name="side") + p = roundtrip(m, tmp_path, engine) + + assert p.variables["side"].indexes["snapshot"].equals(own) + assert p.spec.coords["snapshot"].equals(m.spec.coords["snapshot"]) + assert p.variables["p"].indexes["snapshot"].equals(m.spec.coords["snapshot"]) + + @pytest.mark.parametrize("engine", ENGINES) @pytest.mark.parametrize("frozen", [False, True], ids=["dataset", "csr"]) def test_every_container_shares_the_master_coordinate_dtypes( diff --git a/test/test_spec_operators.py b/test/test_spec_operators.py index 4e470eed..bd9a1b45 100644 --- a/test/test_spec_operators.py +++ b/test/test_spec_operators.py @@ -202,6 +202,29 @@ def test_an_empty_group_on_the_constant_side_is_a_zero_and_not_a_gap() -> None: assert float(m.solution["imports"].sel(bus="south")) == pytest.approx(0.0) +def test_a_lookup_that_maps_nothing_leaves_every_group_at_the_empty_sum() -> None: + """Filtering to the mapped members leaves nothing, and nothing is what xarray will not group.""" + spec = with_( + GROUPED_SPEC, + variables={ + "out": { + "foreach": ["generator"], + "bounds": {"lower": 0, "upper": "capacity"}, + } + }, + expressions={"per_bus": "sum(out, by=gen_bus)"}, + ) + sources = grouped_sources(pd.Series([3.0, 4.0], index=GENS)) + sources["gen_bus"] = pd.Series([], dtype=object) + m = solved(spec, sources) + + assert m.objective.value == pytest.approx(0.0) + per_bus = m.spec.expressions["per_bus"] + assert per_bus.expression.nterm == 0 + assert per_bus.solution.indexes["bus"].tolist() == ["north", "south"] + np.testing.assert_allclose(per_bus.solution.values, [0.0, 0.0]) + + def test_a_member_with_no_value_is_still_refused_through_a_group() -> None: with pytest.raises(SpecDataError, match="parameter 'capacity' covers 1 fewer"): Model.from_spec(GROUPED_SPEC, grouped_sources(pd.Series([3.0], index=GENS[:1])))