diff --git a/benchmarks/models/__init__.py b/benchmarks/models/__init__.py index 66c9a7c76..2b9f7ecac 100644 --- a/benchmarks/models/__init__.py +++ b/benchmarks/models/__init__.py @@ -21,5 +21,6 @@ qp, sos, sparse_network, + spec_pypsa, storage, ) diff --git a/benchmarks/models/spec_pypsa.py b/benchmarks/models/spec_pypsa.py new file mode 100644 index 000000000..fce1719a2 --- /dev/null +++ b/benchmarks/models/spec_pypsa.py @@ -0,0 +1,55 @@ +""" +Model built from math-spec's ``pypsa.yaml`` example (requires math-spec). + +The subject is :meth:`linopy.Model.from_spec`: lowering a spec of PyPSA's full +statement, binding synthetic data to it and building every variable and +constraint it declares. The example lives outside the wheel, so its directory +comes from ``MATH_SPEC_EXAMPLES`` and the case skips without it. A sweep +value is the number of labels per dimension; 40 of them is about 20k +variables. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import TYPE_CHECKING + +from benchmarks.registry import BUILD, FROM_NETCDF, TO_NETCDF, BenchSpec, register + +if TYPE_CHECKING: + import linopy + +SIZES = (5, 40) + +EXAMPLES = os.environ.get("MATH_SPEC_EXAMPLES") +EXAMPLE = Path(EXAMPLES, "pypsa.yaml") if EXAMPLES else None + + +def build_spec_pypsa(n: int) -> linopy.Model: + """Lower ``pypsa.yaml`` and build it with ``n`` labels per dimension.""" + import pytest + + if EXAMPLE is None or not EXAMPLE.exists(): + pytest.skip("set MATH_SPEC_EXAMPLES to a math-spec examples directory") + import math_spec + + import linopy + from linopy.spec.testing import synthetic_sources + + path = str(EXAMPLE) + sources = synthetic_sources(math_spec.to_program(path), n) + with linopy.options as options: + options["semantics"] = "v1" + return linopy.Model.from_spec(path, sources) + + +SPEC = register( + BenchSpec( + name="spec_pypsa", + build=build_spec_pypsa, + sweep=SIZES, + phases=frozenset({BUILD, TO_NETCDF, FROM_NETCDF}), + requires=("math_spec",), + ) +) diff --git a/linopy/io.py b/linopy/io.py index 6ee8b36fd..6c59d0e43 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,10 +1104,19 @@ 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._relaxed_registry: @@ -1260,6 +1275,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 +1412,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 e74e7189d..525e18620 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 000000000..e2214b0e0 --- /dev/null +++ b/linopy/spec/netcdf.py @@ -0,0 +1,216 @@ +""" +Persist the spec of a spec-built model in its netcdf file. + +Variables, constraints and the solution round trip through :mod:`linopy.io` +already. Besides them a spec-built model carries the spec text, the master +coordinates and the lookups; the program is re-lowered from the text on read, +so no lowered ``Program`` ever reaches the file. + +No netcdf type holds a dtype as written. An engine narrows an int64 to +int32, hands a bool back as int8 and a string array back as `` tuple[xr.Dataset, 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. + """ + 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) + 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: + """ + 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) + } + 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 objects.""" + lookups = {name for by_name in spec.lookups.values() for name in by_name} + return [ + str(name) + for name, arr in spec.parameters.items() + if name in lookups or arr.dtype == object + ] + + +def _encode(name: str, arr: xr.DataArray) -> dict[str, xr.DataArray]: + codes, categories = pd.factorize(arr.to_numpy().ravel()) + written = { + CODES + name: _array( + codes.astype(np.int32).reshape(arr.shape), arr.dims, str(arr.dtype) + ) + } + if len(categories): + written[CATEGORIES + name] = _array( + np.asarray(categories), (CATEGORY_DIM + name,) + ) + return written + + +def _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 _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)}) + + +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/spec/testing.py b/linopy/spec/testing.py new file mode 100644 index 000000000..5bec8301f --- /dev/null +++ b/linopy/spec/testing.py @@ -0,0 +1,72 @@ +""" +Synthetic data for a spec, for tests and benchmarks. + +A spec says what data it takes, which is enough to make some up: the shape is +the declaration's, only the values are invented. What comes out builds and +solves, and says nothing about a real system. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pandas as pd +import xarray as xr +from math_spec import program as ms + +_START = "2030-01-01" + + +def synthetic_sources(program: ms.Program, n: int = 3) -> dict[str, Any]: + """ + Dense data for every declaration of *program*, *n* labels per dimension. + + Labels are numbered after their dimension, parameters are a linear ramp, + and each lookup cycles through the labels it maps into. + """ + sources: dict[str, Any] = { + dim: _labels(dim, decl.dtype, n) for dim, decl in program.dimensions.items() + } + for over, lookup in program.lookups: + into = ( + sources[lookup.target] + if lookup.target is not None + else _labels(lookup.name, lookup.dtype, n) + ) + sources[lookup.name] = pd.Series( + [into[i % len(into)] for i in range(n)], index=sources[over] + ) + for name, parameter in program.parameters.items(): + if parameter.derivation is None: + sources[name] = _parameter(name, parameter, sources, n) + return sources + + +def _labels(name: str, dtype: str | None, n: int) -> pd.Index: + """*n* labels of the declared dtype, named after what they label.""" + if dtype == "int": + return pd.Index(range(n), name=name) + if dtype == "datetime": + return pd.date_range(_START, periods=n, freq="h", name=name) + return pd.Index([f"{name}{i}" for i in range(n)], name=name) + + +def _parameter( + name: str, declared: ms.ParameterDeclaration, sources: dict[str, Any], n: int +) -> Any: + dims = declared.dims + shape = [n] * len(dims) + if declared.dtype == "bool": + values: Any = np.ones(shape, dtype=bool) + elif declared.dtype == "int": + values = np.ones(shape, dtype=int) + elif declared.dtype == "str": + values = np.full(shape, "a", dtype=object) + elif dims: + values = np.broadcast_to(1.0 + np.arange(n), shape).copy() + else: + values = np.array(1.0) + if not dims: + return values.item() + return xr.DataArray(values, coords={d: sources[d] for d in dims}, dims=list(dims)) diff --git a/linopy/testing.py b/linopy/testing.py index 5fd167782..59a9bb61e 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) @@ -134,6 +151,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_builder.py b/test/test_spec_builder.py index ad3b5bea5..9c6fd15e8 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 new file mode 100644 index 000000000..4f7911860 --- /dev/null +++ b/test/test_spec_io.py @@ -0,0 +1,246 @@ +""" +Round trips of a spec-built model through netcdf and through ``copy``. + +The spec itself is persisted as its YAML text and lowered again on read, so +what has to survive besides the model is data: the master coordinates, the +lookups and the retained parameters. Labels are the delicate part — a partial +lookup holds NaN in an array of strings — so every lookup shape is checked +value by value and dtype by dtype, on both netcdf engines ``test_io`` uses. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pandas as pd +import pytest +import xarray as xr + +math_spec = pytest.importorskip("math_spec") + +from test_spec_builder import ( # noqa: E402 + DISPATCH_DATA, + EXAMPLE_DISPATCH, + EXAMPLES_DIR, + WHERE_DATA, + WHERE_SPEC, + solved, +) + +import linopy # noqa: E402 +from linopy import Model, read_netcdf # noqa: E402 +from linopy.io import SPEC_ATTR # noqa: E402 +from linopy.spec.testing import synthetic_sources # noqa: E402 +from linopy.testing import assert_model_equal # noqa: E402 + +pytestmark = [ + pytest.mark.v1, + pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs"), +] + +ENGINES = ["netcdf4", "scipy"] + +S1 = pd.Index(["a", "b", "c"], name="s1") +S2 = pd.Index(["p", "q"], name="s2") +I1 = pd.Index([10, 20, 30], name="i1") +I2 = pd.Index([1, 2], name="i2") + +LOOKUP_SPEC: dict[str, Any] = { + "dimensions": { + "s1": {"dtype": "str"}, + "s2": {"dtype": "str"}, + "i1": {"dtype": "int"}, + "i2": {"dtype": "int"}, + }, + "lookups": { + "str_to_str": {"over": "s1", "into": "s2"}, + "str_to_int": {"over": "s1", "into": "i2"}, + "int_to_str": {"over": "i1", "into": "s2"}, + "int_to_int": {"over": "i1", "into": "i2"}, + }, + "parameters": {"cost": {"dims": ["s1"]}}, + "variables": {"x": {"foreach": ["s1"], "bounds": {"lower": 0, "upper": 1}}}, + "objective": {"sense": "minimize", "expression": "sum(x * cost)"}, +} +LOOKUP_OVER = {"str_to_str": S1, "str_to_int": S1, "int_to_str": I1, "int_to_int": I1} +LOOKUP_INTO = {"str_to_str": S2, "str_to_int": I2, "int_to_str": S2, "int_to_int": I2} + +DTYPE_SPEC: dict[str, Any] = { + "dimensions": {"s1": {"dtype": "str"}}, + "parameters": { + "count": {"dims": ["s1"], "dtype": "int"}, + "flag": {"dims": ["s1"], "dtype": "bool"}, + "cost": {"dims": ["s1"]}, + "tag": {"dims": ["s1"], "dtype": "str"}, + }, + "variables": {"x": {"foreach": ["s1"], "bounds": {"lower": 0, "upper": 1}}}, + "objective": {"sense": "minimize", "expression": "sum(x * cost)"}, +} +DTYPE_DATA: dict[str, Any] = { + "s1": S1, + "count": pd.Series([1, 2, 3], index=S1), + "flag": pd.Series([True, False, True], index=S1), + "cost": pd.Series([1.0, 2.0, 3.0], index=S1), + "tag": pd.Series(["u", "v", "w"], index=S1), +} + + +def lookup_sources(mapped: int) -> dict[str, Any]: + """Data for ``LOOKUP_SPEC``, each lookup mapping only its first *mapped* labels.""" + sources: dict[str, Any] = { + "s1": S1, + "s2": S2, + "i1": I1, + "i2": I2, + "cost": pd.Series([1.0, 2.0, 3.0], index=S1), + } + for name, over in LOOKUP_OVER.items(): + into = LOOKUP_INTO[name] + sources[name] = pd.Series( + [into[i % len(into)] for i in range(mapped)], index=over[:mapped] + ) + return sources + + +def roundtrip(m: Model, tmp_path: Path, engine: str) -> Model: + path = tmp_path / f"model-{engine}.nc" + m.to_netcdf(path, engine=engine) + return read_netcdf(path) + + +def assert_arrayequal(a: xr.DataArray, b: xr.DataArray) -> None: + """Assert equal values and dtype — the dtype is what a netcdf type drops.""" + assert a.dtype == b.dtype, f"dtypes differ: {a.dtype} != {b.dtype}" + xr.testing.assert_equal(a, b) + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize("retain", ["report", "all"]) +def test_a_spec_built_model_round_trips( + tmp_path: Path, engine: str, retain: str +) -> None: + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain=retain) + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert p.spec.text == m.spec.text + assert p.spec.program.constraints == m.spec.program.constraints + assert set(p.spec.expressions) == set(m.spec.expressions) + for name in m.spec.expressions: + assert_arrayequal(m.spec.expressions[name], p.spec.expressions[name]) + + +@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 name in p.spec.lookups[over] + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize("name", ["count", "flag", "cost", "tag"]) +def test_a_parameter_keeps_its_dtype(tmp_path: Path, engine: str, name: str) -> None: + m = Model.from_spec(DTYPE_SPEC, DTYPE_DATA, retain="all") + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert p.spec.parameters[name].dtype == m.spec.parameters[name].dtype + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize("frozen", [False, True], ids=["dataset", "csr"]) +def test_every_container_shares_the_master_coordinate_dtypes( + tmp_path: Path, engine: str, frozen: bool +) -> None: + """The master coordinates are canonical: no container may disagree with them.""" + if frozen and engine == "scipy": + pytest.skip( + "netCDF3 holds no unicode-array attr, and a CSR constraint writes one" + ) + m = solved(EXAMPLE_DISPATCH, DISPATCH_DATA, retain="all", freeze_constraints=frozen) + p = roundtrip(m, tmp_path, engine) + + master = {dim: index.dtype for dim, index in p.spec.coords.items()} + holders = [ + *(v.data for _, v in p.variables.items()), + *(c.data for _, c in p.constraints.items()), + p.objective.expression.data, + ] + assert master == {dim: index.dtype for dim, index in m.spec.coords.items()} + for data in holders: + for dim, index in data.indexes.items(): + if str(dim) in master: + assert index.dtype == master[str(dim)], f"{dim} differs on {data}" + + +@pytest.mark.parametrize("engine", ENGINES) +def test_labelled_parameters_and_unreached_coordinates_round_trip( + tmp_path: Path, engine: str +) -> None: + """A str parameter with holes, and a dimension only a lookup reaches.""" + m = Model.from_spec(WHERE_SPEC, WHERE_DATA, retain="all") + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert set(p.spec.coords) == set(m.spec.coords) + + +@pytest.mark.parametrize("deep", [True, False]) +def test_a_copy_carries_the_spec(deep: bool) -> None: + """The copy's spec reads the copy, and only a deep copy owns its buffers.""" + m = Model.from_spec(WHERE_SPEC, WHERE_DATA, retain="all") + p = m.copy(deep=deep) + p.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_model_without_a_spec_carries_none(tmp_path: Path) -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") + m.add_objective(x.sum()) + path = tmp_path / "plain.nc" + m.to_netcdf(path) + + assert SPEC_ATTR not in xr.load_dataset(path).attrs + assert read_netcdf(path)._spec is None + assert m.copy()._spec is None + + +@pytest.mark.skipif( + EXAMPLES_DIR is None, reason="set MATH_SPEC_EXAMPLES to a math-spec examples dir" +) +@pytest.mark.parametrize("engine", ENGINES) +def test_the_pypsa_example_round_trips(tmp_path: Path, engine: str) -> None: + """Nine lookups into one dimension, a datetime axis, bool and str parameters.""" + path = Path(EXAMPLES_DIR or "", "pypsa.yaml") + program = math_spec.to_program(str(path)) + m = Model.from_spec(path, synthetic_sources(program, 3), retain="all") + p = roundtrip(m, tmp_path, engine) + + assert_model_equal(m, p) + assert set(p.spec.coords) == set(m.spec.coords)