feat: multi_start_prodigy point_source image_plane_solved cell (phase 5 of #657)

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,

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"),

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",
+)