fix(core): evolve recomputes defaults computed from the cosmology

docs/core.rst

@@ -443,11 +443,33 @@ 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, 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`` (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
 'sigma_8'?"); an ambiguous one lists the dotted paths to choose from.
@@ -491,13 +513,18 @@ 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``.
-
-    A derived property's input that is not a field of the tree raises a
-    :class:`TypeError` when the table is first built.
+      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. 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, 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): 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

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
@@ -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)
 
@@ -395,10 +397,18 @@ 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 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
         ----------
         given
@@ -412,7 +422,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 = {}

src/hmf/core/routing.py

@@ -158,22 +158,32 @@ class GivenParameters(Mapping[str, Any]):
 
     A read-only mapping from dotted paths to values, passed to
     :meth:`Stage.computed_defaults <hmf.core.stage.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
+    :func:`~hmf.core.stage.undeclared_inputs` can check a stage's
+    :attr:`~hmf.core.stage.Stage.computed_defaults_inputs`.
 
     Parameters
     ----------
     values
         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)
         for i in range(len(path), 0, -1):
             if path[:i] in self._values:
                 try:
@@ -184,10 +194,12 @@ 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)
 
 
@@ -367,6 +379,32 @@ def __init__(self, cls: type[Stage]) -> None:
 
         self._candidates = sorted(set(self.names) | set(self.ambiguous))
 
+        from .stage import Stage
+
+        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", ()))
+            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(
+                        f"{stage_cls.__name__}.computed_defaults_inputs: {_dotted(p)!r} is "
+                        "not a field of the stage tree."
+                    )
+            self.computing[prefix] = inputs
+
     # -- building --------------------------------------------------------------------
 
     def _visit(self, cls: type, prefix: FieldPath, sharers: dict[str, list[FieldPath]]) -> None:
@@ -730,7 +768,59 @@ 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
+    if table.computing and values:
+        values, soft = _recompute_defaults(table, stage, values, soft)
+    return _rebuild(type(stage), stage, values, soft, ())  # 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]]:
+    """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)
+    for prefix, inputs in table.computing.items():
+        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):
+            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
+        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:

src/hmf/core/stage.py

@@ -210,6 +210,13 @@ 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({})
+    #: 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. 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:
         """Return a copy of this stage tree with some parameters changed.
@@ -223,6 +230,14 @@ 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 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
         and this stage, which is immutable, is unchanged.
@@ -242,6 +257,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 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]
 
@@ -325,6 +344,16 @@ 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 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
+        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
         ----------
         given
@@ -523,7 +552,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
@@ -535,17 +564,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
     --------
@@ -595,9 +633,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`.