diff --git a/CHANGELOG.md b/CHANGELOG.md index ce01fadb5..c9dc9b8d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ it releases that version ([RELEASING.md](https://github.com/fluxopt/specsolve/bl ## Upcoming version +- feat(api): an axis is written as plain data with to_dict, and solve_over takes that data in its place ([#1981](https://github.com/fluxopt/specsolve/pull/1981)) - refactor(api): the carry is an argument of the axis rather than of solve_over, so an archive's axis holds its carry ([#1980](https://github.com/fluxopt/specsolve/pull/1980)) ## 0.8.1 (2026-10-11) diff --git a/docs/reference/sweeps.md b/docs/reference/sweeps.md index bc622f02d..6365ad6af 100644 --- a/docs/reference/sweeps.md +++ b/docs/reference/sweeps.md @@ -78,6 +78,22 @@ Solved as a list, the slices key by `key_name=`, and the answer is keyed by slice. Two axes compose as a comprehension over the slices of one, each sliced again by the other. +**An axis is also plain data.** `axis.to_dict()` writes it as a mapping that +`json.dumps` takes, and `solve_over` takes that mapping in the axis's place. An +archive's `axis.json` holds the same mapping: + +```python +axis = sps.EachWindow('snapshot', steps=24, lookahead=24, into='t', carry={'soc_initial': 'soc'}) +axis.to_dict() +# {'each': 'window', 'dim': 'snapshot', 'steps': 24, 'lookahead': 24, 'into': 't', +# 'carry': {'soc_initial': 'soc'}} +sps.solve_over('window.yaml', sources, json.loads(text)) +``` + +A mapping whose `each` names no axis, or that holds a key the axis does not +take or lacks one it needs, is refused before a slice is taken. The message +names the keys the axis takes. + **Sources cross a slice in every shape `build` takes.** A parameter or a relation whose table carries the axis is filtered, as a table or as a parquet path. Every other source passes through as it is. **An index of another diff --git a/src/specsolve/archive_layout.py b/src/specsolve/archive_layout.py index 9b8a6773e..4b3d1d5bd 100644 --- a/src/specsolve/archive_layout.py +++ b/src/specsolve/archive_layout.py @@ -87,7 +87,8 @@ def write_archive( spec: The spec as written, held as ``spec.yaml``. tables: The tidy table each source stands for, keyed as the file declares, each written as ``sources/.parquet``. - axis: The axis manifest, or ``None`` where the sources are not cut. + axis: The axis as its ``to_dict`` writes it, or ``None`` where the + sources are not cut. answer: Lays the answer out in the directory it is given, writing each table with the writer it is given, which adds the run. diff --git a/src/specsolve/axes.py b/src/specsolve/axes.py index 1f7fe3115..3aa71dc78 100644 --- a/src/specsolve/axes.py +++ b/src/specsolve/axes.py @@ -9,7 +9,7 @@ from __future__ import annotations import warnings -from dataclasses import dataclass, field +from dataclasses import MISSING, dataclass, field, fields from functools import partial from typing import TYPE_CHECKING, Any, NamedTuple @@ -31,6 +31,10 @@ from specsolve.inputs import Label, Source +#: The key [`EachCoordinate.to_dict`][] and [`EachWindow.to_dict`][] write an axis's carry under. +CARRY = 'carry' + + class Slice(NamedTuple): """One slice of a sweep: the key, the sources that build it, and what it owns. @@ -117,6 +121,17 @@ class EachCoordinate: def __post_init__(self) -> None: object.__setattr__(self, 'carry', dict(self.carry)) + def to_dict(self) -> dict[str, object]: + """This axis as plain data, which [`solve_over`][specsolve.strategy.solve_over] takes in its place. + + ``json.dumps`` writes it, and an archive's ``axis.json`` is this. It + holds a ``carry`` key only where the axis carries something. + """ + out: dict[str, object] = {'each': 'coordinate', 'dim': self.dim} + if self.carry: + out[CARRY] = dict(self.carry) + return out + def slices(self, sources: Mapping[str, Source]) -> list[tuple[Label, Mapping[str, Source]]]: """The ``(key, sources)`` list this axis would run — what ``axis=`` takes hand-built. @@ -211,6 +226,24 @@ def __post_init__(self) -> None: if self.into == self.dim: raise ValueError(f'into={self.into!r} must differ from dim — the local index replaces the global one') + def to_dict(self) -> dict[str, object]: + """This axis as plain data, which [`solve_over`][specsolve.strategy.solve_over] takes in its place. + + ``json.dumps`` writes it, and an archive's ``axis.json`` is this. It + holds a ``carry`` key only where the axis carries something. + """ + steps = self.steps if isinstance(self.steps, int) else list(self.steps) + out: dict[str, object] = { + 'each': 'window', + 'dim': self.dim, + 'steps': steps, + 'lookahead': self.lookahead, + 'into': self.into, + } + if self.carry: + out[CARRY] = dict(self.carry) + return out + def _key_name(self) -> str: """Where the window started, never ``dim``, which would join silently against data over *dim*.""" return f'{self.dim}_start' @@ -305,6 +338,9 @@ def _blocks(self, total: int) -> list[int]: #: An axis that cuts one set of sources into slices. Axis = EachCoordinate | EachWindow +#: Each axis by the name ``to_dict`` writes under ``each``. +_EACH: dict[str, type[EachCoordinate] | type[EachWindow]] = {'coordinate': EachCoordinate, 'window': EachWindow} + #: Slices written by hand, one ``(key, sources)`` pair each; ``axis=`` takes #: this beside an [`Axis`][]. @@ -331,30 +367,30 @@ def check_no_index_is_cut(program: Program, sources: Mapping[str, Source], axis: ) -def axis_manifest(axis: Axis) -> dict[str, Any]: # pyrefly: ignore[explicit-any] — the archive's own JSON - """*axis* as the JSON an archive carries, read back by [`axis_from`][].""" - carry = {'carry': dict(axis.carry)} if axis.carry else {} - if isinstance(axis, EachCoordinate): - return {'each': 'coordinate', 'dim': axis.dim, **carry} - steps = axis.steps if isinstance(axis.steps, int) else list(axis.steps) - return { - 'each': 'window', - 'dim': axis.dim, - 'steps': steps, - 'lookahead': axis.lookahead, - 'into': axis.into, - **carry, - } - - -def axis_from(manifest: Mapping[str, Any]) -> Axis: # pyrefly: ignore[explicit-any] — the archive's own JSON - """The axis [`axis_manifest`][] wrote.""" - carry = manifest.get('carry', {}) - if manifest['each'] == 'coordinate': - return EachCoordinate(manifest['dim'], carry=carry) - return EachWindow( - manifest['dim'], steps=manifest['steps'], lookahead=manifest['lookahead'], into=manifest['into'], carry=carry - ) +def axis_from(given: Mapping[str, Any]) -> Axis: # pyrefly: ignore[explicit-any] — plain data, as JSON reads it + """The axis *given* as plain data, as [`EachCoordinate.to_dict`][] and [`EachWindow.to_dict`][] write it. + + Raises: + SpecsolveError: An ``each`` that names no axis, or a key the axis + does not take or needs and is not given; the message names the + keys it takes. + """ + each = given.get('each') + if each not in _EACH: + raise SpecsolveError( + f"an axis given as a mapping names its kind under 'each', one of {sorted(_EACH)} (got {each!r})" + ) + kind = _EACH[each] + taken = {f.name: f for f in fields(kind)} + keys = {key: value for key, value in given.items() if key != 'each'} + needed = [name for name, f in taken.items() if f.default is MISSING and f.default_factory is MISSING] + if (unknown := sorted(set(keys) - set(taken))) or (missing := [name for name in needed if name not in keys]): + wrong = f'{unknown} is not one of them' if unknown else f'{missing} is missing' + raise SpecsolveError( + f"{kind.__name__} given as a mapping takes 'each' and {list(taken)}, of which {needed} are " + f'required: {wrong}.' + ) + return kind(**keys) def _one_coordinate(dim: str, key: Label, table: pl.LazyFrame) -> pl.LazyFrame: diff --git a/src/specsolve/strategy.py b/src/specsolve/strategy.py index 0266ce10c..a28ad4509 100644 --- a/src/specsolve/strategy.py +++ b/src/specsolve/strategy.py @@ -17,7 +17,7 @@ import json import shutil import warnings -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from concurrent.futures import Executor, ProcessPoolExecutor, ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, replace @@ -29,7 +29,7 @@ from specsolve.api import build, check from specsolve.archive_layout import check_the_target, write_archive -from specsolve.axes import Axis, EachWindow, HandBuilt, Slice, axis_manifest, check_no_index_is_cut, sources_with_column +from specsolve.axes import Axis, EachWindow, HandBuilt, Slice, axis_from, check_no_index_is_cut, sources_with_column from specsolve.errors import DataError, SpecsolveError, SpecsolveWarning from specsolve.frames import as_frame from specsolve.inputs import declared @@ -62,7 +62,7 @@ ) if TYPE_CHECKING: - from collections.abc import Callable, Generator, Iterable, Mapping, Sequence + from collections.abc import Callable, Generator, Iterable, Sequence from mathspec import Spec from mathspec.program import Program @@ -176,7 +176,7 @@ def _archiving(archive: str | Path | None, axis: Axis | HandBuilt, *, keep_windo def solve_over( spec: Buildable, sources: Mapping[str, Source], - axis: Axis | HandBuilt, + axis: Axis | Mapping[str, Any] | HandBuilt, # pyrefly: ignore[explicit-any] — plain data, as JSON reads it *, key_name: str | None = None, executor: Executor | None = None, @@ -200,9 +200,10 @@ def solve_over( sources: As [`build`][specsolve.api.build] takes them, every shape included; the axis filters the parameters and relations that carry it and passes the rest through. - axis: [`EachCoordinate`][], [`EachWindow`][], or a list of - ``(key, sources)`` written by hand. An axis's ``carry=`` chains - its slices; a hand-built list carries nothing. + axis: [`EachCoordinate`][], [`EachWindow`][], either as the plain + data its ``to_dict`` writes, or a list of ``(key, sources)`` + written by hand. An axis's ``carry=`` chains its slices; a + hand-built list carries nothing. key_name: What to call the slice column; a class axis names its own, a hand-built list has to be told. executor: Any `concurrent.futures.Executor`; ``None`` runs the @@ -261,7 +262,8 @@ def solve_over( The sweep, which reads its answer. Raises: - SpecsolveError: Before a slice is taken: a carry that cannot line up, + SpecsolveError: Before a slice is taken: an axis given as a mapping + whose keys name no axis; a carry that cannot line up, has no seed, collapses a dimension the axis does not advance along, or is asked together with an executor; a key that collides with a column the frames carry; an axis the program does not @@ -281,6 +283,8 @@ def solve_over( the model counts, which every window restarts, or a *start* that leaves a slice no row, which starts that slice from nothing. """ + if isinstance(axis, Mapping): + axis = axis_from(axis) carry: Mapping[str, str] = axis.carry if isinstance(axis, Axis) else {} if carry and executor is not None: raise SpecsolveError( @@ -458,13 +462,14 @@ def _archive_the_sweep( Each source's tidy table comes from *one_slice*, except the ones [`_uncut`][] and [`_spread_over_the_axis`][] give. """ - manifest = axis_manifest(axis) tidied = numbered(program, tidy_sources(program, one_slice)) sliced = sources_with_column(sources, axis.dim) cut = {name: _uncut(program, axis, name, table) for name, table in sliced.items()} held = {**tidied, **_spread_over_the_axis(program, axis, sources, tidied, sliced), **cut} tables = {name: held[name] for name in sources} - write_archive(out, spec, tables, axis=manifest, answer=partial(_the_answer, folded, keep_windows=keep_windows)) + write_archive( + out, spec, tables, axis=axis.to_dict(), answer=partial(_the_answer, folded, keep_windows=keep_windows) + ) def _the_answer(sweep: Sweep, under: Path, write: FrameWriter, *, keep_windows: bool) -> None: diff --git a/tests/test_strategy.py b/tests/test_strategy.py index 8b0a7f5c6..c8f192ed9 100644 --- a/tests/test_strategy.py +++ b/tests/test_strategy.py @@ -26,6 +26,7 @@ import specsolve as sps from specsolve import strategy from specsolve import sweep as sweep_module +from specsolve.axes import axis_from from specsolve.relational.answer_layout import Metrics, Provenance, Record from tests.conftest import DISPATCH_SPEC, override @@ -1137,6 +1138,68 @@ def test_an_axis_that_carries_is_still_hashable(axis): assert axis != replace(axis, carry={}), 'and kept in the equality' +@pytest.mark.parametrize( + 'axis', + [ + pytest.param(sps.EachCoordinate('period'), id='coordinate'), + pytest.param(sps.EachCoordinate('period', carry={'existing': 'total'}), id='coordinate-carrying'), + pytest.param(WINDOW_AXIS, id='window'), + pytest.param(sps.EachWindow('snapshot', steps=(2, 4), lookahead=1, into='t'), id='window-of-listed-steps'), + pytest.param(replace(WINDOW_AXIS, carry={'soc_initial': 'soc'}), id='window-carrying'), + ], +) +def test_an_axis_reads_back_from_the_json_of_its_dict(axis): + """``to_dict`` is plain data: through JSON and back, it is the axis it was.""" + assert axis_from(json.loads(json.dumps(axis.to_dict()))) == axis + + +def test_an_axis_that_carries_nothing_writes_no_carry_key(): + """The ``axis.json`` of an archive whose sweep carried nothing holds no ``carry``, as ``INPUTS_LAYOUT`` 1 has it.""" + assert sps.EachCoordinate('period').to_dict() == {'each': 'coordinate', 'dim': 'period'}, ( + 'no carry key where nothing is carried' + ) + + +def test_solve_over_takes_an_axis_as_its_dict(): + """The dict stands where the axis does, carry included.""" + axis = replace(WINDOW_AXIS, carry={'soc_initial': 'soc'}) + given = sps.solve_over(WINDOW, horizon_sources(), json.loads(json.dumps(axis.to_dict()))) + expected = sps.solve_over(WINDOW, horizon_sources(), axis) + assert given.record['objective'].to_list() == expected.record['objective'].to_list() + assert given.primal('soc').equals(expected.primal('soc')), 'each window started from the one before it' + + +def test_an_archive_holds_the_dict_of_its_axis(tmp_path): + """``axis.json`` is ``to_dict``, so the archive's file and the public form are one.""" + axis = replace(WINDOW_AXIS, carry={'soc_initial': 'soc'}) + sps.solve_over(WINDOW, horizon_sources(), axis, archive=tmp_path / 'roll') + assert json.loads((tmp_path / 'roll' / 'axis.json').read_text()) == axis.to_dict() + + +@pytest.mark.parametrize( + ('given', 'expected'), + [ + pytest.param( + {'dim': 'snapshot'}, "under 'each', one of \\['coordinate', 'window'\\] \\(got None\\)", id='no-each' + ), + pytest.param({'each': 'slice', 'dim': 'snapshot'}, "\\(got 'slice'\\)", id='an-each-that-names-no-axis'), + pytest.param( + {'each': 'coordinate', 'dim': 'period', 'carries': {}}, + "takes 'each' and \\['dim', 'carry'\\].*\\['carries'\\] is not one of them", + id='a-key-the-axis-does-not-take', + ), + pytest.param( + {'each': 'window', 'dim': 'snapshot', 'steps': 4, 'into': 't'}, + "\\['dim', 'steps', 'lookahead', 'into'\\] are required: \\['lookahead'\\] is missing", + id='a-key-the-axis-needs', + ), + ], +) +def test_an_axis_dict_that_names_no_axis_is_refused_with_the_keys_it_takes(given, expected): + with pytest.raises(sps.errors.SpecsolveError, match=expected): + sps.solve_over(WINDOW, horizon_sources(), given) + + def test_carry_and_executor_are_refused_together(): """Sequential by definition, so the combination is a call-time error rather than something discovered at slice two."""