From d248930b4944f88f47beeeb2b3080d129b5e6523 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 14 Jul 2026 19:43:59 +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 lens 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 | 28 +++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/notebooks/guides/modeling/searches.ipynb b/notebooks/guides/modeling/searches.ipynb index 77edcb882..0168b921f 100644 --- a/notebooks/guides/modeling/searches.ipynb +++ b/notebooks/guides/modeling/searches.ipynb @@ -32,6 +32,7 @@ "- **Emcee:** Emcee (https://github.com/dfm/emcee) is an ensemble MCMC sampler that is commonly used in Astronomy.\n", "- **Zeus:** Zeus (https://zeus-mcmc.readthedocs.io/en/latest/) is an ensemble MCMC slice sampler.\n", "- **LBFGS:** LBFGS is a quasi-Newton optimization algorithm from scipy.\n", + "- **MultiStartAdam:** A JAX multi-start gradient maximum a posteriori (MAP) optimizer that works on complex lens parameter spaces.\n", "- **Start Point:** For maximum likelihood estimator (MLE) and Markov Chain Monte Carlo (MCMC) non-linear searches.\n", "- **Search Cookbook:** There are a number of other searches supported by **PyAutoFit** and therefore which can be used.\n", "\n", @@ -280,6 +281,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 lens 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 lens 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 lens 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 2ee7a306d..bfe44a785 100644 --- a/scripts/guides/modeling/searches.py +++ b/scripts/guides/modeling/searches.py @@ -27,6 +27,7 @@ - **Emcee:** Emcee (https://github.com/dfm/emcee) is an ensemble MCMC sampler that is commonly used in Astronomy. - **Zeus:** Zeus (https://zeus-mcmc.readthedocs.io/en/latest/) is an ensemble MCMC slice sampler. - **LBFGS:** LBFGS is a quasi-Newton optimization algorithm from scipy. +- **MultiStartAdam:** A JAX multi-start gradient maximum a posteriori (MAP) optimizer that works on complex lens parameter spaces. - **Start Point:** For maximum likelihood estimator (MLE) and Markov Chain Monte Carlo (MCMC) non-linear searches. - **Search Cookbook:** There are a number of other searches supported by **PyAutoFit** and therefore which can be used. @@ -186,6 +187,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 lens 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 lens 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 lens 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__