diff --git a/autolens/cluster/plot/cluster_plots.py b/autolens/cluster/plot/cluster_plots.py index 26a0e9ea8..32e2aa98b 100644 --- a/autolens/cluster/plot/cluster_plots.py +++ b/autolens/cluster/plot/cluster_plots.py @@ -18,10 +18,25 @@ z=1 and z=2 source planes differ visibly, and plotting only the last plane's curve (the galaxy-scale default) under-reports the lensing structure. -Numerical note: the critical-curve solvers walk a uniform grid with a contour finder -and are numpy-path only (they are not ``jax.vmap``-safe). At arcminute scale a grid -coarser than ~1" per pixel starts missing member-galaxy-scale features; the -docstrings below give concrete grid guidance. +Critical-curve engine: like every other PyAuto plot, the curve and caustic helpers +honour ``general.critical_curves_method`` in the visualize config, read once per plot +call through PyAutoGalaxy's ``autogalaxy.util.plot_utils._critical_curves_method()`` +(which also applies its unknown-value and missing-``jax_zero_contour`` fallbacks to +``marching_squares``): + +- ``marching_squares`` (the default) evaluates the eigen values on the supplied uniform + grid and contours them. At arcminute scale a grid coarser than ~1" per pixel starts + missing member-galaxy-scale features; the docstrings below give concrete grid + guidance. +- ``zero_contour`` traces the zero contour of each eigen value with + ``jax_zero_contour`` via ``LensCalc``'s ``*_via_zero_contour_from`` methods, called + with their defaults exactly as the other PyAuto plots call them. The grid is then + unused, and seed points come from ``LensCalc``'s automatic coarse-grid scan, which + covers only a galaxy-scale (+/-3") box about the origin, so curves outside it are not + found. + +Neither path is claimed to be ``jax.jit`` / ``jax.vmap`` safe here: these are plotting +helpers that run on the numpy path. """ from typing import List, Optional, Sequence, Tuple @@ -171,6 +186,51 @@ def _lens_calc_for_plane(tracer, plane_index: int): return LensCalc.from_tracer(tracer, use_multi_plane=True, plane_j=plane_index) +# LensCalc routine names per engine, keyed by (eigen value kind, curve type). The +# marching-squares routines take the uniform ``grid``; the zero-contour routines take +# no grid and are called with their defaults, mirroring +# ``autogalaxy.util.plot_utils._critical_curves_from`` / ``_caustics_from``. +_CURVE_ROUTINES = { + "marching_squares": { + ("tangential", "critical_curve"): "tangential_critical_curve_list_from", + ("radial", "critical_curve"): "radial_critical_curve_list_from", + ("tangential", "caustic"): "tangential_caustic_list_from", + ("radial", "caustic"): "radial_caustic_list_from", + }, + "zero_contour": { + ( + "tangential", + "critical_curve", + ): "tangential_critical_curve_list_via_zero_contour_from", + ("radial", "critical_curve"): "radial_critical_curve_list_via_zero_contour_from", + ("tangential", "caustic"): "tangential_caustic_list_via_zero_contour_from", + ("radial", "caustic"): "radial_caustic_list_via_zero_contour_from", + }, +} + + +def _critical_curves_method() -> str: + """The configured critical-curve engine, via PyAutoGalaxy's shared selector. + + Returns ``"marching_squares"`` or ``"zero_contour"``; unknown values and a + missing ``jax_zero_contour`` fall back to ``"marching_squares"`` with a warning, + exactly as for every other PyAuto plot. + """ + from autogalaxy.util.plot_utils import _critical_curves_method + + return _critical_curves_method() + + +def _curve_list_from(lens_calc, grid, method: str, kind: str, curve_type: str): + """Compute one plane's ``kind`` (tangential / radial) critical curves or caustics + (``curve_type`` = ``"critical_curve"`` / ``"caustic"``) with the engine *method*. + """ + routine = getattr(lens_calc, _CURVE_ROUTINES[method][(kind, curve_type)]) + if method == "zero_contour": + return routine() + return routine(grid=grid) + + def plot_positions_overlay( positions_list, image=None, @@ -366,16 +426,25 @@ def plot_critical_curves( the positions helpers use for the corresponding source, and labelled with the plane redshift. + The curve engine is ``general.critical_curves_method`` from the visualize config, + read once per call (see the module docstring): ``marching_squares`` (default) + contours the eigen values on *grid*; ``zero_contour`` traces them with + ``jax_zero_contour`` through each plane's ``LensCalc`` + ``*_critical_curve_list_via_zero_contour_from`` methods and ignores *grid*. A + configured ``zero_contour`` without ``jax_zero_contour`` installed falls back to + ``marching_squares`` with a warning. + Parameters ---------- tracer The (multi-plane) tracer whose critical curves are computed via the multi-plane ``LensCalc`` (``use_multi_plane=True, plane_j=j``), numpy path. grid - The uniform grid the marching solver walks. At arcminute scale, resolve to - ~0.5" per pixel or finer over the full field (e.g. ``Grid2D.uniform( - shape_native=(240, 240), pixel_scales=0.5)`` for a 2' field) — coarser grids - miss member-galaxy-scale wiggles in the curves. + The uniform grid the marching-squares solver walks (unused by + ``zero_contour``). At arcminute scale, resolve to ~0.5" per pixel or finer + over the full field (e.g. ``Grid2D.uniform(shape_native=(240, 240), + pixel_scales=0.5)`` for a 2' field) — coarser grids miss + member-galaxy-scale wiggles in the curves. image, pixel_scales Optional background image for the overlay. plane_indices @@ -393,12 +462,16 @@ def plot_critical_curves( if image is not None: _draw_image(ax, image, pixel_scales) + method = _critical_curves_method() + for j in _source_plane_indices(tracer, plane_indices): colour = WONG_PALETTE[(j - 1) % len(WONG_PALETTE)] redshift = float(tracer.planes[j].redshift) lens_calc = _lens_calc_for_plane(tracer, plane_index=j) - tangential_list = lens_calc.tangential_critical_curve_list_from(grid=grid) + tangential_list = _curve_list_from( + lens_calc, grid, method, "tangential", "critical_curve" + ) for k, curve in enumerate(tangential_list): arr = _positions_array(curve) ax.plot( @@ -411,7 +484,9 @@ def plot_critical_curves( ) if include_radial: - radial_list = lens_calc.radial_critical_curve_list_from(grid=grid) + radial_list = _curve_list_from( + lens_calc, grid, method, "radial", "critical_curve" + ) for k, curve in enumerate(radial_list): arr = _positions_array(curve) ax.plot( @@ -458,10 +533,17 @@ def plot_caustics( where each source must sit relative to its own caustic network to produce the observed image multiplicities. + The engine follows ``general.critical_curves_method``, read once per call, exactly + as in ``plot_critical_curves``: ``marching_squares`` (default) uses each plane's + ``LensCalc`` ``*_caustic_list_from(grid=grid)``; ``zero_contour`` uses + ``*_caustic_list_via_zero_contour_from()``. Either way each caustic is that plane's + critical curve ray-traced by that plane's multi-plane deflections. + Parameters ---------- tracer, grid, plane_indices - As in ``plot_critical_curves`` (same grid-resolution guidance applies). + As in ``plot_critical_curves`` (same engine and grid-resolution guidance + applies). include_radial Also draw each plane's radial caustic (dashed); at cluster scale the radial caustic bounds the region producing central demagnified images, so it is on @@ -472,12 +554,16 @@ def plot_caustics( if standalone: fig, ax = subplots(1, 1) + method = _critical_curves_method() + for j in _source_plane_indices(tracer, plane_indices): colour = WONG_PALETTE[(j - 1) % len(WONG_PALETTE)] redshift = float(tracer.planes[j].redshift) lens_calc = _lens_calc_for_plane(tracer, plane_index=j) - tangential_list = lens_calc.tangential_caustic_list_from(grid=grid) + tangential_list = _curve_list_from( + lens_calc, grid, method, "tangential", "caustic" + ) for k, curve in enumerate(tangential_list): arr = _positions_array(curve) ax.plot( @@ -490,7 +576,7 @@ def plot_caustics( ) if include_radial: - radial_list = lens_calc.radial_caustic_list_from(grid=grid) + radial_list = _curve_list_from(lens_calc, grid, method, "radial", "caustic") for k, curve in enumerate(radial_list): arr = _positions_array(curve) ax.plot( diff --git a/test_autolens/cluster/plot/test_cluster_plots.py b/test_autolens/cluster/plot/test_cluster_plots.py index 6361fca96..45c5405e7 100644 --- a/test_autolens/cluster/plot/test_cluster_plots.py +++ b/test_autolens/cluster/plot/test_cluster_plots.py @@ -135,3 +135,241 @@ def test__subplot_cluster_dataset( ) assert str(plot_path / "subplot_cluster_dataset.png") in plot_patch.paths + + +##################################################################### +# Critical-curve engine dispatch (general.critical_curves_method). # +##################################################################### + +import sys + +from autogalaxy.operate.lens_calc import LensCalc + +from autolens.cluster.plot import cluster_plots + +_ROUTINES = { + "marching_squares": [ + "tangential_critical_curve_list_from", + "radial_critical_curve_list_from", + "tangential_caustic_list_from", + "radial_caustic_list_from", + ], + "zero_contour": [ + "tangential_critical_curve_list_via_zero_contour_from", + "radial_critical_curve_list_via_zero_contour_from", + "tangential_caustic_list_via_zero_contour_from", + "radial_caustic_list_via_zero_contour_from", + ], +} + + +@pytest.fixture(name="power_law_multi_plane_tracer") +def make_power_law_multi_plane_tracer(): + """Two-source-plane cluster-like system whose shallower-than-isothermal lens + (slope 1.7, analytic deflections so the test stays fast) has tangential *and* + radial critical curves in both source planes (z=1 and z=2).""" + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.PowerLaw( + centre=(0.0, 0.0), ell_comps=(0.05, 0.0), einstein_radius=1.6, slope=1.7 + ), + ) + return al.Tracer(galaxies=[lens, al.Galaxy(redshift=1.0), al.Galaxy(redshift=2.0)]) + + +@pytest.fixture(name="engine_spy") +def make_engine_spy(monkeypatch, grid): + """Record every LensCalc curve routine called, keyed by the source plane whose + ``LensCalc`` it was called on. + + Marching-squares routines run for real. Zero-contour routines are stubbed (a real + ``jax_zero_contour`` trace costs several seconds per curve) with the matching + marching-squares result *of the same per-plane LensCalc*, so the plotted + geometry still comes from that plane's multi-plane ray-tracing. + """ + calls = [] + depth = [0] + real_lens_calc_for_plane = cluster_plots._lens_calc_for_plane + real_ms = {name: getattr(LensCalc, name) for name in _ROUTINES["marching_squares"]} + + def lens_calc_for_plane(tracer, plane_index): + lens_calc = real_lens_calc_for_plane(tracer, plane_index=plane_index) + lens_calc._spy_plane = plane_index + return lens_calc + + monkeypatch.setattr(cluster_plots, "_lens_calc_for_plane", lens_calc_for_plane) + + def ms_spy(name): + def spy(self, *args, **kwargs): + # Only record calls made by the plot helpers, not the critical-curve + # call a marching-squares caustic routine makes internally. + if depth[0] == 0: + calls.append((self._spy_plane, name, args, kwargs)) + depth[0] += 1 + try: + return real_ms[name](self, *args, **kwargs) + finally: + depth[0] -= 1 + + return spy + + def zc_spy(name, ms_name): + def spy(self, *args, **kwargs): + calls.append((self._spy_plane, name, args, kwargs)) + depth[0] += 1 + try: + return real_ms[ms_name](self, grid=grid) + finally: + depth[0] -= 1 + + return spy + + for name in _ROUTINES["marching_squares"]: + monkeypatch.setattr(LensCalc, name, ms_spy(name)) + for name, ms_name in zip(_ROUTINES["zero_contour"], _ROUTINES["marching_squares"]): + monkeypatch.setattr(LensCalc, name, zc_spy(name, ms_name)) + + return calls + + +def _set_engine(monkeypatch, method): + """Configure ``general.critical_curves_method`` in the visualize config.""" + from autonerves import conf + + monkeypatch.setitem( + conf.instance["visualize"]["general"]["general"], + "critical_curves_method", + method, + ) + + +def _median_radius_by_label(ax): + """``{legend label: median |(y, x)|}`` of the first line drawn per label.""" + return { + line.get_label(): float(np.median(np.hypot(line.get_xdata(), line.get_ydata()))) + for line in ax.get_lines() + if not line.get_label().startswith("_") + } + + +@pytest.mark.parametrize("method", ["marching_squares", "zero_contour"]) +def test__plot_critical_curves__routes_each_plane_to_configured_engine( + method, power_law_multi_plane_tracer, grid, engine_spy, monkeypatch +): + import matplotlib.pyplot as plt + + selector_calls = [] + + def selector(): + selector_calls.append(1) + return method + + monkeypatch.setattr(cluster_plots, "_critical_curves_method", selector) + + fig, ax = plt.subplots() + plot_critical_curves( + power_law_multi_plane_tracer, grid=grid, include_radial=True, ax=ax + ) + + tangential, radial = _ROUTINES[method][0], _ROUTINES[method][1] + assert [(plane, name) for plane, name, _, _ in engine_spy] == [ + (1, tangential), + (1, radial), + (2, tangential), + (2, radial), + ] + for _, _, args, kwargs in engine_spy: + if method == "zero_contour": + assert args == () and kwargs == {} + else: + assert kwargs == {"grid": grid} + assert len(selector_calls) == 1 + + radii = _median_radius_by_label(ax) + assert set(radii) == { + "z=1.00 tangential", + "z=1.00 radial", + "z=2.00 tangential", + "z=2.00 radial", + } + # The z=2 plane sees a stronger lens, so its curves differ from the z=1 plane's. + assert radii["z=2.00 tangential"] > radii["z=1.00 tangential"] * 1.05 + assert not np.isclose(radii["z=2.00 radial"], radii["z=1.00 radial"], rtol=0.02) + plt.close(fig) + + +@pytest.mark.parametrize("method", ["marching_squares", "zero_contour"]) +def test__plot_caustics__routes_each_plane_to_configured_engine( + method, power_law_multi_plane_tracer, grid, engine_spy, monkeypatch +): + import matplotlib.pyplot as plt + + _set_engine(monkeypatch, method) + monkeypatch.setitem(sys.modules, "jax_zero_contour", object()) + + fig, ax = plt.subplots() + plot_caustics(power_law_multi_plane_tracer, grid=grid, ax=ax) + + tangential, radial = _ROUTINES[method][2], _ROUTINES[method][3] + assert [(plane, name) for plane, name, _, _ in engine_spy] == [ + (1, tangential), + (1, radial), + (2, tangential), + (2, radial), + ] + + radii = _median_radius_by_label(ax) + assert set(radii) == { + "z=1.00 tangential", + "z=1.00 radial", + "z=2.00 tangential", + "z=2.00 radial", + } + # Source-plane mapping: each plane's caustics are its own critical curves traced + # through its own deflections, so the two planes' caustics differ too. + assert not np.isclose( + radii["z=2.00 tangential"], radii["z=1.00 tangential"], rtol=0.02 + ) + assert not np.isclose(radii["z=2.00 radial"], radii["z=1.00 radial"], rtol=0.02) + plt.close(fig) + + +def test__plot_critical_curves__tangential_only_by_default( + power_law_multi_plane_tracer, grid, engine_spy, monkeypatch, plot_patch +): + _set_engine(monkeypatch, "zero_contour") + monkeypatch.setitem(sys.modules, "jax_zero_contour", object()) + + plot_critical_curves(power_law_multi_plane_tracer, grid=grid, plane_indices=[2]) + + assert [(plane, name) for plane, name, _, _ in engine_spy] == [ + (2, "tangential_critical_curve_list_via_zero_contour_from") + ] + + +def test__zero_contour_without_dependency__falls_back_to_marching_squares( + power_law_multi_plane_tracer, grid, engine_spy, monkeypatch, caplog +): + import logging + import matplotlib.pyplot as plt + + _set_engine(monkeypatch, "zero_contour") + # A ``None`` entry makes ``import jax_zero_contour`` raise ImportError. + monkeypatch.setitem(sys.modules, "jax_zero_contour", None) + + fig, ax = plt.subplots() + with caplog.at_level(logging.WARNING, logger="autogalaxy.util.plot_utils"): + plot_critical_curves( + power_law_multi_plane_tracer, grid=grid, include_radial=True, ax=ax + ) + plot_caustics(power_law_multi_plane_tracer, grid=grid, ax=ax) + plt.close(fig) + + assert {name for _, name, _, _ in engine_spy} == set(_ROUTINES["marching_squares"]) + fallback_warnings = [ + record + for record in caplog.records + if "jax_zero_contour is not installed" in record.getMessage() + ] + # Read once per plot call, so one warning each. + assert len(fallback_warnings) == 2