diff --git a/autolens/__init__.py b/autolens/__init__.py index b535155a8..66b6c8b6a 100644 --- a/autolens/__init__.py +++ b/autolens/__init__.py @@ -95,6 +95,7 @@ from . import aggregator as agg from .analysis import model_util +from .lens import mappings from .lens import subhalo from .lens.tracer import Tracer from .lens.sensitivity import SubhaloSensitivityResult diff --git a/autolens/config/visualize/plots.yaml b/autolens/config/visualize/plots.yaml index 5b78b606d..5c6b897b7 100644 --- a/autolens/config/visualize/plots.yaml +++ b/autolens/config/visualize/plots.yaml @@ -40,7 +40,7 @@ tracer: # Settings for plots of tracers (e.g. inversion: # Settings for plots of inversions (e.g. InversionPlotter). subplot_inversion: true # Plot subplot of all quantities in each inversion (e.g. reconstrucuted image, reconstruction)? - subplot_mappings: false # Plot subplot of the image-to-source pixels mappings of each pixelization? + subplot_mappings: false # Plot subplot of how the brightest source regions map to their multiple images in the image-plane (works for pixelized and parametric sources)? csv_reconstruction: true # output source_plane_reconstruction_0.csv containing the source-plane mesh y, x, reconstruction and noise map values. adapt: # Settings for plots of adapt images used by adaptive pixelizations. diff --git a/autolens/imaging/model/plotter.py b/autolens/imaging/model/plotter.py index 3f8481bc3..4fad58276 100644 --- a/autolens/imaging/model/plotter.py +++ b/autolens/imaging/model/plotter.py @@ -15,6 +15,7 @@ subplot_fit, subplot_fit_log10, subplot_of_planes, + subplot_mappings, subplot_tracer_from_fit, subplot_fit_combined, subplot_fit_combined_log10, @@ -139,6 +140,21 @@ def should_plot(name): subplot_of_planes(fit, output_path=output_path, output_format=fmt, title_prefix=self.title_prefix) + # The `subplot_mappings` key lives in the `inversion` section of `plots.yaml`, + # where it has sat unread since the mapping visuals were removed in 2025. The + # fit-level subplot it now switches on works for a parametric source too, so it is + # driven from here (which has the fit) rather than from `Plotter.inversion` (which + # has only the inversion, and is never called for a parametric model). + if plot_setting(section="inversion", name="subplot_mappings"): + for plane_index in plane_indexes_to_plot: + subplot_mappings( + fit, plane_index=plane_index, output_path=output_path, + output_format=fmt, + image_plane_lines=ip_lines, image_plane_line_colors=ip_colors, + source_plane_lines=sp_lines, source_plane_line_colors=sp_colors, + title_prefix=self.title_prefix, + ) + if should_plot("fits_fit"): fits_fit(fit=fit, output_path=self.image_path) diff --git a/autolens/imaging/plot/fit_imaging_plots.py b/autolens/imaging/plot/fit_imaging_plots.py index 443eb0743..6fa8bf72a 100644 --- a/autolens/imaging/plot/fit_imaging_plots.py +++ b/autolens/imaging/plot/fit_imaging_plots.py @@ -7,7 +7,7 @@ from autogalaxy.util.plot_utils import plot_array from autoarray.plot.array import _zoom_array_2d -from autoarray.plot.utils import subplots, save_figure, hide_unused_axes, conf_subplot_figsize, tight_layout +from autoarray.plot.utils import subplots, save_figure, hide_unused_axes, conf_subplot_figsize, tight_layout, plot_regions from autoarray.inversion.mappers.abstract import Mapper from autoarray.inversion.plot.mapper_plots import plot_mapper from autogalaxy.util.plot_utils import _critical_curves_from @@ -68,7 +68,9 @@ def _get_source_vmax(fit): def _plot_source_plane(fit, ax, plane_index, zoom_to_brightest=True, colormap=None, use_log10=False, title=None, lines=None, line_colors=None, vmax=None, - zoom_extent_scale: float = 1.0): + zoom_extent_scale: float = 1.0, + regions=None, region_alpha: float = 0.25, + region_labels=None): """ Plot the source-plane image (or a blank inversion placeholder) into an axes. @@ -101,6 +103,15 @@ def _plot_source_plane(fit, ax, plane_index, zoom_to_brightest=True, Matplotlib colormap name. use_log10 : bool, optional If ``True`` the colour scale is applied on a log10 stretch. + regions : list, optional + Source-plane regions to overlay as filled polygons, each a list of ``(N, 2)`` + ``(y, x)`` arrays -- e.g. the ``source_contours`` of a + `~autoarray.inversion.mappings.mapping.Mapping`. Drawn identically on the + parametric and the inversion path, so a mapping keeps its colour across both. + region_alpha : float, optional + The alpha of each region's fill. + region_labels : list of str, optional + A label drawn at the centre of each region, matching its image-plane counterpart. """ tracer = fit.tracer_linear_light_profiles_to_light_profiles if not tracer.planes[plane_index].has(cls=aa.Pixelization): @@ -133,6 +144,13 @@ def _plot_source_plane(fit, ax, plane_index, zoom_to_brightest=True, colormap=colormap, use_log10=use_log10, vmax=vmax, lines=lines, line_colors=line_colors, ) + # `autogalaxy.util.plot_utils.plot_array` is a thin wrapper which does not forward + # `regions=`; the overlay is drawn onto the same axes afterwards instead, using the + # same `plot_regions` the pixelized path below reaches through `plot_mapper`. The + # polygons are arcsec `(y, x)`, so they land correctly whatever zoom the wrapper + # applied. + plot_regions(ax=ax, regions=regions, region_alpha=region_alpha, + region_labels=region_labels) else: # Inversion path: plot the source reconstruction via the mapper. try: @@ -152,6 +170,9 @@ def _plot_source_plane(fit, ax, plane_index, zoom_to_brightest=True, zoom_extent_scale=zoom_extent_scale, lines=lines, line_colors=line_colors, + regions=regions, + region_alpha=region_alpha, + region_labels=region_labels, ) except Exception as exc: logger.warning(f"Could not plot source reconstruction for plane {plane_index}: {exc}") @@ -962,3 +983,213 @@ def _symmetric_vmax(array) -> float: vals = np.asarray(array) finite = vals[np.isfinite(vals)] return float(np.max(np.abs(finite))) if finite.size else 1.0 + + +def subplot_mappings( + fit, + plane_index: int = -1, + shape=None, + output_path: Optional[str] = None, + output_filename: str = "mappings", + output_format: str = None, + colormap: Optional[str] = None, + threshold: Optional[float] = None, + min_pixels: Optional[int] = None, + total_clumps: Optional[int] = None, + pix_indexes: Optional[List] = None, + weight_threshold: float = 0.0, + pixel_scale_precision: Optional[float] = None, + magnification_threshold: float = 0.1, + region_alpha: float = 0.25, + image_plane_lines=None, + image_plane_line_colors=None, + source_plane_lines=None, + source_plane_line_colors=None, + title_prefix: str = None, +): + """ + 2x2 subplot showing, in one look, how the brightest regions of the source map to the + image plane and back. + + The panels are: + + * (0, 0) the lens-light-subtracted data, with the image-plane regions and the critical + curves overlaid; + * (0, 1) the model image of the source plane, with the same regions; + * (1, 0) the source plane zoomed to its brightest region, with the source regions and + the caustics; + * (1, 1) the source plane at full extent, with the same. + + Every source region is filled in its own colour and numbered, and the multiple images + it maps to carry the same colour and the same number in the image-plane panels -- so a + reader can read a merging pair of source galaxies off the figure by colour alone. + + This is the fit-level counterpart of + :func:`~autoarray.inversion.plot.inversion_plots.subplot_mappings`, which works only + for a pixelized source. It works for **either** kind of source: a plane which + ``has(cls=aa.Pixelization)`` is mapped via the inversion's clumps, and a parametric + plane via `ShapeSolver`'s ray-traced triangles. See + :mod:`autolens.lens.mappings` for the two engines. + + Parameters + ---------- + fit : FitImaging + The imaging fit to visualise. + plane_index : int, optional + The index of the source plane whose mappings are drawn. Defaults to the last plane. + shape : Shape, optional + The source-plane region for a parametric source. ``None`` uses a `Circle` at the + first source light profile's centre, with its half-light radius. Ignored for a + pixelized source, which finds its own clumps. + output_path : str, optional + Directory in which to save the figure. ``None`` does not save it. + output_filename : str, optional + The filename stem; ``_{plane_index}`` is appended, as the sibling subplots do. + output_format : str, optional + Image format passed to :func:`~autoarray.plot.utils.save_figure`. + colormap : str, optional + Matplotlib colormap name applied to all image panels. + threshold, min_pixels, total_clumps, pix_indexes, weight_threshold + The clump-finding controls of the pixelized (engine A) path, passed through to + `Inversion.mappings_from`. ``threshold`` is the knob which decides what counts as + a distinct source structure: ~0.5 isolates one smooth source, ~0.2 merges two + nearby galaxies, ~0.8 splits a source into its star-forming knots. + pixel_scale_precision, magnification_threshold + The solver controls of the parametric (engine B) path. + region_alpha : float, optional + The alpha of each drawn region's fill. + image_plane_lines, image_plane_line_colors, source_plane_lines, source_plane_line_colors + Pre-computed critical curves and caustics; computed from the fit when not given. + title_prefix : str, optional + A string prefixed to every panel title. + """ + from autolens.lens import mappings as mappings_module + + plane_index_tag = len(fit.tracer.planes) - 1 if plane_index == -1 else plane_index + + clump_kwargs = {} + + if threshold is not None: + clump_kwargs["threshold"] = threshold + if min_pixels is not None: + clump_kwargs["min_pixels"] = min_pixels + if total_clumps is not None: + clump_kwargs["total_clumps"] = total_clumps + if pix_indexes is not None: + clump_kwargs["pix_indexes"] = pix_indexes + + # No guard around this call. A source plane which cannot be mapped -- a plane with + # neither a pixelization nor a light profile, a profile with no half-light radius, a + # singular inversion -- raises, naming what the caller has to change. Swallowing that + # would draw the four panels with no regions on them, which reads as "this source has + # no multiple images" and is a worse outcome than a traceback. The one case which + # legitimately produces nothing to draw is an empty `mappings` list (below), which is + # a real answer rather than a failure. + mappings = mappings_module.mappings_from_fit( + fit=fit, + plane_index=plane_index, + shape=shape, + pixel_scale_precision=pixel_scale_precision, + magnification_threshold=magnification_threshold, + weight_threshold=weight_threshold, + **clump_kwargs, + ) + + image_regions = [mapping.image_contours for mapping in mappings] + source_regions = [mapping.source_contours for mapping in mappings] + region_labels = [str(i + 1) for i in range(len(mappings))] + + if image_plane_lines is None and source_plane_lines is None: + ( + image_plane_lines, + image_plane_line_colors, + source_plane_lines, + source_plane_line_colors, + ) = _compute_critical_curves_from_fit(fit) + + source_vmax = _get_source_vmax(fit) + + _pf = (lambda t: f"{title_prefix.rstrip()} {t}") if title_prefix else (lambda t: t) + + fig, axes = subplots(2, 2, figsize=conf_subplot_figsize(2, 2)) + axes_flat = list(axes.flatten()) + + # panel (0, 0): the data with everything but this plane subtracted. + try: + subtracted_image = fit.subtracted_images_of_planes_list[plane_index] + except (IndexError, AttributeError): + subtracted_image = None + + if subtracted_image is not None: + plot_array( + array=subtracted_image, + ax=axes_flat[0], + title=_pf("Lens Light Subtracted"), + colormap=colormap, + vmin=0.0 if source_vmax is not None else None, + vmax=source_vmax, + lines=image_plane_lines, + line_colors=image_plane_line_colors, + ) + plot_regions( + ax=axes_flat[0], + regions=image_regions, + region_alpha=region_alpha, + region_labels=region_labels, + ) + else: + axes_flat[0].axis("off") + + # panel (0, 1): the model image of the plane, with the same regions. + try: + model_image = fit.model_images_of_planes_list[plane_index] + except (IndexError, AttributeError): + model_image = None + + if model_image is not None: + plot_array( + array=model_image, + ax=axes_flat[1], + title=_pf("Source Model Image"), + colormap=colormap, + vmax=source_vmax, + lines=image_plane_lines, + line_colors=image_plane_line_colors, + ) + plot_regions( + ax=axes_flat[1], + regions=image_regions, + region_alpha=region_alpha, + region_labels=region_labels, + ) + else: + axes_flat[1].axis("off") + + # panels (1, 0) and (1, 1): the source plane, zoomed and unzoomed. + for ax, zoom, title in ( + (axes_flat[2], True, "Source Plane (Zoom)"), + (axes_flat[3], False, "Source Plane (No Zoom)"), + ): + _plot_source_plane( + fit, + ax, + plane_index, + zoom_to_brightest=zoom, + colormap=colormap, + title=_pf(title), + lines=source_plane_lines, + line_colors=source_plane_line_colors, + vmax=source_vmax, + regions=source_regions, + region_alpha=region_alpha, + region_labels=region_labels, + ) + + hide_unused_axes(axes_flat) + tight_layout() + save_figure( + fig, + path=output_path, + filename=f"{output_filename}_{plane_index_tag}", + format=output_format, + ) diff --git a/autolens/lens/mappings.py b/autolens/lens/mappings.py new file mode 100644 index 000000000..3059ef789 --- /dev/null +++ b/autolens/lens/mappings.py @@ -0,0 +1,475 @@ +""" +How a region of the source plane maps to the image plane, for any kind of source. + +A **mapping** pairs one source-plane region with the image-plane regions (the multiple +images) it maps to. `autoarray.inversion.mappings` builds one for a *pixelized* source +from a `Mapper`'s mapping matrix (engine A); this module builds one for a *parametric* +source from `ShapeSolver`'s ray-traced triangles (engine B), and dispatches between the +two given a `FitImaging`, so a caller does not have to know which kind of source a model +used. + +Both engines return the same `Mapping` and `ImageRegion` objects, so everything +downstream -- `subplot_mappings`, the brightest-pixel positions spectroscopic follow-up +needs, the per-image magnifications -- is written once. + +Everything here is pure numpy and is never jitted: it is diagnostic and visualization +code which runs once per figure or once per result, not inside a likelihood. +""" + +import logging +from typing import List, Optional, Tuple + +import numpy as np + +import autoarray as aa +import autogalaxy as ag + +from autoarray.structures.triangles.shape import Circle, Shape + +from autolens.lens.tracer import Tracer +from autolens.point.solver.shape_solver import ShapeSolver + +logger = logging.getLogger(__name__) + + +def source_mapping_from( + tracer: Tracer, + grid: aa.type.Grid2DLike, + shape: Shape, + plane_index: int = -1, + plane_redshift: Optional[float] = None, + pixel_scale_precision: Optional[float] = None, + magnification_threshold: float = 0.1, +) -> aa.Mapping: + """ + The `Mapping` of a source-plane `Shape` through a tracer, via `ShapeSolver` (engine B). + + This is the thin entry point: it builds the solver for `grid` and delegates to + `ShapeSolver.mapping_from`. + + Parameters + ---------- + tracer + The tracer which ray-traces the image plane to the source plane. + grid + The image-plane grid whose extent is tiled with triangles and whose pixels the + image regions are reported on. + shape + The source-plane region, e.g. a `Circle` at a light profile's centre. + plane_index + The index of the plane `shape` lives in, used only when `plane_redshift` is `None`. + plane_redshift + The redshift of the plane `shape` lives in. `None` derives it from `plane_index`. + pixel_scale_precision + The triangle size the solver refines down to. `None` uses a tenth of the grid's + pixel scale, which resolves the image of a source a few pixels across to well + inside a pixel while keeping the solve fast. + magnification_threshold + Images whose magnification does not exceed this are discarded, which removes the + formally-infinitely-demagnified central image of a singular profile. + + Returns + ------- + The mapping of the shape. + """ + if pixel_scale_precision is None: + pixel_scale_precision = float(grid.pixel_scale) / 10.0 + + if plane_redshift is None: + plane_redshift = float(tracer.plane_redshifts[plane_index]) + + solver = ShapeSolver.for_grid( + grid=grid, + pixel_scale_precision=pixel_scale_precision, + magnification_threshold=magnification_threshold, + ) + + return solver.mapping_from( + tracer=tracer, + shape=shape, + plane_redshift=plane_redshift, + ) + + +def traced_region_from( + tracer: Tracer, + region_grid: aa.type.Grid2DLike, + plane_index: int = -1, +) -> np.ndarray: + """ + The source-plane outline of an image-plane region, as the convex hull of its traced + coordinates. + + An image-plane region is a blob of pixels; ray-tracing it gives a cloud of source-plane + coordinates whose outline is what the source-plane panel of a mapping figure draws. A + convex hull is used rather than an exact outline because the traced cloud of one + multiple image is a small, roughly convex patch of the source plane -- it is the *image* + which is stretched into an arc, not its source. + + Parameters + ---------- + tracer + The tracer which ray-traces the image plane to the source plane. + region_grid + The ``(N, 2)`` ``(y, x)`` image-plane coordinates of the region, e.g. an + `ImageRegion`'s `scaled_coordinates`. + plane_index + The plane the coordinates are traced to. + + Returns + ------- + An ``(M, 2)`` array of ``(y, x)`` source-plane coordinates forming a closed loop + (its first row is repeated as its last). Fewer than three traced coordinates, or a + degenerate (collinear) cloud, returns the traced coordinates themselves. + """ + from scipy.spatial import ConvexHull, QhullError + + if not isinstance(region_grid, aa.Grid2DIrregular): + region_grid = aa.Grid2DIrregular(values=np.asarray(region_grid)) + + traced = np.asarray( + tracer.traced_grid_2d_list_from(grid=region_grid)[plane_index].array + ) + + if traced.shape[0] < 3: + return traced + + try: + hull = ConvexHull(traced) + except QhullError: + return traced + + vertices = traced[hull.vertices] + + return np.vstack([vertices, vertices[:1]]) + + +def _light_profile_shape_from(tracer: Tracer, plane_index: int) -> Circle: + """ + The default source-plane `Circle` for a parametric source: centred on the first light + profile of the plane, with a radius of its half-light radius. + + Parameters + ---------- + tracer + The tracer whose plane holds the source. + plane_index + The index of the source plane. + + Returns + ------- + The circle. + """ + light_profiles = tracer.planes[plane_index].cls_list_from(cls=ag.LightProfile) + + if len(light_profiles) == 0: + raise ValueError( + f"Plane {plane_index} of the tracer has neither a pixelization nor a light " + f"profile, so there is no source region to map. Pass `shape=` explicitly (e.g. " + f"`shape=aa.Circle(y, x, radius=...)`) to say which source-plane region you " + f"want the images of." + ) + + profile = light_profiles[0] + + radius = profile.half_light_radius + + if radius is None: + raise ValueError( + f"The source light profile {type(profile).__name__} has no `half_light_radius` " + f"(it defines no `effective_radius`), so the size of the source region cannot be " + f"inferred from it. Pass `shape=` explicitly, e.g. " + f"`shape=aa.Circle({profile.centre[0]}, {profile.centre[1]}, radius=...)`." + ) + + # `Circle` reads element 0 of a coordinate pair first, and element 0 of every PyAuto + # grid is `y`, so a profile centred at `(y_c, x_c)` gives `Circle(y_c, x_c, radius)` + # positionally. See the `Shape` class docstring in PyAutoArray. + return Circle( + float(profile.centre[0]), float(profile.centre[1]), radius=float(radius) + ) + + +def mappings_from_fit( + fit, + plane_index: int = -1, + shape: Optional[Shape] = None, + grid: Optional[aa.type.Grid2DLike] = None, + pixel_scale_precision: Optional[float] = None, + magnification_threshold: float = 0.1, + **clump_kwargs, +) -> List[aa.Mapping]: + """ + The mappings of a fit's source plane, whichever kind of source the model used. + + A plane which `has(cls=aa.Pixelization)` is mapped by **engine A**: the inversion's + reconstruction is split into clumps and each clump's image-plane regions come from the + mapper's mapping matrix. Any other plane is mapped by **engine B**: a `Circle` at the + first light profile's centre, with its half-light radius, is traced by `ShapeSolver`. + + Parameters + ---------- + fit + A `FitImaging` (or any fit exposing `tracer`, `mask` and, for engine A, `inversion`). + plane_index + The index of the source plane to map. + shape + The source-plane region for engine B. `None` builds it from the plane's first light + profile. Ignored by engine A, which finds its own clumps. + grid + The image-plane grid engine B tiles. `None` uses the fit's unmasked grid. + pixel_scale_precision + The triangle size engine B refines to; `None` uses a tenth of the pixel scale. + magnification_threshold + Engine B discards images below this magnification. + clump_kwargs + `threshold`, `min_pixels`, `total_clumps`, `pix_indexes` and `weight_threshold`, + passed to engine A's `Inversion.mappings_from`. Ignored by engine B. + + Returns + ------- + One `Mapping` per source-plane region: one per clump for engine A, one for engine B. + """ + tracer = fit.tracer_linear_light_profiles_to_light_profiles + + if tracer.planes[plane_index].has(cls=aa.Pixelization): + weight_threshold = clump_kwargs.pop("weight_threshold", 0.0) + + mapper_index = _mapper_index_from(fit=fit, plane_index=plane_index) + + return fit.inversion.mappings_from( + mapper_index=mapper_index, + weight_threshold=weight_threshold, + **clump_kwargs, + ) + + if shape is None: + shape = _light_profile_shape_from(tracer=tracer, plane_index=plane_index) + + if grid is None: + grid = fit.mask.derive_grid.all_false + + return [ + source_mapping_from( + tracer=tracer, + grid=grid, + shape=shape, + plane_index=plane_index, + pixel_scale_precision=pixel_scale_precision, + magnification_threshold=magnification_threshold, + ) + ] + + +def _mapper_index_from(fit, plane_index: int) -> int: + """ + The index into the inversion's mappers of the mapper belonging to `plane_index`. + + A fit has one mapper per pixelized plane and planes are ordered by redshift, so the + mapper of plane ``i`` is mapper ``i - 1`` for the usual lens-plus-source fit. The + normalisation of a negative `plane_index` is done here so the two agree. + + Parameters + ---------- + fit + The fit whose inversion holds the mappers. + plane_index + The index of the source plane. + + Returns + ------- + The mapper index, which is 0 whenever there is only one mapper. + """ + mapper_list = fit.inversion.cls_list_from(cls=aa.Mapper) + + if len(mapper_list) == 1: + return 0 + + total_planes = len(fit.tracer.planes) + + plane_index = plane_index % total_planes + + return max(plane_index - 1, 0) + + +def multiple_image_positions_from( + fit, + plane_index: int = -1, + use_centroid: bool = False, + **mapping_kwargs, +) -> aa.Grid2DIrregular: + """ + The brightest image-plane coordinate of each multiple image of the source. + + This is the quantity spectroscopic follow-up needs: where to point a fibre on each + lensed image. The brightness is read from the *model* image of the source plane, so + the positions describe the lens model rather than the noise realisation of the data. + + Parameters + ---------- + fit + The `FitImaging` whose source is mapped. + plane_index + The index of the source plane. + use_centroid + If `True` each position is the flux-weighted centroid of its image region rather + than its brightest pixel. The centroid is sub-pixel and smooth under small model + changes; the brightest pixel is exactly a data pixel, which is what a fibre is + pointed at. + mapping_kwargs + Passed to `mappings_from_fit`. + + Returns + ------- + A `Grid2DIrregular` of ``(y, x)`` arcsec coordinates, one per multiple image, ordered + as the image regions are (largest image first, mapping by mapping). + """ + model_image = fit.model_images_of_planes_list[plane_index] + + mappings = mappings_from_fit(fit=fit, plane_index=plane_index, **mapping_kwargs) + + positions = [] + + for mapping in mappings: + for region in mapping.image_regions: + if len(region.slim_indexes) == 0: + continue + + if use_centroid: + positions.append(region.centroid_from(array=model_image)) + else: + positions.append(region.brightest_coordinate_from(array=model_image)) + + if len(positions) == 0: + # A shape-(0,) irregular grid would not be indexable as [:, 0], so an empty + # result keeps the (N, 2) contract of a non-empty one. + return aa.Grid2DIrregular(values=np.zeros(shape=(0, 2))) + + return aa.Grid2DIrregular(values=positions) + + +def multiple_image_pixel_coordinates_from( + fit, + plane_index: int = -1, + use_centroid: bool = False, + **mapping_kwargs, +) -> List[Tuple[float, float]]: + """ + The multiple-image positions as pixel coordinates of the fit's data. + + The library deliberately stops at pixel coordinates. Turning them into RA / Dec is a + WCS question, which depends on the header of the FITS the data came from rather than + on anything the lens model knows, and is shown inline in the mappings guide with + astropy. + + Parameters + ---------- + fit + The `FitImaging` whose source is mapped. + plane_index + The index of the source plane. + use_centroid + As `multiple_image_positions_from`. + mapping_kwargs + Passed to `mappings_from_fit`. + + Returns + ------- + One ``(y, x)`` pixel coordinate per multiple image, in the WCS/FITS convention: 1-based, + referred to pixel centres, and continuous (the fractional part is the sub-pixel offset). + """ + positions = multiple_image_positions_from( + fit=fit, + plane_index=plane_index, + use_centroid=use_centroid, + **mapping_kwargs, + ) + + geometry = fit.dataset.data.mask.geometry + + return [ + geometry.pixel_coordinates_wcs_2d_from(scaled_coordinates_2d=tuple(position)) + for position in np.asarray(positions.array) + ] + + +def magnifications_from( + fit, + plane_index: int = -1, + shape: Optional[Shape] = None, + grid: Optional[aa.type.Grid2DLike] = None, + pixel_scale_precision: Optional[float] = None, + magnification_threshold: float = 0.1, + **clump_kwargs, +) -> List[float]: + """ + The magnification of each multiple image of the source. + + The two engines measure it differently, because they know different things: + + - **Engine B** (parametric) measures it geometrically, as + `ShapeSolver.find_magnification(per_image=True)`: the image-plane area the source + shape's images cover, divided by the shape's own area. This is the true + magnification of the lens model, independent of the source's surface brightness. + - **Engine A** (pixelized) has no source-plane *area* to divide by, only a mesh, so it + measures the flux ratio instead: each image region's summed model flux divided by + the total model flux of the whole clump across all of its images. That is the + *relative* magnification of the images of one clump -- it sums to 1 over a clump's + images -- and is not on the same scale as engine B's absolute magnification. + + Parameters + ---------- + fit + The `FitImaging` whose source is mapped. + plane_index + The index of the source plane. + shape, grid, pixel_scale_precision, magnification_threshold + As `mappings_from_fit` (engine B only). + clump_kwargs + As `mappings_from_fit` (engine A only). + + Returns + ------- + One magnification per multiple image, ordered as `multiple_image_positions_from` is. + """ + tracer = fit.tracer_linear_light_profiles_to_light_profiles + + if not tracer.planes[plane_index].has(cls=aa.Pixelization): + if shape is None: + shape = _light_profile_shape_from(tracer=tracer, plane_index=plane_index) + + if grid is None: + grid = fit.mask.derive_grid.all_false + + if pixel_scale_precision is None: + pixel_scale_precision = float(grid.pixel_scale) / 10.0 + + solver = ShapeSolver.for_grid( + grid=grid, + pixel_scale_precision=pixel_scale_precision, + magnification_threshold=magnification_threshold, + ) + + return solver.find_magnification( + tracer=tracer, + shape=shape, + plane_redshift=float(tracer.plane_redshifts[plane_index]), + per_image=True, + ) + + model_image = fit.model_images_of_planes_list[plane_index] + + magnifications = [] + + for mapping in mappings_from_fit(fit=fit, plane_index=plane_index, **clump_kwargs): + fluxes = [ + region.flux_from(array=model_image) for region in mapping.image_regions + ] + + total = float(np.sum(fluxes)) + + magnifications += [ + float(flux / total) if total != 0.0 else 0.0 for flux in fluxes + ] + + return magnifications diff --git a/autolens/plot/__init__.py b/autolens/plot/__init__.py index 7a96de32f..682e61af9 100644 --- a/autolens/plot/__init__.py +++ b/autolens/plot/__init__.py @@ -53,6 +53,7 @@ subplot_fit_x1_plane as subplot_fit_imaging_x1_plane, subplot_fit_log10_x1_plane as subplot_fit_imaging_log10_x1_plane, subplot_of_planes as subplot_fit_imaging_of_planes, + subplot_mappings as subplot_fit_imaging_mappings, subplot_tracer_from_fit as subplot_fit_imaging_tracer, subplot_fit_combined, subplot_fit_combined_log10, diff --git a/autolens/point/solver/shape_solver.py b/autolens/point/solver/shape_solver.py index d9cacf6c1..d7b8f2418 100644 --- a/autolens/point/solver/shape_solver.py +++ b/autolens/point/solver/shape_solver.py @@ -17,7 +17,7 @@ import logging import math -from typing import Tuple, List, Iterator, Optional +from typing import Tuple, List, Iterator, Optional, Union import autoarray as aa @@ -85,6 +85,14 @@ def __init__( self.neighbor_degree = neighbor_degree self.use_jax = use_jax + # The grid ``for_grid`` was built from, when there was one. It is the image-plane + # pixelization ``image_regions_from`` reports its regions on; a solver built via + # ``for_limits_and_scale`` has only an extent and no pixel grid, so it stays `None` + # and the grid must be passed to ``image_regions_from`` explicitly. It is + # deliberately *not* part of the pytree (``tree_flatten``): it is used only by the + # numpy-only, never-jitted mappings layer. + self.grid = None + @property def _xp(self): """The array module the solver runs against by default. ``jnp`` when @@ -155,7 +163,7 @@ def for_grid( y = grid[:, 0] x = grid[:, 1] - return cls.for_limits_and_scale( + solver = cls.for_limits_and_scale( y_min=y.min(), y_max=y.max(), x_min=x.min(), @@ -167,6 +175,13 @@ def for_grid( use_jax=use_jax, ) + # Remembered so ``ShapeSolver.image_regions_from`` can report image-plane regions + # on the same pixelization the caller is plotting, without being handed the grid + # a second time. + solver.grid = grid + + return solver + @classmethod def for_limits_and_scale( cls, @@ -292,7 +307,7 @@ def solve_triangles( self, tracer: Tracer, shape: Shape, - xp, + xp=None, plane_redshift: Optional[float] = None, ): """ @@ -312,7 +327,10 @@ def solve_triangles( shape The shape in the source plane for which we want to identify the image plane coordinates. xp - The array module (``numpy`` or ``jax.numpy``) the solve runs in. + The array module (``numpy`` or ``jax.numpy``) the solve runs in. ``None`` (the + default) falls back to ``self._xp``, which is ``jnp`` when the solver was + constructed with ``use_jax=True`` and ``np`` otherwise -- the same contract as + ``PointSolver.solve``. plane_redshift The redshift of the source plane. @@ -320,6 +338,9 @@ def solve_triangles( ------- A list of image plane coordinates that are traced to the source plane coordinate. """ + if xp is None: + xp = self._xp + # `n_steps` is `ceil(log2(scale / pixel_scale_precision))`, which is <= 0 whenever the # requested precision is coarser than (or equal to) the initial triangle scale. A negative # value used to slip past an `== 0` check and make `steps` an empty list, surfacing as @@ -425,7 +446,7 @@ def steps( self, tracer: Tracer, shape: Shape, - xp, + xp=None, plane_redshift: Optional[float] = None, ) -> Iterator[Step]: """ @@ -439,13 +460,28 @@ def steps( The shape in the source plane for which we want to identify the image plane coordinates. xp The array module (``numpy`` or ``jax.numpy``) the step iteration runs in. + ``None`` (the default) falls back to ``self._xp``. plane_redshift The redshift of the source plane. Returns ------- An iterator over the steps of the triangle solver algorithm. + + Notes + ----- + The area of ``step.filtered_triangles`` is **not** monotone over the steps. Step 0 + tiles the whole image plane at ``self.scale`` and keeps every triangle whose traced + image meets the shape; each later step then filters the *neighbourhood* of the + previous kept set, which is larger than that set, so the kept area oscillates while + it converges. What is monotone is the search envelope: ``step.neighbourhood`` + shrinks at every step, because the rim it adds is one triangle wide and the + triangles halve. The residual on the converged kept area is roughly one triangle + edge per boundary triangle -- see ``ShapeSolver.find_magnification``. """ + if xp is None: + xp = self._xp + initial_triangles = self._initial_triangles(xp) for number in range(self.n_steps): @@ -516,16 +552,141 @@ def tree_unflatten(cls, aux_data, children): class ShapeSolver(AbstractSolver): + """ + The triangle solver for an extended source-plane `Shape`, rather than a point. + + JAX + --- + ``ShapeSolver`` is a NumPy-only solver and rejects ``use_jax=True`` / ``xp=jax.numpy``. + The JAX triangle containers keep static shapes by truncating every refinement step to + ``ArrayTriangles.MAX_CONTAINING_SIZE`` (15) triangles — ample for a ``Point``, which + lies inside a handful of triangles, and meaningless for a shape with area, whose kept + set grows with the size of its images. Before this was found, ``use_jax=True`` was + silently ignored (``find_magnification`` hardcoded ``xp=np``); routing it through + ``self._xp`` instead would have replaced a silently-ignored flag with a silently wrong + number — an Isothermal sphere whose true magnification is 6.86 measures 0.13 under + JAX. Lifting the cap is a redesign of the JAX containers, not a fix here, so the flag + raises instead. See ``test_shape_solver.py::test_jax_and_numpy_kept_triangles_agree``. + """ + + # The one message both rejection routes raise, so a caller sees the same explanation + # whether they set `use_jax=True` or passed `xp=jax.numpy` -- and, since it is a plain + # string, whether or not JAX is installed. + _JAX_REJECTED_MESSAGE = ( + "ShapeSolver does not support JAX. The JAX triangle containers truncate " + "every refinement step to ArrayTriangles.MAX_CONTAINING_SIZE (15) " + "triangles to keep static shapes, which is enough for a Point but not for " + "a Shape with area: the kept triangles of an extended source number in the " + "thousands, so the JAX path silently measures a small fraction of the " + "image and returns a magnification orders of magnitude too small. Use " + "the NumPy path (`use_jax=False`, the default), or `PointSolver` if you " + "want the JAX-differentiable positions of a point source." + ) + + def _xp_from(self, xp): + """ + Resolve the array module for a solve, rejecting JAX. + + The ``use_jax=True`` case is rejected **before** ``self._xp`` is consulted, because + ``self._xp`` imports ``jax.numpy``. Resolving first and checking afterwards made a + JAX-free install raise ``ModuleNotFoundError: No module named 'jax'`` from inside + the solver rather than the explanation below -- the wrong error, and one which + blamed the environment for a decision the solver had already taken. The + explicitly-passed-module check reads only ``__name__``, so it imports nothing + either. + + Parameters + ---------- + xp + The caller's array module, or ``None`` to use ``self._xp``. + + Returns + ------- + The array module, which is always ``numpy``. + """ + if xp is None: + if self.use_jax: + raise NotImplementedError(self._JAX_REJECTED_MESSAGE) + + return self._xp + + if getattr(xp, "__name__", "").startswith("jax"): + raise NotImplementedError(self._JAX_REJECTED_MESSAGE) + + return xp + + def solve_triangles( + self, + tracer: Tracer, + shape: Shape, + xp=None, + plane_redshift: Optional[float] = None, + ): + """ + The kept image-plane triangles whose traced images meet `shape`. + + See ``AbstractSolver.solve_triangles``; this override only rejects the JAX path + (see the class docstring). + """ + return super().solve_triangles( + tracer=tracer, + shape=shape, + xp=self._xp_from(xp), + plane_redshift=plane_redshift, + ) + + def steps( + self, + tracer: Tracer, + shape: Shape, + xp=None, + plane_redshift: Optional[float] = None, + ) -> Iterator[Step]: + """ + Iterate over the refinement steps of the solve. + + See ``AbstractSolver.steps``; this override only rejects the JAX path (see the + class docstring). + """ + return super().steps( + tracer=tracer, + shape=shape, + xp=self._xp_from(xp), + plane_redshift=plane_redshift, + ) + def find_magnification( self, tracer: Tracer, shape: Shape, - xp=np, + xp=None, plane_redshift: Optional[float] = None, - ) -> float: + per_image: bool = False, + ) -> Union[float, List[float]]: """ Find the magnification of the shape in the source plane. + The magnification is the total image-plane area the shape's images cover divided by + the shape's own area, measured on the kept triangles of the finest refinement step. + + Convergence and error + --------------------- + A triangle is kept when its *traced* image meets the shape, so the kept set covers + the true image region plus a rim of boundary triangles which only partly overlap + it. The area therefore approaches the true value **from above**, with an error of + order ``perimeter * side_length`` — one triangle edge per boundary triangle. Since + ``side_length`` at the finest step is ``pixel_scale_precision``, halving + ``pixel_scale_precision`` roughly halves the error. For a source much larger than + ``pixel_scale_precision`` this is a sub-percent effect; for a source comparable to + it the answer is unreliable, and the caller should refine further. + + The other error is physical, not numerical: this measures the magnification of a + *finite* shape, which equals the point magnification only in the limit of a small + shape. An Isothermal sphere with ``einstein_radius=1.0`` and a source at + ``beta=0.3`` has an analytic total magnification ``2 * theta_E / beta = 6.67``; at + ``pixel_scale_precision=0.005`` a ``Circle`` of radius 0.2 gives 7.12 (7% high, + finite-source), radius 0.1 gives 6.76 (1.3%) and radius 0.05 gives 6.66 (0.1%). + Parameters ---------- tracer @@ -533,18 +694,414 @@ def find_magnification( shape The shape of an image plane pixel. xp - The array module (``numpy`` or ``jax.numpy``) the magnification calculation runs in. + The array module the magnification calculation runs in. ``None`` (the default) + falls back to ``self._xp``. This previously defaulted to ``np`` + unconditionally, so a solver built with ``use_jax=True`` silently ran in NumPy; + it now consults the flag and raises, because the JAX path is wrong for a shape + rather than merely unused (see the class docstring). plane_redshift The redshift of the source plane. + per_image + If ``True`` the kept triangles are grouped into their connected components (the + separate multiple images) and one magnification is returned per image, largest + first, dropping any component whose magnification does not exceed + ``self.magnification_threshold``. If ``False`` (the default) the single total + over every kept triangle is returned, unfiltered. Returns ------- - The magnification of the shape in the source plane. + The magnification of the shape in the source plane, or the list of per-image + magnifications when ``per_image=True``. + + Notes + ----- + ``sum(per_image)`` is not identical to the total: the total is unfiltered, so it + also carries any component the threshold rejects. For an Isothermal *sphere* the + rejected component is the demagnified central image the profile's singular centre + produces, worth ~0.004 of a magnification of 6.7. """ - kept_triangles = super().solve_triangles( + kept_triangles = self.solve_triangles( tracer=tracer, shape=shape, xp=xp, plane_redshift=plane_redshift, ) - return kept_triangles.area / shape.area + + if not per_image: + return kept_triangles.area / shape.area + + return [ + magnification + for magnification in sorted( + ( + component.area / shape.area + for component in _connected_components_from(kept_triangles) + ), + reverse=True, + ) + if magnification > self.magnification_threshold + ] + + def image_regions_from( + self, + tracer: Tracer, + shape: Shape, + plane_redshift: Optional[float] = None, + grid: Optional[aa.type.Grid2DLike] = None, + xp=None, + ) -> List[aa.ImageRegion]: + """ + The image-plane regions (the multiple images) a source-plane shape maps to. + + This is the parametric-source counterpart of + ``autoarray.inversion.mappings.image_regions_from``, which does the same job for a + pixelized source via a ``Mapper``. Both return the same ``ImageRegion`` objects, so + ``brightest_coordinate_from`` / ``centroid_from`` / ``flux_from`` / ``area`` behave + identically whichever engine produced them. + + The kept triangles of the finest refinement step are grouped into connected + components by shared vertices — each component is one multiple image — and each + component is converted to an image-plane pixel region by marking every pixel of + `grid` which contains a kept-triangle vertex or centroid. Because the triangles + carry the geometry, a source far smaller than a pixel still produces regions. + + Components whose magnification (their share of the kept area, divided by the + shape's area) does not exceed ``self.magnification_threshold`` are dropped. That + threshold was previously accepted by the constructor and never used by + ``ShapeSolver`` at all; without it the singular centre of an Isothermal sphere + contributes a spurious two-triangle "image" at the lens centre. + + Parameters + ---------- + tracer + The tracer which traces the image plane to the source plane. + shape + The source-plane shape whose images are found. + plane_redshift + The redshift of the plane `shape` lives in. ``None`` is the tracer's last plane. + grid + The image-plane grid whose pixels the regions are reported on. ``None`` uses the + grid the solver was built from by ``for_grid``; a solver built by + ``for_limits_and_scale`` has no grid and must be passed one. + xp + Unused except for symmetry with the rest of the solver: this is diagnostic / + visualization code, so it is numpy-only and never jitted. Passing ``jax.numpy`` + raises. + + Returns + ------- + One ``ImageRegion`` per multiple image, ordered by decreasing area. + """ + if xp is not None and xp is not np: + raise ValueError( + "`ShapeSolver.image_regions_from` is numpy-only and is never jitted: it " + "builds `Mask2D` and contour objects for plotting and diagnostics. Pass " + "`xp=None` (or `xp=numpy`); use `solve_triangles` for the JAX path." + ) + + grid = self._grid_from(grid=grid) + + kept_triangles = self.solve_triangles( + tracer=tracer, + shape=shape, + xp=np, + plane_redshift=plane_redshift, + ) + + components = _connected_components_from(kept_triangles) + + shape_area = shape.area + + regions = [] + + for component in components: + if ( + shape_area > 0.0 + and component.area / shape_area <= self.magnification_threshold + ): + continue + + region = _image_region_from_triangles( + triangles=component.triangles, mask=grid.mask + ) + + if region is not None: + regions.append(region) + + regions.sort(key=lambda region: -len(region.slim_indexes)) + + return regions + + def mapping_from( + self, + tracer: Tracer, + shape: Shape, + plane_redshift: Optional[float] = None, + grid: Optional[aa.type.Grid2DLike] = None, + ) -> aa.Mapping: + """ + The ``Mapping`` pairing a source-plane shape with the image-plane regions it maps to. + + The source-plane side of a shape has no mesh pixels, so ``pix_indexes`` is empty and + ``peak_value`` is ``None``; the shape's own boundary is its one source contour, which + is what the source-plane panels of ``subplot_mappings`` draw. + + Parameters + ---------- + tracer + The tracer which traces the image plane to the source plane. + shape + The source-plane shape which is mapped. + plane_redshift + The redshift of the plane `shape` lives in. ``None`` is the tracer's last plane. + grid + The image-plane grid the regions are reported on; ``None`` uses the solver's own. + + Returns + ------- + The ``Mapping`` of the shape. + """ + return aa.Mapping( + pix_indexes=np.array([], dtype=int), + source_contours=[shape.boundary()], + # `Shape.x` is element 0 of a coordinate pair, which for every PyAuto grid is + # the `y` coordinate — see the `Shape` class docstring. A `Circle` at a light + # profile centre `(y_c, x_c)` is built positionally as `Circle(y_c, x_c, r)`. + source_centre=(float(shape.x), float(shape.y)), + image_regions=self.image_regions_from( + tracer=tracer, + shape=shape, + plane_redshift=plane_redshift, + grid=grid, + ), + peak_value=None, + ) + + def _grid_from(self, grid: Optional[aa.type.Grid2DLike]) -> aa.type.Grid2DLike: + """ + Resolve the image-plane grid the regions are reported on. + + Parameters + ---------- + grid + An explicit grid, or ``None`` to use the one ``for_grid`` remembered. + + Returns + ------- + The grid. + """ + grid = self.grid if grid is None else grid + + if grid is None: + raise ValueError( + "This solver was built with `for_limits_and_scale`, which has an image-plane " + "extent but no pixel grid, so `image_regions_from` does not know which pixels " + "to report its regions on. Pass `grid=` explicitly, or build the solver with " + "`ShapeSolver.for_grid(grid=..., ...)`." + ) + + return grid + + +def _vertex_labels_from(triangles: np.ndarray) -> np.ndarray: + """ + Label the vertices of `triangles` so that coincident vertices share a label. + + Exact float equality (what ``AbstractTriangles.indices`` uses, via ``np.unique``) does + **not** identify coincident vertices here: a lattice point reached from two different + triangles is computed by two different sequences of floating point operations and the + results differ in the last ulp. Grouping the kept triangles of an Isothermal sphere on + exact equality splits its two images into 397 fragments rather than 2. The vertices are + therefore quantised to a thousandth of a triangle edge — far below the spacing of + distinct vertices, far above the ulp — before being deduplicated. + + Parameters + ---------- + triangles + An ``(n, 3, 2)`` array of triangle vertices. + + Returns + ------- + An ``(n, 3)`` integer array of vertex labels. + """ + vertices = np.asarray(triangles).reshape(-1, 2) + + edges = np.linalg.norm( + np.asarray(triangles)[:, 0] - np.asarray(triangles)[:, 1], axis=-1 + ) + + tolerance = float(np.median(edges)) * 1.0e-3 + + if not np.isfinite(tolerance) or tolerance <= 0.0: + tolerance = 1.0e-12 + + _, labels = np.unique(np.round(vertices / tolerance), axis=0, return_inverse=True) + + return np.asarray(labels).reshape(-1, 3) + + +def _connected_components_from(triangles) -> List["_TriangleComponent"]: + """ + Group a set of triangles into connected components which share vertices. + + Two triangles are in the same component when a chain of triangles between them each + share a vertex with the next. For the image-plane triangles of a lensed source, each + component is one multiple image. + + Parameters + ---------- + triangles + An ``AbstractTriangles`` (or anything with a ``triangles`` ``(n, 3, 2)`` array). + + Returns + ------- + One ``_TriangleComponent`` per component. + """ + array = np.asarray(triangles.triangles) + + array = array[~np.isnan(array).any(axis=(1, 2))] + + if array.shape[0] == 0: + return [] + + labels = _vertex_labels_from(array) + + total_triangles = labels.shape[0] + + parent = np.arange(total_triangles) + + def find(index: int) -> int: + while parent[index] != index: + parent[index] = parent[parent[index]] + index = parent[index] + return index + + def union(left: int, right: int): + left, right = find(left), find(right) + if left != right: + parent[left] = right + + # The first triangle seen at each vertex label is unioned with every later triangle + # sharing it, which connects the whole star of triangles around that vertex. + first_at_vertex = {} + + for triangle_index, row in enumerate(labels): + for label in row: + label = int(label) + if label in first_at_vertex: + union(first_at_vertex[label], triangle_index) + else: + first_at_vertex[label] = triangle_index + + roots = np.array([find(index) for index in range(total_triangles)]) + + return [ + _TriangleComponent(triangles=array[roots == root]) for root in np.unique(roots) + ] + + +class _TriangleComponent: + def __init__(self, triangles: np.ndarray): + """ + One connected component of a set of triangles. + + Exposes the same ``triangles`` / ``area`` interface as ``AbstractTriangles`` so a + component can be measured the same way the whole kept set is, without depending on + which triangle container the solve produced. + + Parameters + ---------- + triangles + An ``(n, 3, 2)`` array of the component's triangle vertices. + """ + self.triangles = triangles + + @property + def area(self) -> float: + """ + The total area covered by the component's triangles. + """ + triangles = self.triangles + + return float( + 0.5 + * np.abs( + (triangles[:, 0, 0] * (triangles[:, 1, 1] - triangles[:, 2, 1])) + + (triangles[:, 1, 0] * (triangles[:, 2, 1] - triangles[:, 0, 1])) + + (triangles[:, 2, 0] * (triangles[:, 0, 1] - triangles[:, 1, 1])) + ).sum() + ) + + +def _image_region_from_triangles( + triangles: np.ndarray, mask: aa.Mask2D +) -> Optional[aa.ImageRegion]: + """ + The ``ImageRegion`` covered by a component of image-plane triangles. + + Every pixel of `mask` containing a triangle vertex or centroid is in the region. The + region's boundary contours are the pixel-edge loops of that coverage (the same outline + engine A draws for a pixelized source), so a ring is drawn as an annulus rather than as + the filled disc a convex hull would give. + + Parameters + ---------- + triangles + An ``(n, 3, 2)`` array of ``(y, x)`` image-plane triangle vertices. + mask + The ``Mask2D`` of the image-plane grid, whose ``False`` entries are the data pixels. + + Returns + ------- + The region, or ``None`` when no unmasked pixel of `mask` is covered. + """ + from autoarray.inversion.mappings.mapping import contours_from_bool_native + from autoarray.mask.mask_2d import Mask2D + + points = np.concatenate([triangles.reshape(-1, 2), triangles.mean(axis=1)], axis=0) + + mask_arr = np.asarray(mask) + + geometry = mask.geometry + + rows = np.floor( + (geometry.scaled_maxima[0] - points[:, 0]) / geometry.pixel_scales[0] + ).astype(int) + columns = np.floor( + (points[:, 1] - geometry.scaled_minima[1]) / geometry.pixel_scales[1] + ).astype(int) + + inside = ( + (rows >= 0) + & (rows < mask_arr.shape[0]) + & (columns >= 0) + & (columns < mask_arr.shape[1]) + ) + + region_native = np.zeros(shape=mask_arr.shape, dtype=bool) + region_native[rows[inside], columns[inside]] = True + region_native &= ~mask_arr + + if not region_native.any(): + return None + + slim_index_native = np.full(shape=mask_arr.shape, fill_value=-1, dtype=int) + slim_index_native[~mask_arr] = np.arange(int(np.sum(~mask_arr))) + + centroids = triangles.mean(axis=1) + + return aa.ImageRegion( + slim_indexes=np.sort(slim_index_native[region_native]), + mask=Mask2D( + mask=~region_native, + pixel_scales=mask.pixel_scales, + origin=mask.origin, + ), + contours=contours_from_bool_native( + bool_native=region_native, geometry=geometry + ), + # The mean of the triangle centroids, not of the covered pixel centres: the + # triangles resolve the image far below the pixel scale, so this stays meaningful + # for an image which covers only one or two pixels. + centre=(float(np.mean(centroids[:, 0])), float(np.mean(centroids[:, 1]))), + ) diff --git a/test_autolens/imaging/plot/test_fit_imaging_plots.py b/test_autolens/imaging/plot/test_fit_imaging_plots.py index 0f06bf3e1..490392165 100644 --- a/test_autolens/imaging/plot/test_fit_imaging_plots.py +++ b/test_autolens/imaging/plot/test_fit_imaging_plots.py @@ -11,6 +11,7 @@ subplot_fit_x1_plane, subplot_fit_log10_x1_plane, subplot_of_planes, + subplot_mappings, subplot_tracer_from_fit, subplot_fit_combined, subplot_fit_combined_log10, @@ -216,3 +217,67 @@ def boom(*args, **kwargs): assert record.levelno == logging.WARNING assert record.exc_info is not None, "traceback must be attached" assert isinstance(record.exc_info[1], RuntimeError) + + +def test__subplot_mappings__parametric_source__output_file_created( + fit_imaging_x2_plane_7x7, plot_path, plot_patch +): + """ + The parametric (ShapeSolver) engine. The filename carries the plane index, as the + sibling `mappings_{pixelization_index}` of the inversion subplot does. + """ + subplot_mappings( + fit=fit_imaging_x2_plane_7x7, + output_path=plot_path, + output_format="png", + ) + assert str(plot_path / "mappings_1.png") in plot_patch.paths + + +def test__subplot_mappings__inversion_source__output_file_created( + fit_imaging_x2_plane_inversion_7x7, plot_path, plot_patch +): + """ + The pixelized (inversion) engine, reached through the `pix_indexes` bypass because the + 7x7 fixture reconstructs to all zeros and so has no clumps of its own. + """ + subplot_mappings( + fit=fit_imaging_x2_plane_inversion_7x7, + pix_indexes=[[0, 1, 2]], + output_path=plot_path, + output_format="png", + ) + assert str(plot_path / "mappings_1.png") in plot_patch.paths + + +def test__subplot_mappings__unmappable_plane__raises_rather_than_drawing_nothing( + masked_imaging_7x7, plot_path, plot_patch +): + """ + A source plane which cannot be mapped raises, naming what the caller has to change. It + must not draw four empty panels: a figure with no regions reads as "this source has no + multiple images", which is a wrong answer rather than a missing one -- so + `subplot_mappings` deliberately has no guard around the engine. + + A source plane with neither a pixelization nor a light profile is the case which + cannot be mapped at all. + """ + import autolens as al + + fit = al.FitImaging( + dataset=masked_imaging_7x7, + tracer=al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0), + ), + al.Galaxy(redshift=1.0), + ] + ), + ) + + with pytest.raises(ValueError) as exc_info: + subplot_mappings(fit=fit, output_path=plot_path, output_format="png") + + assert "shape=" in str(exc_info.value) diff --git a/test_autolens/lens/test_mappings.py b/test_autolens/lens/test_mappings.py new file mode 100644 index 000000000..a7b929856 --- /dev/null +++ b/test_autolens/lens/test_mappings.py @@ -0,0 +1,379 @@ +""" +Tests for `autolens.lens.mappings` -- the two engines which pair a source-plane region +with the image-plane regions (the multiple images) it maps to, and the fit-level +quantities built on them. +""" + +import numpy as np +import pytest + +import autoarray as aa +import autolens as al + +from autoarray.structures.triangles.shape import Circle + +from autolens.lens import mappings + + +EINSTEIN_RADIUS = 1.0 + + +@pytest.fixture +def grid(): + return al.Grid2D.uniform(shape_native=(80, 80), pixel_scales=0.05) + + +@pytest.fixture +def tracer(): + return al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph( + centre=(0.0, 0.0), einstein_radius=EINSTEIN_RADIUS + ), + ), + al.Galaxy( + redshift=1.0, + light=al.lp.SersicSph( + centre=(0.2, 0.0), intensity=1.0, effective_radius=0.1 + ), + ), + ] + ) + + +def test__source_inside_the_caustic__maps_to_two_images(tracer, grid): + """ + An isothermal sphere's caustic is the point ``beta = 0``, so any source inside its + Einstein radius is multiply imaged: two images, one on each side of the lens. + """ + mapping = mappings.source_mapping_from( + tracer=tracer, + grid=grid, + shape=Circle(0.2, 0.0, radius=0.05), + pixel_scale_precision=0.005, + ) + + assert len(mapping.image_regions) == 2 + + y_coordinates = sorted(region.centre[0] for region in mapping.image_regions) + + assert y_coordinates[0] == pytest.approx(0.2 - EINSTEIN_RADIUS, abs=0.05) + assert y_coordinates[1] == pytest.approx(0.2 + EINSTEIN_RADIUS, abs=0.05) + + +def test__source_outside_the_einstein_radius__maps_to_one_image(tracer): + """ + A source further from the lens than its Einstein radius has a single image, and it + sits beyond the Einstein radius on the same side as the source. + + The grid is wider than the fixture's: the single image of a source at ``beta = 1.5`` + lies at ``beta + theta_E = 2.5``, which a +/- 2" grid does not contain -- and an image + outside the tiled extent is not found, correctly, since the solver only searches the + image plane it was given. + """ + mapping = mappings.source_mapping_from( + tracer=tracer, + grid=al.Grid2D.uniform(shape_native=(120, 120), pixel_scales=0.05), + shape=Circle(1.5, 0.0, radius=0.05), + pixel_scale_precision=0.005, + ) + + assert len(mapping.image_regions) == 1 + assert mapping.image_regions[0].centre[0] == pytest.approx( + 1.5 + EINSTEIN_RADIUS, abs=0.05 + ) + + +def test__traced_region_is_a_closed_source_plane_loop(tracer, grid): + """ + An image-plane region traced back to the source plane, outlined by its convex hull. + """ + mapping = mappings.source_mapping_from( + tracer=tracer, + grid=grid, + shape=Circle(0.2, 0.0, radius=0.05), + pixel_scale_precision=0.005, + ) + + region = mapping.image_regions[0] + + hull = mappings.traced_region_from( + tracer=tracer, region_grid=region.scaled_coordinates + ) + + assert hull.ndim == 2 + assert hull.shape[1] == 2 + assert hull[0] == pytest.approx(hull[-1]) + + # An image traces back onto its source, so the hull must sit around the shape. + assert np.mean(hull[:, 0]) == pytest.approx(0.2, abs=0.1) + + +def test__traced_region_of_fewer_than_three_points_is_returned_unchanged(tracer): + """ + A convex hull needs three points; two traced coordinates are returned as they are + rather than raising out of a figure. + """ + traced = mappings.traced_region_from( + tracer=tracer, region_grid=np.array([[1.0, 0.0], [-1.0, 0.0]]) + ) + + assert traced.shape == (2, 2) + + +# ------------------------------------------------------------------------------------ # +# engine dispatch +# ------------------------------------------------------------------------------------ # + + +def test__parametric_fit__dispatches_to_the_shape_solver(fit_imaging_x2_plane_7x7): + """ + A plane with no pixelization is mapped by engine B, which returns exactly one + `Mapping` (the one source shape) with no mesh pixel indexes and the shape's own + boundary as its source contour. + """ + results = mappings.mappings_from_fit(fit=fit_imaging_x2_plane_7x7) + + assert len(results) == 1 + + mapping = results[0] + + assert isinstance(mapping, aa.Mapping) + assert len(mapping.pix_indexes) == 0 + assert mapping.peak_value is None + assert len(mapping.source_contours) == 1 + + +def test__pixelized_fit__dispatches_to_the_inversion( + fit_imaging_x2_plane_inversion_7x7, +): + """ + A plane which `has(cls=aa.Pixelization)` is mapped by engine A, whose mappings carry + mesh pixel indexes and a peak reconstructed value -- the two fields engine B leaves + empty. The 7x7 fixture reconstructs to all zeros, so clump finding legitimately finds + nothing at any threshold; `pix_indexes` is the documented bypass (the tutorial path) + and is what exercises the dispatch here. + """ + results = mappings.mappings_from_fit( + fit=fit_imaging_x2_plane_inversion_7x7, pix_indexes=[[0, 1, 2]] + ) + + assert len(results) == 1 + + mapping = results[0] + + assert isinstance(mapping, aa.Mapping) + assert list(mapping.pix_indexes) == [0, 1, 2] + assert mapping.peak_value is not None + + +def test__source_plane_with_no_light_profile__names_the_shape_input( + masked_imaging_7x7, +): + """ + Neither engine can guess a region for a source plane which has neither a pixelization + nor a light profile, so the error must point at `shape=`. + """ + fit = al.FitImaging( + dataset=masked_imaging_7x7, + tracer=al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0), + ), + al.Galaxy(redshift=1.0), + ] + ), + ) + + with pytest.raises(ValueError) as exc_info: + mappings.mappings_from_fit(fit=fit) + + assert "shape=" in str(exc_info.value) + + +# ------------------------------------------------------------------------------------ # +# fit-level quantities +# ------------------------------------------------------------------------------------ # + + +def test__multiple_image_positions__one_per_image_region(fit_imaging_x2_plane_7x7): + positions = mappings.multiple_image_positions_from(fit=fit_imaging_x2_plane_7x7) + + total_regions = sum( + len(mapping.image_regions) + for mapping in mappings.mappings_from_fit(fit=fit_imaging_x2_plane_7x7) + ) + + assert isinstance(positions, aa.Grid2DIrregular) + assert np.asarray(positions.array).shape == (total_regions, 2) + + +def test__multiple_image_positions__centroid_and_brightest_pixel_agree_roughly( + tracer, grid +): + """ + On a real (rather than 7x7 fixture) lens both routes must land on the same image, the + centroid sub-pixel and the brightest pixel snapped to the data grid. + """ + dataset = _simulated_dataset(tracer=tracer, grid=grid) + + fit = al.FitImaging(dataset=dataset, tracer=tracer) + + brightest = np.asarray( + mappings.multiple_image_positions_from( + fit=fit, pixel_scale_precision=0.01 + ).array + ) + centroids = np.asarray( + mappings.multiple_image_positions_from( + fit=fit, use_centroid=True, pixel_scale_precision=0.01 + ).array + ) + + assert brightest.shape == centroids.shape == (2, 2) + assert brightest == pytest.approx(centroids, abs=0.1) + + +def test__multiple_image_pixel_coordinates__are_inside_the_data(tracer, grid): + dataset = _simulated_dataset(tracer=tracer, grid=grid) + + fit = al.FitImaging(dataset=dataset, tracer=tracer) + + pixel_coordinates = mappings.multiple_image_pixel_coordinates_from( + fit=fit, pixel_scale_precision=0.01 + ) + + assert len(pixel_coordinates) == 2 + + shape_native = fit.dataset.data.mask.shape_native + + for y, x in pixel_coordinates: + # WCS pixel coordinates are 1-based and refer to pixel centres. + assert 1.0 <= y <= shape_native[0] + assert 1.0 <= x <= shape_native[1] + + +def test__magnifications__parametric_matches_the_analytic_isothermal_sphere( + tracer, grid +): + """ + Engine B's magnifications are absolute, so they can be checked against the analytic + isothermal sphere: ``1 + theta_E / beta`` and ``theta_E / beta - 1`` for beta = 0.2. + """ + dataset = _simulated_dataset(tracer=tracer, grid=grid) + + fit = al.FitImaging(dataset=dataset, tracer=tracer) + + magnifications = mappings.magnifications_from( + fit=fit, shape=Circle(0.2, 0.0, radius=0.05), pixel_scale_precision=0.005 + ) + + assert len(magnifications) == 2 + + assert magnifications[0] == pytest.approx(1.0 + EINSTEIN_RADIUS / 0.2, rel=0.02) + assert magnifications[1] == pytest.approx(EINSTEIN_RADIUS / 0.2 - 1.0, rel=0.02) + + +def test__magnifications__pixelized_are_flux_shares_summing_to_one(grid): + """ + Engine A has no source-plane area to divide by, so it reports each image's share of + the clump's total model flux -- which sums to one per clump, by construction, and is + documented as being on a different scale from engine B's absolute magnification. + + This needs a reconstruction with something in it, so it fits a real (if small) + simulated lens with a rectangular pixelization rather than using the 7x7 fixture, + whose reconstruction is all zeros. + """ + fit = _pixelized_fit(grid=grid) + + clump_kwargs = dict(threshold=0.5, min_pixels=1) + + magnifications = mappings.magnifications_from(fit=fit, **clump_kwargs) + + total_clumps = len(mappings.mappings_from_fit(fit=fit, **clump_kwargs)) + + assert total_clumps > 0 + assert len(magnifications) >= total_clumps + assert sum(magnifications) == pytest.approx(total_clumps, abs=1.0e-6) + + +def _pixelized_fit(grid): + """ + A `FitImaging` of a simulated lens whose source is a rectangular pixelization, so that + engine A has a non-zero reconstruction to find clumps in. + """ + source_tracer = al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph( + centre=(0.0, 0.0), einstein_radius=EINSTEIN_RADIUS + ), + ), + al.Galaxy( + redshift=1.0, + light=al.lp.SersicSph( + centre=(0.1, 0.0), intensity=1.0, effective_radius=0.2 + ), + ), + ] + ) + + dataset = _simulated_dataset(tracer=source_tracer, grid=grid) + + pixelized_tracer = al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph( + centre=(0.0, 0.0), einstein_radius=EINSTEIN_RADIUS + ), + ), + al.Galaxy( + redshift=1.0, + pixelization=al.Pixelization( + mesh=al.mesh.RectangularUniform(shape=(10, 10)), + regularization=al.reg.Constant(coefficient=1.0), + ), + ), + ] + ) + + return al.FitImaging(dataset=dataset, tracer=pixelized_tracer) + + +def _simulated_dataset(tracer, grid): + """ + A noiseless simulated dataset of `tracer`, masked to the grid's extent. + + Noiseless so that the brightest pixel of a model image is the brightest pixel of the + data, which is what makes the position assertions above deterministic. + """ + simulator = al.SimulatorImaging( + exposure_time=300.0, + psf=al.Convolver.from_gaussian( + shape_native=(3, 3), + pixel_scales=grid.pixel_scales[0], + sigma=0.05, + normalize=True, + ), + add_poisson_noise_to_data=False, + ) + + dataset = simulator.via_tracer_from(tracer=tracer, grid=grid) + + dataset.noise_map = al.Array2D.ones( + shape_native=dataset.data.shape_native, pixel_scales=grid.pixel_scales + ) + + mask = al.Mask2D.circular( + shape_native=dataset.data.shape_native, + pixel_scales=grid.pixel_scales, + radius=1.9, + ) + + return dataset.apply_mask(mask=mask) diff --git a/test_autolens/point/triangles/test_extended.py b/test_autolens/point/triangles/test_extended.py index fde16f2d1..cd63513c0 100644 --- a/test_autolens/point/triangles/test_extended.py +++ b/test_autolens/point/triangles/test_extended.py @@ -1,3 +1,4 @@ +import numpy as np import pytest from autoarray.structures.triangles.shape import Circle @@ -13,12 +14,52 @@ def solver(grid): ) -# def test_solver_basic(solver): -# assert solver.find_magnification( -# tracer=NullTracer(), -# shape=Circle( -# 0.0, -# 0.0, -# radius=0.1, -# ), -# ) == pytest.approx(1.0, abs=0.1) +def test_solver_basic(solver): + """ + The identity lens: `NullTracer` deflects nothing, so the source plane *is* the image + plane and a shape's image has exactly the shape's own area -- a magnification of 1. + + This is the solver's only original test, commented out since ~2024 and revived by the + `image-source-mappings` validation suite. It is the one case where the answer is known + without any lensing at all, so it isolates the tiling and the area accounting from the + ray tracing: a magnification of 1 here means the kept triangles cover the shape and + nothing else. + """ + assert solver.find_magnification( + tracer=NullTracer(), + shape=Circle( + 0.0, + 0.0, + radius=0.1, + ), + ) == pytest.approx(1.0, abs=0.1) + + +def test_solver_basic_is_one_image(solver): + """ + With no deflection there is one image, not several, and its magnification is the total. + """ + magnifications = solver.find_magnification( + tracer=NullTracer(), + shape=Circle(0.0, 0.0, radius=0.1), + per_image=True, + ) + + assert len(magnifications) == 1 + assert magnifications[0] == pytest.approx(1.0, abs=0.1) + + +def test_solver_basic_regions_sit_on_the_shape(solver, grid): + """ + The identity lens again, this time through `image_regions_from`: the one region must + sit on top of the shape it was built from. + """ + regions = solver.image_regions_from( + tracer=NullTracer(), + shape=Circle(0.0, 0.0, radius=0.1), + grid=grid, + ) + + assert len(regions) == 1 + assert regions[0].centre == pytest.approx((0.0, 0.0), abs=0.05) + assert np.asarray(regions[0].slim_indexes).size >= 1 diff --git a/test_autolens/point/triangles/test_shape_solver.py b/test_autolens/point/triangles/test_shape_solver.py new file mode 100644 index 000000000..e90887568 --- /dev/null +++ b/test_autolens/point/triangles/test_shape_solver.py @@ -0,0 +1,641 @@ +""" +The `ShapeSolver` validation suite. + +`ShapeSolver` shipped in ~2024 with one method (`find_magnification`), its only test +commented out, and no consumer other than an unused workspace demo; the audit of +2026-08-19 called it effectively unmaintained. This module is the validation it never +had, written as part of the `image-source-mappings` epic, which reuses the solver as the +source -> image engine for non-pixelized sources. + +Every case here is a statement about the *physics* the solver claims to compute, checked +against something that shares no code with it: the analytic magnification of an +Isothermal sphere, a grid-tracing oracle, the behaviour of shapes of equal area, and the +NumPy path itself. The bugs the suite found are named in the tests that pin them. +""" + +import numpy as np +import pytest +from scipy import ndimage + +import autoarray as aa +import autolens as al + +from autoarray.structures.triangles.shape import Circle, Polygon, Square, Triangle +from autolens.point.solver.shape_solver import ShapeSolver + +EINSTEIN_RADIUS = 1.0 + + +@pytest.fixture +def image_mask(): + return al.Mask2D.circular(shape_native=(100, 100), pixel_scales=0.05, radius=2.4) + + +@pytest.fixture +def image_grid(image_mask): + return al.Grid2D.from_mask(mask=image_mask) + + +@pytest.fixture +def sis_tracer(): + """ + A singular isothermal sphere, whose magnification is analytic. + """ + return al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph( + centre=(0.0, 0.0), einstein_radius=EINSTEIN_RADIUS + ), + ), + al.Galaxy(redshift=1.0), + ] + ) + + +@pytest.fixture +def shape_solver(image_grid): + return ShapeSolver.for_grid(grid=image_grid, pixel_scale_precision=0.005) + + +def sis_total_magnification(beta: float) -> float: + """ + The analytic total magnification of a point source at impact parameter `beta` inside + the Einstein radius of a singular isothermal sphere. + + The two images sit at ``theta = beta +/- theta_E`` with magnifications + ``mu_+ = 1 + theta_E / beta`` and ``mu_- = theta_E / beta - 1``, so the sum of their + absolute magnifications is ``2 * theta_E / beta``. + """ + return 2.0 * EINSTEIN_RADIUS / beta + + +# ---------------------------------------------------------------------------------- # +# (a) the Einstein ring, and the analytic magnification of an off-axis source +# ---------------------------------------------------------------------------------- # + + +def test_on_axis_source_traces_to_an_einstein_ring(shape_solver, sis_tracer): + """ + A small circle at the centre of the source plane of an isothermal *sphere* maps to the + Einstein ring: one connected image whose vertices all sit at the Einstein radius. + + A ring is a single connected component, not two: it is the degenerate case where the + two images of an off-axis source have merged, so the component count is 1 and it is + the *radius* which carries the physics. + """ + # `Circle` takes element 0 first, which for every PyAuto `(y, x)` grid is `y`, so a + # circle at the light profile centre `(y_c, x_c)` is `Circle(y_c, x_c, radius=...)`. + shape = Circle(0.0, 0.0, radius=0.02) + + triangles = shape_solver.solve_triangles(tracer=sis_tracer, shape=shape) + + vertices = np.asarray(triangles.triangles).reshape(-1, 2) + + radii = np.hypot(vertices[:, 0], vertices[:, 1]) + + # The lens centre is itself a (formally infinitely demagnified) image of an on-axis + # source, so it is excluded before the ring is measured -- it is what + # `magnification_threshold` removes from `image_regions_from`. + ring = radii[radii > 0.5 * EINSTEIN_RADIUS] + + assert ring.size > 0.9 * radii.size + assert np.mean(ring) == pytest.approx(EINSTEIN_RADIUS, abs=0.01) + assert np.max(ring) - np.min(ring) < 0.1 * EINSTEIN_RADIUS + + +def test_off_axis_source_has_two_images(shape_solver, sis_tracer): + """ + A source inside the Einstein radius of an isothermal sphere has exactly two images. + """ + beta = 0.3 + + regions = shape_solver.image_regions_from( + tracer=sis_tracer, shape=Circle(beta, 0.0, radius=0.05) + ) + + assert len(regions) == 2 + + # theta = beta + theta_E on one side and beta - theta_E on the other. + y_coordinates = sorted(region.centre[0] for region in regions) + + assert y_coordinates[0] == pytest.approx(beta - EINSTEIN_RADIUS, abs=0.05) + assert y_coordinates[1] == pytest.approx(beta + EINSTEIN_RADIUS, abs=0.05) + + +def test_per_image_magnification_matches_the_analytic_isothermal_sphere( + shape_solver, sis_tracer +): + """ + Each image's magnification, measured as its share of the kept triangle area divided by + the source area, matches the analytic ``1 + theta_E / beta`` and ``theta_E / beta - 1``. + """ + beta = 0.3 + + magnifications = shape_solver.find_magnification( + tracer=sis_tracer, shape=Circle(beta, 0.0, radius=0.05), per_image=True + ) + + assert len(magnifications) == 2 + + assert magnifications[0] == pytest.approx(1.0 + EINSTEIN_RADIUS / beta, rel=0.01) + assert magnifications[1] == pytest.approx(EINSTEIN_RADIUS / beta - 1.0, rel=0.01) + + +def test_total_magnification_converges_as_the_source_shrinks(shape_solver, sis_tracer): + """ + `find_magnification` measures a *finite* source, which equals the point magnification + only in the limit of a small source. The error must therefore shrink as the radius + does -- if it did not, the number would be measuring the tiling rather than the lens. + """ + beta = 0.3 + + analytic = sis_total_magnification(beta=beta) + + errors = [ + abs( + shape_solver.find_magnification( + tracer=sis_tracer, shape=Circle(beta, 0.0, radius=radius) + ) + - analytic + ) + / analytic + for radius in (0.2, 0.1, 0.05) + ] + + assert errors[0] > errors[1] > errors[2] + assert errors[-1] < 0.01 + + +def test_magnification_threshold_removes_the_central_image(image_grid, sis_tracer): + """ + `magnification_threshold` was accepted by the constructor and never used by + `ShapeSolver`: the filter lived only in `PointSolver.solve`. Without it the singular + centre of an isothermal sphere contributes a spurious two-triangle "image" at the lens + centre, which `image_regions_from` would report alongside the two real ones. + """ + shape = Circle(0.3, 0.0, radius=0.05) + + unfiltered = ShapeSolver.for_grid( + grid=image_grid, pixel_scale_precision=0.005, magnification_threshold=0.0 + ) + filtered = ShapeSolver.for_grid( + grid=image_grid, pixel_scale_precision=0.005, magnification_threshold=0.1 + ) + + assert ( + len( + unfiltered.find_magnification( + tracer=sis_tracer, shape=shape, per_image=True + ) + ) + == 3 + ) + assert ( + len(filtered.find_magnification(tracer=sis_tracer, shape=shape, per_image=True)) + == 2 + ) + + +# ---------------------------------------------------------------------------------- # +# (b) the regions against a grid-tracing oracle +# ---------------------------------------------------------------------------------- # + + +def _oracle_regions(tracer, grid, mask, shape): + """ + The image regions found by tracing the image grid and asking the shape which traced + pixel centres it contains. This shares no code with the triangle solver, so it is an + independent check of where the images are -- but it works only at pixel resolution, + which is why it is a test oracle and not the API. + """ + traced = tracer.traced_grid_2d_list_from(grid=grid)[-1] + + return aa.inversion.mappings.mapping.image_regions_from_slim_mask( + mask=mask, + slim_bool=shape.contains(np.asarray(traced.array)), + min_pixels=1, + ) + + +def _native_from_slim(mask, slim_indexes): + native = np.zeros(shape=mask.shape_native, dtype=bool) + + pixels = np.argwhere(~np.asarray(mask)) + + native[pixels[slim_indexes, 0], pixels[slim_indexes, 1]] = True + + return native + + +def test_regions_agree_with_the_grid_tracing_oracle( + shape_solver, sis_tracer, image_grid, image_mask +): + """ + For a source spanning many pixels the solver's regions must contain every pixel the + oracle finds, and must not extend more than one pixel beyond them. + + The solver keeps a triangle when its traced image *meets* the shape, so it marks every + pixel containing any point which traces inside, whereas the oracle marks a pixel only + when its centre does. The solver's region is therefore the oracle's, dilated by up to + a pixel -- and never less than it. + """ + shape = Circle(0.3, 0.0, radius=0.1) + + regions = shape_solver.image_regions_from(tracer=sis_tracer, shape=shape) + oracle = _oracle_regions( + tracer=sis_tracer, grid=image_grid, mask=image_mask, shape=shape + ) + + assert len(regions) == len(oracle) == 2 + + solver_slim = np.unique(np.concatenate([region.slim_indexes for region in regions])) + oracle_slim = np.unique(np.concatenate([region.slim_indexes for region in oracle])) + + # A source spanning >= 5 pixels, as the agreement statement requires. + assert oracle_slim.size >= 5 + + assert set(oracle_slim.tolist()) <= set(solver_slim.tolist()) + + dilated = ndimage.binary_dilation( + _native_from_slim(mask=image_mask, slim_indexes=oracle_slim), + structure=np.ones(shape=(3, 3), dtype=bool), + ) + + assert not ( + _native_from_slim(mask=image_mask, slim_indexes=solver_slim) & ~dilated + ).any() + + +def test_regions_are_image_region_objects_the_phase_1_api_understands( + shape_solver, sis_tracer, image_mask +): + """ + Engine B must return the same object engine A does, so that everything built on + `ImageRegion` works whichever engine produced the region. + """ + regions = shape_solver.image_regions_from( + tracer=sis_tracer, shape=Circle(0.3, 0.0, radius=0.1) + ) + + region = regions[0] + + assert isinstance(region, aa.ImageRegion) + + array = al.Array2D.ones(shape_native=(100, 100), pixel_scales=0.05) + + assert region.flux_from(array=array) == pytest.approx( + float(len(region.slim_indexes)) + ) + assert region.area() == pytest.approx(len(region.slim_indexes) * 0.05 * 0.05) + assert len(region.contours) >= 1 + assert region.contours[0].shape[1] == 2 + + +def test_mapping_from_carries_the_shape_boundary_as_its_source_contour( + shape_solver, sis_tracer +): + mapping = shape_solver.mapping_from( + tracer=sis_tracer, shape=Circle(0.3, 0.0, radius=0.1) + ) + + assert isinstance(mapping, aa.Mapping) + assert len(mapping.image_regions) == 2 + assert mapping.source_centre == pytest.approx((0.3, 0.0)) + assert len(mapping.source_contours) == 1 + + contour = mapping.source_contours[0] + + assert contour[0] == pytest.approx(contour[-1]) + assert np.hypot(contour[:, 0] - 0.3, contour[:, 1] - 0.0) == pytest.approx( + 0.1, abs=1.0e-8 + ) + + +def test_image_regions_from_without_a_grid_names_the_missing_input(): + """ + A solver built from limits has an extent but no pixel grid, so it cannot say which + pixels its regions cover. The error must name the input the caller controls. + """ + solver = ShapeSolver.for_limits_and_scale( + y_min=-1.0, + y_max=1.0, + x_min=-1.0, + x_max=1.0, + scale=0.1, + pixel_scale_precision=0.01, + ) + + assert solver.grid is None + + with pytest.raises(ValueError) as exc_info: + solver.image_regions_from( + tracer=al.Tracer( + galaxies=[al.Galaxy(redshift=0.5), al.Galaxy(redshift=1.0)] + ), + shape=Circle(0.0, 0.0, radius=0.1), + ) + + assert "grid=" in str(exc_info.value) + + +# ---------------------------------------------------------------------------------- # +# (c) the refinement steps +# ---------------------------------------------------------------------------------- # + + +def test_refinement_shrinks_the_search_and_converges(shape_solver, sis_tracer): + """ + The kept area is **not** monotone over the steps, and asserting that it is would pin a + property the algorithm does not have: step 0 keeps whole coarse triangles, and each + later step filters the *neighbourhood* of the previous kept set, which is larger than + that set. Measured over seven steps (``pixel_scale_precision=0.0005``) the kept areas + run 0.2143413, 0.2154238, 0.2126498, 0.2122439, 0.2123877, 0.2124807, 0.2124794 -- + they oscillate at the fourth decimal while converging. + + Two things are true instead, and are what is pinned here. The search envelope + (`neighbourhood`, the set the next step filters) shrinks at every step, since the rim + it adds is one triangle wide and the triangles halve. And the kept area converges: the + final refinement changes it by less than any earlier one (0.0000013 against a first + change of 0.0010825 over those seven steps). + """ + steps = list( + shape_solver.steps(tracer=sis_tracer, shape=Circle(0.3, 0.0, radius=0.1)) + ) + + assert len(steps) == shape_solver.n_steps + + envelopes = [step.neighbourhood.area for step in steps] + + assert envelopes == sorted(envelopes, reverse=True) + + kept = [step.filtered_triangles.area for step in steps] + + changes = [abs(kept[i] - kept[i - 1]) for i in range(1, len(kept))] + + # The final refinement moves the kept area less than any earlier one. + assert changes[-1] == min(changes) + + # And no step is ever far from the converged answer -- the oscillation is small. + assert np.max(np.abs(np.asarray(kept) - kept[-1])) < 0.02 * kept[-1] + + # And the converged answer is the one `solve_triangles` returns. + assert shape_solver.solve_triangles( + tracer=sis_tracer, shape=Circle(0.3, 0.0, radius=0.1) + ).area == pytest.approx(kept[-1]) + + +def test_precision_coarser_than_the_triangle_scale_is_rejected(image_grid): + """ + `n_steps <= 0` used to make `steps` empty and surface as an `IndexError`; it now names + `pixel_scale_precision`. Pinned here because the validation suite is where the + solver's input handling is stated. + """ + with pytest.raises(ValueError) as exc_info: + ShapeSolver.for_grid( + grid=image_grid, pixel_scale_precision=1.0 + ).solve_triangles( + tracer=al.Tracer( + galaxies=[al.Galaxy(redshift=0.5), al.Galaxy(redshift=1.0)] + ), + shape=Circle(0.0, 0.0, radius=0.1), + ) + + assert "pixel_scale_precision" in str(exc_info.value) + + +# ---------------------------------------------------------------------------------- # +# (d) shapes other than a circle +# ---------------------------------------------------------------------------------- # + + +def _equal_area_shapes(centre: float, area: float): + """ + A square, an equilateral triangle and a regular hexagon of the given area, all centred + on ``(centre, 0.0)`` -- element 0 first, as `Shape` requires. + """ + side = np.sqrt(area) + + square = Square( + top=centre - side / 2.0, + bottom=centre + side / 2.0, + left=-side / 2.0, + right=side / 2.0, + ) + + length = np.sqrt(4.0 * area / np.sqrt(3.0)) + height = length * np.sqrt(3.0) / 2.0 + + triangle = Triangle( + (centre + 2.0 * height / 3.0, 0.0), + (centre - height / 3.0, -length / 2.0), + (centre - height / 3.0, length / 2.0), + ) + + radius = np.sqrt(2.0 * area / (3.0 * np.sqrt(3.0))) + angles = np.linspace(0.0, 2.0 * np.pi, 6, endpoint=False) + + hexagon = Polygon( + [(centre + radius * np.cos(angle), radius * np.sin(angle)) for angle in angles] + ) + + return {"square": square, "triangle": triangle, "hexagon": hexagon} + + +def test_shapes_of_equal_area_give_the_same_images(shape_solver, sis_tracer): + """ + Magnification is an area ratio, so it must depend on the source's area and position + and not on which `Shape` subclass expresses them. Every shape's `mask` reads the + triangle vertices in the same element order, so a disagreement here would be a + coordinate-order bug in one of them (`Triangle.triangle_contains_mask` had exactly + that bug before phase 2a). + """ + beta = 0.3 + + circle = Circle(beta, 0.0, radius=0.1) + + reference = shape_solver.find_magnification( + tracer=sis_tracer, shape=circle, per_image=True + ) + + assert len(reference) == 2 + + for name, shape in _equal_area_shapes(centre=beta, area=circle.area).items(): + assert shape.area == pytest.approx(circle.area, rel=1.0e-6), name + + magnifications = shape_solver.find_magnification( + tracer=sis_tracer, shape=shape, per_image=True + ) + + assert len(magnifications) == 2, name + assert magnifications == pytest.approx(reference, rel=0.02), name + + assert ( + len(shape_solver.image_regions_from(tracer=sis_tracer, shape=shape)) == 2 + ), name + + +# ---------------------------------------------------------------------------------- # +# (e) a source on an intermediate plane +# ---------------------------------------------------------------------------------- # + + +@pytest.fixture +def multi_plane_tracer(): + return al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.2), + ), + al.Galaxy( + redshift=1.0, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.2), + ), + al.Galaxy(redshift=2.0), + ] + ) + + +def test_intermediate_plane_regions_back_trace_into_the_shape( + shape_solver, multi_plane_tracer +): + """ + A source on the *middle* plane of a three-plane tracer. Solving to the last plane by + mistake gives regions which are not images of this source at all, so the check is not + that regions exist but that they trace back into the shape at the plane asked for. + """ + shape = Circle(0.2, 0.0, radius=0.1) + + regions = shape_solver.image_regions_from( + tracer=multi_plane_tracer, shape=shape, plane_redshift=1.0 + ) + + assert len(regions) == 2 + + plane_index = multi_plane_tracer.plane_index_via_redshift_from(redshift=1.0) + + traced = multi_plane_tracer.traced_grid_2d_list_from( + grid=al.Grid2DIrregular([region.centre for region in regions]) + )[plane_index] + + assert shape.contains(np.asarray(traced.array)).all() + + +def test_intermediate_plane_differs_from_the_last_plane( + shape_solver, multi_plane_tracer +): + """ + `plane_redshift` must actually change the answer, otherwise the test above would pass + for a solver which ignored it. + """ + shape = Circle(0.2, 0.0, radius=0.1) + + at_middle = shape_solver.image_regions_from( + tracer=multi_plane_tracer, shape=shape, plane_redshift=1.0 + ) + at_last = shape_solver.image_regions_from(tracer=multi_plane_tracer, shape=shape) + + assert len(at_middle) != len(at_last) + + +# ---------------------------------------------------------------------------------- # +# (f) the JAX path +# ---------------------------------------------------------------------------------- # + + +def test_use_jax_solver_is_rejected_rather_than_silently_wrong(image_grid, sis_tracer): + """ + `use_jax=True` used to be silently ignored by `find_magnification`, which hardcoded + `xp=np`. Consulting `self._xp` instead would have made it silently *wrong*: see + `test_jax_and_numpy_kept_triangles_agree`. It therefore raises, naming the cap. + + **No `importorskip` here, deliberately.** The rejection must fire on an install with + no JAX at all -- that is the install most likely to hit it -- so this test runs on the + `unittest-nojax` CI leg, where it caught the solver resolving `self._xp` (which + imports `jax.numpy`) before the check, so the user got `ModuleNotFoundError` instead + of the explanation. + """ + solver = ShapeSolver.for_grid( + grid=image_grid, pixel_scale_precision=0.02, use_jax=True + ) + + for call in ( + lambda: solver.find_magnification( + tracer=sis_tracer, shape=Circle(0.3, 0.0, radius=0.1) + ), + lambda: solver.solve_triangles( + tracer=sis_tracer, shape=Circle(0.3, 0.0, radius=0.1) + ), + lambda: list( + solver.steps(tracer=sis_tracer, shape=Circle(0.3, 0.0, radius=0.1)) + ), + ): + with pytest.raises(NotImplementedError) as exc_info: + call() + + assert "MAX_CONTAINING_SIZE" in str(exc_info.value) + + +def test_explicit_jax_module_is_rejected_rather_than_silently_wrong( + image_grid, sis_tracer +): + """ + The other route to the JAX path: a NumPy-configured solver handed ``xp=jax.numpy`` + explicitly. It must give the same error as `use_jax=True`, so neither route can reach + the truncated containers. + + This one needs JAX to have a module to pass, hence the guard -- unlike the test above, + which is exactly the case that must work *without* it. + """ + jnp = pytest.importorskip("jax.numpy") + + solver = ShapeSolver.for_grid(grid=image_grid, pixel_scale_precision=0.02) + + with pytest.raises(NotImplementedError) as exc_info: + solver.find_magnification( + tracer=sis_tracer, shape=Circle(0.3, 0.0, radius=0.1), xp=jnp + ) + + assert "MAX_CONTAINING_SIZE" in str(exc_info.value) + + +@pytest.mark.xfail( + strict=True, + reason=( + "DEFERRED: the JAX triangle containers truncate every refinement step to " + "ArrayTriangles.MAX_CONTAINING_SIZE (15) triangles to keep static shapes. That is " + "enough for a Point but not for a Shape with area, so the JAX path keeps 15 of the " + "~800 triangles the NumPy path keeps and measures a magnification of 0.13 where " + "the truth is 6.86. Lifting the cap is a redesign of the JAX containers (the cap " + "is what makes their shapes static, and an extended source has no static bound), " + "not a fix in ShapeSolver, so it is deferred and ShapeSolver raises on use_jax " + "instead. Remove this xfail when the containers grow a dynamic kept set." + ), +) +def test_jax_and_numpy_kept_triangles_agree(image_grid, sis_tracer): + jnp = pytest.importorskip("jax.numpy") + + shape = Circle(0.3, 0.0, radius=0.1) + + numpy_triangles = ShapeSolver.for_grid( + grid=image_grid, pixel_scale_precision=0.02 + ).solve_triangles(tracer=sis_tracer, shape=shape, xp=np) + + jax_solver = ShapeSolver.for_grid( + grid=image_grid, pixel_scale_precision=0.02, use_jax=True + ) + + # `_xp_from` is what refuses JAX; the parity this test asserts is of the underlying + # algorithm, so it is bypassed deliberately here rather than the guard being weakened. + jax_triangles = super(ShapeSolver, jax_solver).solve_triangles( + tracer=sis_tracer, shape=shape, xp=jnp + ) + + def vertex_set(triangles): + vertices = np.asarray(triangles.triangles).reshape(-1, 2) + vertices = vertices[~np.isnan(vertices).any(axis=1)] + return set(map(tuple, np.unique(np.round(vertices, 8), axis=0))) + + assert vertex_set(numpy_triangles) == vertex_set(jax_triangles)