searches: register af.MultiStartAdam as a first-class profiling sampler (imaging/mge)

searches/README.md

@@ -20,7 +20,7 @@ Auto-generated by `scripts/build_readme.py` from the versioned artifacts under `
 
 | Dimension      | Values                                                                    |
 |----------------|---------------------------------------------------------------------------|
-| Sampler        | `nautilus` (more to come via `_samplers.SAMPLER_BUILDERS`)                 |
+| Sampler        | `nautilus`, `multi_start_adam` (more via `_samplers.SAMPLER_BUILDERS`)     |
 | Dataset class  | `imaging`, `interferometer`, `point_source`, `datacube`                   |
 | Model type     | `mge`, `pixelization`, `delaunay`, `image_plane`, `source_plane`          |
 | Instrument     | per-dataset-class (HST/Euclid/JWST/AO; SMA/ALMA/ALMA-high/JVLA; simple)   |
@@ -47,6 +47,14 @@ searches/
 
 ## Key design choices
 
+**MAP optimizers alongside samplers.** `multi_start_adam` (`af.MultiStartAdam`,
+a JAX/optax multi-start gradient MAP optimizer) is registered as a first-class
+search too, but only for the `imaging/mge` cell — the benchmark-proven cell where
+a gradient MAP optimizer is meaningful (pixelization/Delaunay/interferometer/
+point-source are outside its use case). It is JAX-only (a pure-NumPy config
+raises) and has no `n_live` (it records `n_starts`/`n_steps`; the JSON stores
+`n_live: null`).
+
 **First-class only.** No more wrapping `nautilus.Sampler` directly. The
 old `simple.py` / `jax.py` scripts are deleted. Every cell goes through
 `af.Nautilus.fit(model, analysis)`, so visualization, output writes,

searches/_runner.py

@@ -39,9 +39,14 @@
 from searches._metrics import attach_viz_timer, collect_metrics  # noqa: E402
 from searches._samplers import (  # noqa: E402
     SAMPLER_BUILDERS,
+    multi_start_settings,
     n_live_for,
     vmap_batch_for_cell,
 )
+
+# Samplers that have an ``n_live`` (nested sampling). MAP optimizers such as
+# ``multi_start_adam`` do not, and record ``null`` rather than a misleading value.
+_SAMPLERS_WITH_N_LIVE = frozenset({"nautilus"})
 from searches._setup import build_for_cell, format_best_fit  # noqa: E402
 
 _DEFAULT_INSTRUMENTS: dict[str, str] = {
@@ -78,12 +83,15 @@ def run_search(
     config_name = cli.config_name or "default"
     use_jax = _decide_use_jax()
 
+    uses_n_live = sampler in _SAMPLERS_WITH_N_LIVE
+    n_live = n_live_for(dataset_class, model_type) if uses_n_live else None
+
     print(
         f"\n--- searches/{sampler}/{dataset_class}/{model_type}"
         f" [{instrument}, {config_name}, use_jax={use_jax},"
         f" mp={cli.use_mixed_precision}] ---"
     )
-    print(f"  n_live: {n_live_for(dataset_class, model_type)}")
+    print(f"  n_live: {n_live if n_live is not None else 'n/a (MAP optimizer)'}")
 
     print("  Building dataset / model / analysis...")
     dataset, model, analysis = build_for_cell(
@@ -141,7 +149,7 @@ def run_search(
         cli=cli,
         use_jax=use_jax,
         n_free_params=int(model.total_free_parameters),
-        n_live=n_live_for(dataset_class, model_type),
+        n_live=n_live,
         metrics=metrics,
         viz_n_calls=viz_timer.n_calls,
         best_fit=best_fit,
@@ -167,7 +175,7 @@ def _sampler_config_dict(
     dataset_class: str,
     model_type: str,
     instrument: str,
-    n_live: int,
+    n_live: int | None,
     use_jax: bool,
 ) -> dict:
     """Return the JSON-friendly sampler config block for the metric write.
@@ -176,8 +184,8 @@ def _sampler_config_dict(
     actually constructs the search with — so the JSON faithfully
     records what was run, including the per-cell vmap batch cap.
     """
-    batch = vmap_batch_for_cell(dataset_class, model_type, instrument)
     if sampler == "nautilus":
+        batch = vmap_batch_for_cell(dataset_class, model_type, instrument)
         return {
             "n_live": n_live,
             "n_batch": batch,
@@ -186,6 +194,9 @@ def _sampler_config_dict(
             "force_x1_cpu": use_jax,
             "iterations_per_update": 3 * n_live,
         }
+    if sampler == "multi_start_adam":
+        # MAP optimizer: no n_live; records its own multi-start knobs.
+        return {**multi_start_settings(), "number_of_cores": 1}
     return {"n_live": n_live, "_note": f"unknown sampler {sampler!r}"}
 
 
@@ -210,7 +221,7 @@ def _build_summary(
     cli: Any,
     use_jax: bool,
     n_free_params: int,
-    n_live: int,
+    n_live: int | None,
     metrics: Any,
     viz_n_calls: int,
     best_fit: str,

searches/_samplers.py

@@ -122,7 +122,62 @@ def build_nautilus(
     )
 
 
+# MultiStartAdam profiling settings. Single-sourced here so the builder and the
+# JSON config block (``_runner._sampler_config_dict``) record identical values.
+# These are illustrative profiling values, not the A100 scaling run (the
+# GIGA-Lens recipe uses hundreds of starts); ``n_starts=64`` is a representative
+# multi-start batch for a local/A100 profile.
+_MULTI_START_N_STARTS = 64
+_MULTI_START_N_STEPS = 300
+_MULTI_START_LEARNING_RATE = 0.01
+
+
+def multi_start_settings() -> dict:
+    """The knobs ``build_multi_start_adam`` constructs the search with.
+
+    Exposed so ``_sampler_config_dict`` records exactly what was run.
+    """
+    return {
+        "n_starts": _MULTI_START_N_STARTS,
+        "n_steps": _MULTI_START_N_STEPS,
+        "learning_rate": _MULTI_START_LEARNING_RATE,
+    }
+
+
+def build_multi_start_adam(
+    *,
+    sampler: str,
+    dataset_class: str,
+    model_type: str,
+    instrument: str,
+    config_name: str,
+    use_jax: bool,
+) -> af.MultiStartAdam:
+    """Construct a first-class ``af.MultiStartAdam`` search for one profiling cell.
+
+    ``MultiStartAdam`` is a JAX / ``optax`` multi-start first-order gradient MAP
+    optimizer: it runs ``n_starts`` broad starts in parallel (its own ``jax.vmap``)
+    and returns the best-basin point. Unlike ``af.Nautilus`` it:
+
+    - is JAX-native and **requires** a JAX-traceable analysis (``use_jax=True``);
+      a pure-NumPy config will raise. The sweep runs JAX-on by default.
+    - has no ``n_live`` (it uses ``n_starts`` / ``n_steps``), and
+    - does not use the ``use_jax_vmap`` / ``force_x1_cpu`` ``Fitness`` path — it
+      builds its own batched ``value_and_grad``.
+
+    ``number_of_cores=1`` is kept for consistency with the profile convention
+    (it is metadata here; the search runs a single-process vmap loop).
+    """
+    return af.MultiStartAdam(
+        name=config_name,
+        path_prefix=f"searches/{sampler}/{dataset_class}/{model_type}/{instrument}",
+        number_of_cores=1,
+        **multi_start_settings(),
+    )
+
+
 SamplerBuilder = Callable[..., af.NonLinearSearch]
 SAMPLER_BUILDERS: dict[str, SamplerBuilder] = {
     "nautilus": build_nautilus,
+    "multi_start_adam": build_multi_start_adam,
 }

searches/multi_start_adam/imaging/mge.py

@@ -0,0 +1,30 @@
+"""First-class af.MultiStartAdam search profiling — imaging MGE.
+
+Drives a full ``af.MultiStartAdam`` fit on an MGE lens + MGE source imaging model
+across the canonical instruments (hst / euclid / jwst / ao). MultiStartAdam is a
+JAX / optax multi-start gradient MAP optimizer; this is the cell where the search
+is meaningful and benchmark-proven (the GPU MAP-optimizer benchmark recovered the
+truth basin on the HST MGE lens likelihood). See ``searches/README.md`` for the
+design and the sweep workflow.
+
+Note: MultiStartAdam is JAX-native and requires ``use_jax=True`` (the sweep runs
+JAX-on by default); a pure-NumPy config will raise.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]  # autolens_profiling/
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from searches._runner import run_search  # noqa: E402
+
+run_search(
+    sampler="multi_start_adam",
+    dataset_class="imaging",
+    model_type="mge",
+    default_instrument="hst",
+)

searches/sweep.py

@@ -71,6 +71,9 @@
     ("nautilus", "interferometer", "pixelization"),
     ("nautilus", "interferometer", "delaunay"),
     ("nautilus", "datacube", "delaunay"),
+    # MAP optimizer — only the MGE cell, where a multi-start gradient optimizer
+    # is meaningful and benchmark-proven (JAX-only).
+    ("multi_start_adam", "imaging", "mge"),
 ]
 
 

simulators/README.md

@@ -46,11 +46,15 @@ Auto-generated by `scripts/build_readme.py` from the latest `*_summary_v<version
 | Script | Total wall time | PyAutoLens version |
 |--------|-----------------|--------------------|
 | `cluster.py` | 234.99 s | v2026.7.6.649 |
+| `group.py` | 38.11 s | v2026.5.14.2 |
+| `imaging.py` | 3.12 s | v2026.5.14.2 |
 | `imaging_ao.py` | 186.06 s | v2026.7.6.649 |
 | `imaging_euclid.py` | 2.28 s | v2026.7.6.649 |
 | `imaging_jwst.py` | 20.33 s | v2026.7.6.649 |
+| `interferometer.py` | 25.52 s | v2026.5.14.2 |
 | `interferometer_alma.py` | 339.65 s | v2026.7.6.649 |
 | `interferometer_alma_high.py` | 397.46 s | v2026.7.6.649 |
+| `multi.py` | 10.33 s | v2026.5.14.2 |
 <!-- END auto-table:simulators -->
 
 ## Running a script