From b744d222c3792c59a39e7664a62ff98a6b1de65e Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 8 Nov 2026 08:42:42 +0000 Subject: [PATCH 1/4] fix(core): evolve re-derives defaults computed from the cosmology `build` takes sigma_8 and n_s from the cosmology's meta when they are not given, but `evolve(cosmology=...)` kept the old values, so a cosmology scan done with `evolve` silently gave hybrid models (e.g. Planck15 with Planck18's sigma_8), and evolving to a cosmology without meta silently kept sigma_8 where `build` raises. Each stage now records which parameters below it took a default from its `computed_defaults`, and which paths of the given parameters the hook looked up (`ComputedRecord`, in an init=False, eq=False, physics=False field, so it is not part of the value, the hash or the content hash). `evolve` makes every parameter it sets explicit, and when it changes a path the hook read, calls the hook again on the new tree for the parameters still computed. So `build(**a).evolve(**b) == build(**a, **b)`. A pickle or copy keeps the record; a configuration does not, so every value a configuration gives is explicit when read. `Stage.computed_parameters()` lists the computed ones. The missing-meta error now names the cosmology's class when it has no name, and asks to set the parameter explicitly. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_019U8u44xFH87Am2mNHVrimP --- docs/core.rst | 34 ++- src/hmf/core/config.py | 54 +++- src/hmf/core/mass_function.py | 15 +- src/hmf/core/routing.py | 296 +++++++++++++++++++++- src/hmf/core/stage.py | 79 +++++- tests/core/test_core_computed_defaults.py | 212 ++++++++++++++++ tests/core/test_core_routing.py | 30 ++- 7 files changed, 689 insertions(+), 31 deletions(-) create mode 100644 tests/core/test_core_computed_defaults.py diff --git a/docs/core.rst b/docs/core.rst index 5f28b673..eddc95aa 100644 --- a/docs/core.rst +++ b/docs/core.rst @@ -443,11 +443,35 @@ So the flat names of a mass function are exactly :meth:`build mf = MassFunction.build(**{"transfer.model": "EH"}, fit="ST", sigma_8=0.8) mf.evolve(sigma_8=0.85, n_s=0.95) - mf.evolve(cosmology=Planck15) # the cosmology of the whole tree + mf.evolve(cosmology=Planck15) # the cosmology of the whole tree, and + # n_s from Planck15 (sigma_8 is set above) mf.evolve(**{"fit.a": 0.75}) # a parameter of the fit mf.evolve(**{"linear_power.transfer.n_s": 0.95}) MassFunction.from_flat({"transfer.model": "EH", "linear_power": {"sigma_8": 0.8}}) +**Computed defaults.** ``build`` takes σ8 and n_s, when they are not given, from the +cosmology's ``meta`` (a stage's :meth:`~hmf.core.stage.Stage.computed_defaults`). A +parameter that took such a default stays *computed* until it is set explicitly, and +``evolve`` computes it again whenever a parameter it is computed from changes: so +``build(**a).evolve(**b)`` equals ``build(**a, **b)``, and a cosmology scan done with +``evolve`` gives the same models as one done with ``build``:: + + mf = MassFunction.build() # sigma_8 = 0.8102, n_s = 0.9665 (Planck18) + mf.evolve(cosmology=Planck15) # sigma_8 = 0.8159, n_s = 0.9667 (Planck15) + MassFunction.build(sigma_8=0.8).evolve(cosmology=Planck15) # sigma_8 stays 0.8 + mf.evolve(sigma_8=0.8).evolve(cosmology=Planck15) # so does this one + +A value set explicitly, in ``build`` or in any ``evolve``, is kept from then on, and so +is every value below a stage given whole (``evolve(linear_power=lp)``). +:meth:`~hmf.core.stage.Stage.computed_parameters` lists the parameters still computed. +If the new cosmology does not give a computed parameter (its ``meta`` has no +``sigma8``, as for a cosmology made by hand), ``evolve`` raises the +:class:`ValueError` ``build`` raises, which names the parameter: set it explicitly. +Which parameters are computed is not part of a stage's value: it does not change +``==``, the hash or the content hash. A pickle or a copy keeps it; a configuration +does not record it, so a stage read from one has every value the configuration gives +explicit. + A name that is not a parameter raises a :class:`TypeError` at once, before anything is built, with the closest names (``mf.evolve(sigma8=0.8)``: "Did you mean 'sigma_8'?"); an ambiguous one lists the dotted paths to choose from. @@ -491,7 +515,9 @@ Extension stages * the fields whose names are too generic to stand alone in ``qualified_parameters``, listed with the stage's name; * and defaults that depend on other parameters (as ``build`` takes σ8 and n_s from - the cosmology) in the class method ``computed_defaults``. + the cosmology) in the class method ``computed_defaults``. It must compute them + from the parameters it is given alone, looking up each one it depends on: + ``evolve`` calls it again when one of those it looked up changes. A derived property's input that is not a field of the tree raises a :class:`TypeError` when the table is first built. @@ -569,7 +595,9 @@ TOML file. The disk cache's keys use the same canonical form **Pickling.** A stage pickles, and deep-copies, as its configuration, with the Boltzmann-code runs its tree has already used, as plain arrays. Unpickling rebuilds the tree and reuses those runs, so it does not run CAMB or CLASS again, and gives -bit-identical results. No CAMB or CLASS object is ever pickled. +bit-identical results. No CAMB or CLASS object is ever pickled. The pickle also records +which parameters hold computed defaults (see "Computed defaults" above), so a copy +evolves as the original does. API --- diff --git a/src/hmf/core/config.py b/src/hmf/core/config.py index a649a408..85f2e03f 100644 --- a/src/hmf/core/config.py +++ b/src/hmf/core/config.py @@ -99,6 +99,13 @@ sigma_8 = 0.8 # (at the top, before [hmf], in a file) fit = "ST" +A configuration does not record which values were computed defaults (see +:meth:`Stage.computed_parameters `): every +value it gives is explicit in the stage read from it, so a later +``evolve(cosmology=...)`` keeps it. Only a value a partial configuration leaves out, +and so takes a computed default, is computed again. A configuration written by +:func:`to_config` gives every value, so the stage read from it has none computed. + An unknown field raises a :class:`TypeError` with the closest names; an unknown model a :class:`~hmf.core.model.ModelNotFoundError` with suggestions; a wrong unit an :class:`astropy.units.UnitConversionError`. The error notes where in the @@ -142,7 +149,9 @@ the Boltzmann-code runs its tree has already used, as plain arrays (see :func:`hmf.core._boltzmann.export_runs`). Unpickling rebuilds the tree from the configuration, and puts the runs in the in-process memo, so it does not run CAMB or -CLASS again; results are bit-identical. No CAMB or CLASS object is ever pickled. +CLASS again; results are bit-identical. No CAMB or CLASS object is ever pickled. The +pickle also records which parameters hold computed defaults, which the configuration +does not, so a pickled or copied stage evolves as the original does. """ from __future__ import annotations @@ -1342,12 +1351,45 @@ def walk(obj: Any) -> None: return runs -def _pickle_state(stage: Stage) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]: - """What a stage pickles as: its configuration, and the runs its tree has used.""" - return to_config(stage), export_runs(_computed_runs(stage)) +def _pickle_state( + stage: Stage, +) -> tuple[dict[str, Any], dict[str, dict[str, Any]], dict[str, tuple[list[str], list[str]]]]: + """What a stage pickles as: its configuration, its runs and its computed defaults. + The runs are those its tree has used; the computed defaults, which parameters hold + one (see :meth:`Stage.computed_parameters + `), which the configuration does not + record. + """ + records = { + ".".join(prefix): ( + sorted(".".join(p) for p in record.paths), + sorted(".".join(p) for p in record.reads), + ) + for prefix, record in routing.computed_records(stage).items() + } + return to_config(stage), export_runs(_computed_runs(stage)), records -def _unpickle(config: Mapping[str, Any], runs: Mapping[str, Mapping[str, Any]]) -> Stage: + +def _unpickle( + config: Mapping[str, Any], + runs: Mapping[str, Mapping[str, Any]], + computed: Mapping[str, tuple[list[str], list[str]]] | None = None, +) -> Stage: """Rebuild a pickled stage (see :meth:`Stage.__reduce__ `).""" import_runs(runs) - return from_config(config) + stage = from_config(config) + if computed: + + def split(name: str) -> tuple[str, ...]: + return tuple(name.split(".")) if name else () + + records = { + split(prefix): routing.ComputedRecord( + paths=frozenset(split(p) for p in paths), + reads=frozenset(split(r) for r in reads), + ) + for prefix, (paths, reads) in computed.items() + } + routing.with_records(stage, records) + return stage diff --git a/src/hmf/core/mass_function.py b/src/hmf/core/mass_function.py index 261698c7..8dd35c95 100644 --- a/src/hmf/core/mass_function.py +++ b/src/hmf/core/mass_function.py @@ -208,9 +208,11 @@ def _cosmology_parameter(cosmology: FLRW, key: str, name: str) -> float: """A parameter from an astropy cosmology's ``meta`` (e.g. Planck18's ``sigma8``).""" value = cosmology.meta.get(key) if value is None: + label = cosmology.name or type(cosmology).__name__ raise ValueError( - f"MassFunction.build: the cosmology {cosmology.name!r} does not give {name} " - f"(its meta has no {key!r}, as astropy's realisations do): pass {name}=..." + f"MassFunction: the cosmology {label!r} does not give {name}, its default " + f"(its meta has no {key!r}, as astropy's realisations do): set {name} " + f"explicitly ({name}=...)." ) return float(value) @@ -399,6 +401,12 @@ def build(cls, **parameters: Any) -> MassFunction: def computed_defaults(cls, given: GivenParameters) -> Mapping[str, Any]: """The defaults of n_s and sigma_8: the cosmology's (see :meth:`build`). + They are computed again when :meth:`~hmf.core.stage.Stage.evolve` changes the + cosmology (or the transfer or linear-power stage), for those of them that were + not set explicitly, so that ``build().evolve(cosmology=c)`` equals + ``build(cosmology=c)``. One that was set, when building or in an ``evolve``, + keeps its value. + Parameters ---------- given @@ -412,7 +420,8 @@ def computed_defaults(cls, given: GivenParameters) -> Mapping[str, Any]: Raises ------ ValueError - If one of them is not given and the cosmology's ``meta`` does not have it. + If one of them is not given and the cosmology's ``meta`` does not have it + (the message names it, and asks to set it explicitly). """ cosmology = given.get("linear_power.transfer.cosmology", Planck18) out = {} diff --git a/src/hmf/core/routing.py b/src/hmf/core/routing.py index 87dfb689..6c0a892e 100644 --- a/src/hmf/core/routing.py +++ b/src/hmf/core/routing.py @@ -56,7 +56,7 @@ import difflib import types import typing -from collections.abc import Callable, Collection, Iterator, Mapping +from collections.abc import Callable, Collection, Iterable, Iterator, Mapping from functools import cached_property from typing import TYPE_CHECKING, Any, cast @@ -68,6 +68,7 @@ from .stage import Stage __all__ = [ + "ComputedRecord", "GivenParameters", "ParameterInfo", "Router", @@ -158,7 +159,9 @@ class GivenParameters(Mapping[str, Any]): A read-only mapping from dotted paths to values, passed to :meth:`Stage.computed_defaults `. Looking - up a path below a stage given whole returns the value in that stage. + up a path below a stage given whole returns the value in that stage. It records + the paths looked up (:attr:`reads`), so that :meth:`Stage.evolve + ` knows which changes call for the defaults again. Parameters ---------- @@ -166,14 +169,25 @@ class GivenParameters(Mapping[str, Any]): The given values, by path. """ - __slots__ = ("_values",) + __slots__ = ("_reads", "_values") def __init__(self, values: Mapping[FieldPath, Any]) -> None: self._values = dict(values) + self._reads: set[FieldPath] = set() + + @property + def reads(self) -> frozenset[FieldPath]: + """The paths looked up so far (``()`` if the mapping was iterated: every path).""" + return frozenset(self._reads) def __getitem__(self, key: str) -> Any: """The value at a dotted path (a ``KeyError`` if it is not given).""" path = _split(key) + self._reads.add(path) + return self._lookup(path, key) + + def _lookup(self, path: FieldPath, key: str) -> Any: + """The value at ``path`` (a ``KeyError`` naming ``key`` if it is not given).""" for i in range(len(path), 0, -1): if path[:i] in self._values: try: @@ -184,13 +198,110 @@ def __getitem__(self, key: str) -> Any: def __iter__(self) -> Iterator[str]: """The given dotted paths.""" + self._reads.add(()) return (_dotted(p) for p in self._values) def __len__(self) -> int: """The number of given paths.""" + self._reads.add(()) return len(self._values) +class _TreeParameters(GivenParameters): + """The parameters of a stage tree being evolved, as :class:`GivenParameters`. + + What :meth:`Stage.computed_defaults ` sees + when :func:`evolve` computes its defaults again: every parameter of the new tree + (the changes, and the old stage's values of the rest) except those that are still + computed, as :func:`from_flat` would have been given them. A stage path is given + only if it is changed whole. + """ + + __slots__ = ("_computed", "_leaves", "_old") + + def __init__( + self, + old: Stage, + changes: Mapping[FieldPath, Any], + computed: frozenset[FieldPath], + ) -> None: + super().__init__(changes) + self._old = old + self._computed = computed + self._leaves = router(type(old)).leaves + + def _lookup(self, path: FieldPath, key: str) -> Any: + values, leaves = self._values, self._leaves + is_parameter = False + for i in range(1, len(path) + 1): + if path[:i] in values: + return super()._lookup(path, key) + is_parameter = is_parameter or path[:i] in leaves + if not is_parameter or path in self._computed: + raise KeyError(key) + try: + return _walk(self._old, path) + except AttributeError: + raise KeyError(key) from None + + def _paths(self) -> list[FieldPath]: + changed = [p for p in self._values if p not in self._leaves] + leaves = [p for p in self._leaves if p not in self._computed] + return [*changed, *leaves] + + def __iter__(self) -> Iterator[str]: + self._reads.add(()) + return (_dotted(p) for p in self._paths()) + + def __len__(self) -> int: + self._reads.add(()) + return len(self._paths()) + + +@attrs.frozen +class ComputedRecord: + """Which parameters below a stage hold a default computed by its ``computed_defaults``. + + A stage records it (see :meth:`Stage.computed_parameters + `) so that :meth:`Stage.evolve + ` computes those defaults again when a parameter + they are computed from changes. It is not part of the stage's value. + + Parameters + ---------- + paths + The paths, below the stage, of the parameters that took a computed default. + reads + The paths, below the stage, that ``computed_defaults`` looked up (``()``: it + read every parameter). + """ + + paths: frozenset[FieldPath] = frozenset() + reads: frozenset[FieldPath] = frozenset() + + def explicit(self, changed: Iterable[FieldPath]) -> ComputedRecord: + """The record with the parameters at or below the changed paths made explicit. + + Parameters + ---------- + changed + Paths, below the stage, that are set. + + Returns + ------- + ComputedRecord + """ + changed = tuple(changed) + dropped = [p for p in self.paths for c in changed if _overlaps(p, c)] + if not dropped: + return self + return ComputedRecord(paths=self.paths.difference(dropped), reads=self.reads) + + +#: The record of a stage with no computed defaults. +NO_COMPUTED = ComputedRecord() + + @attrs.frozen class _Node: """What evolve and from_flat need to know of one stage class.""" @@ -367,6 +478,16 @@ def __init__(self, cls: type[Stage]) -> None: self._candidates = sorted(set(self.names) | set(self.ambiguous)) + from .stage import Stage + + #: The stages of the tree whose class computes defaults (from the top). + self.computing: tuple[FieldPath, ...] = tuple( + prefix + for prefix, stage_cls in self.stages.items() + if getattr(getattr(stage_cls, "computed_defaults", None), "__func__", None) + is not getattr(Stage.computed_defaults, "__func__", None) + ) + # -- building -------------------------------------------------------------------- def _visit(self, cls: type, prefix: FieldPath, sharers: dict[str, list[FieldPath]]) -> None: @@ -730,7 +851,139 @@ def evolve(stage: Stage, changes: Mapping[str, Any]) -> Stage: values = resolved.values if resolved.models: values = _apply_models(table, resolved, lambda p: _walk(stage, p)) - return _rebuild(type(stage), stage, values, resolved.soft, ()) # type: ignore[no-any-return] + soft = resolved.soft + records: dict[FieldPath, ComputedRecord] = {} + if table.computing and values: + values, soft, records = _recompute_defaults(table, stage, values, soft) + new = _rebuild(type(stage), stage, values, soft, ()) + _record(new, records, stage) + return new # type: ignore[no-any-return] + + +def _recompute_defaults( + table: Router, + stage: Stage, + values: Mapping[FieldPath, Any], + soft: frozenset[FieldPath], +) -> tuple[dict[FieldPath, Any], frozenset[FieldPath], dict[FieldPath, ComputedRecord]]: + """The changes of an evolve, with the computed defaults it changes. + + For each stage of the tree that computes defaults and is not changed whole, the + parameters that are changed become explicit, and if a change is one its + ``computed_defaults`` read, the defaults of the parameters still computed are + computed again from the new tree, and those that change are added. Stages are + visited from the top, as in :func:`_with_computed_defaults`. + + Returns the changes, the paths set by a shared name, and the new record of each + stage with a record that has a change below it, by its path. + """ + out = dict(values) + soft_out = set(soft) + records: dict[FieldPath, ComputedRecord] = {} + for prefix in table.computing: + if any(prefix[:i] in values for i in range(len(prefix) + 1)): + continue # given whole, or below a stage given whole + n = len(prefix) + below = {p[n:]: v for p, v in out.items() if p[:n] == prefix} + if not below: + continue + try: + old = _walk(stage, prefix) + except AttributeError: # pragma: no cover - below a stage that is None + continue + if old is None: + continue + record: ComputedRecord = old._computed + if not record.paths: + continue + new_record = record.explicit(below) + if new_record.paths and any(_overlaps(r, c) for r in record.reads for c in below): + given = _TreeParameters(old, below, new_record.paths) + computed = type(old).computed_defaults(given) + extra = router(type(old)).resolve(computed) if computed else None + sub_table = router(type(old)) + for path in new_record.paths: + if extra is not None and path in extra.values: + value = extra.values[path] + if path in extra.soft: + soft_out.add(prefix + path) + else: + # As from_flat: a parameter with no computed default takes its own. + value = _default(sub_table.leaves[path], repr(sub_table.preferred[path])) + if not _same_value(value, _walk(old, path)): + out[prefix + path] = value + new_record = ComputedRecord(paths=new_record.paths, reads=given.reads) + records[prefix] = new_record + return out, frozenset(soft_out), records + + +def _same_value(a: Any, b: Any) -> bool: + """Whether two values are the same object or compare equal (False if eq fails).""" + if a is b: + return True + try: + return bool(a == b) + except Exception: # noqa: BLE001 - an eq that can't answer means "not the same" + return False + + +def _record( + stage: Stage, records: Mapping[FieldPath, ComputedRecord], old: Stage | None = None +) -> None: + """Set the computed-defaults record of the stages at the given paths of a new tree. + + Only for stages just built: a stage shared with the old tree ``old`` keeps its own. + """ + for prefix, record in records.items(): + node = _walk(stage, prefix) + if old is not None and node is _walk(old, prefix): + continue + if node._computed is not record: + object.__setattr__(node, "_computed", record) + + +def computed_records(stage: Stage) -> dict[FieldPath, ComputedRecord]: + """The computed-defaults record of each stage of a tree that has one. + + Parameters + ---------- + stage + The stage at the root of the tree. + + Returns + ------- + dict + The non-empty records, by the path of their stage. + """ + out: dict[FieldPath, ComputedRecord] = {} + for prefix in router(type(stage)).computing: + try: + node = _walk(stage, prefix) + except AttributeError: # pragma: no cover - below a stage that is None + continue + if node is not None and node._computed.paths: + out[prefix] = node._computed + return out + + +def with_records(stage: Stage, records: Mapping[FieldPath, ComputedRecord]) -> Stage: + """Set the computed-defaults records of a tree that was just built (and return it). + + Parameters + ---------- + stage + The stage at the root of a tree that nothing else holds yet (e.g. one just + unpickled). + records + Records by the path of their stage, as :func:`computed_records` gives them. + + Returns + ------- + Stage + ``stage``. + """ + _record(stage, records) + return stage def from_flat(cls: type[Stage], mapping: Mapping[str, Any]) -> Stage: @@ -748,7 +1001,7 @@ def from_flat(cls: type[Stage], mapping: Mapping[str, Any]) -> Stage: Stage """ table = router(cls) - resolved = _with_computed_defaults(table, table.resolve(mapping)) + resolved, records = _with_computed_defaults(table, table.resolve(mapping)) def base(path: FieldPath) -> Any: return _default(table.leaves[path], repr(table.preferred[path])) @@ -756,7 +1009,9 @@ def base(path: FieldPath) -> Any: values = resolved.values if resolved.models: values = _apply_models(table, resolved, base) - return _rebuild(cls, None, values, resolved.soft, ()) # type: ignore[no-any-return] + stage = _rebuild(cls, None, values, resolved.soft, ()) + _record(stage, records) + return stage # type: ignore[no-any-return] def from_paths( @@ -809,40 +1064,55 @@ def from_paths( f"{cls.__name__}: {_dotted(path)!r} is given, and so is the stage " f"{_dotted(path[:i])!r} above it. Give one or the other." ) - resolved = _with_computed_defaults( + resolved, records = _with_computed_defaults( table, _Resolved(values=paths, soft=frozenset(soft & set(paths)), models={}) ) - return _rebuild(cls, None, resolved.values, resolved.soft, ()) # type: ignore[no-any-return] + stage = _rebuild(cls, None, resolved.values, resolved.soft, ()) + _record(stage, records) + return stage # type: ignore[no-any-return] -def _with_computed_defaults(table: Router, resolved: _Resolved) -> _Resolved: +def _with_computed_defaults( + table: Router, resolved: _Resolved +) -> tuple[_Resolved, dict[FieldPath, ComputedRecord]]: """Add the computed defaults of each stage of the tree that is not given whole. Each stage's :meth:`~hmf.core.stage.Stage.computed_defaults` sees the parameters given below it; stages are visited from the top, and a default is used only for a field that is not given (or set by a default from a stage above), and not below a given stage. + + Returns the values with the defaults added, and the record of the defaults each + stage computed (the parameters that took one, and the paths it read), by the + stage's path. """ values = dict(resolved.values) soft = set(resolved.soft) given = [*resolved.values, *resolved.models] - for prefix, stage_cls in table.stages.items(): + records: dict[FieldPath, ComputedRecord] = {} + for prefix in table.computing: if any(prefix[:i] in values for i in range(len(prefix) + 1)): continue + stage_cls = table.stages[prefix] n = len(prefix) below = {p[n:]: v for p, v in values.items() if p[:n] == prefix} - computed = stage_cls.computed_defaults(GivenParameters(below)) # type: ignore[attr-defined] + parameters = GivenParameters(below) + computed = stage_cls.computed_defaults(parameters) # type: ignore[attr-defined] if not computed: continue extra = router(stage_cls).resolve(computed) + used = set() for path, value in extra.values.items(): full = prefix + path if not any(_overlaps(full, g) for g in given): values[full] = value given.append(full) + used.add(path) if path in extra.soft: soft.add(full) - return _Resolved(values=values, soft=frozenset(soft), models=resolved.models) + if used: + records[prefix] = ComputedRecord(paths=frozenset(used), reads=parameters.reads) + return _Resolved(values=values, soft=frozenset(soft), models=resolved.models), records def parameter_defaults(cls: type[Stage]) -> dict[str, Any]: @@ -860,7 +1130,7 @@ def parameter_defaults(cls: type[Stage]) -> dict[str, Any]: :meth:`Stage.parameter_names`. """ table = router(cls) - computed = _with_computed_defaults(table, _Resolved(values={}, soft=frozenset(), models={})) + computed, _ = _with_computed_defaults(table, _Resolved(values={}, soft=frozenset(), models={})) out: dict[str, Any] = {} for name, info in table.parameter_info.items(): path = _split(info.path) diff --git a/src/hmf/core/stage.py b/src/hmf/core/stage.py index eb127dbb..5053c6d7 100644 --- a/src/hmf/core/stage.py +++ b/src/hmf/core/stage.py @@ -210,6 +210,19 @@ class Stage(Documented): #: recorded when the class is created. _derived_properties: ClassVar[Mapping[str, DerivedProperty[Any]]] = MappingProxyType({}) + _computed: routing.ComputedRecord = field( + default=routing.NO_COMPUTED, + init=False, + eq=False, + repr=False, + physics=False, + doc=( + "The parameters below the stage that hold a default computed by its " + ":meth:`computed_defaults` (see :meth:`computed_parameters`). Not part of " + "the stage's value." + ), + ) + def evolve(self, **changes: Any) -> Self: """Return a copy of this stage tree with some parameters changed. @@ -222,6 +235,13 @@ def evolve(self, **changes: Any) -> Self: :meth:`evolve_own`; every other stage is shared, with its cached results, and so is every :func:`derived_property` whose inputs do not change. + A parameter that holds a default computed from other parameters (see + :meth:`computed_defaults`, e.g. a mass function's sigma_8 and n_s, from the + cosmology's ``meta``) is computed again when a parameter it is computed from + changes, so that ``build(**a).evolve(**b)`` equals ``build(**a, **b)``. A + parameter set explicitly, here or when the tree was built, keeps its value + (see :meth:`computed_parameters`). + Every name is checked before anything is built. The new stages are built, and so validated, by their constructors; if that fails, the error propagates and this stage, which is immutable, is unchanged. @@ -241,6 +261,10 @@ def evolve(self, **changes: Any) -> Self: TypeError If a name is not a parameter (with suggestions), is ambiguous (with the dotted paths to choose from), is a derived property, or is given twice. + ValueError + If a default that is computed again cannot be computed from the new + parameters (e.g. the new cosmology does not give sigma_8), or for an + invalid value. """ return routing.evolve(self, changes) # type: ignore[return-value] @@ -250,7 +274,9 @@ def evolve_own(self, **changes: Any) -> Self: No routing: each name is a field of this stage, and nothing else changes. It is the step :meth:`evolve` takes at each stage it rebuilds. The value of each :func:`derived_property` already computed is carried to the new stage if none - of its inputs changes; every other cached result is computed again. + of its inputs changes; every other cached result is computed again. No + default is computed again: a field given here, and every parameter below it, + is explicit in the new stage (see :meth:`computed_parameters`). Parameters ---------- @@ -278,6 +304,12 @@ def evolve_own(self, **changes: Any) -> Self: msg += f" Fields: {names}." raise TypeError(msg) new = attrs.evolve(self, **changes) + record = self._computed + if record.paths: + # A field given here is explicit, with every parameter below it. + if any(p[0] in changes for p in record.paths): + record = record.explicit((n,) for n in changes) + object.__setattr__(new, "_computed", record) for name, prop in self._derived_properties.items(): try: value = _computed_value(self, name) @@ -324,6 +356,15 @@ def computed_defaults(cls, given: routing.GivenParameters) -> Mapping[str, Any]: the cosmology); the base stage has none. Defaults for parameters (or below stages) that are given are ignored. + The stage records which parameters took a computed default, and which paths of + ``given`` the hook looked up (see :meth:`computed_parameters`). When + :meth:`evolve` changes one of those paths, it calls the hook again, with every + parameter of the new tree except the ones still computed, and the computed + parameters take the new defaults. A parameter set explicitly, in + :meth:`from_flat` or in :meth:`evolve`, is never computed again. So the hook + must compute its defaults from ``given`` alone, and look up (with ``in``, + ``[]`` or ``get``) every parameter it depends on. + Parameters ---------- given @@ -336,6 +377,40 @@ def computed_defaults(cls, given: routing.GivenParameters) -> Mapping[str, Any]: """ return {} + def computed_parameters(self) -> tuple[str, ...]: + """The parameters of the stage tree that hold a default computed from others. + + Those that took a default from a stage's :meth:`computed_defaults` when the + tree was built (by :meth:`from_flat`, or :meth:`build` where a stage has one) + and have not been set explicitly since. :meth:`evolve` computes them again when + a parameter they are computed from changes (see :meth:`computed_defaults`). + Which parameters they are is not part of the stage's value: it does not change + ``==``, the hash or the content hash, and a configuration does not record it + (a stage read from a configuration has every value it gives explicit), but a + pickle or a copy keeps it. + + Returns + ------- + tuple of str + Their dotted paths, sorted. + + Examples + -------- + >>> from hmf.core.mass_function import MassFunction + >>> mf = MassFunction.build(**{"transfer.model": "EH"}) + >>> mf.computed_parameters() + ('linear_power.sigma_8', 'linear_power.transfer.n_s') + >>> mf.evolve(sigma_8=0.8).computed_parameters() + ('linear_power.transfer.n_s',) + """ + return tuple( + sorted( + ".".join((*prefix, *path)) + for prefix, record in routing.computed_records(self).items() + for path in record.paths + ) + ) + def to_flat(self) -> dict[str, Any]: """The value of every parameter of the stage tree, by the name it is routed by. @@ -463,7 +538,7 @@ def _walk_path(obj: Any, dotted: str) -> Any: def _fresh(stage: Stage) -> Stage: """A copy of a stage tree with nothing computed: every stage of it is a new one.""" - changes = {} + changes: dict[str, Any] = {} for a in attrs.fields(type(stage)): value = getattr(stage, a.name) if a.init and isinstance(value, Stage): diff --git a/tests/core/test_core_computed_defaults.py b/tests/core/test_core_computed_defaults.py new file mode 100644 index 00000000..ff606599 --- /dev/null +++ b/tests/core/test_core_computed_defaults.py @@ -0,0 +1,212 @@ +"""Tests of computed defaults under evolve: build(**a).evolve(**b) == build(**a, **b). + +A mass function's sigma_8 and n_s default to the cosmology's (its ``meta``). A +parameter that took such a default is computed again when ``evolve`` changes what it is +computed from; one set explicitly keeps its value. +""" + +import copy +import pickle + +import attrs +import numpy as np +import pytest +from astropy.cosmology import WMAP7, WMAP9, FlatLambdaCDM, Planck13, Planck15, Planck18 + +from hmf.core import config, field +from hmf.core.mass_function import MassFunction +from hmf.core.routing import ComputedRecord +from hmf.core.stage import Stage +from hmf.core.transfer import Transfer +from hmf.core.units import Msun_h + +EH = {"transfer.model": "EH"} +M = np.logspace(10, 15, 11) * Msun_h +Z = np.array([0.0, 1.0])[:, None] +COMPUTED = ("linear_power.sigma_8", "linear_power.transfer.n_s") +#: A cosmology whose meta has neither sigma8 nor n. +BARE = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.05) + + +@attrs.frozen(kw_only=True) +class HaloModel(Stage): + """A toy extension stage holding a mass function.""" + + mass_function: MassFunction = field( + factory=lambda: MassFunction.build(**EH), doc="The mass function." + ) + bias: str = field(default="Tinker10", doc="The bias model.") + + +@pytest.fixture(scope="module") +def default(): + """The default mass function, with the fast transfer model.""" + return MassFunction.build(**EH) + + +def _dndm(mf: MassFunction) -> np.ndarray: + return mf.dndm(m=M[None, :], z=Z).value + + +@pytest.mark.parametrize("cosmology", [Planck13, Planck15, WMAP7, WMAP9]) +def test_evolve_cosmology_equals_build(default, cosmology): + built = MassFunction.build(cosmology=cosmology, **EH) + evolved = default.evolve(cosmology=cosmology) + flat = MassFunction.from_flat({**EH, "cosmology": cosmology}) + # The defaults are the cosmology's own (astropy's realisations give them). + assert built.linear_power.sigma_8 == cosmology.meta["sigma8"] + assert built.linear_power.transfer.n_s == cosmology.meta["n"] + assert evolved == built + assert flat == built + assert evolved.computed_parameters() == built.computed_parameters() == COMPUTED + assert np.array_equal(_dndm(evolved), _dndm(built)) + assert np.array_equal(_dndm(flat), _dndm(built)) + # And from a non-default cosmology, and back. + other = MassFunction.build(cosmology=Planck15, **EH).evolve(cosmology=cosmology) + assert other == built + assert evolved.evolve(cosmology=Planck18) == default + + +def test_evolve_transfer_stage_given_whole_equals_build(default): + transfer = Transfer(model="EH", cosmology=Planck15, n_s=0.95) + evolved = default.evolve(**{"linear_power.transfer": transfer}) + built = MassFunction.from_flat({"linear_power.transfer": transfer}) + assert evolved == built + assert evolved.linear_power.sigma_8 == Planck15.meta["sigma8"] + assert evolved.linear_power.transfer.n_s == 0.95 + # n_s, in the stage given whole, is explicit; sigma_8 is still computed. + assert evolved.computed_parameters() == built.computed_parameters() + assert evolved.computed_parameters() == ("linear_power.sigma_8",) + + +@pytest.mark.parametrize( + "start", + [ + lambda: MassFunction.build(sigma_8=0.7, **EH), + lambda: MassFunction.build(**EH).evolve(sigma_8=0.7), + lambda: MassFunction.build(**EH).evolve(sigma_8=0.7).evolve(cosmology=WMAP9), + ], +) +def test_explicit_sigma_8_survives_a_cosmology_change(start): + new = start().evolve(cosmology=Planck15) + assert new.linear_power.sigma_8 == 0.7 + # n_s was not set, so it is still the cosmology's. + assert new.linear_power.transfer.n_s == Planck15.meta["n"] + assert new.computed_parameters() == ("linear_power.transfer.n_s",) + assert new == MassFunction.build(sigma_8=0.7, cosmology=Planck15, **EH) + + +def test_explicit_value_equal_to_the_default_is_kept(default): + # Set to the value it had: still explicit from then on. + pinned = default.evolve(sigma_8=default.linear_power.sigma_8) + assert pinned.evolve(cosmology=Planck15).linear_power.sigma_8 == Planck18.meta["sigma8"] + + +def test_a_stage_given_whole_is_explicit(default): + lp = default.linear_power + replaced = default.evolve(linear_power=lp) + assert replaced.computed_parameters() == () + new = replaced.evolve(cosmology=Planck15) + assert new.linear_power.sigma_8 == Planck18.meta["sigma8"] + assert new.linear_power.transfer.n_s == Planck18.meta["n"] + + +def test_evolve_without_a_source_raises_when_computed(default): + with pytest.raises(ValueError, match=r"'FlatLambdaCDM' does not give n_s.*set n_s explicitly"): + default.evolve(cosmology=BARE) + with pytest.raises(ValueError, match=r"does not give sigma_8.*set sigma_8 explicitly"): + default.evolve(n_s=0.96).evolve(cosmology=BARE) + # build raises the same. + with pytest.raises(ValueError, match=r"does not give sigma_8.*set sigma_8 explicitly"): + MassFunction.build(cosmology=BARE, n_s=0.96, **EH) + named = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.05, name="mine") + with pytest.raises(ValueError, match="the cosmology 'mine' does not give n_s"): + default.evolve(cosmology=named) + # The original is unchanged. + assert default.linear_power.sigma_8 == Planck18.meta["sigma8"] + + +def test_evolve_without_a_source_works_when_explicit(): + mf = MassFunction.build(sigma_8=0.8, n_s=0.96, **EH) + new = mf.evolve(cosmology=BARE) + assert new.linear_power.sigma_8 == 0.8 + assert new.linear_power.transfer.n_s == 0.96 + assert new == MassFunction.build(cosmology=BARE, sigma_8=0.8, n_s=0.96, **EH) + + +def test_provenance_is_not_part_of_the_value(default): + explicit = MassFunction.build( + sigma_8=default.linear_power.sigma_8, n_s=default.linear_power.transfer.n_s, **EH + ) + assert explicit.computed_parameters() == () + assert default.computed_parameters() == COMPUTED + assert explicit == default + assert hash(explicit) == hash(default) + assert config.content_hash(explicit) == config.content_hash(default) + assert config.to_config(explicit) == config.to_config(default) + assert "_computed" not in repr(default) + assert "_computed" not in MassFunction.parameter_names() + + +@pytest.mark.parametrize("copier", [lambda s: pickle.loads(pickle.dumps(s)), copy.deepcopy]) +def test_pickling_keeps_provenance(copier): + mf = MassFunction.build(sigma_8=0.7, **EH) + same = copier(mf) + assert same == mf + assert same.computed_parameters() == mf.computed_parameters() + assert same.linear_power._computed == mf.linear_power._computed + assert same._computed == mf._computed + assert same.evolve(cosmology=Planck15) == mf.evolve(cosmology=Planck15) + + +def test_a_configuration_records_no_provenance(default): + read = config.from_config(config.to_config(default)) + assert read == default + assert read.computed_parameters() == () + # So every value it gives is kept. + assert read.evolve(cosmology=Planck15).linear_power.sigma_8 == Planck18.meta["sigma8"] + # A value a partial configuration leaves out takes a computed default, and stays one. + partial = config.from_config( + {"hmf": {"schema_version": 1}, "sigma_8": 0.8, **EH}, stage=MassFunction + ) + assert partial.computed_parameters() == ("linear_power.transfer.n_s",) + + +def test_evolve_own_makes_the_fields_it_sets_explicit(default): + assert default.evolve_own(delta_c=1.7).computed_parameters() == COMPUTED + assert default.evolve_own(linear_power=default.linear_power).computed_parameters() == () + + +def test_unrelated_changes_share_the_stages(default): + new = default.evolve(fit="ST") + assert new.linear_power is default.linear_power + assert new.computed_parameters() == COMPUTED + lower = default.evolve(sigma_8=0.7) + assert lower.linear_power.transfer is default.linear_power.transfer + assert lower.computed_parameters() == ("linear_power.transfer.n_s",) + # A cosmology change with the same sigma_8 and n_s rebuilds only what it must. + same_meta = Planck18.clone(H0=Planck18.H0) + assert default.evolve(cosmology=same_meta) == default + + +def test_extension_stage_recomputes_the_mass_functions_defaults(default): + built = HaloModel.from_flat({**EH, "cosmology": Planck15}) + assert built.mass_function == MassFunction.build(cosmology=Planck15, **EH) + evolved = HaloModel.from_flat(EH).evolve(cosmology=Planck15) + assert evolved == built + # A mass function given whole brings its own record. + held = HaloModel(mass_function=default).evolve(cosmology=Planck15) + assert held == built + assert held.computed_parameters() == ( + "mass_function.linear_power.sigma_8", + "mass_function.linear_power.transfer.n_s", + ) + explicit = HaloModel(mass_function=default.evolve(sigma_8=0.7)).evolve(cosmology=Planck15) + assert explicit.mass_function.linear_power.sigma_8 == 0.7 + + +def test_computed_record(): + record = ComputedRecord(paths=frozenset({("a", "b"), ("c",)}), reads=frozenset({("d",)})) + assert record.explicit([("a",)]).paths == {("c",)} + assert record.explicit([("e",)]) is record + assert record.explicit([("c",)]).reads == record.reads diff --git a/tests/core/test_core_routing.py b/tests/core/test_core_routing.py index cc52f0b2..7cf08b7f 100644 --- a/tests/core/test_core_routing.py +++ b/tests/core/test_core_routing.py @@ -137,7 +137,7 @@ def test_every_field_is_routable(): for path, cls in table.stages.items(): for a in attrs.fields(cls): full = (*path, a.name) - if full not in table.stages: + if a.init and full not in table.stages: fields.add(".".join(full)) assert fields == routed assert set(table.stages) == {(), ("linear_power",), ("linear_power", "transfer")} @@ -200,6 +200,17 @@ def test_changes_cover_every_parameter(): assert set(CHANGES) == set(MassFunction.parameter_names()) +def _explicit_flat(mf): + """``mf.to_flat()`` without the parameters that hold computed defaults. + + So that ``from_flat({**_explicit_flat(mf), name: value})`` is ``mf.evolve(name=value)``: + the computed defaults are computed again, from the new value. + """ + computed = set(mf.computed_parameters()) + info = MassFunction.parameter_info() + return {k: v for k, v in mf.to_flat().items() if info[k].path not in computed} + + def _two_species(mf): """``mf`` with a tabulated transfer function whose total matter is suppressed at small scales. @@ -227,7 +238,7 @@ def test_evolve_every_parameter(mf, tmp_path, name): got, old = getattr(got, part), getattr(old, part) assert got is not old, path # A fresh tree, with every other parameter the same. - fresh = MassFunction.from_flat({**mf.to_flat(), name: value}) + fresh = MassFunction.from_flat({**_explicit_flat(mf), name: value}) assert new == fresh if quantity is not None: method, z = quantity @@ -301,7 +312,8 @@ def test_cosmology_change_is_consistent(mf): assert stage.cosmology is Planck15 assert lp.growth.transfer is lp.transfer assert new.variance.power.transfer is lp.transfer - assert new == MassFunction.from_flat({**mf.to_flat(), "cosmology": Planck15}) + assert new == MassFunction.from_flat({**_explicit_flat(mf), "cosmology": Planck15}) + assert new == MassFunction.build(**{"transfer.model": "EH"}, cosmology=Planck15) def test_k_accuracy_change_is_consistent(mf): @@ -549,6 +561,13 @@ def test_given_parameters(): assert "linear_power.transfer.nope" not in given assert sorted(given) == ["delta_c", "linear_power.transfer"] assert len(given) == 2 + # It records what was looked up (() for iterating: everything). + fresh = GivenParameters({("delta_c",): 1.6}) + assert "linear_power.sigma_8" not in fresh + assert fresh.get("delta_c") == 1.6 + assert fresh.reads == {("linear_power", "sigma_8"), ("delta_c",)} + list(fresh) + assert () in fresh.reads # --------------------------------------------------------------------------------- @@ -852,7 +871,10 @@ def test_toy_shared_parameter_and_derived_properties(): double = built.double assert built.evolve(**{"right.count": 2}).double is double assert built.evolve(acc=3).double == 18 - assert built.evolve(base=4).double == 16 + # base is what acc's default is computed from: given, acc takes its own default. + assert built.evolve(base=4) == Holder.from_flat({"base": 4}) + assert built.evolve(base=4).double == 8 + assert built.evolve(acc=4).evolve(base=4).double == 32 # acc set: kept assert Holder.invalidated_by("acc") == ("", "double", "left", "right") assert Holder.invalidated_by("base") == ("", "double", "unannotated", "unresolved") # Without the composing stage's declaration, the name is ambiguous. From b7e2655cfa7aa66c2c2c9ae831de52c9e11ac3ad Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 8 Nov 2026 08:50:21 +0000 Subject: [PATCH 2/4] fix(core): no computed-defaults record on an optional stage left out from_flat on a tree with an optional stage that stays None (only shared names reach it) tried to set that stage's record on None. Skip it. Cover the remaining branches of the re-derivation: a shared computed default, a stage given whole inside an extension, and the evolving tree's parameter view. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_019U8u44xFH87Am2mNHVrimP --- src/hmf/core/routing.py | 5 +- tests/core/test_core_computed_defaults.py | 100 +++++++++++++++++++++- 2 files changed, 102 insertions(+), 3 deletions(-) diff --git a/src/hmf/core/routing.py b/src/hmf/core/routing.py index 6c0a892e..009e7b2d 100644 --- a/src/hmf/core/routing.py +++ b/src/hmf/core/routing.py @@ -932,11 +932,12 @@ def _record( ) -> None: """Set the computed-defaults record of the stages at the given paths of a new tree. - Only for stages just built: a stage shared with the old tree ``old`` keeps its own. + Only for stages just built: a stage shared with the old tree ``old`` keeps its own, + and an optional stage left out (None) has none. """ for prefix, record in records.items(): node = _walk(stage, prefix) - if old is not None and node is _walk(old, prefix): + if node is None or (old is not None and node is _walk(old, prefix)): continue if node._computed is not record: object.__setattr__(node, "_computed", record) diff --git a/tests/core/test_core_computed_defaults.py b/tests/core/test_core_computed_defaults.py index ff606599..fe3a540b 100644 --- a/tests/core/test_core_computed_defaults.py +++ b/tests/core/test_core_computed_defaults.py @@ -7,13 +7,14 @@ import copy import pickle +from typing import ClassVar import attrs import numpy as np import pytest from astropy.cosmology import WMAP7, WMAP9, FlatLambdaCDM, Planck13, Planck15, Planck18 -from hmf.core import config, field +from hmf.core import config, field, routing from hmf.core.mass_function import MassFunction from hmf.core.routing import ComputedRecord from hmf.core.stage import Stage @@ -210,3 +211,100 @@ def test_computed_record(): assert record.explicit([("a",)]).paths == {("c",)} assert record.explicit([("e",)]) is record assert record.explicit([("c",)]).reads == record.reads + + +def test_extension_stage_with_a_mass_function_given_whole(default): + explicit = default.evolve(sigma_8=0.7) + built = HaloModel.from_flat({"mass_function": default}) + assert built.computed_parameters() == tuple(f"mass_function.{p}" for p in COMPUTED) + # Replacing the stage whole brings the new stage's own record, untouched. + replaced = built.evolve(mass_function=explicit) + assert replaced.mass_function is explicit + assert replaced.computed_parameters() == ("mass_function.linear_power.transfer.n_s",) + + +# --------------------------------------------------------------------------------- +# Toy trees: shared and optional stages with computed defaults +# --------------------------------------------------------------------------------- + + +@attrs.frozen(kw_only=True) +class Part(Stage): + """A toy stage with a setting.""" + + acc: int = field(default=1, doc="A setting.") + + +@attrs.frozen(kw_only=True) +class Pair(Stage): + """A toy stage sharing its parts' setting, whose default is computed from ``base``.""" + + shared_parameters: ClassVar = frozenset({"acc"}) + + left: Part = field(factory=Part, doc="One.") + right: Part = field(factory=Part, doc="Two.") + base: int = field(default=3, doc="What acc's default is computed from.") + + @classmethod + def computed_defaults(cls, given): + return {"acc": 2 * given.get("base", 3)} + + +@attrs.frozen(kw_only=True) +class Top(Stage): + """A toy stage holding an optional Pair.""" + + pair: Pair | None = field(default=None, doc="An optional pair.") + other: int = field(default=0, doc="Something else.") + + +def test_toy_shared_computed_default_is_recomputed(): + built = Pair.from_flat({}) + assert built.left.acc == built.right.acc == 6 + new = built.evolve(base=4) + assert new.left.acc == new.right.acc == 8 + assert new == Pair.from_flat({"base": 4}) + assert new.computed_parameters() == ("left.acc", "right.acc") + # Set by its shared name: explicit for both holders. + pinned = built.evolve(acc=5) + assert pinned.computed_parameters() == () + assert pinned.evolve(base=4).left.acc == 5 + + +def test_toy_optional_stage_left_out_has_no_record(): + top = Top.from_flat({}) + assert top.pair is None + assert top.computed_parameters() == () + # A shared name that only reaches the absent stage changes nothing. + assert top.evolve(acc=5).pair is None + with_pair = Top.from_flat({"pair.base": 2}) + assert with_pair.pair.left.acc == 4 + assert with_pair.evolve(**{"pair.base": 5}).pair == Pair.from_flat({"base": 5}) + + +def test_tree_parameters_view_of_an_evolving_tree(default): + computed = frozenset({("linear_power", "sigma_8")}) + given = routing._TreeParameters(default, {("delta_c",): 1.7}, computed) + assert given["delta_c"] == 1.7 + assert given["linear_power.transfer.n_s"] == default.linear_power.transfer.n_s + assert given["fit.A_200"] == default.fit.A_200 # into a model field + assert "linear_power.sigma_8" not in given # still computed + assert "linear_power" not in given # a stage, not changed whole + assert "fit.nope" not in given + names = list(given) + assert "delta_c" in names + assert "linear_power.sigma_8" not in names + assert len(given) == len(names) + assert () in given.reads + + +def test_same_value_when_eq_cannot_answer(): + a = np.array([1.0, 2.0]) + assert routing._same_value(a, a) + assert not routing._same_value(a, np.array([1.0, 3.0])) + assert routing._same_value(1, 1.0) + + +def test_base_stage_computes_no_defaults(): + assert Stage.computed_defaults(routing.GivenParameters({})) == {} + assert routing.router(Part).computing == () From 1d760a92a57e8611a6b13a0d2f636039974f628f Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 8 Nov 2026 10:26:44 +0000 Subject: [PATCH 3/4] refactor(core): a cosmology change always recomputes its defaults Replace the provenance design with a rule that depends only on the stage's value: when evolve changes a parameter a stage's defaults are computed from, every parameter computed_defaults gives a default for takes the new default, unless it is given in the same call. So mf.evolve(cosmology=c) takes sigma_8 and n_s from c, and mf.evolve(cosmology=c, sigma_8=0.8) keeps 0.8; any other change computes no default again. A stage declares what its defaults are computed from in computed_defaults_inputs (MassFunction: the transfer stage's cosmology); the router checks each is a field. Tracking what the hook looks up would also count its "is it given?" checks, so evolve(sigma_8=...) would reset an n_s set earlier. undeclared_inputs now also checks the hook against its declared inputs, through the lookups GivenParameters records. Removed: the init=False provenance field on Stage, ComputedRecord, computed_parameters(), and the computed-defaults element of the pickle state. A stage read from a pickle or a configuration evolves as the original does. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_019U8u44xFH87Am2mNHVrimP --- docs/core.rst | 57 ++-- src/hmf/core/config.py | 54 +--- src/hmf/core/mass_function.py | 14 +- src/hmf/core/routing.py | 318 +++++---------------- src/hmf/core/stage.py | 136 ++++----- tests/core/test_core_computed_defaults.py | 331 +++++++++++----------- tests/core/test_core_routing.py | 26 +- 7 files changed, 342 insertions(+), 594 deletions(-) diff --git a/docs/core.rst b/docs/core.rst index 677f9fa2..09694cc2 100644 --- a/docs/core.rst +++ b/docs/core.rst @@ -443,34 +443,32 @@ So the flat names of a mass function are exactly :meth:`build mf = MassFunction.build(**{"transfer.model": "EH"}, fit="ST", sigma_8=0.8) mf.evolve(sigma_8=0.85, n_s=0.95) - mf.evolve(cosmology=Planck15) # the cosmology of the whole tree, and - # n_s from Planck15 (sigma_8 is set above) + mf.evolve(cosmology=Planck15) # the cosmology of the whole tree, with + # sigma_8 and n_s from Planck15 mf.evolve(**{"fit.a": 0.75}) # a parameter of the fit mf.evolve(**{"linear_power.transfer.n_s": 0.95}) MassFunction.from_flat({"transfer.model": "EH", "linear_power": {"sigma_8": 0.8}}) **Computed defaults.** ``build`` takes σ8 and n_s, when they are not given, from the -cosmology's ``meta`` (a stage's :meth:`~hmf.core.stage.Stage.computed_defaults`). A -parameter that took such a default stays *computed* until it is set explicitly, and -``evolve`` computes it again whenever a parameter it is computed from changes: so -``build(**a).evolve(**b)`` equals ``build(**a, **b)``, and a cosmology scan done with -``evolve`` gives the same models as one done with ``build``:: - - mf = MassFunction.build() # sigma_8 = 0.8102, n_s = 0.9665 (Planck18) - mf.evolve(cosmology=Planck15) # sigma_8 = 0.8159, n_s = 0.9667 (Planck15) - MassFunction.build(sigma_8=0.8).evolve(cosmology=Planck15) # sigma_8 stays 0.8 - mf.evolve(sigma_8=0.8).evolve(cosmology=Planck15) # so does this one - -A value set explicitly, in ``build`` or in any ``evolve``, is kept from then on, and so -is every value below a stage given whole (``evolve(linear_power=lp)``). -:meth:`~hmf.core.stage.Stage.computed_parameters` lists the parameters still computed. -If the new cosmology does not give a computed parameter (its ``meta`` has no -``sigma8``, as for a cosmology made by hand), ``evolve`` raises the -:class:`ValueError` ``build`` raises, which names the parameter: set it explicitly. -Which parameters are computed is not part of a stage's value: it does not change -``==``, the hash or the content hash. A pickle or a copy keeps it; a configuration -does not record it, so a stage read from one has every value the configuration gives -explicit. +cosmology's ``meta`` (the mass function's +:meth:`~hmf.core.stage.Stage.computed_defaults`). They are computed from the cosmology, +so ``evolve`` computes them again whenever it changes the cosmology (or a stage holding +it, such as ``"linear_power.transfer"``): each takes the new cosmology's value unless it +is given in the same call. Any other change leaves them as they are:: + + mf = MassFunction.build(sigma_8=0.8, n_s=0.95) + mf.evolve(cosmology=Planck15) # sigma_8 = 0.8159, n_s = 0.9667 + mf.evolve(cosmology=Planck15, sigma_8=0.8) # sigma_8 = 0.8, n_s = 0.9667 + mf.evolve(sigma_8=0.85) # n_s stays 0.95 + +So ``mf.evolve(cosmology=c)`` equals ``build`` with ``c`` and every other parameter of +``mf`` except σ8 and n_s, and a cosmology scan done with ``evolve`` gives the same models +as one done with ``build``. To keep σ8 fixed through a scan, pass it with each +cosmology. The result depends only on the stage and the changes: a stage read back from +a configuration or a pickle evolves as the original does. If the new cosmology does not +give a parameter that is not passed (its ``meta`` has no ``sigma8``, as for a +cosmology made by hand), ``evolve`` raises the :class:`ValueError` ``build`` raises, +which names the parameter: set it explicitly. A name that is not a parameter raises a :class:`TypeError` at once, before anything is built, with the closest names (``mf.evolve(sigma8=0.8)``: "Did you mean @@ -515,15 +513,16 @@ Extension stages * the fields whose names are too generic to stand alone in ``qualified_parameters``, listed with the stage's name; * and defaults that depend on other parameters (as ``build`` takes σ8 and n_s from - the cosmology) in the class method ``computed_defaults``. It must compute them - from the parameters it is given alone, looking up each one it depends on: - ``evolve`` calls it again when one of those it looked up changes. + the cosmology) in the class method ``computed_defaults``, with the parameters + they are computed from in ``computed_defaults_inputs`` (dotted paths below the + stage), so that ``evolve`` computes them again when one of those changes. A derived property's input that is not a field of the tree raises a :class:`TypeError` when the table is first built. :func:`~hmf.core.stage.undeclared_inputs` checks, in a test, that each derived - property declares every field it reads (and every stage it holds): run on a stage - of the extension, it returns the fields each one reads but does not declare. + property declares every field it reads (and every stage it holds), and that + ``computed_defaults`` reads only its declared inputs: run on a stage of the + extension, it returns the fields each one reads but does not declare. Declaring an input that is read only on some branches is safe. Configurations @@ -599,8 +598,6 @@ tree has already used, as plain arrays. Unpickling rebuilds the tree and reuses runs, so it does not run CAMB or CLASS again, and gives bit-identical results. No CAMB or CLASS object is ever pickled. A stage is immutable, so ``copy.copy`` and ``copy.deepcopy`` return the stage itself, with its caches. -The pickle also records which parameters hold computed defaults (see "Computed -defaults" above), so an unpickled stage evolves as the original does. API --- diff --git a/src/hmf/core/config.py b/src/hmf/core/config.py index 8d4b0a76..0a74dd2e 100644 --- a/src/hmf/core/config.py +++ b/src/hmf/core/config.py @@ -104,13 +104,6 @@ sigma_8 = 0.8 # (at the top, before [hmf], in a file) fit = "ST" -A configuration does not record which values were computed defaults (see -:meth:`Stage.computed_parameters `): every -value it gives is explicit in the stage read from it, so a later -``evolve(cosmology=...)`` keeps it. Only a value a partial configuration leaves out, -and so takes a computed default, is computed again. A configuration written by -:func:`to_config` gives every value, so the stage read from it has none computed. - An unknown field raises a :class:`TypeError` with the closest names; an unknown model a :class:`~hmf.core.model.ModelNotFoundError` with suggestions; a wrong unit an :class:`astropy.units.UnitConversionError`. The error notes where in the @@ -154,9 +147,7 @@ runs its tree has already used, as plain arrays (see :func:`hmf.core._boltzmann.export_runs`). Unpickling rebuilds the tree from the configuration, and puts the runs in the in-process memo, so it does not run CAMB or -CLASS again; results are bit-identical. No CAMB or CLASS object is ever pickled. The -pickle also records which parameters hold computed defaults, which the configuration -does not, so an unpickled stage evolves as the original does. +CLASS again; results are bit-identical. No CAMB or CLASS object is ever pickled. :func:`copy.copy` and :func:`copy.deepcopy` of a stage return the stage itself (it is immutable), with everything it has computed. """ @@ -1433,45 +1424,12 @@ def walk(obj: Any) -> None: return runs -def _pickle_state( - stage: Stage, -) -> tuple[dict[str, Any], dict[str, dict[str, Any]], dict[str, tuple[list[str], list[str]]]]: - """What a stage pickles as: its configuration, its runs and its computed defaults. +def _pickle_state(stage: Stage) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]: + """What a stage pickles as: its configuration, and the runs its tree has used.""" + return to_config(stage), export_runs(_computed_runs(stage)) - The runs are those its tree has used; the computed defaults, which parameters hold - one (see :meth:`Stage.computed_parameters - `), which the configuration does not - record. - """ - records = { - ".".join(prefix): ( - sorted(".".join(p) for p in record.paths), - sorted(".".join(p) for p in record.reads), - ) - for prefix, record in routing.computed_records(stage).items() - } - return to_config(stage), export_runs(_computed_runs(stage)), records - -def _unpickle( - config: Mapping[str, Any], - runs: Mapping[str, Mapping[str, Any]], - computed: Mapping[str, tuple[list[str], list[str]]] | None = None, -) -> Stage: +def _unpickle(config: Mapping[str, Any], runs: Mapping[str, Mapping[str, Any]]) -> Stage: """Rebuild a pickled stage (see :meth:`Stage.__reduce__ `).""" import_runs(runs) - stage = from_config(config) - if computed: - - def split(name: str) -> tuple[str, ...]: - return tuple(name.split(".")) if name else () - - records = { - split(prefix): routing.ComputedRecord( - paths=frozenset(split(p) for p in paths), - reads=frozenset(split(r) for r in reads), - ) - for prefix, (paths, reads) in computed.items() - } - routing.with_records(stage, records) - return stage + return from_config(config) diff --git a/src/hmf/core/mass_function.py b/src/hmf/core/mass_function.py index 8dd35c95..b8afc0c6 100644 --- a/src/hmf/core/mass_function.py +++ b/src/hmf/core/mass_function.py @@ -137,7 +137,7 @@ import math from collections.abc import Mapping from functools import cached_property -from typing import Any +from typing import Any, ClassVar import astropy.units as u import attrs @@ -397,15 +397,17 @@ def build(cls, **parameters: Any) -> MassFunction: """ return cls.from_flat(parameters) + computed_defaults_inputs: ClassVar[tuple[str, ...]] = ("linear_power.transfer.cosmology",) + @classmethod def computed_defaults(cls, given: GivenParameters) -> Mapping[str, Any]: """The defaults of n_s and sigma_8: the cosmology's (see :meth:`build`). - They are computed again when :meth:`~hmf.core.stage.Stage.evolve` changes the - cosmology (or the transfer or linear-power stage), for those of them that were - not set explicitly, so that ``build().evolve(cosmology=c)`` equals - ``build(cosmology=c)``. One that was set, when building or in an ``evolve``, - keeps its value. + They are computed from the cosmology alone (its + :attr:`~hmf.core.stage.Stage.computed_defaults_inputs`), so + :meth:`~hmf.core.stage.Stage.evolve` computes them again whenever it changes the + cosmology (or a stage holding it): ``mf.evolve(cosmology=c)`` takes both from + ``c``, and ``mf.evolve(cosmology=c, sigma_8=0.8)`` takes only n_s from it. Parameters ---------- diff --git a/src/hmf/core/routing.py b/src/hmf/core/routing.py index 009e7b2d..ce4086c3 100644 --- a/src/hmf/core/routing.py +++ b/src/hmf/core/routing.py @@ -56,7 +56,7 @@ import difflib import types import typing -from collections.abc import Callable, Collection, Iterable, Iterator, Mapping +from collections.abc import Callable, Collection, Iterator, Mapping from functools import cached_property from typing import TYPE_CHECKING, Any, cast @@ -68,7 +68,6 @@ from .stage import Stage __all__ = [ - "ComputedRecord", "GivenParameters", "ParameterInfo", "Router", @@ -160,8 +159,9 @@ class GivenParameters(Mapping[str, Any]): A read-only mapping from dotted paths to values, passed to :meth:`Stage.computed_defaults `. Looking up a path below a stage given whole returns the value in that stage. It records - the paths looked up (:attr:`reads`), so that :meth:`Stage.evolve - ` knows which changes call for the defaults again. + the paths looked up (:attr:`reads`), so that + :func:`~hmf.core.stage.undeclared_inputs` can check a stage's + :attr:`~hmf.core.stage.Stage.computed_defaults_inputs`. Parameters ---------- @@ -184,10 +184,6 @@ def __getitem__(self, key: str) -> Any: """The value at a dotted path (a ``KeyError`` if it is not given).""" path = _split(key) self._reads.add(path) - return self._lookup(path, key) - - def _lookup(self, path: FieldPath, key: str) -> Any: - """The value at ``path`` (a ``KeyError`` naming ``key`` if it is not given).""" for i in range(len(path), 0, -1): if path[:i] in self._values: try: @@ -207,101 +203,6 @@ def __len__(self) -> int: return len(self._values) -class _TreeParameters(GivenParameters): - """The parameters of a stage tree being evolved, as :class:`GivenParameters`. - - What :meth:`Stage.computed_defaults ` sees - when :func:`evolve` computes its defaults again: every parameter of the new tree - (the changes, and the old stage's values of the rest) except those that are still - computed, as :func:`from_flat` would have been given them. A stage path is given - only if it is changed whole. - """ - - __slots__ = ("_computed", "_leaves", "_old") - - def __init__( - self, - old: Stage, - changes: Mapping[FieldPath, Any], - computed: frozenset[FieldPath], - ) -> None: - super().__init__(changes) - self._old = old - self._computed = computed - self._leaves = router(type(old)).leaves - - def _lookup(self, path: FieldPath, key: str) -> Any: - values, leaves = self._values, self._leaves - is_parameter = False - for i in range(1, len(path) + 1): - if path[:i] in values: - return super()._lookup(path, key) - is_parameter = is_parameter or path[:i] in leaves - if not is_parameter or path in self._computed: - raise KeyError(key) - try: - return _walk(self._old, path) - except AttributeError: - raise KeyError(key) from None - - def _paths(self) -> list[FieldPath]: - changed = [p for p in self._values if p not in self._leaves] - leaves = [p for p in self._leaves if p not in self._computed] - return [*changed, *leaves] - - def __iter__(self) -> Iterator[str]: - self._reads.add(()) - return (_dotted(p) for p in self._paths()) - - def __len__(self) -> int: - self._reads.add(()) - return len(self._paths()) - - -@attrs.frozen -class ComputedRecord: - """Which parameters below a stage hold a default computed by its ``computed_defaults``. - - A stage records it (see :meth:`Stage.computed_parameters - `) so that :meth:`Stage.evolve - ` computes those defaults again when a parameter - they are computed from changes. It is not part of the stage's value. - - Parameters - ---------- - paths - The paths, below the stage, of the parameters that took a computed default. - reads - The paths, below the stage, that ``computed_defaults`` looked up (``()``: it - read every parameter). - """ - - paths: frozenset[FieldPath] = frozenset() - reads: frozenset[FieldPath] = frozenset() - - def explicit(self, changed: Iterable[FieldPath]) -> ComputedRecord: - """The record with the parameters at or below the changed paths made explicit. - - Parameters - ---------- - changed - Paths, below the stage, that are set. - - Returns - ------- - ComputedRecord - """ - changed = tuple(changed) - dropped = [p for p in self.paths for c in changed if _overlaps(p, c)] - if not dropped: - return self - return ComputedRecord(paths=self.paths.difference(dropped), reads=self.reads) - - -#: The record of a stage with no computed defaults. -NO_COMPUTED = ComputedRecord() - - @attrs.frozen class _Node: """What evolve and from_flat need to know of one stage class.""" @@ -480,13 +381,22 @@ def __init__(self, cls: type[Stage]) -> None: from .stage import Stage - #: The stages of the tree whose class computes defaults (from the top). - self.computing: tuple[FieldPath, ...] = tuple( - prefix - for prefix, stage_cls in self.stages.items() - if getattr(getattr(stage_cls, "computed_defaults", None), "__func__", None) - is not getattr(Stage.computed_defaults, "__func__", None) - ) + base_hook = getattr(Stage.computed_defaults, "__func__", None) + #: The stages of the tree that compute defaults (from the top), and the paths + #: below each that its defaults are computed from. + self.computing: dict[FieldPath, tuple[FieldPath, ...]] = {} + for prefix, stage_cls in self.stages.items(): + hook = getattr(getattr(stage_cls, "computed_defaults", None), "__func__", None) + if hook is None or hook is base_hook: + continue + inputs = tuple(_split(p) for p in getattr(stage_cls, "computed_defaults_inputs", ())) + for p in inputs: + if not _is_field(stage_cls, p): + raise TypeError( + f"{stage_cls.__name__}.computed_defaults_inputs: {_dotted(p)!r} is " + "not a field of the stage tree." + ) + self.computing[prefix] = inputs # -- building -------------------------------------------------------------------- @@ -852,12 +762,9 @@ def evolve(stage: Stage, changes: Mapping[str, Any]) -> Stage: if resolved.models: values = _apply_models(table, resolved, lambda p: _walk(stage, p)) soft = resolved.soft - records: dict[FieldPath, ComputedRecord] = {} if table.computing and values: - values, soft, records = _recompute_defaults(table, stage, values, soft) - new = _rebuild(type(stage), stage, values, soft, ()) - _record(new, records, stage) - return new # type: ignore[no-any-return] + values, soft = _recompute_defaults(table, stage, values, soft) + return _rebuild(type(stage), stage, values, soft, ()) # type: ignore[no-any-return] def _recompute_defaults( @@ -865,126 +772,48 @@ def _recompute_defaults( stage: Stage, values: Mapping[FieldPath, Any], soft: frozenset[FieldPath], -) -> tuple[dict[FieldPath, Any], frozenset[FieldPath], dict[FieldPath, ComputedRecord]]: - """The changes of an evolve, with the computed defaults it changes. - - For each stage of the tree that computes defaults and is not changed whole, the - parameters that are changed become explicit, and if a change is one its - ``computed_defaults`` read, the defaults of the parameters still computed are - computed again from the new tree, and those that change are added. Stages are - visited from the top, as in :func:`_with_computed_defaults`. - - Returns the changes, the paths set by a shared name, and the new record of each - stage with a record that has a change below it, by its path. +) -> tuple[dict[FieldPath, Any], frozenset[FieldPath]]: + """The changes of an evolve, with the computed defaults they change. + + For each stage of the tree that computes defaults, is not changed whole, and has a + change to one of its :attr:`~hmf.core.stage.Stage.computed_defaults_inputs`, its + ``computed_defaults`` is called with the changes below it and the new values of its + inputs. Each default it gives for a parameter that is not changed (nor below a + changed stage) is added, if it differs from the old tree's value. Stages are + visited from the top, as in :func:`_with_computed_defaults`, so a default computed + above is seen below. + + Returns the changes and the paths set by a shared name. """ + from .stage import _same + out = dict(values) soft_out = set(soft) - records: dict[FieldPath, ComputedRecord] = {} - for prefix in table.computing: - if any(prefix[:i] in values for i in range(len(prefix) + 1)): - continue # given whole, or below a stage given whole + for prefix, inputs in table.computing.items(): + if not inputs or any(prefix[:i] in values for i in range(len(prefix) + 1)): + continue # nothing to recompute from, or given whole (or below a stage given whole) n = len(prefix) below = {p[n:]: v for p, v in out.items() if p[:n] == prefix} - if not below: - continue - try: - old = _walk(stage, prefix) - except AttributeError: # pragma: no cover - below a stage that is None + if not any(_overlaps(c, i) for c in below for i in inputs): continue + old = _walk(stage, prefix) if old is None: + continue # an optional stage that is absent: nothing to recompute + given = GivenParameters(below) + at = {p: v for p, v in below.items() if p not in inputs} + at.update({i: given.get(_dotted(i), _walk(old, i)) for i in inputs}) + computed = type(old).computed_defaults(GivenParameters(at)) + if not computed: continue - record: ComputedRecord = old._computed - if not record.paths: - continue - new_record = record.explicit(below) - if new_record.paths and any(_overlaps(r, c) for r in record.reads for c in below): - given = _TreeParameters(old, below, new_record.paths) - computed = type(old).computed_defaults(given) - extra = router(type(old)).resolve(computed) if computed else None - sub_table = router(type(old)) - for path in new_record.paths: - if extra is not None and path in extra.values: - value = extra.values[path] - if path in extra.soft: - soft_out.add(prefix + path) - else: - # As from_flat: a parameter with no computed default takes its own. - value = _default(sub_table.leaves[path], repr(sub_table.preferred[path])) - if not _same_value(value, _walk(old, path)): - out[prefix + path] = value - new_record = ComputedRecord(paths=new_record.paths, reads=given.reads) - records[prefix] = new_record - return out, frozenset(soft_out), records - - -def _same_value(a: Any, b: Any) -> bool: - """Whether two values are the same object or compare equal (False if eq fails).""" - if a is b: - return True - try: - return bool(a == b) - except Exception: # noqa: BLE001 - an eq that can't answer means "not the same" - return False - - -def _record( - stage: Stage, records: Mapping[FieldPath, ComputedRecord], old: Stage | None = None -) -> None: - """Set the computed-defaults record of the stages at the given paths of a new tree. - - Only for stages just built: a stage shared with the old tree ``old`` keeps its own, - and an optional stage left out (None) has none. - """ - for prefix, record in records.items(): - node = _walk(stage, prefix) - if node is None or (old is not None and node is _walk(old, prefix)): - continue - if node._computed is not record: - object.__setattr__(node, "_computed", record) - - -def computed_records(stage: Stage) -> dict[FieldPath, ComputedRecord]: - """The computed-defaults record of each stage of a tree that has one. - - Parameters - ---------- - stage - The stage at the root of the tree. - - Returns - ------- - dict - The non-empty records, by the path of their stage. - """ - out: dict[FieldPath, ComputedRecord] = {} - for prefix in router(type(stage)).computing: - try: - node = _walk(stage, prefix) - except AttributeError: # pragma: no cover - below a stage that is None - continue - if node is not None and node._computed.paths: - out[prefix] = node._computed - return out - - -def with_records(stage: Stage, records: Mapping[FieldPath, ComputedRecord]) -> Stage: - """Set the computed-defaults records of a tree that was just built (and return it). - - Parameters - ---------- - stage - The stage at the root of a tree that nothing else holds yet (e.g. one just - unpickled). - records - Records by the path of their stage, as :func:`computed_records` gives them. - - Returns - ------- - Stage - ``stage``. - """ - _record(stage, records) - return stage + new = router(type(old)).resolve(computed) + for path, value in new.values.items(): + if any(_overlaps(path, c) for c in below): + continue # given in this evolve + if not _same(value, _walk(old, path)): + out[prefix + path] = value + if path in new.soft: + soft_out.add(prefix + path) + return out, frozenset(soft_out) def from_flat(cls: type[Stage], mapping: Mapping[str, Any]) -> Stage: @@ -1002,7 +831,7 @@ def from_flat(cls: type[Stage], mapping: Mapping[str, Any]) -> Stage: Stage """ table = router(cls) - resolved, records = _with_computed_defaults(table, table.resolve(mapping)) + resolved = _with_computed_defaults(table, table.resolve(mapping)) def base(path: FieldPath) -> Any: return _default(table.leaves[path], repr(table.preferred[path])) @@ -1010,9 +839,7 @@ def base(path: FieldPath) -> Any: values = resolved.values if resolved.models: values = _apply_models(table, resolved, base) - stage = _rebuild(cls, None, values, resolved.soft, ()) - _record(stage, records) - return stage # type: ignore[no-any-return] + return _rebuild(cls, None, values, resolved.soft, ()) # type: ignore[no-any-return] def from_paths( @@ -1065,55 +892,40 @@ def from_paths( f"{cls.__name__}: {_dotted(path)!r} is given, and so is the stage " f"{_dotted(path[:i])!r} above it. Give one or the other." ) - resolved, records = _with_computed_defaults( + resolved = _with_computed_defaults( table, _Resolved(values=paths, soft=frozenset(soft & set(paths)), models={}) ) - stage = _rebuild(cls, None, resolved.values, resolved.soft, ()) - _record(stage, records) - return stage # type: ignore[no-any-return] + return _rebuild(cls, None, resolved.values, resolved.soft, ()) # type: ignore[no-any-return] -def _with_computed_defaults( - table: Router, resolved: _Resolved -) -> tuple[_Resolved, dict[FieldPath, ComputedRecord]]: +def _with_computed_defaults(table: Router, resolved: _Resolved) -> _Resolved: """Add the computed defaults of each stage of the tree that is not given whole. Each stage's :meth:`~hmf.core.stage.Stage.computed_defaults` sees the parameters given below it; stages are visited from the top, and a default is used only for a field that is not given (or set by a default from a stage above), and not below a given stage. - - Returns the values with the defaults added, and the record of the defaults each - stage computed (the parameters that took one, and the paths it read), by the - stage's path. """ values = dict(resolved.values) soft = set(resolved.soft) given = [*resolved.values, *resolved.models] - records: dict[FieldPath, ComputedRecord] = {} - for prefix in table.computing: + for prefix, stage_cls in table.stages.items(): if any(prefix[:i] in values for i in range(len(prefix) + 1)): continue - stage_cls = table.stages[prefix] n = len(prefix) below = {p[n:]: v for p, v in values.items() if p[:n] == prefix} - parameters = GivenParameters(below) - computed = stage_cls.computed_defaults(parameters) # type: ignore[attr-defined] + computed = stage_cls.computed_defaults(GivenParameters(below)) # type: ignore[attr-defined] if not computed: continue extra = router(stage_cls).resolve(computed) - used = set() for path, value in extra.values.items(): full = prefix + path if not any(_overlaps(full, g) for g in given): values[full] = value given.append(full) - used.add(path) if path in extra.soft: soft.add(full) - if used: - records[prefix] = ComputedRecord(paths=frozenset(used), reads=parameters.reads) - return _Resolved(values=values, soft=frozenset(soft), models=resolved.models), records + return _Resolved(values=values, soft=frozenset(soft), models=resolved.models) def parameter_defaults(cls: type[Stage]) -> dict[str, Any]: @@ -1131,7 +943,7 @@ def parameter_defaults(cls: type[Stage]) -> dict[str, Any]: :meth:`Stage.parameter_names`. """ table = router(cls) - computed, _ = _with_computed_defaults(table, _Resolved(values={}, soft=frozenset(), models={})) + computed = _with_computed_defaults(table, _Resolved(values={}, soft=frozenset(), models={})) out: dict[str, Any] = {} for name, info in table.parameter_info.items(): path = _split(info.path) diff --git a/src/hmf/core/stage.py b/src/hmf/core/stage.py index c5757077..4b99ed33 100644 --- a/src/hmf/core/stage.py +++ b/src/hmf/core/stage.py @@ -210,19 +210,11 @@ class Stage(Documented): #: The derived properties of the class (see :func:`derived_property`), by name, #: recorded when the class is created. _derived_properties: ClassVar[Mapping[str, DerivedProperty[Any]]] = MappingProxyType({}) - - _computed: routing.ComputedRecord = field( - default=routing.NO_COMPUTED, - init=False, - eq=False, - repr=False, - physics=False, - doc=( - "The parameters below the stage that hold a default computed by its " - ":meth:`computed_defaults` (see :meth:`computed_parameters`). Not part of " - "the stage's value." - ), - ) + #: The dotted paths, below the stage, of the parameters its + #: :meth:`computed_defaults` are computed from (e.g. the cosmology, for sigma_8 and + #: n_s). :meth:`evolve` computes the defaults again when one of them changes. A + #: parameter the hook only checks for (whether it is given) is not one of them. + computed_defaults_inputs: ClassVar[tuple[str, ...]] = () def evolve(self, **changes: Any) -> Self: """Return a copy of this stage tree with some parameters changed. @@ -236,12 +228,13 @@ def evolve(self, **changes: Any) -> Self: :meth:`evolve_own`; every other stage is shared, with its cached results, and so is every :func:`derived_property` whose inputs do not change. - A parameter that holds a default computed from other parameters (see - :meth:`computed_defaults`, e.g. a mass function's sigma_8 and n_s, from the - cosmology's ``meta``) is computed again when a parameter it is computed from - changes, so that ``build(**a).evolve(**b)`` equals ``build(**a, **b)``. A - parameter set explicitly, here or when the tree was built, keeps its value - (see :meth:`computed_parameters`). + A change to a parameter that a stage's defaults are computed from (its + :attr:`computed_defaults_inputs`, e.g. a mass function's cosmology) computes + them again: every parameter that :meth:`computed_defaults` gives a default for + takes the new default, unless it is given in the same call. So + ``mf.evolve(cosmology=Planck15)`` takes sigma_8 and n_s from Planck15, and + ``mf.evolve(cosmology=Planck15, sigma_8=0.8)`` keeps sigma_8 = 0.8. Other + changes compute no default again. Every name is checked before anything is built. The new stages are built, and so validated, by their constructors; if that fails, the error propagates @@ -263,9 +256,9 @@ def evolve(self, **changes: Any) -> Self: If a name is not a parameter (with suggestions), is ambiguous (with the dotted paths to choose from), is a derived property, or is given twice. ValueError - If a default that is computed again cannot be computed from the new - parameters (e.g. the new cosmology does not give sigma_8), or for an - invalid value. + If a default computed again cannot be computed from the new parameters + (e.g. the new cosmology does not give sigma_8, and sigma_8 is not given), + or for an invalid value. """ return routing.evolve(self, changes) # type: ignore[return-value] @@ -275,9 +268,7 @@ def evolve_own(self, **changes: Any) -> Self: No routing: each name is a field of this stage, and nothing else changes. It is the step :meth:`evolve` takes at each stage it rebuilds. The value of each :func:`derived_property` already computed is carried to the new stage if none - of its inputs changes; every other cached result is computed again. No - default is computed again: a field given here, and every parameter below it, - is explicit in the new stage (see :meth:`computed_parameters`). + of its inputs changes; every other cached result is computed again. Parameters ---------- @@ -305,12 +296,6 @@ def evolve_own(self, **changes: Any) -> Self: msg += f" Fields: {names}." raise TypeError(msg) new = attrs.evolve(self, **changes) - record = self._computed - if record.paths: - # A field given here is explicit, with every parameter below it. - if any(p[0] in changes for p in record.paths): - record = record.explicit((n,) for n in changes) - object.__setattr__(new, "_computed", record) for name, prop in self._derived_properties.items(): try: value = _computed_value(self, name) @@ -357,14 +342,14 @@ def computed_defaults(cls, given: routing.GivenParameters) -> Mapping[str, Any]: the cosmology); the base stage has none. Defaults for parameters (or below stages) that are given are ignored. - The stage records which parameters took a computed default, and which paths of - ``given`` the hook looked up (see :meth:`computed_parameters`). When - :meth:`evolve` changes one of those paths, it calls the hook again, with every - parameter of the new tree except the ones still computed, and the computed - parameters take the new defaults. A parameter set explicitly, in - :meth:`from_flat` or in :meth:`evolve`, is never computed again. So the hook - must compute its defaults from ``given`` alone, and look up (with ``in``, - ``[]`` or ``get``) every parameter it depends on. + A stage that defines it lists the parameters its defaults are computed from in + :attr:`computed_defaults_inputs`. :meth:`evolve` calls the hook again when one + of them changes, with the parameters changed in that call and the new values of + the inputs, and every parameter it gives a default for, and that the call does + not give, takes that default. So the hook must compute its defaults from its + inputs alone, and give a default for each parameter it computes whenever that + parameter is not given. :func:`undeclared_inputs` checks, in a test, that it + reads no other parameter. Parameters ---------- @@ -378,40 +363,6 @@ def computed_defaults(cls, given: routing.GivenParameters) -> Mapping[str, Any]: """ return {} - def computed_parameters(self) -> tuple[str, ...]: - """The parameters of the stage tree that hold a default computed from others. - - Those that took a default from a stage's :meth:`computed_defaults` when the - tree was built (by :meth:`from_flat`, or :meth:`build` where a stage has one) - and have not been set explicitly since. :meth:`evolve` computes them again when - a parameter they are computed from changes (see :meth:`computed_defaults`). - Which parameters they are is not part of the stage's value: it does not change - ``==``, the hash or the content hash, and a configuration does not record it - (a stage read from a configuration has every value it gives explicit), but a - pickle or a copy keeps it. - - Returns - ------- - tuple of str - Their dotted paths, sorted. - - Examples - -------- - >>> from hmf.core.mass_function import MassFunction - >>> mf = MassFunction.build(**{"transfer.model": "EH"}) - >>> mf.computed_parameters() - ('linear_power.sigma_8', 'linear_power.transfer.n_s') - >>> mf.evolve(sigma_8=0.8).computed_parameters() - ('linear_power.transfer.n_s',) - """ - return tuple( - sorted( - ".".join((*prefix, *path)) - for prefix, record in routing.computed_records(self).items() - for path in record.paths - ) - ) - def to_flat(self) -> dict[str, Any]: """The value of every parameter of the stage tree, by the name it is routed by. @@ -556,7 +507,7 @@ def _walk_path(obj: Any, dotted: str) -> Any: def _fresh(stage: Stage) -> Stage: """A copy of a stage tree with nothing computed: every stage of it is a new one.""" - changes: dict[str, Any] = {} + changes = {} for a in attrs.fields(type(stage)): value = getattr(stage, a.name) if a.init and isinstance(value, Stage): @@ -598,7 +549,7 @@ def _captured( def undeclared_inputs(stage: Stage) -> dict[str, tuple[str, ...]]: - """The fields each derived property of a stage reads, but does not declare as inputs. + """The fields each derived property (or computed default) reads but does not declare. A check for tests, of a stage's :func:`derived_property` declarations. Each derived property is evaluated on a fresh copy of the stage's tree (nothing computed), while @@ -610,17 +561,26 @@ def undeclared_inputs(stage: Stage) -> dict[str, tuple[str, ...]]: too much is safe, declaring too little is not, because :meth:`Stage.evolve_own` would carry a stale value. + It also checks the stage's :meth:`~Stage.computed_defaults`, under the key + ``"computed_defaults"``: the hook is called with no parameters given, and every + parameter it looks up must be one of its + :attr:`~Stage.computed_defaults_inputs` (or below one), or one it gives a default + for (checking whether that is given). Iterating the parameters reads them all + (``"*"``). An undeclared input would leave a default stale after an + :meth:`~Stage.evolve` that changes it. + Parameters ---------- stage - A stage whose derived properties to check (those of the stages it holds are - not checked: call it on them too). + A stage whose derived properties and computed defaults to check (those of the + stages it holds are not checked: call it on them too). Returns ------- dict - For each derived property that reads an undeclared field, the dotted paths - (below the stage) of those fields. Empty if every read is declared. + For each derived property (and ``"computed_defaults"``) that reads an + undeclared field, the dotted paths (below the stage) of those fields. Empty if + every read is declared. Examples -------- @@ -670,9 +630,27 @@ def tracing( missing = sorted(".".join(r) for r in reads if not any(r[: len(i)] == i for i in inputs)) if missing: out[name] = tuple(missing) + hook_missing = _undeclared_default_inputs(type(stage)) + if hook_missing: + out["computed_defaults"] = hook_missing return out +def _undeclared_default_inputs(cls: type[Stage]) -> tuple[str, ...]: + """The parameters a stage class's computed_defaults reads but does not declare.""" + inputs = routing.router(cls).computing.get(()) + if inputs is None: + return () + given = routing.GivenParameters({}) + outputs = set(routing.router(cls).resolve(cls.computed_defaults(given)).values) + if () in given.reads: + return ("*",) + covered = [*inputs, *outputs] + return tuple( + sorted(".".join(r) for r in given.reads if not any(r[: len(c)] == c for c in covered)) + ) + + def _cosmology_field(*, also: str = "") -> Any: """The ``cosmology`` field of a :class:`CosmologyStage`. diff --git a/tests/core/test_core_computed_defaults.py b/tests/core/test_core_computed_defaults.py index fe3a540b..bcdacc9d 100644 --- a/tests/core/test_core_computed_defaults.py +++ b/tests/core/test_core_computed_defaults.py @@ -1,8 +1,9 @@ -"""Tests of computed defaults under evolve: build(**a).evolve(**b) == build(**a, **b). +"""Tests of computed defaults under evolve. -A mass function's sigma_8 and n_s default to the cosmology's (its ``meta``). A -parameter that took such a default is computed again when ``evolve`` changes what it is -computed from; one set explicitly keeps its value. +A mass function's sigma_8 and n_s default to the cosmology's (its ``meta``). They are +computed from the cosmology alone (``computed_defaults_inputs``), so an ``evolve`` that +changes the cosmology takes both from the new one, unless they are given in the same +call; any other change leaves them as they are. """ import copy @@ -16,15 +17,15 @@ from hmf.core import config, field, routing from hmf.core.mass_function import MassFunction -from hmf.core.routing import ComputedRecord -from hmf.core.stage import Stage +from hmf.core.stage import Stage, undeclared_inputs from hmf.core.transfer import Transfer from hmf.core.units import Msun_h EH = {"transfer.model": "EH"} M = np.logspace(10, 15, 11) * Msun_h Z = np.array([0.0, 1.0])[:, None] -COMPUTED = ("linear_power.sigma_8", "linear_power.transfer.n_s") +#: The flat names of the parameters a mass function computes from the cosmology. +DEFAULTED = ("sigma_8", "n_s") #: A cosmology whose meta has neither sigma8 nor n. BARE = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.05) @@ -49,182 +50,167 @@ def _dndm(mf: MassFunction) -> np.ndarray: return mf.dndm(m=M[None, :], z=Z).value +def _sigma_8(mf: MassFunction) -> float: + return mf.linear_power.sigma_8 + + +def _n_s(mf: MassFunction) -> float: + return mf.linear_power.transfer.n_s + + +def _rebuilt(mf: MassFunction, **changes) -> MassFunction: + """``mf``'s parameters, without those computed from the cosmology, and ``changes``.""" + flat = {k: v for k, v in mf.to_flat().items() if k not in DEFAULTED} + return MassFunction.from_flat({**flat, **changes}) + + @pytest.mark.parametrize("cosmology", [Planck13, Planck15, WMAP7, WMAP9]) def test_evolve_cosmology_equals_build(default, cosmology): built = MassFunction.build(cosmology=cosmology, **EH) evolved = default.evolve(cosmology=cosmology) - flat = MassFunction.from_flat({**EH, "cosmology": cosmology}) # The defaults are the cosmology's own (astropy's realisations give them). - assert built.linear_power.sigma_8 == cosmology.meta["sigma8"] - assert built.linear_power.transfer.n_s == cosmology.meta["n"] + assert _sigma_8(built) == cosmology.meta["sigma8"] + assert _n_s(built) == cosmology.meta["n"] assert evolved == built - assert flat == built - assert evolved.computed_parameters() == built.computed_parameters() == COMPUTED + assert MassFunction.from_flat({**EH, "cosmology": cosmology}) == built assert np.array_equal(_dndm(evolved), _dndm(built)) - assert np.array_equal(_dndm(flat), _dndm(built)) - # And from a non-default cosmology, and back. - other = MassFunction.build(cosmology=Planck15, **EH).evolve(cosmology=cosmology) - assert other == built + # From any starting cosmology, and back. + assert MassFunction.build(cosmology=Planck15, **EH).evolve(cosmology=cosmology) == built assert evolved.evolve(cosmology=Planck18) == default -def test_evolve_transfer_stage_given_whole_equals_build(default): - transfer = Transfer(model="EH", cosmology=Planck15, n_s=0.95) - evolved = default.evolve(**{"linear_power.transfer": transfer}) - built = MassFunction.from_flat({"linear_power.transfer": transfer}) - assert evolved == built - assert evolved.linear_power.sigma_8 == Planck15.meta["sigma8"] - assert evolved.linear_power.transfer.n_s == 0.95 - # n_s, in the stage given whole, is explicit; sigma_8 is still computed. - assert evolved.computed_parameters() == built.computed_parameters() - assert evolved.computed_parameters() == ("linear_power.sigma_8",) - - @pytest.mark.parametrize( "start", [ - lambda: MassFunction.build(sigma_8=0.7, **EH), - lambda: MassFunction.build(**EH).evolve(sigma_8=0.7), - lambda: MassFunction.build(**EH).evolve(sigma_8=0.7).evolve(cosmology=WMAP9), + lambda: MassFunction.build(sigma_8=0.7, n_s=0.95, **EH), + lambda: MassFunction.build(**EH).evolve(sigma_8=0.7, n_s=0.95), + lambda: MassFunction.build(cosmology=WMAP9, **EH).evolve(sigma_8=0.7), ], ) -def test_explicit_sigma_8_survives_a_cosmology_change(start): - new = start().evolve(cosmology=Planck15) - assert new.linear_power.sigma_8 == 0.7 - # n_s was not set, so it is still the cosmology's. - assert new.linear_power.transfer.n_s == Planck15.meta["n"] - assert new.computed_parameters() == ("linear_power.transfer.n_s",) - assert new == MassFunction.build(sigma_8=0.7, cosmology=Planck15, **EH) +def test_a_cosmology_change_takes_the_new_cosmologys_defaults(start): + """However sigma_8 and n_s were set before, the new cosmology gives them.""" + mf = start() + new = mf.evolve(cosmology=Planck15) + assert _sigma_8(new) == Planck15.meta["sigma8"] + assert _n_s(new) == Planck15.meta["n"] + assert new == _rebuilt(mf, cosmology=Planck15) + assert new == MassFunction.build(cosmology=Planck15, **EH) + + +def test_values_given_with_the_cosmology_are_kept(default): + new = default.evolve(cosmology=Planck15, sigma_8=0.7) + assert _sigma_8(new) == 0.7 + assert _n_s(new) == Planck15.meta["n"] + assert new == MassFunction.build(cosmology=Planck15, sigma_8=0.7, **EH) + both = default.evolve(cosmology=Planck15, sigma_8=0.7, n_s=0.95) + assert (_sigma_8(both), _n_s(both)) == (0.7, 0.95) + + +def test_other_changes_compute_no_default(default): + """Only the cosmology's change recomputes: n_s, set before, survives a sigma_8 change.""" + mf = default.evolve(n_s=0.95) + lower = mf.evolve(sigma_8=0.8) + assert _n_s(lower) == 0.95 + assert lower.linear_power.transfer is mf.linear_power.transfer + assert _sigma_8(mf.evolve(n_s=0.9)) == _sigma_8(mf) + new = default.evolve(fit="ST") + assert new.linear_power is default.linear_power + + +def test_a_cosmology_with_the_same_defaults_rebuilds_only_the_cosmology(default): + # Same parameters and meta: an equal tree. + assert default.evolve(cosmology=Planck18.clone(H0=Planck18.H0)) == default + # The defaults are recomputed but equal the old ones: only the cosmology changes. + other = Planck18.clone(Om0=0.32) + new = default.evolve(cosmology=other) + assert (_sigma_8(new), _n_s(new)) == (_sigma_8(default), _n_s(default)) -def test_explicit_value_equal_to_the_default_is_kept(default): - # Set to the value it had: still explicit from then on. - pinned = default.evolve(sigma_8=default.linear_power.sigma_8) - assert pinned.evolve(cosmology=Planck15).linear_power.sigma_8 == Planck18.meta["sigma8"] +def test_a_transfer_stage_given_whole(default): + """A stage holding the cosmology is a change to it; n_s is the stage's own.""" + transfer = Transfer(model="EH", cosmology=Planck15, n_s=0.95) + evolved = default.evolve(**{"linear_power.transfer": transfer}) + assert evolved == MassFunction.from_flat({"linear_power.transfer": transfer}) + assert _sigma_8(evolved) == Planck15.meta["sigma8"] + assert _n_s(evolved) == 0.95 -def test_a_stage_given_whole_is_explicit(default): - lp = default.linear_power - replaced = default.evolve(linear_power=lp) - assert replaced.computed_parameters() == () - new = replaced.evolve(cosmology=Planck15) - assert new.linear_power.sigma_8 == Planck18.meta["sigma8"] - assert new.linear_power.transfer.n_s == Planck18.meta["n"] +def test_a_linear_power_given_whole_is_kept(default): + lp = default.evolve(sigma_8=0.7, cosmology=Planck15).linear_power + new = default.evolve(linear_power=lp) + assert new.linear_power is lp + assert _sigma_8(new) == 0.7 -def test_evolve_without_a_source_raises_when_computed(default): +def test_evolve_without_a_source_raises(default): with pytest.raises(ValueError, match=r"'FlatLambdaCDM' does not give n_s.*set n_s explicitly"): default.evolve(cosmology=BARE) with pytest.raises(ValueError, match=r"does not give sigma_8.*set sigma_8 explicitly"): - default.evolve(n_s=0.96).evolve(cosmology=BARE) - # build raises the same. - with pytest.raises(ValueError, match=r"does not give sigma_8.*set sigma_8 explicitly"): - MassFunction.build(cosmology=BARE, n_s=0.96, **EH) + default.evolve(cosmology=BARE, n_s=0.96) + # Values set before do not help: the new cosmology gives the defaults. + with pytest.raises(ValueError, match="does not give n_s"): + default.evolve(sigma_8=0.8, n_s=0.96).evolve(cosmology=BARE) named = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.05, name="mine") with pytest.raises(ValueError, match="the cosmology 'mine' does not give n_s"): default.evolve(cosmology=named) - # The original is unchanged. - assert default.linear_power.sigma_8 == Planck18.meta["sigma8"] + # build raises the same, and the original is unchanged. + with pytest.raises(ValueError, match=r"does not give sigma_8.*set sigma_8 explicitly"): + MassFunction.build(cosmology=BARE, n_s=0.96, **EH) + assert _sigma_8(default) == Planck18.meta["sigma8"] -def test_evolve_without_a_source_works_when_explicit(): - mf = MassFunction.build(sigma_8=0.8, n_s=0.96, **EH) - new = mf.evolve(cosmology=BARE) - assert new.linear_power.sigma_8 == 0.8 - assert new.linear_power.transfer.n_s == 0.96 +def test_evolve_without_a_source_works_when_given(): + mf = MassFunction.build(**EH) + new = mf.evolve(cosmology=BARE, sigma_8=0.8, n_s=0.96) assert new == MassFunction.build(cosmology=BARE, sigma_8=0.8, n_s=0.96, **EH) + # And from it, any other change computes nothing (BARE gives no defaults). + assert _sigma_8(new.evolve(fit="ST")) == 0.8 + # Back to a cosmology with defaults. + assert new.evolve(cosmology=Planck15) == MassFunction.build(cosmology=Planck15, **EH) -def test_provenance_is_not_part_of_the_value(default): - explicit = MassFunction.build( - sigma_8=default.linear_power.sigma_8, n_s=default.linear_power.transfer.n_s, **EH - ) - assert explicit.computed_parameters() == () - assert default.computed_parameters() == COMPUTED - assert explicit == default - assert hash(explicit) == hash(default) - assert config.content_hash(explicit) == config.content_hash(default) - assert config.to_config(explicit) == config.to_config(default) - assert "_computed" not in repr(default) - assert "_computed" not in MassFunction.parameter_names() - - -@pytest.mark.parametrize("copier", [lambda s: pickle.loads(pickle.dumps(s)), copy.deepcopy]) -def test_pickling_keeps_provenance(copier): +@pytest.mark.parametrize( + "copier", + [ + lambda s: pickle.loads(pickle.dumps(s)), + copy.deepcopy, + lambda s: config.from_config(config.to_config(s)), + ], +) +def test_evolve_depends_only_on_the_value(copier): + """No hidden state: a copy, a pickle or a configuration evolves as the original.""" mf = MassFunction.build(sigma_8=0.7, **EH) same = copier(mf) assert same == mf - assert same.computed_parameters() == mf.computed_parameters() - assert same.linear_power._computed == mf.linear_power._computed - assert same._computed == mf._computed assert same.evolve(cosmology=Planck15) == mf.evolve(cosmology=Planck15) - - -def test_a_configuration_records_no_provenance(default): - read = config.from_config(config.to_config(default)) - assert read == default - assert read.computed_parameters() == () - # So every value it gives is kept. - assert read.evolve(cosmology=Planck15).linear_power.sigma_8 == Planck18.meta["sigma8"] - # A value a partial configuration leaves out takes a computed default, and stays one. - partial = config.from_config( - {"hmf": {"schema_version": 1}, "sigma_8": 0.8, **EH}, stage=MassFunction - ) - assert partial.computed_parameters() == ("linear_power.transfer.n_s",) - - -def test_evolve_own_makes_the_fields_it_sets_explicit(default): - assert default.evolve_own(delta_c=1.7).computed_parameters() == COMPUTED - assert default.evolve_own(linear_power=default.linear_power).computed_parameters() == () - - -def test_unrelated_changes_share_the_stages(default): - new = default.evolve(fit="ST") - assert new.linear_power is default.linear_power - assert new.computed_parameters() == COMPUTED - lower = default.evolve(sigma_8=0.7) - assert lower.linear_power.transfer is default.linear_power.transfer - assert lower.computed_parameters() == ("linear_power.transfer.n_s",) - # A cosmology change with the same sigma_8 and n_s rebuilds only what it must. - same_meta = Planck18.clone(H0=Planck18.H0) - assert default.evolve(cosmology=same_meta) == default + assert same.evolve(sigma_8=0.8) == mf.evolve(sigma_8=0.8) def test_extension_stage_recomputes_the_mass_functions_defaults(default): built = HaloModel.from_flat({**EH, "cosmology": Planck15}) assert built.mass_function == MassFunction.build(cosmology=Planck15, **EH) - evolved = HaloModel.from_flat(EH).evolve(cosmology=Planck15) - assert evolved == built - # A mass function given whole brings its own record. - held = HaloModel(mass_function=default).evolve(cosmology=Planck15) - assert held == built - assert held.computed_parameters() == ( - "mass_function.linear_power.sigma_8", - "mass_function.linear_power.transfer.n_s", - ) - explicit = HaloModel(mass_function=default.evolve(sigma_8=0.7)).evolve(cosmology=Planck15) - assert explicit.mass_function.linear_power.sigma_8 == 0.7 - - -def test_computed_record(): - record = ComputedRecord(paths=frozenset({("a", "b"), ("c",)}), reads=frozenset({("d",)})) - assert record.explicit([("a",)]).paths == {("c",)} - assert record.explicit([("e",)]) is record - assert record.explicit([("c",)]).reads == record.reads + assert HaloModel.from_flat(EH).evolve(cosmology=Planck15) == built + held = HaloModel(mass_function=default.evolve(sigma_8=0.7)) + assert held.evolve(cosmology=Planck15) == built + assert _sigma_8(held.evolve(bias="SMT").mass_function) == 0.7 def test_extension_stage_with_a_mass_function_given_whole(default): - explicit = default.evolve(sigma_8=0.7) - built = HaloModel.from_flat({"mass_function": default}) - assert built.computed_parameters() == tuple(f"mass_function.{p}" for p in COMPUTED) - # Replacing the stage whole brings the new stage's own record, untouched. - replaced = built.evolve(mass_function=explicit) + explicit = default.evolve(sigma_8=0.7, cosmology=Planck15) + hm = HaloModel.from_flat({"mass_function": default}) + assert hm.mass_function is default + # Replacing the stage whole computes nothing in it. + replaced = hm.evolve(mass_function=explicit) assert replaced.mass_function is explicit - assert replaced.computed_parameters() == ("mass_function.linear_power.transfer.n_s",) + + +def test_mass_function_declares_its_inputs(default): + assert MassFunction.computed_defaults_inputs == ("linear_power.transfer.cosmology",) + assert undeclared_inputs(default) == {} # --------------------------------------------------------------------------------- -# Toy trees: shared and optional stages with computed defaults +# Toy trees: shared and optional stages with computed defaults, and bad declarations # --------------------------------------------------------------------------------- @@ -240,6 +226,7 @@ class Pair(Stage): """A toy stage sharing its parts' setting, whose default is computed from ``base``.""" shared_parameters: ClassVar = frozenset({"acc"}) + computed_defaults_inputs: ClassVar = ("base",) left: Part = field(factory=Part, doc="One.") right: Part = field(factory=Part, doc="Two.") @@ -264,47 +251,63 @@ def test_toy_shared_computed_default_is_recomputed(): new = built.evolve(base=4) assert new.left.acc == new.right.acc == 8 assert new == Pair.from_flat({"base": 4}) - assert new.computed_parameters() == ("left.acc", "right.acc") - # Set by its shared name: explicit for both holders. - pinned = built.evolve(acc=5) - assert pinned.computed_parameters() == () - assert pinned.evolve(base=4).left.acc == 5 + # Set by its shared name, then recomputed by a change of base. + assert built.evolve(acc=5).left.acc == 5 + assert built.evolve(acc=5).evolve(base=4) == new + assert built.evolve(acc=5, base=4).right.acc == 5 + assert undeclared_inputs(built) == {} -def test_toy_optional_stage_left_out_has_no_record(): +def test_toy_optional_stage(): top = Top.from_flat({}) assert top.pair is None - assert top.computed_parameters() == () - # A shared name that only reaches the absent stage changes nothing. - assert top.evolve(acc=5).pair is None + assert top.evolve(acc=5).pair is None # a shared name reaching only an absent stage + with pytest.raises(ValueError, match="pair is None"): + top.evolve(**{"pair.base": 5}) # an input of an absent stage: nothing to recompute with_pair = Top.from_flat({"pair.base": 2}) assert with_pair.pair.left.acc == 4 assert with_pair.evolve(**{"pair.base": 5}).pair == Pair.from_flat({"base": 5}) -def test_tree_parameters_view_of_an_evolving_tree(default): - computed = frozenset({("linear_power", "sigma_8")}) - given = routing._TreeParameters(default, {("delta_c",): 1.7}, computed) - assert given["delta_c"] == 1.7 - assert given["linear_power.transfer.n_s"] == default.linear_power.transfer.n_s - assert given["fit.A_200"] == default.fit.A_200 # into a model field - assert "linear_power.sigma_8" not in given # still computed - assert "linear_power" not in given # a stage, not changed whole - assert "fit.nope" not in given - names = list(given) - assert "delta_c" in names - assert "linear_power.sigma_8" not in names - assert len(given) == len(names) - assert () in given.reads +def test_undeclared_inputs_of_computed_defaults(): + @attrs.frozen(kw_only=True) + class Undeclared(Stage): + a: int = field(default=1, doc="Read, not declared.") + b: int = field(default=2, doc="Computed.") + + @classmethod + def computed_defaults(cls, given): + return {} if "b" in given else {"b": 2 * given.get("a", 1)} + + @attrs.frozen(kw_only=True) + class Iterates(Stage): + a: int = field(default=1, doc="A parameter.") + + @classmethod + def computed_defaults(cls, given): + return {"a": len(given)} + + assert undeclared_inputs(Undeclared()) == {"computed_defaults": ("a",)} + assert undeclared_inputs(Iterates()) == {"computed_defaults": ("*",)} + # Without declared inputs, evolve recomputes nothing. + assert Undeclared.from_flat({}).evolve(a=5).b == 2 + + +def test_declared_input_must_be_a_field(): + @attrs.frozen(kw_only=True) + class BadInput(Stage): + a: int = field(default=1, doc="A parameter.") + computed_defaults_inputs: ClassVar = ("nope",) + @classmethod + def computed_defaults(cls, given): + return {} -def test_same_value_when_eq_cannot_answer(): - a = np.array([1.0, 2.0]) - assert routing._same_value(a, a) - assert not routing._same_value(a, np.array([1.0, 3.0])) - assert routing._same_value(1, 1.0) + with pytest.raises(TypeError, match=r"BadInput.computed_defaults_inputs: 'nope'"): + BadInput.parameter_names() def test_base_stage_computes_no_defaults(): assert Stage.computed_defaults(routing.GivenParameters({})) == {} - assert routing.router(Part).computing == () + assert routing.router(Part).computing == {} + assert undeclared_inputs(Part()) == {} diff --git a/tests/core/test_core_routing.py b/tests/core/test_core_routing.py index 7cf08b7f..cc731e2c 100644 --- a/tests/core/test_core_routing.py +++ b/tests/core/test_core_routing.py @@ -137,7 +137,7 @@ def test_every_field_is_routable(): for path, cls in table.stages.items(): for a in attrs.fields(cls): full = (*path, a.name) - if a.init and full not in table.stages: + if full not in table.stages: fields.add(".".join(full)) assert fields == routed assert set(table.stages) == {(), ("linear_power",), ("linear_power", "transfer")} @@ -200,15 +200,16 @@ def test_changes_cover_every_parameter(): assert set(CHANGES) == set(MassFunction.parameter_names()) -def _explicit_flat(mf): - """``mf.to_flat()`` without the parameters that hold computed defaults. +def _rebuilt_flat(mf, name): + """``mf.to_flat()``, for ``from_flat({**_rebuilt_flat(mf, name), name: value})``. - So that ``from_flat({**_explicit_flat(mf), name: value})`` is ``mf.evolve(name=value)``: - the computed defaults are computed again, from the new value. + That is ``mf.evolve(name=value)``: a change of the cosmology computes sigma_8 and + n_s again, from the new one, so they are left out for it. """ - computed = set(mf.computed_parameters()) - info = MassFunction.parameter_info() - return {k: v for k, v in mf.to_flat().items() if info[k].path not in computed} + flat = mf.to_flat() + if name == "cosmology": + del flat["sigma_8"], flat["n_s"] + return flat def _two_species(mf): @@ -238,7 +239,7 @@ def test_evolve_every_parameter(mf, tmp_path, name): got, old = getattr(got, part), getattr(old, part) assert got is not old, path # A fresh tree, with every other parameter the same. - fresh = MassFunction.from_flat({**_explicit_flat(mf), name: value}) + fresh = MassFunction.from_flat({**_rebuilt_flat(mf, name), name: value}) assert new == fresh if quantity is not None: method, z = quantity @@ -312,7 +313,7 @@ def test_cosmology_change_is_consistent(mf): assert stage.cosmology is Planck15 assert lp.growth.transfer is lp.transfer assert new.variance.power.transfer is lp.transfer - assert new == MassFunction.from_flat({**_explicit_flat(mf), "cosmology": Planck15}) + assert new == MassFunction.from_flat({**_rebuilt_flat(mf, "cosmology"), "cosmology": Planck15}) assert new == MassFunction.build(**{"transfer.model": "EH"}, cosmology=Planck15) @@ -871,10 +872,7 @@ def test_toy_shared_parameter_and_derived_properties(): double = built.double assert built.evolve(**{"right.count": 2}).double is double assert built.evolve(acc=3).double == 18 - # base is what acc's default is computed from: given, acc takes its own default. - assert built.evolve(base=4) == Holder.from_flat({"base": 4}) - assert built.evolve(base=4).double == 8 - assert built.evolve(acc=4).evolve(base=4).double == 32 # acc set: kept + assert built.evolve(base=4).double == 16 # no declared inputs: acc is kept assert Holder.invalidated_by("acc") == ("", "double", "left", "right") assert Holder.invalidated_by("base") == ("", "double", "unannotated", "unresolved") # Without the composing stage's declaration, the name is ambiguous. From 5e155f151b679b6c2ed788a5c85e0338993f08d2 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 8 Nov 2026 15:27:33 +0000 Subject: [PATCH 4/4] feat(core): require computed_defaults_inputs on a stage with computed_defaults A stage that defines computed_defaults without declaring what its defaults are computed from would never have them recomputed by evolve, which is the bug this PR fixes. Building its routing table now raises a TypeError that asks for the inputs. No stage outside the tests defines the hook except MassFunction, which declares its cosmology. The toy Holder stage in the routing tests declares base, with a hook computed from base's value, so evolve(base=4) now equals from_flat({"base": 4}). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_019U8u44xFH87Am2mNHVrimP --- docs/core.rst | 8 +++++--- src/hmf/core/routing.py | 11 +++++++++-- src/hmf/core/stage.py | 9 ++++++--- tests/core/test_core_computed_defaults.py | 20 +++++++++++++++++++- tests/core/test_core_routing.py | 8 ++++++-- 5 files changed, 45 insertions(+), 11 deletions(-) diff --git a/docs/core.rst b/docs/core.rst index 09694cc2..3f8178cd 100644 --- a/docs/core.rst +++ b/docs/core.rst @@ -515,10 +515,12 @@ Extension stages * and defaults that depend on other parameters (as ``build`` takes σ8 and n_s from the cosmology) in the class method ``computed_defaults``, with the parameters they are computed from in ``computed_defaults_inputs`` (dotted paths below the - stage), so that ``evolve`` computes them again when one of those changes. + stage), so that ``evolve`` computes them again when one of those changes. The + inputs are required: a stage that defines ``computed_defaults`` without them + raises a :class:`TypeError` when its routing table is first built. - A derived property's input that is not a field of the tree raises a - :class:`TypeError` when the table is first built. + A derived property's input, or a computed default's, that is not a field of the + tree raises a :class:`TypeError` when the table is first built. :func:`~hmf.core.stage.undeclared_inputs` checks, in a test, that each derived property declares every field it reads (and every stage it holds), and that ``computed_defaults`` reads only its declared inputs: run on a stage of the diff --git a/src/hmf/core/routing.py b/src/hmf/core/routing.py index ce4086c3..0013f241 100644 --- a/src/hmf/core/routing.py +++ b/src/hmf/core/routing.py @@ -390,6 +390,13 @@ def __init__(self, cls: type[Stage]) -> None: if hook is None or hook is base_hook: continue inputs = tuple(_split(p) for p in getattr(stage_cls, "computed_defaults_inputs", ())) + if not inputs: + raise TypeError( + f"{stage_cls.__name__} defines computed_defaults but no " + "computed_defaults_inputs: list the parameters its defaults are computed " + "from (dotted paths below the stage), so that evolve computes them again " + "when one of them changes." + ) for p in inputs: if not _is_field(stage_cls, p): raise TypeError( @@ -790,8 +797,8 @@ def _recompute_defaults( out = dict(values) soft_out = set(soft) for prefix, inputs in table.computing.items(): - if not inputs or any(prefix[:i] in values for i in range(len(prefix) + 1)): - continue # nothing to recompute from, or given whole (or below a stage given whole) + if any(prefix[:i] in values for i in range(len(prefix) + 1)): + continue # given whole, or below a stage given whole n = len(prefix) below = {p[n:]: v for p, v in out.items() if p[:n] == prefix} if not any(_overlaps(c, i) for c in below for i in inputs): diff --git a/src/hmf/core/stage.py b/src/hmf/core/stage.py index 4b99ed33..10b626db 100644 --- a/src/hmf/core/stage.py +++ b/src/hmf/core/stage.py @@ -213,7 +213,9 @@ class Stage(Documented): #: The dotted paths, below the stage, of the parameters its #: :meth:`computed_defaults` are computed from (e.g. the cosmology, for sigma_8 and #: n_s). :meth:`evolve` computes the defaults again when one of them changes. A - #: parameter the hook only checks for (whether it is given) is not one of them. + #: parameter the hook only checks for (whether it is given) is not one of them. A + #: stage that defines :meth:`computed_defaults` must declare at least one: its + #: routing table raises a :class:`TypeError` otherwise. computed_defaults_inputs: ClassVar[tuple[str, ...]] = () def evolve(self, **changes: Any) -> Self: @@ -342,8 +344,9 @@ def computed_defaults(cls, given: routing.GivenParameters) -> Mapping[str, Any]: the cosmology); the base stage has none. Defaults for parameters (or below stages) that are given are ignored. - A stage that defines it lists the parameters its defaults are computed from in - :attr:`computed_defaults_inputs`. :meth:`evolve` calls the hook again when one + A stage that defines it must list the parameters its defaults are computed from + in :attr:`computed_defaults_inputs` (a :class:`TypeError` is raised when its + routing table is first built if it lists none). :meth:`evolve` calls the hook again when one of them changes, with the parameters changed in that call and the new values of the inputs, and every parameter it gives a default for, and that the call does not give, takes that default. So the hook must compute its defaults from its diff --git a/tests/core/test_core_computed_defaults.py b/tests/core/test_core_computed_defaults.py index bcdacc9d..69850a43 100644 --- a/tests/core/test_core_computed_defaults.py +++ b/tests/core/test_core_computed_defaults.py @@ -274,6 +274,8 @@ def test_undeclared_inputs_of_computed_defaults(): class Undeclared(Stage): a: int = field(default=1, doc="Read, not declared.") b: int = field(default=2, doc="Computed.") + c: int = field(default=0, doc="Declared, not read.") + computed_defaults_inputs: ClassVar = ("c",) @classmethod def computed_defaults(cls, given): @@ -282,6 +284,7 @@ def computed_defaults(cls, given): @attrs.frozen(kw_only=True) class Iterates(Stage): a: int = field(default=1, doc="A parameter.") + computed_defaults_inputs: ClassVar = ("a",) @classmethod def computed_defaults(cls, given): @@ -289,10 +292,25 @@ def computed_defaults(cls, given): assert undeclared_inputs(Undeclared()) == {"computed_defaults": ("a",)} assert undeclared_inputs(Iterates()) == {"computed_defaults": ("*",)} - # Without declared inputs, evolve recomputes nothing. + # An undeclared input leaves the default stale: what the check is for. assert Undeclared.from_flat({}).evolve(a=5).b == 2 +def test_computed_defaults_need_declared_inputs(): + @attrs.frozen(kw_only=True) + class NoInputs(Stage): + a: int = field(default=1, doc="A parameter.") + + @classmethod + def computed_defaults(cls, given): + return {"a": 2} + + with pytest.raises(TypeError, match="NoInputs defines computed_defaults but no"): + NoInputs.from_flat({}) + with pytest.raises(TypeError, match="NoInputs defines computed_defaults but no"): + NoInputs().evolve(a=3) + + def test_declared_input_must_be_a_field(): @attrs.frozen(kw_only=True) class BadInput(Stage): diff --git a/tests/core/test_core_routing.py b/tests/core/test_core_routing.py index cc731e2c..5dd07599 100644 --- a/tests/core/test_core_routing.py +++ b/tests/core/test_core_routing.py @@ -837,6 +837,7 @@ class Holder(Stage): """A toy stage composing two Leafy stages, sharing their setting, with a computed default.""" shared_parameters: ClassVar = frozenset({"acc"}) + computed_defaults_inputs: ClassVar = ("base",) left: Leafy = field(factory=Leafy, doc="One.") right: Leafy = field(factory=Leafy, doc="Two.") @@ -859,7 +860,7 @@ def unresolved(self) -> "NotAType": # noqa: F821 - an annotation that can't be @classmethod def computed_defaults(cls, given): - return {"acc": 2} if "base" not in given else {} + return {"acc": 2 if given.get("base", 3) == 3 else 1} def test_toy_shared_parameter_and_derived_properties(): @@ -872,7 +873,10 @@ def test_toy_shared_parameter_and_derived_properties(): double = built.double assert built.evolve(**{"right.count": 2}).double is double assert built.evolve(acc=3).double == 18 - assert built.evolve(base=4).double == 16 # no declared inputs: acc is kept + # acc's default is computed from base: a change of base computes it again. + assert built.evolve(base=4) == Holder.from_flat({"base": 4}) + assert built.evolve(base=4).double == 8 + assert built.evolve(base=4, acc=3).double == 24 assert Holder.invalidated_by("acc") == ("", "double", "left", "right") assert Holder.invalidated_by("base") == ("", "double", "unannotated", "unresolved") # Without the composing stage's declaration, the name is ambiguous.