diff --git a/searches/README.md b/searches/README.md index 4a85ab8..63c367a 100644 --- a/searches/README.md +++ b/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, diff --git a/searches/_runner.py b/searches/_runner.py index fca4f28..bc741c6 100644 --- a/searches/_runner.py +++ b/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, diff --git a/searches/_samplers.py b/searches/_samplers.py index 6360f76..c9aac19 100644 --- a/searches/_samplers.py +++ b/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, } diff --git a/searches/multi_start_adam/imaging/mge.py b/searches/multi_start_adam/imaging/mge.py new file mode 100644 index 0000000..7f64b96 --- /dev/null +++ b/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", +) diff --git a/searches/sweep.py b/searches/sweep.py index 3081382..4a409e3 100644 --- a/searches/sweep.py +++ b/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"), ] diff --git a/simulators/README.md b/simulators/README.md index 532d22e..888bc06 100644 --- a/simulators/README.md +++ b/simulators/README.md @@ -46,11 +46,15 @@ Auto-generated by `scripts/build_readme.py` from the latest `*_summary_v ## Running a script