From 91420035df88f4211ea2d8e982d4dc876b669c13 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 14:56:22 +0200 Subject: [PATCH 1/2] 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 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..5d3fac635 --- /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 6ee8b36fd..f64b90b51 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 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..de2559292 --- /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 5fd167782..c51d3d343 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 000000000..16007306f --- /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 7fbfd881c21472e0f1be1e5e8da3c7b8d4692928 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 3 Sep 2026 15:35:19 +0200 Subject: [PATCH 2/2] 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 5d3fac635..fce1719a2 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 f64b90b51..6c59d0e43 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 de2559292..e2214b0e0 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 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 c51d3d343..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) 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 index 16007306f..4f7911860 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)