From a6b47b6c31c83b60884cce6b87ed15f4120d0e36 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 14 Jul 2026 19:44:22 +0100 Subject: [PATCH] docs(searches): add multi-start gradient searches to the modeling searches guide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a MultiStartAdam section to scripts/guides/modeling/searches.py documenting af.MultiStartAdam (with MultiStartADABelief / MultiStartLion as alternatives) — the JAX multi-start gradient MAP optimizer that works on the complex model parameter spaces where single-start LBFGS gets stuck. Config-only, mirroring the LBFGS section; regenerate the paired notebook. Co-Authored-By: Claude Opus 4.8 --- notebooks/guides/modeling/searches.ipynb | 39 ++++++++++++++++++++++++ scripts/guides/modeling/searches.py | 30 +++++++++++++++++- 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/notebooks/guides/modeling/searches.ipynb b/notebooks/guides/modeling/searches.ipynb index 16c5cd83..5bd1bf86 100644 --- a/notebooks/guides/modeling/searches.ipynb +++ b/notebooks/guides/modeling/searches.ipynb @@ -32,6 +32,7 @@ "- **Emcee**: An ensemble MCMC sampler that is commonly used in Astronomy and Astrophysics.\n", "- **Zeus**: An ensemble MCMC slice sampler that is the most effective MCMC method for modeling.\n", "- **LBFGS**: A quasi-Newton optimization algorithm that is a maximum likelihood estimator (MLE) method.\n", + "- **MultiStartAdam**: A JAX multi-start gradient maximum a posteriori (MAP) optimizer that works on complex model parameter spaces.\n", "- **Start Point**: An API that allows the user to specify the start-point of a model-fit, which is useful for MCMC and MLE methods.\n", "- **Search Cookbook**: A cookbook that documents all searches available in **PyAutoFit**, including those not documented here.\n", "\n", @@ -278,6 +279,44 @@ "outputs": [], "execution_count": null }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MultiStartAdam__\n", + "\n", + "`MultiStartAdam` is a JAX / `optax` multi-start first-order gradient optimizer, and a maximum a posteriori (MAP)\n", + "estimator. It directly addresses the weakness of a single-start optimizer like `LBFGS` noted above: instead of\n", + "descending from a single starting point (which, for the complex parameter spaces of galaxy models, frequently gets\n", + "stuck in a local maximum), it launches `n_starts` independent optimizations from broad starting points in parallel\n", + "via `jax.vmap` and returns the best one. This wide population of starts reliably finds the global maximum-likelihood\n", + "basin, making it a robust and fast optimizer even for models where single-start optimizers fail.\n", + "\n", + "Because it is gradient-based, it requires a JAX-traceable analysis (created via `use_jax=True`). It also manages its\n", + "own broad starting points, so unlike `LBFGS` it does not use the start-point API described below.\n", + "\n", + "`MultiStartADABelief` and `MultiStartLion` are drop-in alternatives that use a different `optax` update rule (`Lion`\n", + "is sign-based and therefore prefers a ~10x smaller `learning_rate`).\n", + "\n", + "Like all optimizers it returns a single best-fit model, not a posterior with errors, so `Nautilus` remains the\n", + "default recommendation when parameter uncertainties are required." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.MultiStartAdam(\n", + " path_prefix=Path(\"imaging\", \"searches\"),\n", + " name=\"MultiStartAdam\",\n", + " n_starts=50,\n", + " n_steps=500,\n", + " learning_rate=0.01,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, { "cell_type": "markdown", "metadata": {}, diff --git a/scripts/guides/modeling/searches.py b/scripts/guides/modeling/searches.py index 37a1573e..eb54c9b1 100644 --- a/scripts/guides/modeling/searches.py +++ b/scripts/guides/modeling/searches.py @@ -27,6 +27,7 @@ - **Emcee**: An ensemble MCMC sampler that is commonly used in Astronomy and Astrophysics. - **Zeus**: An ensemble MCMC slice sampler that is the most effective MCMC method for modeling. - **LBFGS**: A quasi-Newton optimization algorithm that is a maximum likelihood estimator (MLE) method. +- **MultiStartAdam**: A JAX multi-start gradient maximum a posteriori (MAP) optimizer that works on complex model parameter spaces. - **Start Point**: An API that allows the user to specify the start-point of a model-fit, which is useful for MCMC and MLE methods. - **Search Cookbook**: A cookbook that documents all searches available in **PyAutoFit**, including those not documented here. @@ -35,7 +36,7 @@ If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook. """ -# from autoconf import setup_notebook; setup_notebook() +# from autoconf import setup_notebook; setup_notebook() from pathlib import Path import autofit as af @@ -184,6 +185,33 @@ unique_tag="example", ) +""" +__MultiStartAdam__ + +`MultiStartAdam` is a JAX / `optax` multi-start first-order gradient optimizer, and a maximum a posteriori (MAP) +estimator. It directly addresses the weakness of a single-start optimizer like `LBFGS` noted above: instead of +descending from a single starting point (which, for the complex parameter spaces of galaxy models, frequently gets +stuck in a local maximum), it launches `n_starts` independent optimizations from broad starting points in parallel +via `jax.vmap` and returns the best one. This wide population of starts reliably finds the global maximum-likelihood +basin, making it a robust and fast optimizer even for models where single-start optimizers fail. + +Because it is gradient-based, it requires a JAX-traceable analysis (created via `use_jax=True`). It also manages its +own broad starting points, so unlike `LBFGS` it does not use the start-point API described below. + +`MultiStartADABelief` and `MultiStartLion` are drop-in alternatives that use a different `optax` update rule (`Lion` +is sign-based and therefore prefers a ~10x smaller `learning_rate`). + +Like all optimizers it returns a single best-fit model, not a posterior with errors, so `Nautilus` remains the +default recommendation when parameter uncertainties are required. +""" +search = af.MultiStartAdam( + path_prefix=Path("imaging", "searches"), + name="MultiStartAdam", + n_starts=50, + n_steps=500, + learning_rate=0.01, +) + """ __Start Point__