From 5b8b55c77b4d155e2e6e008821ac62ccc89d5982 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 14 Jul 2026 19:01:57 +0100 Subject: [PATCH] test(searches): validate MultiStartAdam recovers truth basin under JAX MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds scripts/searches/MultiStartAdam.py — the JAX cross-backend end-to-end validation for af.MultiStartAdam (PyAutoFit#1369), asserting truth-basin recovery on the 1D Gaussian dataset (centre/normalization/sigma = 50/25/10) via the standard search.fit() API with a use_jax=True analysis. Kept out of PyAutoFit's NumPy-only unit suite by design. Co-Authored-By: Claude Opus 4.8 --- scripts/searches/MultiStartAdam.py | 100 +++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 scripts/searches/MultiStartAdam.py diff --git a/scripts/searches/MultiStartAdam.py b/scripts/searches/MultiStartAdam.py new file mode 100644 index 0000000..d85034f --- /dev/null +++ b/scripts/searches/MultiStartAdam.py @@ -0,0 +1,100 @@ +""" +Searches: MultiStartAdam (JAX-jitted gradient MAP optimizer) +============================================================ + +End-to-end validation of ``af.MultiStartAdam`` — the multi-start first-order +gradient MAP search promoted from the GPU MAP-optimizer benchmark (PyAutoFit +#1369). It runs ``N`` broad starts in parallel via ``jax.vmap``, taking a fixed +self-normalised Adam step per start on the unconstrained parameterization, and +returns the best-basin start as the maximum-log-posterior point. + +This is the JAX cross-backend check that deliberately lives here rather than in +PyAutoFit's unit suite (library unit tests stay NumPy-only). It asserts that the +search recovers the known truth basin of the 1D Gaussian dataset +(``centre, normalization, sigma = 50, 25, 10``). + +Two pytree-registration calls (as in ``Nautilus_jax.py``) let +``model.instance_from_vector`` flow through ``jax.jit``: + + - ``enable_pytrees()`` registers ``Model`` / ``Collection`` / ``ModelInstance`` + and the prior classes once per process. + - ``register_model(model)`` registers each concrete ``cls`` in the model + (here ``af.ex.Gaussian``) so its instances become traceable pytrees. +""" + +import numpy as np +from os import path + +import autofit as af +from autofit.jax.pytrees import enable_pytrees, register_model + +enable_pytrees() + +""" +__Data__ + +Load the 1D Gaussian dataset (truth: centre=50, normalization=25, sigma=10). +Simulate it first if it is not already on disk, so this script is self-contained. +""" +dataset_path = path.join("dataset", "example_1d", "gaussian_x1") + +if not path.exists(dataset_path): + import subprocess + import sys + + subprocess.run( + [sys.executable, "scripts/simulators/simulators.py"], + check=True, + ) + +data = af.util.numpy_array_from_json(file_path=path.join(dataset_path, "data.json")) +noise_map = af.util.numpy_array_from_json( + file_path=path.join(dataset_path, "noise_map.json") +) + +""" +__Model + Analysis__ + +The standard N=3 ``Gaussian`` model with a JAX-traceable analysis +(``use_jax=True`` routes the likelihood maths through ``jax.numpy``). +""" +model = af.Model(af.ex.Gaussian) + +model.centre = af.UniformPrior(lower_limit=0.0, upper_limit=100.0) +model.normalization = af.LogUniformPrior(lower_limit=1e-2, upper_limit=1e2) +model.sigma = af.UniformPrior(lower_limit=0.0, upper_limit=30.0) + +register_model(model) + +analysis = af.ex.Analysis(data=data, noise_map=noise_map, use_jax=True) + +""" +__Search__ + +Run ``MultiStartAdam`` with ``NullPaths`` (no ``name`` / ``path_prefix``) so +nothing is written to disk. A modest start count and step budget keep the run +cheap — the goal is to verify the JAX gradient path executes end-to-end and +lands in the truth basin, not to profile it. +""" +search = af.MultiStartAdam(n_starts=16, n_steps=500, learning_rate=0.5) + +result = search.fit(model=model, analysis=analysis) + +""" +__Assertion__ + +The best-basin instance must recover the truth to within a loose tolerance +(the MAP sits near truth, offset only by the dataset noise). +""" +instance = result.samples.max_log_likelihood() + +print( + f"Recovered: centre={instance.centre:.3f}, " + f"normalization={instance.normalization:.3f}, sigma={instance.sigma:.3f}" +) + +assert abs(instance.centre - 50.0) < 2.0, instance.centre +assert abs(instance.normalization - 25.0) < 3.0, instance.normalization +assert abs(instance.sigma - 10.0) < 2.0, instance.sigma + +print("MultiStartAdam recovered the truth basin of the 1D Gaussian dataset.")