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/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/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 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..f893bf43 --- /dev/null +++ b/linopy/spec/binder.py @@ -0,0 +1,618 @@ +""" +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, field +from typing import Any, Literal, get_args + +import numpy as np +import pandas as pd +import xarray as xr +from math_spec import did_you_mean +from math_spec import program as ms + +from linopy.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", + "i": "int", + "u": "int", + "b": "bool", + "O": "str", + "U": "str", + "S": "str", + "M": "datetime", +} +_EMPTY_DTYPES: dict[str, Any] = { + "float": float, + "int": int, + "bool": bool, + "str": object, +} +_SCALARS = (bool, int, float, str, np.number, np.bool_) +_DIMENSION_SHAPES = "a pandas Index, a list, a tuple, a 1-D numpy array, a pandas Series or a 1-D DataArray" +_LOOKUP_SHAPES = "a pandas Series indexed by '{over}', a dict keyed by '{over}' labels, or a 1-D DataArray over '{over}'" +_PARAMETER_SHAPES = ( + "a DataArray over {dims}, a pandas Series whose (Multi)Index levels are {dims}, " + "a DataFrame with columns {columns} or in wide form, a dict keyed by label, or one number" +) + + +def 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 ``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()) + _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, eq=False) +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. + """ + + program: ms.Program + coords: Mapping[str, pd.Index] + lookups: Mapping[str, Mapping[str, xr.DataArray]] + retain: Retain + sources: Mapping[str, Any] + _keys: frozenset[str] = field(repr=False) + + def parameter(self, name: str) -> xr.DataArray: + """ + The parameter *name* resolved from ``sources`` and aligned to ``coords``. + + Resolved on every call and never cached. An already aligned array is + returned without a copy; a mismatching one is reindexed onto the + master coordinates, leaving NaN (``False`` for ``bool``) where no row + was supplied. + + Raises: + SpecDataError: No data, a shape no reader accepts, a rank other + than declared, a label its dimension lacks, two rows for one + coordinate, a null value in a row, or values of another type + than declared. + """ + declared = self._declaration(name) + if name not in self._keys: + raise SpecDataError(f"no data provided for parameter '{name}'") + arr = _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 [] + parameters = self.program.parameters + keep = ( + set(parameters) if self.retain == "all" else _report_closure(self.program) + ) + return [n for n, p in parameters.items() if p.derivation is None and n in keep] + + +def _report_closure(program: ms.Program) -> set[str]: + """Every parameter a named expression reads, by node or by name.""" + bodies = tuple(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 + one = len(unknown) == 1 + lead = f"source key {unknown[0]!r} names" if one else f"source keys {unknown} name" + raise SpecDataError( + f"{lead} neither a parameter, a dimension nor a lookup this spec declares. " + f"{did_you_mean(unknown[0], known)} Pass only what the spec takes." + ) + + +# --------------------------------------------------------------------------- +# dimensions +# --------------------------------------------------------------------------- + + +def _reached(program: ms.Program) -> set[str]: + dims: set[str] = set() + for declared in (program.parameters, program.variables, program.constraints): + dims.update(d for decl in declared.values() for d in decl.dims) + dims.update(pw.over for pw in program.piecewise.values()) + for over, lk in program.lookups: + dims.add(over) + if lk.target is not None: + dims.add(lk.target) + return dims + + +def _master_coords( + program: ms.Program, sources: Mapping[str, Any], keys: frozenset[str] +) -> dict[str, pd.Index]: + reached = _reached(program) + coords: dict[str, pd.Index] = {} + for dim in program.dimensions: + if dim in keys: + coords[dim] = _index(dim, sources[dim]) + elif dim in reached: + raise SpecDataError( + f"dimension '{dim}' has no index: pass its labels under key '{dim}' as " + f"{_DIMENSION_SHAPES}. The index is what says which labels exist, and without " + f"one a mistyped label is indistinguishable from a new one." + ) + return coords + + +def _index(dim: str, obj: Any) -> pd.Index: + if isinstance(obj, (pd.Series, xr.DataArray, np.ndarray)): + if obj.ndim != 1: + raise SpecDataError( + f"index for dimension '{dim}' is {obj.ndim}-dimensional; pass {_DIMENSION_SHAPES}." + ) + values: Any = np.asarray(obj) + elif isinstance(obj, (pd.Index, list, tuple)): + values = obj + else: + raise SpecDataError( + f"index for dimension '{dim}': cannot read labels out of {type(obj).__name__}; pass {_DIMENSION_SHAPES}." + ) + index = pd.Index(values, name=dim) + if index.has_duplicates: + twice = index[index.duplicated()].unique().tolist() + raise SpecDataError( + f"dimension '{dim}' lists {_shown(twice)} more than once. A dimension's members are a set: " + f"each label appears once, in the order the source gives it." + ) + return index + + +# --------------------------------------------------------------------------- +# lookups +# --------------------------------------------------------------------------- + + +def _lookups( + program: ms.Program, + sources: Mapping[str, Any], + keys: frozenset[str], + coords: Mapping[str, pd.Index], +) -> dict[str, dict[str, xr.DataArray]]: + out: dict[str, dict[str, xr.DataArray]] = {} + for over, lk in program.lookups: + space = lk.target or lk.name + if lk.name not in keys: + raise SpecDataError( + f"no data provided for lookup '{lk.name}'. Pass it under key '{lk.name}' as " + f"{_LOOKUP_SHAPES.format(over=over)}, holding a '{space}' value for each " + f"'{over}' label it maps and nothing for a label it does not." + ) + series = _lookup_series(lk.name, over, sources[lk.name]) + _check_lookup(series, lk, over, coords) + padded = series.reindex(coords[over]) + out.setdefault(over, {})[lk.name] = xr.DataArray(padded, name=lk.name) + 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.dtype is not None: + _check_value_dtype(lk.name, lk.dtype, series.dtype, kind="lookup") + if lk.target is None: + return + values = pd.Index(series.to_numpy()) + foreign = values[~values.isin(coords[lk.target])].unique().tolist() + if foreign: + raise SpecDataError( + f"dimension '{over}' lookup '{lk.name}' has value(s) that are not '{lk.target}' labels: " + f"{_shown(foreign)}. Every value must be a declared '{lk.target}' label, otherwise " + f"sum(by={lk.name}) drops those terms in the join that places them, and the model " + f"builds and solves without them." + ) + + +# --------------------------------------------------------------------------- +# parameters +# --------------------------------------------------------------------------- + + +def _fill(declared: ms.ParameterDeclaration) -> Any: + return False if declared.dtype == "bool" else np.nan + + +def _as_array( + name: str, + declared: ms.ParameterDeclaration, + obj: Any, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + if isinstance(obj, xr.DataArray): + return _from_dense(name, declared, obj) + if isinstance(obj, pd.DataFrame): + return _from_frame(name, declared, obj, coords) + if isinstance(obj, pd.Series): + return _from_rows(name, declared, obj, coords) + if isinstance(obj, Mapping): + return _from_rows(name, declared, pd.Series(dict(obj)), coords) + if isinstance(obj, _SCALARS): + return _from_scalar(name, declared, obj, coords) + dims = declared.dims + raise SpecDataError( + f"parameter '{name}': cannot adapt {type(obj).__name__} to an array over {list(dims)}; " + f"pass {_parameter_shapes(dims)}." + ) + + +def _parameter_shapes(dims: Sequence[str]) -> str: + return _PARAMETER_SHAPES.format(dims=list(dims), columns=[*dims, "value"]) + + +def _from_scalar( + name: str, + declared: ms.ParameterDeclaration, + obj: Any, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + if pd.isna(obj): + raise SpecDataError( + f"parameter '{name}' is one value and that value is a hole (null or NaN). " + f"A number was meant, or the parameter has no data and should not be passed." + ) + value = np.asarray(obj) + _check_value_dtype(name, declared.dtype, value.dtype) + if declared.dtype == "float": + value = value.astype(float) + arr = xr.DataArray(value, name=name) + if declared.dims: + arr = arr.expand_dims({d: coords[d] for d in declared.dims}) + return arr + + +def _from_dense( + name: str, declared: ms.ParameterDeclaration, arr: xr.DataArray +) -> xr.DataArray: + dims = declared.dims + _check_value_dtype(name, declared.dtype, arr.dtype) + if set(arr.dims) != set(dims) or len(arr.dims) != len(dims): + raise SpecDataError( + f"parameter '{name}' arrived as a DataArray over {list(arr.dims)}, and '{name}' is over " + f"{list(dims)}. The dims must be the declared ones, in any order." + ) + for d in dims: + if d not in arr.indexes: + raise SpecDataError( + f"parameter '{name}' has no coordinate labels along '{d}'. A parameter is read for " + f"values against its labels, so every dimension needs an index coordinate." + ) + _refuse_duplicate_coordinates(name, (d,), arr.indexes[d]) + return arr.transpose(*dims) + + +def _from_frame( + name: str, + declared: ms.ParameterDeclaration, + df: pd.DataFrame, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + dims = declared.dims + tidy = df + if not set(dims) <= set(df.columns) and set(dims) <= _headers(df): + tidy = df.reset_index() + if "value" in tidy.columns and set(dims) <= set(tidy.columns): + indexed = tidy.set_index(list(dims)) if dims else tidy + return _from_rows(name, declared, indexed["value"], coords) + if len(dims) == 2: + return _from_dense(name, declared, xr.DataArray(_wide(name, dims, df))) + raise SpecDataError( + f"parameter '{name}' arrived as a DataFrame with columns {list(df.columns)}; a table for " + f"'{name}' carries columns {[*dims, 'value']}." + ) + + +def _headers(df: pd.DataFrame) -> set[Any]: + return set(df.columns) | set(df.index.names) + + +def _wide(name: str, dims: tuple[str, ...], df: pd.DataFrame) -> pd.DataFrame: + names = (df.index.name, df.columns.name) + if names == (None, None): + return df.rename_axis(index=dims[0], columns=dims[1]) + if set(names) == set(dims): + return df + raise SpecDataError( + f"parameter '{name}' arrived as a wide DataFrame with index '{names[0]}' and columns " + f"'{names[1]}', and '{name}' is over {list(dims)}. Name the index and columns after the " + f"two dims, or pass a table with columns {[*dims, 'value']}." + ) + + +def _from_rows( + name: str, + declared: ms.ParameterDeclaration, + series: pd.Series, + coords: Mapping[str, pd.Index], +) -> xr.DataArray: + dims = declared.dims + if not dims: + if len(series) != 1: + raise SpecDataError( + f"parameter '{name}' is declared with no dims, which means one value broadcast " + f"everywhere, but its source has {len(series)} rows. Declare the dims it is indexed " + f"by, or pass one number." + ) + return _from_scalar(name, declared, series.iloc[0], coords) + series = _with_dims(name, dims, series) + if series.empty: + series = series.astype(_EMPTY_DTYPES[declared.dtype]) + holes = series.isna() + if holes.any(): + raise SpecDataError( + f"parameter '{name}' carries {int(holes.sum())} row(s) with no value, null or NaN: " + f"{_coordinates_shown(dims, series.index[holes][:3])}. In a table the absence of a " + f"value is the absence of the row, and such a row says the coordinate exists and denies " + f"it in the same breath. Drop those rows, or supply the values." + ) + _check_value_dtype(name, declared.dtype, series.dtype) + _refuse_duplicate_coordinates(name, dims, series.index) + for d in dims: + _refuse_strangers(name, d, series.index.get_level_values(d), coords[d]) + onto = [coords[d] for d in dims] + full = onto[0] if len(dims) == 1 else pd.MultiIndex.from_product(onto, names=dims) + values = series.reindex(full, fill_value=_fill(declared)).to_numpy() + dense = values.reshape(tuple(len(index) for index in onto)) + return xr.DataArray(dense, dims=dims, coords=dict(zip(dims, onto)), name=name) + + +def _with_dims(name: str, dims: tuple[str, ...], series: pd.Series) -> pd.Series: + index = series.index + if index.nlevels != len(dims): + raise SpecDataError( + f"parameter '{name}': a Series or dict carries one label per level, and its index has " + f"{index.nlevels} level(s) where '{name}' is over {list(dims)}. " + f"Pass {_parameter_shapes(dims)}." + ) + names = list(index.names) + if all(n is None for n in names): + return series.set_axis(index.set_names(list(dims))) + if set(names) != set(dims): + raise SpecDataError( + f"parameter '{name}' is indexed by {names}, and '{name}' is over {list(dims)}. " + f"Name the index levels after the declared dims." + ) + if tuple(names) != dims: + series = series.reorder_levels(list(dims)) + return series + + +def _refuse_duplicate_coordinates( + name: str, dims: tuple[str, ...], index: pd.Index +) -> None: + duplicated = index.duplicated() + if not duplicated.any(): + return + counts = index[duplicated].value_counts() + shown = "; ".join( + f"{_coordinate(dims, key)} ({n + 1} rows)" for key, n in counts.iloc[:3].items() + ) + raise SpecDataError( + f"parameter '{name}' has more than one row for a coordinate: {shown}. A parameter is a " + f"function of its dims, so which value applies is undefined; aggregate the source to one " + f"row per {list(dims)} before attaching it." + ) + + +def _refuse_strangers(name: str, dim: str, labels: pd.Index, known: pd.Index) -> None: + strangers = labels[~labels.isin(known)].unique().tolist() + if not strangers: + return + raise SpecDataError( + f"parameter '{name}' has label(s) in dimension '{dim}' that are not coordinates of it: " + f"{_shown(strangers)}.\n {dim} has: {_shown(known.tolist(), 10)}\n" + f"A 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()): + stale = {d: i for d, i in onto.items() if arr.indexes[d].dtype != i.dtype} + return arr.assign_coords(stale) if stale else arr + for d, index in onto.items(): + _refuse_strangers(name, d, arr.indexes[d], index) + return arr.reindex(onto, fill_value=fill) + + +def _check_value_dtype( + name: str, declared: str, dtype: Any, kind: str = "parameter" +) -> None: + if str(dtype.kind) in _ACCEPTED_KINDS[declared]: + return + arrived = _KIND_NAMES.get(str(dtype.kind), str(dtype)) + raise SpecDataError( + f"{kind} '{name}' is declared '{declared}' and its values arrived as '{arrived}'. " + f"A declared dtype is a claim about the values, and it is checked here: the file says what " + f"the values are, or the values are not attached.\n" + f" Cast the values to {declared}, if the declaration is what you meant\n" + f" Or declare what the data has: {{dtype: {arrived}}}" + ) + + +# --------------------------------------------------------------------------- +# wording +# --------------------------------------------------------------------------- + + +def _shown(labels: Sequence[Any], limit: int = 5) -> str: + head = ", ".join(repr(x) for x in labels[:limit]) + return head + (f" (and {len(labels) - limit} more)" if len(labels) > limit else "") + + +def _coordinate(dims: Sequence[str], key: Hashable) -> str: + row = key if isinstance(key, tuple) else (key,) + return ", ".join(f"{d}={v!r}" for d, v in zip(dims, row)) + + +def _coordinates_shown(dims: Sequence[str], rows: Iterable[Hashable]) -> str: + return "; ".join(_coordinate(dims, row) for row in rows) diff --git a/linopy/spec/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/pyproject.toml b/pyproject.toml index b9c0c051..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. @@ -161,7 +169,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 new file mode 100644 index 00000000..c791eb6e --- /dev/null +++ b/test/test_spec_binder.py @@ -0,0 +1,807 @@ +"""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: dict[str, Any] = { + "dimensions": {"f": {"dtype": "str"}, "t": {"dtype": "int"}, "g": {"dtype": "str"}}, + "lookups": {"grp": {"over": "f", "into": "g"}}, + "parameters": { + "cost": {"dims": ["f"]}, + "cap": {"dims": ["f", "t"]}, + "flag": {"dims": ["f"], "dtype": "bool"}, + "rate": {"dims": []}, + "lead": {"dims": ["f"], "dtype": "int"}, + }, + "variables": { + "x": { + "foreach": ["f", "t"], + "where": "flag", + "bounds": {"lower": 0, "upper": "cap"}, + } + }, + "constraints": { + "k": {"foreach": ["g", "t"], "expression": "sum(x, by=grp) <= 10"}, + "s": { + "foreach": ["f", "t"], + "expression": "shift(x, over=t, offset=lead, edge=0) >= 0", + }, + }, + "objective": {"sense": "maximize", "expression": "sum(x * cost)"}, + "expressions": { + "spend": "sum(x * cost, over=t)", + "total": "sum(spend, over=f) * rate", + }, +} + +F = pd.Index(["b", "a", "c"], name="f") +T = pd.Index([0, 1, 2], name="t") +CAP = xr.DataArray(np.arange(9.0).reshape(3, 3), coords={"f": F, "t": T}, name="cap") +COST = pd.Series([1.0, 2.0, 3.0], index=F) + +DUP_ROWS = pd.Series([1.0, 9.0, 2.0], index=pd.Index(["a", "a", "b"], name="f")) +NULL_ROW = pd.Series({"a": 1.0, "b": None}) +NAN_ROW = pd.Series({"a": 1.0, "b": np.nan}) +NULL_FRAME = pd.DataFrame({"f": ["a", "b"], "value": [1.0, None]}) +STRAY_ROW = pd.Series({"a": 1.0, "zz": 2.0}) +DEEP_INDEX = pd.MultiIndex.from_tuples([("a", 0), ("b", 0)], names=["f", "k"]) +DEEP_ROWS = pd.Series([5.0, 5.0], index=DEEP_INDEX) + + +def sources_from( + base: Mapping[str, Any], override: Mapping[str, Any] +) -> dict[str, Any]: + """*base* with *override* applied; a ``None`` value drops the key instead.""" + merged = {**base, **override} + for key, value in override.items(): + if value is None: + merged.pop(key) + return merged + + +@pytest.fixture(scope="module") +def program() -> Any: + return math_spec.to_program(SPEC) + + +@pytest.fixture +def good() -> dict[str, Any]: + return { + "f": list(F), + "t": list(T), + "g": ["n", "e"], + "cost": COST, + "cap": CAP, + "flag": pd.Series([True, False, True], index=F), + "rate": 0.5, + "lead": pd.Series([1, 0, 1], index=F), + "grp": pd.Series(["n", "e", "n"], index=F), + } + + +def read_all(program: Any, sources: Mapping[str, Any]) -> list[xr.DataArray]: + 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"), + "indexed-frame": COST.to_frame("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}), + "lead": pd.Series({"a": 1}), + "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() + 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] + 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 + + +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: + good.pop("cost") + bound = bind(program, good) + with pytest.raises(SpecDataError, match="no data provided for parameter 'cost'"): + 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"]}, + r"dimension 'f' lists 'a' more than once", + id="duplicate-member", + ), + pytest.param( + {"cost": STRAY_ROW}, r"parameter 'cost'.*'f'.*'zz'", id="unknown-label" + ), + pytest.param( + {"cap": CAP.assign_coords(t=[0, 1, 9])}, + r"parameter 'cap'.*'t'.*\b9\b", + id="unknown-label-dense", + ), + pytest.param( + {"cap": CAP.assign_coords(t=[9, 0, 7])}, + r"not coordinates of it: 9, 7\.", + id="unknown-labels-dense-in-source-order", + ), + pytest.param( + {"cap": CAP.to_series().reset_index(name="value").assign(t=[9, 0, 7] * 3)}, + r"not coordinates of it: 9, 7\.", + id="unknown-labels-rows-in-source-order", + ), + pytest.param( + {"cost": DUP_ROWS}, + r"parameter 'cost' has more than one row for a coordinate: f='a' \(2 rows\)", + id="duplicated-coordinate-row", + ), + pytest.param( + {"cap": xr.DataArray([1.0, 2.0], coords={"f": ["a", "a"]})}, + r"parameter 'cap' arrived as a DataArray over \['f'\]", + id="dense-wrong-dims", + ), + pytest.param( + {"cap": xr.DataArray(np.ones((3, 3)), coords={"f": list(F), "t": [0, 0, 1]})}, + r"parameter 'cap' has more than one row for a coordinate: t=0 \(2 rows\)", + id="dense-duplicate-coordinate", + ), + pytest.param( + {"cap": COST}, + r"parameter 'cap'.*1 level\(s\) where 'cap' is over \['f', 't'\]", + id="wrong-rank", + ), + pytest.param( + {"cap": DEEP_ROWS.rename_axis(["f", "q"])}, + r"parameter 'cap' is indexed by \['f', 'q'\]", + id="wrong-level-names", + ), + pytest.param( + {"rate": COST}, + r"parameter 'rate' is declared with no dims.*3 rows", + id="rows-for-scalar", + ), + pytest.param( + {"cost": {"a", "b"}}, + r"parameter 'cost': cannot adapt set", + id="unsupported-shape", + ), + pytest.param( + {"cost": pd.DataFrame({"f": ["a"], "amount": [1.0]})}, + r"parameter 'cost' arrived as a DataFrame with columns \['f', 'amount'\]", + id="frame-without-value-column", + ), + pytest.param( + {"cap": CAP.to_pandas().rename_axis(index="f", columns="q")}, + r"parameter 'cap' arrived as a wide DataFrame with index 'f' and columns 'q'", + id="wide-frame-wrong-axis-names", + ), + pytest.param( + {"cap": xr.DataArray(np.ones((3, 3)), dims=["f", "t"])}, + r"parameter 'cap' has no coordinate labels along 'f'", + id="dense-without-labels", + ), + pytest.param( + {"cost": NULL_ROW}, r"parameter 'cost' carries 1 row.*f='b'", id="null-row" + ), + pytest.param({"cost": NAN_ROW}, r"parameter 'cost' carries 1 row", id="nan-row"), + pytest.param( + {"rate": float("nan")}, + r"parameter 'rate' is one value and that value is a hole", + id="nan-scalar", + ), + pytest.param( + {"rate": pd.DataFrame({"value": [None]})}, + r"parameter 'rate' is one value and that value is a hole", + id="nan-scalar-frame", + ), + pytest.param( + {"lead": pd.Series([1.5, 0.0, 1.0], index=F)}, + r"'lead' is declared 'int'.*'float'", + id="float-for-int", + ), + pytest.param( + {"flag": pd.Series([1, 0, 1], index=F)}, + r"'flag' is declared 'bool'.*'int'", + id="int-for-bool", + ), + pytest.param( + {"flag": 1.0}, r"'flag' is declared 'bool'.*'float'", id="float-scalar-for-bool" + ), + pytest.param( + {"rate": "1.5"}, r"'rate' is declared 'float'.*'str'", id="numeric-str-scalar" + ), + pytest.param( + {"rate": True}, + r"'rate' is declared 'float'.*'bool'", + id="bool-scalar-for-float", + ), + pytest.param( + {"rate": "abc"}, r"'rate' is declared 'float'.*'str'", id="str-scalar-for-float" + ), + pytest.param( + {"cost": pd.Series(["x", "y", "z"], index=F)}, + r"'cost' is declared 'float'.*'str'", + id="str-for-float", + ), + pytest.param( + {"csot": COST}, r"source key 'csot'.*Did you mean 'cost'", id="unknown-key" + ), + pytest.param({"f": None}, r"dimension 'f' has no index", id="missing-dimension"), + pytest.param( + {"f": {"a": 1}}, + r"index for dimension 'f': cannot read labels out of dict", + id="dimension-shape", + ), + pytest.param( + {"f": np.ones((2, 2))}, + r"index for dimension 'f' is 2-dimensional", + id="dimension-rank", + ), + pytest.param( + {"grp": xr.DataArray(["n"], coords={"t": [0]})}, + r"lookup 'grp' arrived as a DataArray over \['t'\]", + id="lookup-wrong-dataarray-dim", + ), + pytest.param( + {"grp": None}, r"no data provided for lookup 'grp'", id="missing-lookup" + ), + pytest.param( + {"grp": {"zz": "n"}}, + r"lookup 'grp' maps 'zz', which are not labels of 'f'", + id="lookup-stray-key", + ), + pytest.param( + {"grp": {"a": "zz"}}, + r"lookup 'grp' has value\(s\) that are not 'g' labels: 'zz'", + id="lookup-stray-value", + ), + pytest.param( + {"grp": pd.Series(["n", "e"], index=pd.Index(["a", "a"], name="f"))}, + r"lookup 'grp' maps 1 'f' label\(s\) more than once: 'a'", + id="lookup-two-values", + ), + pytest.param( + {"grp": {"a": None, "b": "n"}}, + r"lookup 'grp' carries 1 row\(s\) with a null in 'g': f='a'", + id="lookup-null", + ), + pytest.param( + {"grp": pd.DataFrame({"f": ["a"], "g": ["n"]})}, + r"lookup 'grp': cannot adapt DataFrame", + id="lookup-shape", + ), + pytest.param( + {"grp": pd.Series(["n"], index=pd.Index(["a"], name="t"))}, + r"lookup 'grp' is a Series indexed by 't'", + id="lookup-wrong-index", + ), +] + + +@pytest.mark.parametrize(("override", "match"), REFUSALS) +def test_malformed_data_is_refused_naming_the_symbol( + program: Any, good: dict[str, Any], override: dict[str, Any], match: str +) -> None: + with pytest.raises(SpecDataError, match=match): + read_all(program, sources_from(good, override)) + + +def test_int_labels_are_shown_as_written(program: Any, good: dict[str, Any]) -> None: + with pytest.raises(SpecDataError, match=r"\b99\b") as error: + read_all(program, {**good, "cap": CAP.assign_coords(t=[0, 1, 99])}) + assert "int64" not in str(error.value) + + +def test_dataset_is_a_source(program: Any, good: dict[str, Any]) -> None: + dims = {"f": F, "t": T, "g": ["n", "e"]} + values = {k: xr.DataArray(v) for k, v in good.items() if k not in dims} + from_dataset = 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: + 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"}, + "span": {"dims": ["f"], "dtype": "int"}, + "on": {"dims": ["f"], "dtype": "bool"}, + "other": {"dims": ["f"]}, + }, + "variables": {"x": {"foreach": ["f", "t"], "bounds": {"lower": 0, "upper": 1}}}, + "objective": {"sense": "maximize", "expression": "sum(x * other)"}, + "expressions": { + "recent": "sum_back(x, over=t, within=span)", + "late": { + "foreach": ["f", "t"], + "cases": { + "active": { + "when": "on", + "expression": "shift(x, over=t, offset=lag, edge=0)", + } + }, + "otherwise": "x * cost", + }, + }, + } + program = math_spec.to_program(spec) + f = pd.Index(["a"], name="f") + sources = { + "f": f, + "t": [0, 1], + "cost": pd.Series([1.0], index=f), + "lag": pd.Series([1], index=f), + "span": pd.Series([2], index=f), + "on": pd.Series([True], index=f), + "other": pd.Series([2.0], index=f), + } + retained = 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"]) +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_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"}}, + "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: dict[str, Any] = { + "dimensions": {"f": {"dtype": "str"}}, + "parameters": {"cost": {"dims": ["f"]}, "cap": {"dims": ["f"]}}, + "variables": {"x": {"foreach": ["f"], "bounds": {"lower": 0, "upper": "cap"}}}, + "constraints": {"k": {"foreach": ["f"], "expression": "x <= cap"}}, + "objective": {"sense": "maximize", "expression": "sum(x * cost)"}, +} + +GOOD = { + "f": ["a", "b"], + "cost": pd.Series({"a": 1.0, "b": 2.0}), + "cap": pd.Series({"a": 5.0, "b": 5.0}), +} +ACCEPTED = "accepted" + +PARITY_CASES = [ + pytest.param({}, ACCEPTED, id="valid"), + pytest.param({"cap": None}, SpecDataError, id="parameter-missing-entirely"), + pytest.param({"cost": pd.Series({"a": 1.0})}, ACCEPTED, id="coefficient-sparse"), + pytest.param({"cost": DUP_ROWS}, SpecDataError, id="duplicated-coordinate-row"), + pytest.param({"cost": STRAY_ROW}, SpecDataError, id="label-not-in-the-dimension"), + pytest.param({"cost": NULL_ROW}, SpecDataError, id="a-null-value"), + pytest.param({"cost": NAN_ROW}, SpecDataError, id="a-nan-value"), + pytest.param({"cap": NULL_ROW}, SpecDataError, id="a-hole-in-a-bound"), + pytest.param({"cost": float("nan")}, SpecDataError, id="a-hole-as-a-scalar"), + pytest.param({"cost": [1.0, np.nan]}, SpecDataError, id="a-hole-in-a-sequence"), + pytest.param({"cost": {"a": 1.0, "b": None}}, SpecDataError, id="a-hole-in-a-dict"), + pytest.param({"cost": NULL_FRAME}, SpecDataError, id="a-hole-in-a-tidy-frame"), + pytest.param({"cost": pd.Series({"a": 1, "b": 2})}, ACCEPTED, id="whole-numbers"), + pytest.param({"csot": COST}, SpecDataError, id="an-undeclared-source-key"), + pytest.param({"cost": DEEP_ROWS}, SpecDataError, id="a-series-too-deep"), +] + + +@pytest.mark.parametrize(("override", "verdict"), PARITY_CASES) +def test_parity_with_lpspec_data_verdicts( + override: dict[str, Any], verdict: Any +) -> None: + program = math_spec.to_program(PARITY_SPEC) + sources = sources_from(GOOD, override) + if verdict is ACCEPTED: + read_all(program, sources) + return + with pytest.raises(verdict): + read_all(program, sources) + + +def test_a_hole_is_named_where_it_sits() -> None: + program = math_spec.to_program(PARITY_SPEC) + with pytest.raises(SpecDataError, match="parameter 'cost'") as error: + read_all(program, {**GOOD, "cost": NULL_ROW}) + assert "divisor" not in str(error.value) + assert "f='b'" in str(error.value) + + +FLAG_SPEC = { + "dimensions": {"g": {"dtype": "str"}}, + "parameters": {"active": {"dims": ["g"], "dtype": "bool"}}, + "variables": { + "x": {"foreach": ["g"], "where": "active", "bounds": {"lower": 0, "upper": 1}} + }, + "objective": {"sense": "maximize", "expression": "sum(x)"}, +} + + +@pytest.mark.parametrize( + ("column", "verdict"), + [ + pytest.param(pd.Series({"a": True, "b": False}), ACCEPTED, id="a-bool-column"), + pytest.param(pd.Series({"a": 1, "b": 0}), SpecDataError, id="a-1-0-int-column"), + pytest.param( + pd.Series({"a": 1.0, "b": 0.0}), SpecDataError, id="a-1-0-float-column" + ), + ], +) +def test_a_flag_binds_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 + return + with pytest.raises(SpecDataError, match="declared 'bool'"): + read_all(program, sources) + + +LOOKUP_SPEC = { + "dimensions": {"g": {}, "b": {"dtype": "str"}}, + "lookups": {"gen_bus": {"over": "g", "into": "b"}}, + "parameters": {"p_max": {"dims": ["g"]}}, + "variables": {"x": {"foreach": ["g"], "bounds": {"lower": 0, "upper": "p_max"}}}, + "constraints": {"k": {"foreach": ["b"], "expression": "sum(x, by=gen_bus) <= 10"}}, + "objective": {"sense": "maximize", "expression": "sum(x)"}, +} +G_TWICE = pd.Index(["w", "w", "s"], name="g") +LOOKUP_GOOD = { + "p_max": pd.Series({"w": 5.0, "s": 5.0}), + "g": ["w", "s"], + "b": ["n", "e"], + "gen_bus": pd.Series({"w": "n", "s": "e"}), +} + + +@pytest.mark.parametrize( + ("override", "match"), + [ + pytest.param( + {"g": None, "b": None}, "dimension 'g' has no index", id="a-map-no-labels" + ), + pytest.param( + {"gen_bus": None}, "no data provided for lookup", id="an-index-no-map" + ), + pytest.param( + {"gen_bus": pd.Series({"w": "n", "s": "zz"})}, + "not 'b' labels", + id="a-stray-value", + ), + pytest.param( + {"gen_bus": pd.Series(["n", "e", "e"], index=G_TWICE)}, + "more than once", + id="two-values-for-one-label", + ), + pytest.param( + {"gen_bus": pd.Series({"w": None, "s": "e"})}, + "null in 'b'", + id="mapping-a-label-to-nothing", + ), + pytest.param( + {"gen_bus": pd.Series([None, "n", "n"], index=G_TWICE)}, + "null in 'b'", + id="a-label-held-twice-with-a-null", + ), + ], +) +def test_a_lookup_defect_is_refused(override: dict[str, Any], match: str) -> None: + program = math_spec.to_program(LOOKUP_SPEC) + with pytest.raises(SpecDataError, match=match): + read_all(program, sources_from(LOOKUP_GOOD, override)) + + +TAG_SPEC = { + **LOOKUP_SPEC, + "lookups": {"tag": {"over": "g", "dtype": "int"}}, + "constraints": {"k": {"foreach": ["g"], "expression": "x <= 10"}}, +} +TAG_GOOD = sources_from(LOOKUP_GOOD, {"gen_bus": None, "b": None}) + + +def test_a_label_space_lookup_is_padded_onto_the_dimension() -> None: + program = math_spec.to_program(TAG_SPEC) + 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"}}} + ) + 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)