diff --git a/docs/core.rst b/docs/core.rst index bfc7c9f1..3f8178cd 100644 --- a/docs/core.rst +++ b/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 diff --git a/src/hmf/core/mass_function.py b/src/hmf/core/mass_function.py index 261698c7..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 @@ -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 = {} diff --git a/src/hmf/core/routing.py b/src/hmf/core/routing.py index 87dfb689..0013f241 100644 --- a/src/hmf/core/routing.py +++ b/src/hmf/core/routing.py @@ -158,7 +158,10 @@ 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 + :func:`~hmf.core.stage.undeclared_inputs` can check a stage's + :attr:`~hmf.core.stage.Stage.computed_defaults_inputs`. Parameters ---------- @@ -166,14 +169,21 @@ 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) 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: diff --git a/src/hmf/core/stage.py b/src/hmf/core/stage.py index c2f9d4a5..10b626db 100644 --- a/src/hmf/core/stage.py +++ b/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`. diff --git a/tests/core/test_core_computed_defaults.py b/tests/core/test_core_computed_defaults.py new file mode 100644 index 00000000..69850a43 --- /dev/null +++ b/tests/core/test_core_computed_defaults.py @@ -0,0 +1,331 @@ +"""Tests of computed defaults under evolve. + +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 +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, routing +from hmf.core.mass_function import MassFunction +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] +#: 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) + + +@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 + + +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) + # The defaults are the cosmology's own (astropy's realisations give them). + assert _sigma_8(built) == cosmology.meta["sigma8"] + assert _n_s(built) == cosmology.meta["n"] + assert evolved == built + assert MassFunction.from_flat({**EH, "cosmology": cosmology}) == built + assert np.array_equal(_dndm(evolved), _dndm(built)) + # From any starting cosmology, and back. + assert MassFunction.build(cosmology=Planck15, **EH).evolve(cosmology=cosmology) == built + assert evolved.evolve(cosmology=Planck18) == default + + +@pytest.mark.parametrize( + "start", + [ + 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_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_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_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(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(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) + # 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_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) + + +@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.evolve(cosmology=Planck15) == mf.evolve(cosmology=Planck15) + 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) + 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, 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 + + +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, and bad declarations +# --------------------------------------------------------------------------------- + + +@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"}) + computed_defaults_inputs: ClassVar = ("base",) + + 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}) + # 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(): + top = Top.from_flat({}) + assert top.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_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.") + c: int = field(default=0, doc="Declared, not read.") + computed_defaults_inputs: ClassVar = ("c",) + + @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.") + computed_defaults_inputs: ClassVar = ("a",) + + @classmethod + def computed_defaults(cls, given): + return {"a": len(given)} + + assert undeclared_inputs(Undeclared()) == {"computed_defaults": ("a",)} + assert undeclared_inputs(Iterates()) == {"computed_defaults": ("*",)} + # 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): + a: int = field(default=1, doc="A parameter.") + computed_defaults_inputs: ClassVar = ("nope",) + + @classmethod + def computed_defaults(cls, given): + return {} + + 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 undeclared_inputs(Part()) == {} diff --git a/tests/core/test_core_routing.py b/tests/core/test_core_routing.py index cc52f0b2..5dd07599 100644 --- a/tests/core/test_core_routing.py +++ b/tests/core/test_core_routing.py @@ -200,6 +200,18 @@ def test_changes_cover_every_parameter(): assert set(CHANGES) == set(MassFunction.parameter_names()) +def _rebuilt_flat(mf, name): + """``mf.to_flat()``, for ``from_flat({**_rebuilt_flat(mf, name), name: 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. + """ + flat = mf.to_flat() + if name == "cosmology": + del flat["sigma_8"], flat["n_s"] + return flat + + def _two_species(mf): """``mf`` with a tabulated transfer function whose total matter is suppressed at small scales. @@ -227,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({**mf.to_flat(), name: value}) + fresh = MassFunction.from_flat({**_rebuilt_flat(mf, name), name: value}) assert new == fresh if quantity is not None: method, z = quantity @@ -301,7 +313,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({**_rebuilt_flat(mf, "cosmology"), "cosmology": Planck15}) + assert new == MassFunction.build(**{"transfer.model": "EH"}, cosmology=Planck15) def test_k_accuracy_change_is_consistent(mf): @@ -549,6 +562,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 # --------------------------------------------------------------------------------- @@ -817,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.") @@ -839,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(): @@ -852,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 + # 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.