From 33c2e9d3188c6acbf50abaf2c788787b71e006b0 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sun, 27 Sep 2026 18:52:02 +0100 Subject: [PATCH] feat: analysis-declared gradient_mode for Fitness.grad and MultiStartGradient An af.Analysis now declares how gradient searches differentiate its likelihood via the class attribute `gradient_mode` ("reverse" default, or "forward" = jax.jacfwd over the flat parameter vector). New helper module autofit/jax/gradient.py (GRADIENT_MODES, resolve_gradient_mode, value_and_grad_from, grad_from) resolves the declaration against an optional override and builds the transform; JAX is imported lazily. - Fitness: optional `gradient_mode` override; `_grad` built through grad_from(resolved mode); survives pickling, old pickles default None. - MultiStartGradient: `gradient_mode` keyword (validated at construction, dict/pickle round-trip), resolved against the analysis at fit time; both value_and_grad sites (physical and scaler/bijector stepped objective) go through value_and_grad_from; resolved mode is logged at search start, persisted in search_internal and reported in samples_info. - Defaults unchanged: every analysis stays "reverse" unless it declares otherwise (PyAutoLens AnalysisPoint will declare "forward", per autolens_profiling #327/#331). Refs #1648 Co-Authored-By: Claude Opus 5.5 --- autofit/jax/gradient.py | 118 ++++++ autofit/non_linear/analysis/analysis.py | 26 ++ autofit/non_linear/fitness.py | 31 +- .../search/mle/multi_start_gradient/search.py | 71 +++- test_autofit/non_linear/test_gradient_mode.py | 361 ++++++++++++++++++ 5 files changed, 594 insertions(+), 13 deletions(-) create mode 100644 autofit/jax/gradient.py create mode 100644 test_autofit/non_linear/test_gradient_mode.py diff --git a/autofit/jax/gradient.py b/autofit/jax/gradient.py new file mode 100644 index 000000000..97700510a --- /dev/null +++ b/autofit/jax/gradient.py @@ -0,0 +1,118 @@ +""" +Analysis-declared gradient mode. + +An ``af.Analysis`` declares how its likelihood is best differentiated through the class attribute +``gradient_mode``: ``"reverse"`` (the default, ``jax.value_and_grad``) or ``"forward"`` +(``jax.jacfwd`` over the flat parameter vector). A gradient search may override the declaration +(e.g. ``af.MultiStartAdam(gradient_mode="reverse")``). + +Both modes compute the same gradient; they differ only in cost. Reverse mode costs a small constant +multiple of one likelihood evaluation whatever the number of parameters, so it is the right default. +Forward mode costs one tangent per free parameter, but wins when the likelihood itself contains an +inner forward-mode derivative (reverse mode must then run reverse-over-forward through it) and the +parameter count is modest -- e.g. the source-plane point-source likelihood, whose lensing Hessian is +``jax.jacfwd``-built (autolens_profiling #327/#331: 2-4.5x faster, up to 8x faster to compile). + +JAX is imported lazily inside the builders so that ``autofit`` imports without JAX. +""" + +from typing import Callable, Optional + +GRADIENT_MODES = ("reverse", "forward") + + +def _check_mode(mode: str, source: str) -> str: + if mode not in GRADIENT_MODES: + raise ValueError( + f"Unknown gradient_mode {mode!r} ({source}). Valid modes are " + f"{', '.join(repr(m) for m in GRADIENT_MODES)}." + ) + return mode + + +def validate_gradient_mode(mode: Optional[str]) -> Optional[str]: + """ + Return ``mode`` unchanged if it is ``None`` or a valid mode, else raise ``ValueError``. + + Used by constructors (``Fitness``, ``MultiStartGradient``) that accept an optional override, + so a mistyped mode fails at construction rather than at the first gradient evaluation. + """ + if mode is None: + return None + return _check_mode(mode, "override") + + +def resolve_gradient_mode(analysis, override: Optional[str] = None) -> str: + """ + The gradient mode to use for ``analysis``. + + ``override`` (a search- or ``Fitness``-level setting) wins when it is not ``None``; otherwise + the analysis's ``gradient_mode`` class attribute is used, defaulting to ``"reverse"`` for + objects that do not declare one. An unknown value from either source raises ``ValueError``. + """ + if override is not None: + return _check_mode(override, "override") + + mode = getattr(analysis, "gradient_mode", "reverse") + return _check_mode(mode, f"declared by {type(analysis).__name__}") + + +def _forward_value_and_grad(func: Callable) -> Callable: + import jax + import jax.numpy as jnp + + def value_twice(vector): + value = func(vector) + if jnp.ndim(value) != 0: + raise ValueError( + "gradient_mode='forward' requires a scalar objective; the function returned " + f"an array of shape {jnp.shape(value)}." + ) + return value, value + + jac = jax.jacfwd(value_twice, has_aux=True) + + def value_and_grad(vector): + grad, value = jac(vector) + return value, grad + + return value_and_grad + + +def value_and_grad_from(func: Callable, mode: str) -> Callable: + """ + A callable returning ``(value, grad)`` of the scalar ``func`` of a flat parameter vector, + with the same contract as ``jax.value_and_grad(func)``. + + - ``"reverse"``: ``jax.value_and_grad(func)``; + - ``"forward"``: ``jax.jacfwd`` over the vector, with the value carried as auxiliary output so + the likelihood is traced once. + """ + _check_mode(mode, "value_and_grad_from") + + if mode == "reverse": + import jax + + return jax.value_and_grad(func) + + return _forward_value_and_grad(func) + + +def grad_from(func: Callable, mode: str) -> Callable: + """ + The gradient-only counterpart of ``value_and_grad_from``: ``jax.grad(func)`` in reverse mode, + the ``jax.jacfwd`` gradient of the scalar ``func`` in forward mode. + """ + _check_mode(mode, "grad_from") + + if mode == "reverse": + import jax + + return jax.grad(func) + + value_and_grad = _forward_value_and_grad(func) + + def grad(vector): + return value_and_grad(vector)[1] + + return grad diff --git a/autofit/non_linear/analysis/analysis.py b/autofit/non_linear/analysis/analysis.py index 6a4df2af0..47046d179 100644 --- a/autofit/non_linear/analysis/analysis.py +++ b/autofit/non_linear/analysis/analysis.py @@ -59,6 +59,32 @@ class Analysis(ABC): # `ZeroSolver`. LATENT_BATCH_MODE = "vmap" + # How gradient-based searches differentiate this analysis's likelihood + # (resolved by `autofit.jax.gradient.resolve_gradient_mode`; used by + # `Fitness.grad` and the `MultiStartGradient` searches). + # + # - "reverse" (default): `jax.value_and_grad`. Its cost is a small constant + # multiple of one likelihood evaluation whatever the number of free + # parameters, so it is right for almost every analysis. + # + # - "forward": `jax.jacfwd` over the flat parameter vector. Declare it when + # the likelihood itself contains an inner forward-mode derivative (reverse + # mode then runs reverse-over-forward through it) and the model has a + # modest number of free parameters. The source-plane point-source + # likelihood is the reference case: its lensing Hessian is built with + # `jax.jacfwd`, and forward mode measured 2-4.5x faster and up to 8x + # faster to compile through 24 free parameters (autolens_profiling + # #327/#331), so `autolens.AnalysisPoint` declares it. + # + # Memory caveat: forward mode carries one tangent per free parameter, so + # under `jax.vmap` (e.g. `MultiStartGradient`'s starts) its memory scales + # with `n_starts * n_params`. + # + # A search may override the declaration, e.g. + # `af.MultiStartAdam(gradient_mode="reverse")`. Both modes compute the same + # gradient; only speed and memory differ. + gradient_mode = "reverse" + def __init__( self, use_jax: bool = False, diff --git a/autofit/non_linear/fitness.py b/autofit/non_linear/fitness.py index 42452e13e..b448e6618 100644 --- a/autofit/non_linear/fitness.py +++ b/autofit/non_linear/fitness.py @@ -15,6 +15,7 @@ from autofit.mapper.prior_model.abstract import AbstractPriorModel +from autofit.jax.gradient import validate_gradient_mode from autofit.non_linear.jax_compile import log_on_first_compile from autofit.non_linear.paths.abstract import AbstractPaths from autofit.non_linear.analysis import Analysis @@ -182,6 +183,7 @@ def __init__( iterations_per_quick_update: Optional[int] = None, background_quick_update: bool = False, live_visual_update: bool = False, + gradient_mode: Optional[str] = None, ): """ Interfaces with any non-linear search to fit the model to the data and return a log likelihood via @@ -240,6 +242,10 @@ def __init__( chi-squared value that is minimized. store_history If `True`, the parameters and log likelihood values of every model-fit are stored in lists. + gradient_mode + How `grad` differentiates the likelihood: `"reverse"` (`jax.grad`) or `"forward"` (`jax.jacfwd` + over the flat parameter vector). `None` (default) uses the analysis's declared + `Analysis.gradient_mode`. See `autofit.jax.gradient`. """ self.analysis = analysis @@ -264,6 +270,7 @@ def __init__( self.use_jax_vmap = use_jax_vmap self.use_jax_jit = use_jax_jit + self.gradient_mode = validate_gradient_mode(gradient_mode) if getattr(self.analysis, "_use_jax", False): from autofit.jax.pytrees import enable_pytrees, register_model @@ -847,6 +854,8 @@ def __setstate__(self, state): or "_apply_assertions_traced" not in self.__dict__ ): self._set_traced_assertions() + # `Fitness` objects pickled before the gradient-mode override existed defer to the analysis. + self.__dict__.setdefault("gradient_mode", None) self._call = self.call if getattr(self, "use_jax_vmap", False): self._call = self._vmap @@ -916,22 +925,22 @@ def _jit(self): @cached_property def _grad(self): """ - Gradient of the JIT-compiled likelihood function. + Gradient of the likelihood function (`self.call`) with respect to the flat parameter vector. - This wraps the JIT-compiled likelihood function (`self._call`) with - `jax.grad`, returning a function that computes gradients of the - likelihood with respect to its input parameters. Useful for gradient- - based optimization and inference methods. + The transform is chosen by `autofit.jax.gradient.resolve_gradient_mode`: this `Fitness`'s + `gradient_mode` override if set, otherwise the analysis's declared `Analysis.gradient_mode` + -- `jax.grad` in reverse mode, `jax.jacfwd` in forward mode. Both give the same gradient. - Since this is a `cached_property`, the gradient function is compiled - and cached on first access, ensuring that expensive setup is done - only once. + Since this is a `cached_property`, the gradient function is built and cached on first + access, ensuring that expensive setup is done only once. """ - import jax + from autofit.jax import gradient + + mode = gradient.resolve_gradient_mode(self.analysis, self.gradient_mode) return log_on_first_compile( - jax.grad(self.call), - "likelihood function gradient", + gradient.grad_from(self.call, mode), + f"likelihood function gradient ({mode} mode)", ) def grad(self, *args, **kwargs): diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 374f269df..0d2792ecf 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -12,6 +12,11 @@ from autofit.non_linear.search.mle.abstract_mle import AbstractMLE from autofit.non_linear.analysis import Analysis from autofit.non_linear.fitness import Fitness +from autofit.jax.gradient import ( + resolve_gradient_mode, + validate_gradient_mode, + value_and_grad_from, +) from autofit.non_linear.bijector import AbstractBijector, BijectorNone from autofit.non_linear.clipper import ( AbstractClipper, @@ -58,6 +63,7 @@ def __init__( n_steps: int = 300, learning_rate: Optional[float] = None, batch_size: Optional[int] = None, + gradient_mode: Optional[str] = None, max_consecutive_nan: int = 8, start_lower_limit: float = 0.15, start_upper_limit: float = 0.85, @@ -142,6 +148,23 @@ def __init__( algorithmic knob (how many points it proposes per iteration) that autofit merely forwards. Here ``n_starts`` is the algorithm; this only decides how many of those starts are evaluated at a time. + gradient_mode + How the objective is differentiated: ``"reverse"`` + (``jax.value_and_grad``) or ``"forward"`` (``jax.jacfwd`` over the + flat parameter vector). ``None`` (default) uses the analysis's + declared ``Analysis.gradient_mode`` -- ``"reverse"`` for every + analysis except those that declare otherwise, e.g. + ``autolens.AnalysisPoint`` declares ``"forward"`` because its + source-plane likelihood carries an inner forward-mode lensing + Hessian (autolens_profiling #327/#331). Pass a value to override + the declaration for this search. Both modes give the same gradient + and so the same fit; only speed, compile time and memory differ. + + Memory caveat, alongside ``batch_size`` above: forward mode carries + one tangent per free parameter, so under the vmap over starts its + memory scales with ``n_starts * n_params`` (a ``batch_size``-wide + chunk's worth when batching). If a forward-mode fit exhausts memory, + set ``batch_size`` or override with ``gradient_mode="reverse"``. start_lower_limit, start_upper_limit The unit-cube bounds broad starts are drawn uniformly from. The interior default ``(0.15, 0.85)`` avoids the prior edges where many @@ -371,6 +394,10 @@ def __init__( self.n_starts = n_starts self.n_steps = n_steps self.batch_size = batch_size + # Validated at construction so a mistyped mode fails here rather than + # at the first compile; resolved against the analysis at fit time + # (``_resolved_gradient_mode``), since no analysis exists yet. + self.gradient_mode = validate_gradient_mode(gradient_mode) self.max_consecutive_nan = max_consecutive_nan self.learning_rate = ( learning_rate if learning_rate is not None else self._default_learning_rate @@ -436,6 +463,17 @@ def __init__( self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search") + def _resolved_gradient_mode(self, analysis) -> str: + """ + The gradient mode this fit uses: the search's ``gradient_mode`` if set, + otherwise the analysis's declaration. ``getattr`` because a search + unpickled from before the keyword existed carries no attribute and must + defer to the analysis. + """ + return resolve_gradient_mode( + analysis, override=getattr(self, "gradient_mode", None) + ) + def _is_final_boundary(self, converged: bool, total_steps: int) -> bool: """ Whether the chunk boundary just reached ends the search — either the @@ -971,7 +1009,23 @@ def _fit( # compile) while the chunk alone compiles in minutes. The tiling is # numerically identical to the vmap; `batch_size` never changes # results, it only bounds memory and compile. - _value_and_grad = jax.value_and_grad(fitness.call) + # + # The AD mode (``gradient_mode``) is resolved once here, against the + # analysis's declaration and this search's override, and used for both + # objectives below. Forward and reverse give the same ``(value, grad)``; + # everything downstream of the two builders is mode-agnostic. + gradient_mode = self._resolved_gradient_mode(analysis) + if not self.silence: + source = ( + "search override" + if getattr(self, "gradient_mode", None) is not None + else f"declared by {type(analysis).__name__}" + ) + self.logger.info( + f"MultiStartGradient gradient mode: {gradient_mode} ({source})." + ) + + _value_and_grad = value_and_grad_from(fitness.call, gradient_mode) # Per-start gradient finiteness is reduced *inside* the jitted call, as a # third output, rather than by a separate op on its result. The reduction @@ -1069,7 +1123,9 @@ def _to_physical(vector): # objective: ``_broad_starts`` below jits it directly to filter candidate # draws, and those draws are, and stay, physical. _value_and_grad_stepped = ( - jax.value_and_grad(lambda phi: fitness.call(_to_physical(phi))) + value_and_grad_from( + lambda phi: fitness.call(_to_physical(phi)), gradient_mode + ) if (has_scaler or has_bijector) else _value_and_grad ) @@ -1609,6 +1665,9 @@ def batched_value_and_grad(params): "params": params_physical, "scale": scale, "bijector": self.bijector.kinds if has_bijector else None, + # The AD mode this process ran with, resolved against the + # analysis (which ``samples_via_internal_from`` cannot see). + "gradient_mode": gradient_mode, "opt_state": opt_state, "best_params": best_params, "best_fom": best_fom, @@ -1983,6 +2042,14 @@ def samples_via_internal_from( "n_starts": self.n_starts, "n_steps": self.n_steps, "batch_size": self.batch_size, + # The RESOLVED AD mode (search override, else the analysis's + # declaration), read from ``search_internal`` because the analysis + # is not available here. Falls back to the search's own setting + # (possibly ``None`` = "the analysis's declaration, unrecorded") for + # a ``search_internal`` written before the mode was persisted. + "gradient_mode": search_internal.get( + "gradient_mode", getattr(self, "gradient_mode", None) + ), "total_steps": total_steps, "optax_method": self.optax_method, "learning_rate": self.learning_rate, diff --git a/test_autofit/non_linear/test_gradient_mode.py b/test_autofit/non_linear/test_gradient_mode.py new file mode 100644 index 000000000..6e91df30d --- /dev/null +++ b/test_autofit/non_linear/test_gradient_mode.py @@ -0,0 +1,361 @@ +""" +Analysis-declared gradient mode (PyAutoFit#1648). + +An ``Analysis`` declares how its likelihood is best differentiated through the class attribute +``gradient_mode`` ("reverse" by default, "forward" for likelihoods like the source-plane point-source +chi-squared, which carry an inner forward-mode derivative). ``autofit.jax.gradient`` resolves that +declaration against an optional search-level override and builds the matching JAX transform; +``Fitness.grad`` and ``MultiStartGradient`` both go through it. + +The two modes are the same mathematical gradient, so every test below pins forward against reverse +rather than against a hand-written derivative: what this change may alter is speed, never numbers. +""" + +import pickle +import shutil +import uuid + +import numpy as np +import pytest + +import autofit as af +from autofit.non_linear.fitness import Fitness +from autonerves.dictable import from_dict, to_dict + +jax = pytest.importorskip("jax") + +from autofit.jax.gradient import ( # noqa: E402 + GRADIENT_MODES, + grad_from, + resolve_gradient_mode, + value_and_grad_from, +) + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +def _gaussian_data(): + xvalues = np.arange(40, dtype=float) + truth = af.ex.Gaussian(centre=20.0, normalization=5.0, sigma=6.0) + data = truth.model_data_from(xvalues=xvalues) + return data, np.full_like(data, 0.5) + + +class ForwardAnalysis(af.ex.Analysis): + gradient_mode = "forward" + + +def _fitness(analysis_cls=af.ex.Analysis, **kwargs): + model = af.Model(af.ex.Gaussian) + data, noise_map = _gaussian_data() + analysis = analysis_cls(data=data, noise_map=noise_map, use_jax=True) + return Fitness( + model=model, + analysis=analysis, + fom_is_log_likelihood=False, + convert_to_chi_squared=True, + **kwargs, + ) + + +def _vectors(model, n=5, seed=3): + rng = np.random.default_rng(seed) + vectors = [np.asarray(model.physical_values_from_prior_medians, dtype=float)] + for _ in range(n - 1): + unit = rng.uniform(0.2, 0.8, size=model.prior_count) + vectors.append(np.asarray(model.vector_from_unit_vector(list(unit)), dtype=float)) + return vectors + + +# -------------------------------------------------------------------------- +# resolve_gradient_mode +# -------------------------------------------------------------------------- + + +def test__modes_are_reverse_and_forward(): + assert GRADIENT_MODES == ("reverse", "forward") + + +def test__analysis_default_is_reverse(): + assert af.Analysis.gradient_mode == "reverse" + assert resolve_gradient_mode(af.ex.Analysis(data=np.ones(3), noise_map=np.ones(3))) == "reverse" + + +def test__class_level_declaration_is_resolved(): + analysis = ForwardAnalysis(data=np.ones(3), noise_map=np.ones(3)) + + assert resolve_gradient_mode(analysis) == "forward" + + +def test__override_wins_over_the_analysis_declaration(): + analysis = ForwardAnalysis(data=np.ones(3), noise_map=np.ones(3)) + + assert resolve_gradient_mode(analysis, override="reverse") == "reverse" + assert ( + resolve_gradient_mode( + af.ex.Analysis(data=np.ones(3), noise_map=np.ones(3)), override="forward" + ) + == "forward" + ) + + +def test__an_object_without_the_attribute_resolves_to_reverse(): + assert resolve_gradient_mode(object()) == "reverse" + + +@pytest.mark.parametrize("where", ["override", "analysis"]) +def test__an_unknown_mode_raises_naming_the_valid_ones(where): + class Bad(af.ex.Analysis): + gradient_mode = "sideways" + + if where == "override": + analysis, override = af.ex.Analysis(data=np.ones(3), noise_map=np.ones(3)), "sideways" + else: + analysis, override = Bad(data=np.ones(3), noise_map=np.ones(3)), None + + with pytest.raises(ValueError) as exc_info: + resolve_gradient_mode(analysis, override=override) + + message = str(exc_info.value) + assert "sideways" in message + assert "reverse" in message and "forward" in message + + +def test__an_unknown_mode_raises_in_the_builders(): + with pytest.raises(ValueError): + value_and_grad_from(lambda v: v.sum(), "sideways") + with pytest.raises(ValueError): + grad_from(lambda v: v.sum(), "sideways") + + +# -------------------------------------------------------------------------- +# The transforms: forward is the reverse gradient, to round-off +# -------------------------------------------------------------------------- + + +def test__value_and_grad_from__forward_matches_reverse_on_the_fitness_objective(): + fitness = _fitness() + + reverse = jax.jit(value_and_grad_from(fitness.call, "reverse")) + forward = jax.jit(value_and_grad_from(fitness.call, "forward")) + + for vector in _vectors(fitness.model): + value_r, grad_r = reverse(vector) + value_f, grad_f = forward(vector) + + assert np.shape(value_f) == () + assert np.shape(grad_f) == np.shape(vector) + assert np.all(np.asarray(grad_r) != 0.0) + np.testing.assert_allclose(value_f, value_r, rtol=1e-10) + np.testing.assert_allclose(grad_f, grad_r, rtol=1e-10) + + +def test__value_and_grad_from__reverse_is_jax_value_and_grad(): + fitness = _fitness() + vector = _vectors(fitness.model, n=1)[0] + + value, grad = value_and_grad_from(fitness.call, "reverse")(vector) + value_ref, grad_ref = jax.value_and_grad(fitness.call)(vector) + + np.testing.assert_array_equal(value, value_ref) + np.testing.assert_array_equal(grad, grad_ref) + + +def test__grad_from__forward_matches_reverse(): + fitness = _fitness() + + for vector in _vectors(fitness.model): + np.testing.assert_allclose( + grad_from(fitness.call, "forward")(vector), + grad_from(fitness.call, "reverse")(vector), + rtol=1e-10, + ) + + +def test__forward_mode_rejects_a_non_scalar_objective(): + with pytest.raises(ValueError, match="scalar"): + value_and_grad_from(lambda v: v * 2.0, "forward")(np.ones(3)) + + +# -------------------------------------------------------------------------- +# Fitness.grad honours the mode +# -------------------------------------------------------------------------- + + +def test__fitness_gradient_mode_defaults_to_the_analysis_declaration(): + assert _fitness().gradient_mode is None + assert _fitness(gradient_mode="forward").gradient_mode == "forward" + + +def test__fitness_grad_forward_matches_reverse(): + reverse = _fitness() + forward_by_override = _fitness(gradient_mode="forward") + forward_by_analysis = _fitness(analysis_cls=ForwardAnalysis) + + for vector in _vectors(reverse.model): + grad_r = np.asarray(reverse.grad(vector)) + assert np.all(grad_r != 0.0) + np.testing.assert_allclose(forward_by_override.grad(vector), grad_r, rtol=1e-10) + np.testing.assert_allclose(forward_by_analysis.grad(vector), grad_r, rtol=1e-10) + + +def test__fitness_grad_builds_the_resolved_mode(monkeypatch): + """The parity above cannot tell WHICH transform ran; this can.""" + import autofit.jax.gradient as gradient + + seen = [] + real = gradient.grad_from + + def spy(func, mode): + seen.append(mode) + return real(func, mode) + + monkeypatch.setattr(gradient, "grad_from", spy) + + fitness = _fitness(analysis_cls=ForwardAnalysis) + fitness.grad(_vectors(fitness.model, n=1)[0]) + overridden = _fitness(analysis_cls=ForwardAnalysis, gradient_mode="reverse") + overridden.grad(_vectors(overridden.model, n=1)[0]) + + assert seen == ["forward", "reverse"] + + +def test__fitness_invalid_mode_raises_at_construction(): + with pytest.raises(ValueError): + _fitness(gradient_mode="sideways") + + +def test__fitness_pickle_round_trip_keeps_the_mode(): + fitness = _fitness(gradient_mode="forward") + vector = _vectors(fitness.model, n=1)[0] + before = np.asarray(fitness.grad(vector)) + + restored = pickle.loads(pickle.dumps(fitness)) + + assert restored.gradient_mode == "forward" + np.testing.assert_allclose(restored.grad(vector), before, rtol=1e-12) + + +def test__fitness_pickled_before_the_mode_existed_still_loads(): + fitness = _fitness() + state = fitness.__getstate__() + state.pop("gradient_mode") + + restored = Fitness.__new__(Fitness) + restored.__setstate__(state) + + assert restored.gradient_mode is None + vector = _vectors(restored.model, n=1)[0] + np.testing.assert_allclose(restored.grad(vector), fitness.grad(vector), rtol=1e-12) + + +# -------------------------------------------------------------------------- +# MultiStartGradient +# -------------------------------------------------------------------------- + + +def test__multi_start_gradient_mode_defaults_to_none_and_is_carried(): + assert af.MultiStartAdam().gradient_mode is None + for cls in ( + af.MultiStartAdam, + af.MultiStartADABelief, + af.MultiStartLion, + af.MultiStartProdigy, + ): + assert cls(gradient_mode="forward").gradient_mode == "forward" + + +def test__multi_start_invalid_gradient_mode_raises_at_construction(): + with pytest.raises(ValueError) as exc_info: + af.MultiStartAdam(gradient_mode="sideways") + + assert "reverse" in str(exc_info.value) + + +def test__multi_start_dict_round_trip_keeps_the_mode(): + restored = from_dict(to_dict(af.MultiStartAdam(gradient_mode="forward", n_starts=5))) + + assert isinstance(restored, af.MultiStartAdam) + assert restored.gradient_mode == "forward" + + +def test__multi_start_pickle_round_trip_keeps_the_mode(): + restored = pickle.loads(pickle.dumps(af.MultiStartAdam(gradient_mode="forward"))) + + assert restored.gradient_mode == "forward" + + +def test__multi_start_pickled_before_the_mode_existed_resolves_the_analysis(): + """A search unpickled from before this attribute existed carries no + ``gradient_mode``; the fit must fall back to the analysis declaration.""" + search = af.MultiStartAdam() + del search.__dict__["gradient_mode"] + + assert search._resolved_gradient_mode(ForwardAnalysis(data=np.ones(3), noise_map=np.ones(3))) == "forward" + + +def _multi_start_fit(gradient_mode, analysis_cls=af.ex.Analysis, **search_kwargs): + pytest.importorskip("optax") + + model = af.Model(af.ex.Gaussian) + model.centre = af.UniformPrior(lower_limit=0.0, upper_limit=40.0) + model.normalization = af.UniformPrior(lower_limit=0.1, upper_limit=20.0) + model.sigma = af.UniformPrior(lower_limit=1.0, upper_limit=15.0) + + data, noise_map = _gaussian_data() + analysis = analysis_cls(data=data, noise_map=noise_map, use_jax=True) + + search = af.MultiStartAdam( + name=f"gradient_mode_{gradient_mode}_{uuid.uuid4().hex}", + n_starts=4, + n_steps=25, + seed=11, + convergence=af.MultiStartGradientConvergence(check_for_convergence=False), + gradient_mode=gradient_mode, + **search_kwargs, + ) + + try: + result = search.fit(model=model, analysis=analysis) + finally: + shutil.rmtree(search.paths.output_path, ignore_errors=True) + + return result + + +@pytest.mark.parametrize( + "search_kwargs", + [ + {}, + {"scaler": af.ScalerPriorWidth()}, + {"bijector": af.BijectorAuto()}, + ], + ids=["physical", "scaler", "bijector"], +) +def test__multi_start_forward_and_reverse_give_the_same_fit(search_kwargs): + reverse = _multi_start_fit("reverse", **search_kwargs) + forward = _multi_start_fit("forward", **search_kwargs) + + np.testing.assert_allclose( + forward.samples.max_log_likelihood(as_instance=False), + reverse.samples.max_log_likelihood(as_instance=False), + rtol=1e-8, + ) + np.testing.assert_allclose( + forward.log_likelihood, reverse.log_likelihood, rtol=1e-8 + ) + + assert reverse.samples.samples_info["gradient_mode"] == "reverse" + assert forward.samples.samples_info["gradient_mode"] == "forward" + + +def test__multi_start_resolves_the_analysis_declaration_when_unset(): + result = _multi_start_fit(None, analysis_cls=ForwardAnalysis) + + assert result.samples.samples_info["gradient_mode"] == "forward" + + +def test__multi_start_keyword_overrides_the_analysis_declaration(): + result = _multi_start_fit("reverse", analysis_cls=ForwardAnalysis) + + assert result.samples.samples_info["gradient_mode"] == "reverse"