diff --git a/scripts/misc/searches/_setup.py b/scripts/misc/searches/_setup.py index 42a4b4aeb..42f707c04 100644 --- a/scripts/misc/searches/_setup.py +++ b/scripts/misc/searches/_setup.py @@ -523,6 +523,8 @@ def _build_model(dataset_class: str, model_type: str, *, mask_radius: float) -> return _delaunay_matern_model(mask_radius=mask_radius) if model_type in ("image_plane", "source_plane"): return _point_source_model() + if model_type in ("image_plane_solved", "source_plane_solved"): + return _point_source_model(solved=True) raise ValueError(f"Unknown model_type: {model_type!r}") @@ -744,10 +746,12 @@ def _delaunay_matern_model(*, mask_radius: float) -> af.Collection: return af.Collection(galaxies=af.Collection(lens=lens, source=source)) -def _point_source_model() -> af.Collection: +def _point_source_model(solved: bool = False) -> af.Collection: mass, _ = _lens_mass_and_shear() # No shear for the point-source profile. lens = af.Model(al.Galaxy, redshift=0.5, mass=mass) - point_0 = af.Model(al.ps.PointFlux) + # solved: parameter-free PointSolved — the *Solved fit classes solve the source + # centre analytically (#657), dropping 3 free parameters vs PointFlux. + point_0 = af.Model(al.ps.PointSolved) if solved else af.Model(al.ps.PointFlux) source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0) return af.Collection(galaxies=af.Collection(lens=lens, source=source)) @@ -853,9 +857,12 @@ def _build_analysis( pixel_scale_precision=solver_kwargs["pixel_scale_precision"], magnification_threshold=solver_kwargs["magnification_threshold"], ) - fit_positions_cls = ( - al.FitPositionsImagePairAll if model_type == "image_plane" else al.FitPositionsSource - ) + fit_positions_cls = { + "image_plane": al.FitPositionsImagePairAll, + "source_plane": al.FitPositionsSource, + "image_plane_solved": al.FitPositionsImagePairAllSolved, + "source_plane_solved": al.FitPositionsSourceSolved, + }[model_type] return al.AnalysisPoint( dataset=dataset, solver=solver, diff --git a/scripts/misc/searches/sweep.py b/scripts/misc/searches/sweep.py index 545438720..4f3c4be32 100644 --- a/scripts/misc/searches/sweep.py +++ b/scripts/misc/searches/sweep.py @@ -82,6 +82,9 @@ def _profiling_root() -> _Path: CELLS: list[tuple[str, str, str]] = [ ("nautilus", "point_source", "image_plane"), ("nautilus", "point_source", "source_plane"), + # Gradient search on the solver-chained solved-centre likelihood — enabled by the + # PointSolver implicit-diff custom_jvp (#657 phase 5). JAX-only. + ("multi_start_prodigy", "point_source", "image_plane_solved"), ("nautilus", "imaging", "mge"), ("nautilus", "imaging", "pixelization"), ("nautilus", "imaging", "delaunay"), diff --git a/scripts/point_source/searches/multi_start_prodigy/image_plane_solved.py b/scripts/point_source/searches/multi_start_prodigy/image_plane_solved.py new file mode 100644 index 000000000..3ec1ee53e --- /dev/null +++ b/scripts/point_source/searches/multi_start_prodigy/image_plane_solved.py @@ -0,0 +1,44 @@ +"""af.MultiStartProdigy search profiling — point-source image-plane solved centre. + +JAX multi-start Prodigy gradient MAP optimizer on the solver-chained +``FitPositionsImagePairAll``-family likelihood with an ``al.ps.PointSolved`` +source (``FitPositionsImagePairAllSolved``): the source centre is solved +analytically and the lens-equation solve carries implicit-diff gradients +(``autolens.point.solver.implicit_diff``, #657 phase 5) — the first gradient +search able to run on an image-plane point-source likelihood. Prodigy +self-tunes its learning rate. See ``searches/README.md``. +""" + +from __future__ import annotations + +import sys as _sys +from pathlib import Path as _Path + + +def _profiling_root() -> _Path: + for _p in _Path(__file__).resolve().parents: + if (_p / "ruff.toml").exists(): + return _p + raise RuntimeError("autolens_profiling root (ruff.toml) not found") + + +_misc_dir = str(_profiling_root() / "scripts" / "misc") +if _misc_dir not in _sys.path: + _sys.path.insert(0, _misc_dir) + + +import sys +from pathlib import Path + +_REPO_ROOT = _profiling_root() # 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_prodigy", + dataset_class="point_source", + model_type="image_plane_solved", + default_instrument="simple", +)