diff --git a/autolens/point/plot/fit_point_plots.py b/autolens/point/plot/fit_point_plots.py index af87e2ac6..9fb44651b 100644 --- a/autolens/point/plot/fit_point_plots.py +++ b/autolens/point/plot/fit_point_plots.py @@ -1,8 +1,12 @@ +import logging + import numpy as np from typing import Optional from autoarray.plot.utils import subplots, save_figure, conf_subplot_figsize, tight_layout +logger = logging.getLogger(__name__) + def subplot_fit( fit, @@ -18,30 +22,34 @@ def subplot_fit( """ Produce a subplot summarising a `FitPointDataset`. - The subplot contains one or two panels depending on whether flux - measurements are present in the dataset: + The subplot is one row of two or three panels, depending on whether flux measurements are present in the + dataset: * **Positions panel** (always shown), titled ``Fit Positions ()``. Drawn on the same frame as the dataset positions panel (the ``PointSolver`` grid extent, no axis labels), it shows: - the observed positions as large filled markers with their noise-map - circles and the sigma in the legend; - - the model positions on top as red crosses; - - a thin ``Residual`` line from each observed position to the model + circles and the sigma in the legend, numbered 1..N in the order the + ``PointSolver`` figure numbers its images, each number beside its image; + - for an image-plane fit, the model positions on top as red crosses, a + thin ``Residual`` line from each observed position to the model position the fit's own pairing matches it to (Hungarian for ``FitPositionsImagePair``, nearest for ``FitPositionsImagePairRepeat`` - and the dominant nearest term of ``FitPositionsImagePairAll``); - - model images no observed position is paired with as hollow grey + and the dominant nearest term of ``FitPositionsImagePairAll``), and + model images no observed position is paired with as hollow grey crosses (``Unmatched model image``); - the tangential / radial critical curves, thin and light, clipped at the panel edge where they extend beyond the solver grid; - the centre of every mass profile in the tracer (``Mass centre``). - For a source-plane fit (e.g. ``FitPositionsSource``) the model data are - the observed positions ray-traced to the source plane: they are drawn as - ``Model (source plane)`` crosses and no residual lines or unmatched images - are drawn, as there is no image-plane pairing. + A source-plane fit (e.g. ``FitPositionsSource``) draws no model positions here: its model data lie in the + source plane and are drawn on the source-plane panel. + * **Source-plane panel** (always shown), titled ``Fit Source Plane ()``: the observed positions + ray-traced to the point source's plane, numbered as on the positions panel, the source centre and the + caustics of that plane, zoomed onto the traced positions with an inset of the whole caustic. What it draws is + paired to the fit class (see :mod:`autolens.point.plot.source_plane_panel`): context only for an image-plane + chi-squared; the noise footprint and chi-squared of the fit's weighting for a source-plane chi-squared. * **Fluxes panel** (shown only when ``fit.dataset.fluxes`` is not ``None``): a bar/line plot of the observed flux values. @@ -63,7 +71,10 @@ def subplot_fit( Colours of each image-plane line (``"white"`` tangential, ``"yellow"`` radial, as returned by the critical-curve helper). source_plane_lines, source_plane_line_colors - Accepted for a uniform plotter interface; not drawn. + Caustics drawn on the source-plane panel when the point source is in the tracer's last plane (the plane + these lines are computed for). For a point source in an earlier plane of a multi-plane tracer, or when no + lines are passed, the caustics of the point source's own plane are computed on + ``critical_curve_grid_from``. Pass ``[]`` for both line arguments to draw no curves on either panel. title_prefix : str, optional Text prepended to every panel title. solver : PointSolver, optional @@ -71,40 +82,49 @@ def subplot_fit( ``fit.solver``. """ from autoarray.plot.yx import plot_yx + from autolens.lens.plot.critical_curves import critical_curve_grid_from from autolens.point.plot.positions_panel import ( finish_positions_panel, + image_numbers_from, plot_critical_curves, plot_fit_model_positions, + plot_image_numbers, plot_mass_centres, plot_observed_positions, positions_extent_from, ) + from autolens.point.plot import source_plane_panel has_fluxes = fit.dataset.fluxes is not None - n = 2 if has_fluxes else 1 + n = 3 if has_fluxes else 2 if solver is None: solver = getattr(fit, "solver", None) - if image_plane_lines is None and source_plane_lines is None: - from autolens.lens.plot.critical_curves import ( - _compute_critical_curve_lines, - critical_curve_grid_from, - ) + lines_given = image_plane_lines is not None or source_plane_lines is not None + + if not lines_given: + from autolens.lens.plot.critical_curves import _compute_critical_curve_lines - image_plane_lines, image_plane_line_colors, _, _ = ( - _compute_critical_curve_lines( - fit.tracer, critical_curve_grid_from(dataset=fit.dataset, solver=solver) - ) + ( + image_plane_lines, + image_plane_line_colors, + source_plane_lines, + source_plane_line_colors, + ) = _compute_critical_curve_lines( + fit.tracer, critical_curve_grid_from(dataset=fit.dataset, solver=solver) ) fig, axes = subplots(1, n, figsize=conf_subplot_figsize(1, n)) - axes_flat = [axes] if n == 1 else list(np.array(axes).flatten()) + axes_flat = list(np.array(axes).flatten()) _prefix = f"{title_prefix.rstrip()} " if title_prefix else "" + numbers = image_numbers_from(fit.dataset.positions) + # Positions panel ax = axes_flat[0] + extent = positions_extent_from(positions=fit.dataset.positions, solver=solver) plot_critical_curves(ax, image_plane_lines, image_plane_line_colors) plot_mass_centres(ax, fit.tracer) plot_observed_positions( @@ -116,17 +136,48 @@ def subplot_fit( finish_positions_panel( ax=ax, title=f"{_prefix}Fit Positions ({fit.dataset.name})", - extent=positions_extent_from(positions=fit.dataset.positions, solver=solver), + extent=extent, ) + plot_image_numbers(ax, fit.dataset.positions, numbers, extent) + + # Source-plane panel + source_plane_title = f"{_prefix}Fit Source Plane ({fit.dataset.name})" + try: + data = source_plane_panel.source_plane_data_from( + fit.positions, numbers=numbers + ) + tangential, radial = source_plane_caustics_from( + fit=fit, + solver=solver, + lines_given=lines_given, + source_plane_lines=source_plane_lines, + source_plane_line_colors=source_plane_line_colors, + ) + source_plane_panel.plot_source_plane_panel( + ax=axes_flat[1], + data=data, + title=source_plane_title, + tangential_caustics=tangential, + radial_caustics=radial, + ) + except Exception: + # A visualization must never stop a model-fit: report the failure and leave the panel empty. + logger.warning( + "The source-plane panel of fit.png could not be drawn; it is left empty.", + exc_info=True, + ) + axes_flat[1].cla() + axes_flat[1].set_title(source_plane_title) + axes_flat[1].set_axis_off() # Fluxes panel - if has_fluxes and n > 1: + if has_fluxes: y = np.array(fit.dataset.fluxes) x = np.arange(len(y)) plot_yx( y=y, x=x, - ax=axes_flat[1], + ax=axes_flat[2], title=f"{_prefix}Fit Fluxes ({fit.dataset.name})", output_path=None, output_filename="fit_point_fluxes", @@ -137,6 +188,43 @@ def subplot_fit( save_figure(fig, path=output_path, filename="fit", format=output_format) +def source_plane_caustics_from( + fit, + solver=None, + lines_given: bool = False, + source_plane_lines=None, + source_plane_line_colors=None, +): + """ + Returns the ``(tangential, radial)`` caustics drawn on the source-plane panel of ``fit``: those of the plane + the point source is in. + + The pre-computed ``source_plane_lines`` are the caustics of the tracer's last plane, so they are reused when + the point source is in that plane. Otherwise (a point in an earlier plane of a multi-plane tracer, or no lines + passed) the caustics of the point's own plane are computed on ``critical_curve_grid_from``. If lines were given + and both are empty (``[]``), no caustics are drawn. + """ + from autolens.lens.plot.critical_curves import critical_curve_grid_from + from autolens.point.plot import source_plane_panel + + plane_index = fit.positions.plane_index + last_plane = len(fit.tracer.planes) - 1 + in_last_plane = plane_index in (-1, last_plane) + + if lines_given and source_plane_lines is not None and in_last_plane: + return source_plane_panel.caustics_from_lines( + source_plane_lines, source_plane_line_colors + ) + if lines_given and source_plane_lines is not None and len(source_plane_lines) == 0: + return [], [] + + return source_plane_panel.point_plane_caustics_from( + tracer=fit.tracer, + plane_index=plane_index, + grid=critical_curve_grid_from(dataset=fit.dataset, solver=solver), + ) + + def plot_image_with_points( image, dataset, diff --git a/autolens/point/plot/point_solver_plots.py b/autolens/point/plot/point_solver_plots.py index a825bfb63..caa9696e7 100644 --- a/autolens/point/plot/point_solver_plots.py +++ b/autolens/point/plot/point_solver_plots.py @@ -37,6 +37,15 @@ import numpy as np +from autolens.point.plot.positions_panel import ( + MARKER_RADIUS_POINTS, + NUMBER_PLACEMENTS, + adjacent_number_placements_from, + box_overlaps_circles, + boxes_overlap, + image_order_from, +) + logger = logging.getLogger(__name__) TESTED_COLOR = "#8a8f95" @@ -284,7 +293,7 @@ def point_solver_data_from( plane_redshift=plane_redshift, ) ) - solved = solved[np.argsort(-np.arctan2(solved[:, 0], solved[:, 1]))] + solved = solved[image_order_from(solved)] positions = None if positions is None else _yx(positions) extent = _extent_from(solver=solver, positions=positions) @@ -356,19 +365,13 @@ def panel_title_from(step, data: PointSolverFigureData, short: bool = False) -> return f"{title} ({counts})" -# Candidate placements of an image number relative to its image, as (dx, dy) offsets in points plus the text's -# horizontal and vertical alignment: the number's nearest corner sits a few points off the image marker, upper-left -# first (the teaching figure's look), so a number is always immediately adjacent to its own image. -_NUMBER_PLACEMENTS = ( - (-6, 5, "right", "bottom"), - (6, 5, "left", "bottom"), - (-6, -5, "right", "top"), - (6, -5, "left", "top"), -) +# The placement of image numbers is shared with the point-source fit panels (``positions_panel``), so an image's +# number sits beside it the same way in every point-source figure. +_NUMBER_PLACEMENTS = NUMBER_PLACEMENTS NUMBER_FONTSIZE = 14 # Radius (points) treated as occupied around an image marker when placing labels. -_MARKER_RADIUS_POINTS = 6.0 +_MARKER_RADIUS_POINTS = MARKER_RADIUS_POINTS # Candidate placements of the rejected-candidate label next to a single candidate, as (x, y) offsets in points and # horizontal alignment. @@ -402,17 +405,8 @@ def _to_points(data: PointSolverFigureData, xy: np.ndarray, panel_inches: float) return (xy - np.array([data.extent[0], data.extent[2]])) * scale -def _box_overlaps_circles(box, centres, radius) -> np.ndarray: - x0, y0, x1, y1 = box - if len(centres) == 0: - return np.zeros(0, dtype=bool) - gap_x = np.maximum(np.maximum(x0 - centres[:, 0], centres[:, 0] - x1), 0.0) - gap_y = np.maximum(np.maximum(y0 - centres[:, 1], centres[:, 1] - y1), 0.0) - return np.hypot(gap_x, gap_y) < radius - - -def _boxes_overlap(a, b) -> bool: - return not (a[2] <= b[0] or b[2] <= a[0] or a[3] <= b[1] or b[3] <= a[1]) +_box_overlaps_circles = box_overlaps_circles +_boxes_overlap = boxes_overlap def _panel_points(data: PointSolverFigureData, panel_inches: float): @@ -439,31 +433,13 @@ def number_placements_from( panel_inches, ) width, height = _panel_points(data, panel_inches) - placed_boxes, placements = [], [] - for index, (x, y) in enumerate(markers[: len(data.solved)]): - text_width = 0.65 * NUMBER_FONTSIZE * len(str(index + 1)) + 2.0 - text_height = NUMBER_FONTSIZE + 2.0 - others = np.delete(markers, index, axis=0) - best, best_score = _NUMBER_PLACEMENTS[0], np.inf - for dx, dy, ha, va in _NUMBER_PLACEMENTS: - x0 = x + dx - (text_width if ha == "right" else 0.0) - y0 = y + dy - (text_height if va == "top" else 0.0) - box = (x0, y0, x0 + text_width, y0 + text_height) - inside = ( - box[0] >= 0 and box[1] >= 0 and box[2] <= width and box[3] <= height - ) - score = ( - int(_box_overlaps_circles(box, others, _MARKER_RADIUS_POINTS).sum()) - + sum(_boxes_overlap(box, other) for other in placed_boxes) - + (0 if inside else 100) - ) - if score < best_score: - best, best_score, best_box = (dx, dy, ha, va), score, box - if score == 0: - break - placements.append(best) - placed_boxes.append(best_box) - return placements + return adjacent_number_placements_from( + markers=markers, + labels=[str(index + 1) for index in range(len(data.solved))], + width=width, + height=height, + fontsize=NUMBER_FONTSIZE, + ) def _rejected_label_placement( diff --git a/autolens/point/plot/positions_panel.py b/autolens/point/plot/positions_panel.py index 5b5066913..36eccf718 100644 --- a/autolens/point/plot/positions_panel.py +++ b/autolens/point/plot/positions_panel.py @@ -13,7 +13,12 @@ - **No axis labels**, so the arc-second tick labels are not overlapped. The fit panel additionally draws the model positions, the residual lines of the fit's own pairing, unmatched model -images, critical curves and the centres of the tracer's mass profiles. +images, critical curves, the centres of the tracer's mass profiles and the image numbers. + +**Image numbers.** Images are numbered 1..N in the order of :func:`image_order_from` (descending angle about the +origin, the order the ``PointSolver`` figure numbers its solved images in), and each number is placed immediately +beside its own image by :func:`adjacent_number_placements_from`, which the ``PointSolver`` figure shares. The fit's +image-plane and source-plane panels use the same numbers, so image ``i`` can be followed from one to the other. """ import numpy as np @@ -36,6 +41,23 @@ CRITICAL_CURVE_COLOR_MAP = {"white": "0.6", "yellow": "tab:olive"} +# Candidate placements of an image number relative to its image, as (dx, dy) offsets in points plus the text's +# horizontal and vertical alignment: the number's nearest corner sits a few points off the image marker, upper-left +# first, so a number is always immediately adjacent to its own image. +NUMBER_PLACEMENTS = ( + (-6, 5, "right", "bottom"), + (6, 5, "left", "bottom"), + (-6, -5, "right", "top"), + (6, -5, "left", "top"), +) + +# Radius (points) treated as occupied around an image marker when placing labels. +MARKER_RADIUS_POINTS = 6.0 + +# Font size of the image numbers on the fit panels. +FIT_NUMBER_FONTSIZE = 12 + + def _np_grid(grid) -> np.ndarray: """ Returns a (N, 2) numpy array of (y,x) coordinates from a grid-like input. @@ -299,6 +321,40 @@ def unmatched_model_mask_from(model: np.ndarray, indexes: np.ndarray) -> np.ndar return mask +def residual_style_from(fit_positions) -> Tuple[str, str]: + """ + Returns the ``(linestyle, legend label)`` of an image-plane fit's residual lines, by its pairing. + + Hard-pairing fits (``FitPositionsImagePair``, ``FitPositionsImagePairRepeat`` and its ``*Solved`` variant) pair + each observed position with one model position: solid ``"Residual"`` lines. ``FitPositionsImagePairAll`` (and + its ``*Solved`` variant) uses an all-pairs mixture likelihood with no hard pairing, so the line drawn to the + nearest model position is only its dominant term: dashed ``"Residual (nearest, soft pairing)"`` lines. The class + is looked up along its method resolution order, so a subclass is styled as its parent. + """ + from autolens.point.fit.positions.image.pair_all import FitPositionsImagePairAll + + soft = ("--", "Residual (nearest, soft pairing)") + table = {FitPositionsImagePairAll: soft} + for cls in type(fit_positions).__mro__: + if cls in table: + return table[cls] + return "-", "Residual" + + +def model_label_from(fit_positions) -> str: + """ + Returns the legend label of an image-plane fit's model positions: ``"Model"``, or + ``"Model (from solved β*)"`` for a ``*Solved`` fit, whose model images are solved from the analytically solved + source centre rather than a free ``centre`` parameter (the same ``β*`` the source-plane panel marks with a + diamond). + """ + from autolens.point.fit.solved import SolvedCentre + + if isinstance(fit_positions, SolvedCentre): + return "Model (from solved β*)" + return "Model" + + def plot_fit_model_positions(ax, fit_positions): """ Draws the model positions of a positions fit on top of the observed positions. @@ -307,29 +363,16 @@ def plot_fit_model_positions(ax, fit_positions): each observed position to the model position the fit pairs it with, and model images no observed position is paired with are drawn as hollow grey crosses (``Unmatched model image``). - For a source-plane fit (e.g. ``FitPositionsSource``) the model data are the observed positions ray-traced to - the source plane: they are drawn as red crosses labelled ``Model (source plane)`` and no residual lines are - drawn, as there is no image-plane pairing. + For a source-plane fit (e.g. ``FitPositionsSource``) nothing is drawn: its model data are the observed positions + ray-traced to the source plane, which belong on the fit's source-plane panel, not the image plane. """ pairing = model_pairing_from(fit_positions) if pairing is None: - model = _np_grid(fit_positions.model_data) - model = model[np.isfinite(model).all(axis=1)] - if len(model): - ax.scatter( - model[:, 1], - model[:, 0], - s=MODEL_MARKER_SIZE, - c=MODEL_COLOR, - marker="x", - linewidths=MODEL_MARKER_LINEWIDTH, - zorder=7, - label="Model (source plane)", - ) return observed, model, indexes = pairing + linestyle, label = residual_style_from(fit_positions) labelled = False for obs, index in zip(observed, indexes): @@ -340,9 +383,10 @@ def plot_fit_model_positions(ax, fit_positions): [obs[0], model[index, 0]], color=MODEL_COLOR, linewidth=1.0, + linestyle=linestyle, alpha=0.7, zorder=6, - label=None if labelled else "Residual", + label=None if labelled else label, ) labelled = True @@ -357,7 +401,7 @@ def plot_fit_model_positions(ax, fit_positions): marker="x", linewidths=MODEL_MARKER_LINEWIDTH, zorder=7, - label="Model", + label=model_label_from(fit_positions), ) if np.any(unmatched): ax.scatter( @@ -384,3 +428,155 @@ def finish_positions_panel(ax, title: str, extent: List[float]): apply_extent(ax, extent) ax.set_aspect("equal", adjustable="box") ax.legend(loc="best", fontsize=9, framealpha=0.85, markerscale=0.7) + + +def image_order_from(positions) -> np.ndarray: + """ + Returns the indexes of the positions in numbering order: descending angle about the origin, the order in which + the ``PointSolver`` figure numbers its solved images. Image ``k`` (1-based) is ``positions[order[k - 1]]``. + """ + positions = _np_grid(positions) + return np.argsort(-np.arctan2(positions[:, 0], positions[:, 1]), kind="stable") + + +def image_numbers_from(positions) -> np.ndarray: + """ + Returns the 1-based number of every position, in the positions' own (dataset) order, following + :func:`image_order_from`. + """ + order = image_order_from(positions) + numbers = np.empty(len(order), dtype=int) + numbers[order] = np.arange(1, len(order) + 1) + return numbers + + +def box_overlaps_circles(box, centres, radius) -> np.ndarray: + """ + Returns, for every circle of ``radius`` around ``centres``, whether it overlaps the ``(x0, y0, x1, y1)`` box. + """ + x0, y0, x1, y1 = box + if len(centres) == 0: + return np.zeros(0, dtype=bool) + gap_x = np.maximum(np.maximum(x0 - centres[:, 0], centres[:, 0] - x1), 0.0) + gap_y = np.maximum(np.maximum(y0 - centres[:, 1], centres[:, 1] - y1), 0.0) + return np.hypot(gap_x, gap_y) < radius + + +def boxes_overlap(a, b) -> bool: + return not (a[2] <= b[0] or b[2] <= a[0] or a[3] <= b[1] or b[3] <= a[1]) + + +def adjacent_number_placements_from( + markers: np.ndarray, + labels: List[str], + width: float, + height: float, + fontsize: float, +) -> List[Tuple[float, float, str, str]]: + """ + Returns, for each of the first ``len(labels)`` markers, the ``(dx, dy, ha, va)`` offset (points) of its label. + + Each label sits a few points off its own marker (one of ``NUMBER_PLACEMENTS``): the first placement which stays + inside the panel and overlaps neither another marker nor an already-placed label, else the one overlapping + least. A label is therefore never moved away from its marker to find space. + + Parameters + ---------- + markers + The (x, y) positions of every marker in points from the panel's lower-left corner; markers after the first + ``len(labels)`` are obstacles only (e.g. rejected candidates). + labels + The text of each numbered marker's label. + width, height + The size of the panel in points. + fontsize + The label font size in points. + """ + markers = np.asarray(markers, dtype=float).reshape(-1, 2) + placed_boxes, placements = [], [] + for index, (x, y) in enumerate(markers[: len(labels)]): + text_width = 0.65 * fontsize * len(labels[index]) + 2.0 + text_height = fontsize + 2.0 + others = np.delete(markers, index, axis=0) + best, best_score, best_box = NUMBER_PLACEMENTS[0], np.inf, None + for dx, dy, ha, va in NUMBER_PLACEMENTS: + x0 = x + dx - (text_width if ha == "right" else 0.0) + y0 = y + dy - (text_height if va == "top" else 0.0) + box = (x0, y0, x0 + text_width, y0 + text_height) + inside = ( + box[0] >= 0 and box[1] >= 0 and box[2] <= width and box[3] <= height + ) + score = ( + int(box_overlaps_circles(box, others, MARKER_RADIUS_POINTS).sum()) + + sum(boxes_overlap(box, other) for other in placed_boxes) + + (0 if inside else 100) + ) + if score < best_score: + best, best_score, best_box = (dx, dy, ha, va), score, box + if score == 0: + break + placements.append(best) + placed_boxes.append(best_box) + return placements + + +def _panel_size_points(ax, extent) -> Tuple[float, float, float]: + """ + Returns the ``(width, height, points_per_arcsec)`` of an equal-aspect panel showing ``extent``, from the axes' + position in its figure. + """ + fig_width, fig_height = ax.figure.get_size_inches() + position = ax.get_position() + box_width = position.width * fig_width * 72.0 + box_height = position.height * fig_height * 72.0 + span_x = max(extent[1] - extent[0], 1e-30) + span_y = max(extent[3] - extent[2], 1e-30) + scale = min(box_width / span_x, box_height / span_y) + return span_x * scale, span_y * scale, scale + + +def plot_image_numbers( + ax, + positions, + numbers, + extent: List[float], + fontsize: float = FIT_NUMBER_FONTSIZE, +): + """ + Labels every position with its image number, placed immediately beside it by + :func:`adjacent_number_placements_from`. + + Parameters + ---------- + positions + The (y,x) positions which are labelled. + numbers + The number of each position (see :func:`image_numbers_from`). + extent + The ``[x_min, x_max, y_min, y_max]`` extent of the panel, used to convert positions to points. + """ + positions = _np_grid(positions) + if len(positions) == 0: + return + width, height, scale = _panel_size_points(ax, extent) + markers = (positions[:, ::-1] - np.array([extent[0], extent[2]])) * scale + placements = adjacent_number_placements_from( + markers=markers, + labels=[str(number) for number in numbers], + width=width, + height=height, + fontsize=fontsize, + ) + for (y, x), number, (dx, dy, ha, va) in zip(positions, numbers, placements): + ax.annotate( + f"{number}", + (x, y), + xytext=(dx, dy), + textcoords="offset points", + ha=ha, + va=va, + fontsize=fontsize, + weight="bold", + zorder=10, + bbox=dict(fc="w", ec="none", alpha=0.8, pad=0.5), + ) diff --git a/autolens/point/plot/source_plane_panel.py b/autolens/point/plot/source_plane_panel.py new file mode 100644 index 000000000..55c8005d1 --- /dev/null +++ b/autolens/point/plot/source_plane_panel.py @@ -0,0 +1,708 @@ +""" +The "Fit Source Plane" panel of the point-source ``fit.png``: the observed positions ray-traced to the source plane, +showing how closely they trace back together. + +What the panel draws is paired to the positions fit class, so that it never shows a chi-squared the fit does not +compute (:func:`presentation_from`, one entry per class of ``autolens.point.fit.positions``): + +- **Image-plane chi-squared** (``FitPositionsImagePair``, ``FitPositionsImagePairRepeat``, + ``FitPositionsImagePairAll`` and their ``*Solved`` variants): the panel is context only. It shows the traced + positions, the source centre and the caustics, with no noise footprints and no chi-squared box, as the fit's + chi-squared is measured in the image plane. +- **Source-plane chi-squared** (``FitPositionsSource``, ``FitPositionsSourceSolved``): the footprint the + chi-squared actually uses is drawn on every traced position: + + - ``weighting = "magnification"``: a circle of radius ``σ/|μ|`` (the scalar noise of + ``χ²ᵢ = |β̂ᵢ − c|² µᵢ² / σᵢ²``), plus the ``A·σ`` ellipse dashed as the propagated noise (reference); + - ``weighting = "jacobian"``: the ``A·σ`` ellipse, the source-plane covariance ``σ² A Aᵀ = Wᵢ⁻¹`` of the + precision tensor ``Wᵢ = Aᵢ⁻ᵀΘᵢAᵢ⁻¹`` the chi-squared uses; + + and the legend gives the chi-squared the fit uses (only that one), as a handle-less entry ``χ² = ``. + +The source centre is a red star for a free-centre fit (the point-source profile's ``centre``) and a red diamond for +a ``*Solved`` fit (the analytically solved ``β*``). + +Every fit-type cue lives in the legend, as marker symbols and labels (the centre marker, the weighting-specific +noise label, the chi-squared entry): the panel carries no free-floating annotation text about the fit type. + +The panel is zoomed onto the traced positions and the centre with a margin; an inset in the emptiest corner shows the +whole tangential caustic with a box marking the zoom. The caustics are those of the plane the point source is in +(``fit.plane_index``), so a multi-plane (e.g. cluster) point is drawn against its own caustics. The images are +numbered as on the image-plane panel. +""" + +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import numpy as np + +from autolens.point.plot import positions_panel as pp + +IMAGE_PLANE = "image" +SOURCE_PLANE = "source" + +CAUSTIC_COLOR = "0.6" +RADIAL_CAUSTIC_COLOR = "tab:olive" +CENTRE_COLOR = pp.MODEL_COLOR +MAGNIFICATION_NOISE_COLOR = "tab:purple" +JACOBIAN_NOISE_COLOR = pp.NOISE_COLOR + +# The zoom extends this fraction of its half-width beyond the traced positions and centre, on every side. +ZOOM_MARGIN = 0.3 + +# Size of the caustic inset, as a fraction of the panel. +INSET_SIZE = 0.3 + +# The propagated-noise reference ellipses under magnification weighting: dashed and darker than the primary +# footprint's colour so they read on white, but thinner than the primary σ/|μ| circles. +REFERENCE_NOISE_COLOR = "#c25400" + + +@dataclass(frozen=True) +class SourcePlanePresentation: + """ + How the source-plane panel presents a positions fit (see :func:`presentation_from`). + + Parameters + ---------- + chi_squared_plane + ``"image"`` if the fit's chi-squared is measured in the image plane (context-only panel), ``"source"`` if + it is measured in the source plane. + solved_centre + Whether the source centre is the analytically solved ``β*`` (a ``*Solved`` fit) rather than a free + ``centre`` parameter. + weighting + The source-plane chi-squared weighting (``"magnification"`` or ``"jacobian"``); ``None`` for image-plane + fits. + """ + + chi_squared_plane: str + solved_centre: bool + weighting: Optional[str] = None + + @property + def is_source_plane(self) -> bool: + return self.chi_squared_plane == SOURCE_PLANE + + @property + def centre_marker(self) -> str: + return "D" if self.solved_centre else "*" + + @property + def centre_label(self) -> str: + return ( + "Solved source centre (β*)" if self.solved_centre else "Model source centre" + ) + + +def presentation_table(): + """ + Returns the presentation of every positions fit class in ``autolens.point.fit.positions``, as a dict mapping + the class to ``(chi_squared_plane, solved_centre)``. + """ + from autolens.point.fit.positions.image.pair import FitPositionsImagePair + from autolens.point.fit.positions.image.pair_all import ( + FitPositionsImagePairAll, + FitPositionsImagePairAllSolved, + ) + from autolens.point.fit.positions.image.pair_repeat import ( + FitPositionsImagePairRepeat, + FitPositionsImagePairRepeatSolved, + ) + from autolens.point.fit.positions.source.separations import ( + FitPositionsSource, + FitPositionsSourceSolved, + ) + + return { + FitPositionsImagePair: (IMAGE_PLANE, False), + FitPositionsImagePairRepeat: (IMAGE_PLANE, False), + FitPositionsImagePairRepeatSolved: (IMAGE_PLANE, True), + FitPositionsImagePairAll: (IMAGE_PLANE, False), + FitPositionsImagePairAllSolved: (IMAGE_PLANE, True), + FitPositionsSource: (SOURCE_PLANE, False), + FitPositionsSourceSolved: (SOURCE_PLANE, True), + } + + +def presentation_from(fit_positions) -> SourcePlanePresentation: + """ + Returns how the source-plane panel presents a positions fit, paired to its class. + + The class is looked up in :func:`presentation_table` along its method resolution order, so a user subclass of a + fit class (e.g. ``FitPositionsSource`` with ``weighting = "jacobian"``) is presented as its parent. The weighting + of a source-plane fit is read from the fit itself. A class with no entry is presented as context only, which + never draws a chi-squared the fit may not compute. + """ + table = presentation_table() + for cls in type(fit_positions).__mro__: + if cls in table: + plane, solved = table[cls] + break + else: + from autolens.point.fit.solved import SolvedCentre + + plane, solved = IMAGE_PLANE, isinstance(fit_positions, SolvedCentre) + + weighting = ( + getattr(fit_positions, "weighting", None) if plane == SOURCE_PLANE else None + ) + return SourcePlanePresentation( + chi_squared_plane=plane, solved_centre=solved, weighting=weighting + ) + + +def _np(values) -> np.ndarray: + values = values.array if hasattr(values, "array") else values + return np.asarray(values, dtype=float) + + +def traced_positions_from(fit_positions) -> np.ndarray: + """ + Returns the observed positions ray-traced to the plane of the point source, as a (N, 2) (y,x) array. + + This is the same ray-tracing as ``FitPositionsSource.model_data`` (multi-plane deflections to + ``fit.plane_index`` when the tracer has more than two planes), applied to any positions fit. + """ + import autoarray as aa + + tracer = fit_positions.tracer + # A NumPy copy of the positions: the figure is NumPy-only, even for a fit built with `use_jax=True` (whose data + # grid is JAX-backed). + positions = aa.Grid2DIrregular(values=pp._np_grid(fit_positions.data)) + + if len(tracer.planes) <= 2: + deflections = tracer.deflections_yx_2d_from(grid=positions, xp=np) + else: + deflections = tracer.deflections_between_planes_from( + grid=positions, xp=np, plane_i=0, plane_j=fit_positions.plane_index + ) + return pp._np_grid( + positions.grid_2d_via_deflection_grid_from(deflection_grid=deflections, xp=np) + ) + + +def jacobian_covariances_from(fit_positions) -> np.ndarray: + """ + Returns the (N, 2, 2) source-plane covariance ``Cᵢ = Wᵢ⁻¹ = σᵢ² Aᵢ Aᵢᵀ`` of every traced position, in (y,x) + order, where ``Wᵢ`` is the jacobian precision tensor of + :func:`~autolens.point.fit.solved.precision_tensor_components_from`. A position whose tensor is singular (on + a critical curve) has a ``NaN`` covariance. + """ + from autolens.point.fit.solved import precision_tensor_components_from + + w11, w12, w21, w22 = ( + _np(w) for w in precision_tensor_components_from(fit_positions, "jacobian") + ) + precision = np.stack( + [np.stack([w11, w12], axis=-1), np.stack([w21, w22], axis=-1)], axis=-2 + ) + covariances = np.full_like(precision, np.nan) + for i, w in enumerate(precision): + if np.all(np.isfinite(w)) and abs(np.linalg.det(w)) > 0: + covariances[i] = np.linalg.inv(w) + return covariances + + +def ellipse_from(covariance: np.ndarray) -> Optional[Tuple[float, float, float]]: + """ + Returns the ``(semi_major, semi_minor, angle_deg)`` of the 1-sigma ellipse of a (y,x) covariance, the angle + measured anticlockwise from the x axis, or ``None`` for a non-finite covariance. + """ + if not np.all(np.isfinite(covariance)): + return None + values, vectors = np.linalg.eigh(0.5 * (covariance + covariance.T)) + values = np.clip(values, 0.0, None) + major = vectors[:, 1] + angle = float(np.degrees(np.arctan2(major[0], major[1]))) + return float(np.sqrt(values[1])), float(np.sqrt(values[0])), angle + + +def chi_squared_from(residuals: np.ndarray, covariances: np.ndarray) -> float: + """ + Returns ``Σᵢ rᵢᵀ Cᵢ⁻¹ rᵢ`` for (y,x) residuals ``rᵢ`` and covariances ``Cᵢ``, skipping non-finite terms. + """ + total = 0.0 + for r, c in zip(residuals, covariances): + if np.all(np.isfinite(c)) and np.all(np.isfinite(r)): + total += float(r @ np.linalg.solve(c, r)) + return total + + +@dataclass +class SourcePlaneData: + """ + Everything the source-plane panel draws, read off a positions fit. + """ + + presentation: SourcePlanePresentation + traced: np.ndarray + centre: np.ndarray + numbers: np.ndarray + sigma: np.ndarray + magnification_radii: np.ndarray + covariances: np.ndarray + chi_squared: Optional[float] + + +def source_plane_data_from(fit_positions, numbers=None) -> SourcePlaneData: + """ + Returns the source-plane quantities of a positions fit, computed with the fit's own machinery so the drawn + noise is exactly that of its chi-squared. + + Parameters + ---------- + fit_positions + Any positions fit (``fit.positions`` of a ``FitPointDataset``). + numbers + The image number of each observed position; ``None`` uses + :func:`~autolens.point.plot.positions_panel.image_numbers_from`. + """ + presentation = presentation_from(fit_positions) + + traced = traced_positions_from(fit_positions) + centre = np.asarray(fit_positions.source_plane_coordinate, dtype=float).ravel()[:2] + sigma = np.broadcast_to(_np(fit_positions.noise_map).ravel(), (len(traced),)).copy() + + try: + mu = np.abs(_np(fit_positions.magnifications_at_positions).ravel()) + with np.errstate(divide="ignore"): + radii = sigma / mu + except Exception: + radii = np.full(len(traced), np.nan) + + try: + covariances = jacobian_covariances_from(fit_positions) + except Exception: + covariances = np.full((len(traced), 2, 2), np.nan) + + chi_squared = None + if presentation.is_source_plane: + chi_squared = float(np.sum(_np(fit_positions.chi_squared_map))) + + if numbers is None: + numbers = pp.image_numbers_from(fit_positions.data) + + return SourcePlaneData( + presentation=presentation, + traced=traced, + centre=centre, + numbers=np.asarray(numbers), + sigma=sigma, + magnification_radii=radii, + covariances=covariances, + chi_squared=chi_squared, + ) + + +def point_plane_caustics_from( + tracer, plane_index: int, grid +) -> Tuple[List[np.ndarray], List[np.ndarray]]: + """ + Returns the ``(tangential, radial)`` caustics of the plane ``plane_index`` the point source is in, as lists of + (N, 2) (y,x) arrays, computed by marching squares on the image-plane ``grid``. + + For a two-plane tracer this is the tracer's own source plane; for a multi-plane tracer the caustics are those of + the deflections from the image plane to ``plane_index`` (``LensCalc.from_tracer(..., plane_j=plane_index)``), + the same deflections the fit traces the positions with. Radial caustics are only computed when a radial critical + curve is resolved by the grid. Returns empty lists when ``PYAUTO_FAST_PLOTS=1`` or the model has no caustics. + """ + import os + + if os.environ.get("PYAUTO_FAST_PLOTS") == "1": + return [], [] + + from autogalaxy.operate.lens_calc import LensCalc + + def _lines(curves): + lines = [] + for curve in curves or []: + line = pp._np_grid(curve) + line = line[np.isfinite(line).all(axis=1)] + if len(line) > 1: + lines.append(line) + return lines + + use_multi_plane = len(tracer.planes) > 2 + try: + lens_calc = LensCalc.from_tracer( + tracer=tracer, + use_multi_plane=use_multi_plane, + plane_i=0, + plane_j=plane_index if use_multi_plane else -1, + ) + tangential = _lines(lens_calc.tangential_caustic_list_from(grid=grid)) + except Exception: + return [], [] + + radial = [] + try: + areas = lens_calc.radial_critical_curve_area_list_from(grid=grid) + if any(area > grid.pixel_scale for area in areas): + radial = _lines(lens_calc.radial_caustic_list_from(grid=grid)) + except Exception: + radial = [] + + return tangential, radial + + +def caustics_from_lines( + lines, line_colors +) -> Tuple[List[np.ndarray], List[np.ndarray]]: + """ + Splits pre-computed source-plane lines into ``(tangential, radial)`` caustics by their colour (``"white"`` + tangential, ``"yellow"`` radial, as returned by the critical-curve helper). + """ + tangential, radial = [], [] + for i, line in enumerate(lines or []): + line = np.asarray(line, dtype=float) + if line.ndim != 2 or len(line) < 2: + continue + color = line_colors[i] if line_colors and i < len(line_colors) else "white" + (radial if color == "yellow" else tangential).append(line) + return tangential, radial + + +def zoom_extent_from(data: SourcePlaneData, margin: float = ZOOM_MARGIN) -> List[float]: + """ + Returns the ``[x_min, x_max, y_min, y_max]`` square zoom of the panel: the traced positions and the centre, + with ``margin`` of the half-width added on every side. + + The half-width is at least three times the largest ``σ/|μ|`` (so a near-perfect fit is not zoomed to a point) + and, for a source-plane fit, at least the largest footprint its chi-squared uses, so every footprint is visible. + """ + points = np.vstack([data.traced, data.centre[None, :]]) + points = points[np.isfinite(points).all(axis=1)] + lo, hi = points.min(axis=0), points.max(axis=0) + half = 0.5 * float(np.max(hi - lo)) + + radii = data.magnification_radii[np.isfinite(data.magnification_radii)] + floor = 3.0 * float(radii.max()) if len(radii) else 0.0 + + if data.presentation.is_source_plane: + if data.presentation.weighting == "magnification" and len(radii): + floor = max(floor, float(radii.max())) + elif data.presentation.weighting == "jacobian": + majors = [e[0] for e in map(ellipse_from, data.covariances) if e] + if majors: + floor = max(floor, max(majors)) + + if floor <= 0.0: + floor = 0.05 * max(float(np.max(np.abs(points))), 1e-3) + + half = max(half, floor) * (1.0 + 2.0 * margin) + centre_y, centre_x = 0.5 * (lo + hi) + return [centre_x - half, centre_x + half, centre_y - half, centre_y + half] + + +CORNER_LOCATIONS = { + (0, 0): "lower left", + (1, 0): "lower right", + (0, 1): "upper left", + (1, 1): "upper right", +} + + +def ranked_corners_from( + extent: List[float], points: np.ndarray, size: float = INSET_SIZE, pad: float = 0.02 +) -> List[Tuple[Tuple[int, int], Tuple[float, float, float, float]]]: + """ + Returns the panel's corners from emptiest to fullest, as ``((i, j), (x0, y0, width, height))`` with the box in + axes fraction: ranked by how few of the drawn ``points`` (y,x) a ``size`` box in the corner holds, then by how + far the nearest point is from its centre. + """ + x_min, x_max, y_min, y_max = extent + points = np.asarray(points, dtype=float).reshape(-1, 2) + fx = (points[:, 1] - x_min) / (x_max - x_min) + fy = (points[:, 0] - y_min) / (y_max - y_min) + ranked = [] + for i, x0 in enumerate((pad, 1.0 - pad - size)): + for j, y0 in enumerate((pad, 1.0 - pad - size)): + inside = int( + np.sum( + (fx > x0 - 0.03) + & (fx < x0 + size + 0.03) + & (fy > y0 - 0.03) + & (fy < y0 + size + 0.03) + ) + ) + distance = ( + float(np.min(np.hypot(fx - (x0 + size / 2), fy - (y0 + size / 2)))) + if len(points) + else 0.0 + ) + ranked.append(((inside, -distance), (i, j), (x0, y0, size, size))) + ranked.sort(key=lambda r: r[0]) + return [(r[1], r[2]) for r in ranked] + + +def _footprint_points(data: SourcePlaneData) -> np.ndarray: + """ + The traced positions, the centre and the ends of every drawn ellipse, used to find empty corners. + """ + points = [data.traced, data.centre[None, :]] + covariances = data.covariances if data.presentation.is_source_plane else [] + for (y, x), covariance in zip(data.traced, covariances): + ellipse = ellipse_from(covariance) + if ellipse is None: + continue + semi_major, _, angle = ellipse + t = np.radians(angle) + d = semi_major * np.array([np.sin(t), np.cos(t)]) + points.append(np.array([[y + d[0], x + d[1]], [y - d[0], x - d[1]]])) + points = np.vstack(points) + return points[np.isfinite(points).all(axis=1)] + + +def _plot_caustics(ax, tangential, radial, linewidth=0.9, label=True): + labelled = False + for line, color in [(c, CAUSTIC_COLOR) for c in tangential] + [ + (c, RADIAL_CAUSTIC_COLOR) for c in radial + ]: + ax.plot( + line[:, 1], + line[:, 0], + color=color, + linewidth=linewidth, + alpha=0.8, + zorder=2, + label="Caustic" if (label and not labelled) else None, + ) + labelled = True + + +def _plot_centre(ax, data: SourcePlaneData, size=180, label=True): + presentation = data.presentation + ax.scatter( + [data.centre[1]], + [data.centre[0]], + s=size if presentation.centre_marker == "*" else 0.5 * size, + c=CENTRE_COLOR, + marker=presentation.centre_marker, + linewidths=0.6, + edgecolors="white", + zorder=8, + label=presentation.centre_label if label else None, + ) + + +def _milliarcsec_range(values: np.ndarray) -> str: + values = 1e3 * values[np.isfinite(values)] + if len(values) == 0: + return "n/a" + if np.allclose(values, values[0]): + return f"{values[0]:.2g} mas" + return f"{values.min():.2g}–{values.max():.2g} mas" + + +def _format_chi_squared(value: float) -> str: + return f"{value:,.0f}" if abs(value) >= 1000 else f"{value:.3g}" + + +def chi_squared_label_from(data: SourcePlaneData) -> Optional[str]: + """ + Returns the legend entry of the chi-squared the fit uses, e.g. ``"χ² = 10,023"``, or ``None`` for an + image-plane fit (which has no source-plane chi-squared to show). + """ + if not data.presentation.is_source_plane or data.chi_squared is None: + return None + return f"χ² = {_format_chi_squared(data.chi_squared)}" + + +def _plot_footprints(ax, data: SourcePlaneData): + """ + Draws the noise footprint the fit's chi-squared uses on every traced position (source-plane fits only). + """ + from matplotlib.patches import Circle, Ellipse + + presentation = data.presentation + if not presentation.is_source_plane: + return + + def _ellipses(**kwargs): + drawn = False + for (y, x), covariance in zip(data.traced, data.covariances): + ellipse = ellipse_from(covariance) + if ellipse is None: + continue + semi_major, semi_minor, angle = ellipse + ax.add_patch( + Ellipse( + (x, y), + 2.0 * semi_major, + 2.0 * semi_minor, + angle=angle, + fill=False, + **kwargs, + ) + ) + drawn = True + return drawn + + if presentation.weighting == "magnification": + reference = dict( + edgecolor=REFERENCE_NOISE_COLOR, + linewidth=1.0, + linestyle="--", + alpha=0.9, + zorder=4.8, + ) + if _ellipses(**reference): + ax.plot( + [], + [], + color=REFERENCE_NOISE_COLOR, + linewidth=1.0, + linestyle="--", + alpha=0.9, + label="Propagated noise A·σ (reference)", + ) + for (y, x), radius in zip(data.traced, data.magnification_radii): + if np.isfinite(radius) and radius > 0: + ax.add_patch( + Circle( + (x, y), + radius, + fill=False, + edgecolor=MAGNIFICATION_NOISE_COLOR, + linewidth=1.2, + zorder=5.5, + ) + ) + ax.plot( + [], + [], + color=MAGNIFICATION_NOISE_COLOR, + linewidth=1.2, + label=f"χ² noise σ/|μ| ({_milliarcsec_range(data.magnification_radii)})", + ) + else: + if _ellipses(edgecolor=JACOBIAN_NOISE_COLOR, linewidth=1.2, zorder=5.5): + ax.plot( + [], + [], + color=JACOBIAN_NOISE_COLOR, + linewidth=1.2, + label="χ² noise A·σ (jacobian)", + ) + + +def _add_inset(ax, data: SourcePlaneData, tangential, radial, zoom, bounds): + """ + Adds the context inset: the whole tangential caustic (radial caustics are drawn but not framed, as for a + cluster they can be far larger), the centre and a box marking the main panel's zoom. + """ + from matplotlib.patches import Rectangle + + framed = list(tangential) or list(radial) + if not framed: + return None + + inset = ax.inset_axes(bounds) + _plot_caustics(inset, tangential, radial, linewidth=0.8, label=False) + _plot_centre(inset, data, size=50, label=False) + + points = np.vstack(framed + [data.traced, data.centre[None, :]]) + points = points[np.isfinite(points).all(axis=1)] + lo, hi = points.min(axis=0), points.max(axis=0) + half = 0.55 * float(np.max(hi - lo)) + centre_y, centre_x = 0.5 * (lo + hi) + inset.set_xlim(centre_x - half, centre_x + half) + inset.set_ylim(centre_y - half, centre_y + half) + inset.set_aspect("equal", adjustable="box") + + x_min, x_max, y_min, y_max = zoom + # The zoom is often far smaller than an inset pixel: keep its box visible. + minimum = 0.04 * 2.0 * half + width, height = max(x_max - x_min, minimum), max(y_max - y_min, minimum) + x_mid, y_mid = 0.5 * (x_min + x_max), 0.5 * (y_min + y_max) + inset.add_patch( + Rectangle( + (x_mid - width / 2.0, y_mid - height / 2.0), + width, + height, + fill=False, + edgecolor="k", + linewidth=1.0, + zorder=6, + ) + ) + inset.set_xticks([]) + inset.set_yticks([]) + for spine in inset.spines.values(): + spine.set_edgecolor("0.5") + inset.patch.set_alpha(0.92) + return inset + + +def plot_source_plane_panel( + ax, + data: SourcePlaneData, + title: str, + tangential_caustics=None, + radial_caustics=None, +): + """ + Draws the "Fit Source Plane" panel of a positions fit on ``ax`` (see the module docstring). + + Parameters + ---------- + ax + The matplotlib axes drawn on. + data + The panel's quantities, from :func:`source_plane_data_from`. + title + The panel title, e.g. ``"Fit Source Plane (point_0)"``. + tangential_caustics, radial_caustics + The caustics of the point source's plane, as lists of (N, 2) (y,x) arrays. + """ + from autoarray.plot.utils import apply_extent, apply_labels + + tangential = list(tangential_caustics or []) + radial = list(radial_caustics or []) + + zoom = zoom_extent_from(data) + corners = ranked_corners_from(zoom, _footprint_points(data)) + + has_inset = bool(tangential or radial) + inset_corner, inset_bounds = corners[0] + remaining = corners[1:] if has_inset else corners + legend_corner = remaining[0][0] + + _plot_caustics(ax, tangential, radial) + _plot_footprints(ax, data) + + traced = data.traced + ax.scatter( + traced[:, 1], + traced[:, 0], + s=110, + c=pp.OBSERVED_COLOR, + marker="o", + zorder=4, + label=f"Observed, traced ({pp.sigma_label_from(data.sigma)})", + ) + _plot_centre(ax, data) + + apply_labels(ax, title=title, xlabel="", ylabel="") + apply_extent(ax, zoom) + ax.set_aspect("equal", adjustable="box") + + pp.plot_image_numbers(ax, traced, data.numbers, zoom) + + chi_squared_label = chi_squared_label_from(data) + if chi_squared_label is not None: + # A handle-less entry: the chi-squared the fit uses, in the legend with the other fit-type cues. + ax.plot([], [], linestyle="none", marker="none", label=chi_squared_label) + + ax.legend( + loc=CORNER_LOCATIONS[legend_corner], + fontsize=8, + framealpha=0.85, + markerscale=0.6, + ) + + if has_inset: + _add_inset(ax, data, tangential, radial, zoom, inset_bounds) + + return zoom diff --git a/test_autolens/point/plot/test_positions_panel.py b/test_autolens/point/plot/test_positions_panel.py index 04a933ed8..1698ac71f 100644 --- a/test_autolens/point/plot/test_positions_panel.py +++ b/test_autolens/point/plot/test_positions_panel.py @@ -132,7 +132,8 @@ def test__subplot_fit__title_extent_legend_mass_centres_and_curves(figures): ax = figures[0].axes[0] assert ax.get_title() == "Fit Positions (point_0)" - assert figures[0].axes[1].get_title() == "Fit Fluxes (point_0)" + assert figures[0].axes[1].get_title() == "Fit Source Plane (point_0)" + assert figures[0].axes[2].get_title() == "Fit Fluxes (point_0)" assert ax.get_xlabel() == "" and ax.get_ylabel() == "" assert ax.get_xlim() == pytest.approx((fit.solver.x_min, fit.solver.x_max)) @@ -171,13 +172,14 @@ def fake_lines(tracer, grid): assert "Critical curve" not in _legend_labels(figures[0].axes[0]) -def test__subplot_fit__source_plane_fit_has_no_residual_lines(figures): +def test__subplot_fit__source_plane_fit_draws_no_model_on_image_plane(figures): fit_point_plots.subplot_fit( fit=_fit(fit_positions_cls=al.FitPositionsSource), image_plane_lines=[] ) labels = _legend_labels(figures[0].axes[0]) - assert "Model (source plane)" in labels + assert "Model (source plane)" not in labels + assert "Model" not in labels assert "Residual" not in labels assert "Unmatched model image" not in labels diff --git a/test_autolens/point/plot/test_source_plane_panel.py b/test_autolens/point/plot/test_source_plane_panel.py new file mode 100644 index 000000000..88b359a93 --- /dev/null +++ b/test_autolens/point/plot/test_source_plane_panel.py @@ -0,0 +1,445 @@ +import importlib +import inspect +import pkgutil + +import matplotlib.pyplot as plt +import numpy as np +import pytest +from matplotlib.patches import Circle, Ellipse + +import autolens as al +from autolens.point.fit.positions.abstract import AbstractFitPositions +from autolens.point.plot import fit_point_plots +from autolens.point.plot import source_plane_panel as spp +from autolens.point.plot.positions_panel import image_numbers_from, image_order_from + +SOURCE_CENTRE = (0.07, 0.07) + + +class FitPositionsSourceJacobian(al.FitPositionsSource): + weighting = "jacobian" + + +# Every positions fit class and the presentation its source-plane panel must use. +PRESENTATIONS = [ + (al.FitPositionsImagePair, "image", False, None), + (al.FitPositionsImagePairRepeat, "image", False, None), + (al.FitPositionsImagePairRepeatSolved, "image", True, None), + (al.FitPositionsImagePairAll, "image", False, None), + (al.FitPositionsImagePairAllSolved, "image", True, None), + (al.FitPositionsSource, "source", False, "magnification"), + (FitPositionsSourceJacobian, "source", False, "jacobian"), + (al.FitPositionsSourceSolved, "source", True, "jacobian"), +] + + +def _tracer(solved: bool, einstein_radius: float = 1.6): + point = al.ps.PointSolved() if solved else al.ps.Point(centre=SOURCE_CENTRE) + return al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), + einstein_radius=einstein_radius, + ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0), + ), + ), + al.Galaxy(redshift=1.0, point_0=point), + ] + ) + + +@pytest.fixture(name="solver", scope="module") +def make_solver(): + return al.PointSolver.for_grid( + grid=al.Grid2D.uniform(shape_native=(11, 11), pixel_scales=0.5), + pixel_scale_precision=0.01, + magnification_threshold=0.1, + ) + + +@pytest.fixture(name="dataset", scope="module") +def make_dataset(solver): + """The images of the true model, so every fit class (incl. the solved ones) has a well-posed fit.""" + positions = solver.solve( + tracer=_tracer(solved=False), source_plane_coordinate=SOURCE_CENTRE + ) + positions = np.asarray(positions.array) + positions = positions[np.isfinite(positions).all(axis=1)] + return al.PointDataset( + name="point_0", + positions=al.Grid2DIrregular(positions + 0.002), + positions_noise_map=0.005, + ) + + +def _fit(cls, dataset, solver, einstein_radius=1.62): + return al.FitPointDataset( + dataset=dataset, + tracer=_tracer( + solved="Solved" in cls.__name__, einstein_radius=einstein_radius + ), + solver=solver, + fit_positions_cls=cls, + ) + + +def _positions_fit_classes(): + import autolens.point.fit.positions as package + + classes = set() + for module_info in pkgutil.walk_packages(package.__path__, package.__name__ + "."): + module = importlib.import_module(module_info.name) + for _, cls in inspect.getmembers(module, inspect.isclass): + if ( + issubclass(cls, AbstractFitPositions) + and cls.__module__ == module.__name__ + and not cls.__name__.startswith("Abstract") + ): + classes.add(cls) + return classes + + +def test__presentation_table__covers_every_positions_fit_class(): + classes = _positions_fit_classes() + + assert len(classes) == 7 + assert classes == set(spp.presentation_table()) + + +@pytest.mark.parametrize("cls, plane, solved, weighting", PRESENTATIONS) +def test__presentation_from__paired_to_fit_class( + cls, plane, solved, weighting, dataset, solver +): + presentation = spp.presentation_from(_fit(cls, dataset, solver).positions) + + assert presentation.chi_squared_plane == plane + assert presentation.solved_centre is solved + assert presentation.weighting == weighting + + assert presentation.centre_marker == ("D" if solved else "*") + assert presentation.centre_label == ( + "Solved source centre (β*)" if solved else "Model source centre" + ) + + +def _draw(fit_positions, caustics=True): + data = spp.source_plane_data_from(fit_positions) + tangential = [ + np.array([[0.0, 0.2], [0.2, 0.0], [0.0, -0.2], [-0.2, 0.0], [0.0, 0.2]]) + ] + fig, ax = plt.subplots(figsize=(6, 6)) + zoom = spp.plot_source_plane_panel( + ax, + data, + title="Fit Source Plane (point_0)", + tangential_caustics=tangential if caustics else [], + ) + return fig, ax, data, zoom + + +def _texts(ax): + return [text.get_text() for text in ax.texts] + + +def _legend_labels(ax): + return [text.get_text() for text in ax.get_legend().get_texts()] + + +@pytest.mark.parametrize("cls, plane, solved, weighting", PRESENTATIONS) +def test__plot_source_plane_panel__draws_what_the_fit_class_uses( + cls, plane, solved, weighting, dataset, solver +): + fit_positions = _fit(cls, dataset, solver).positions + fig, ax, data, zoom = _draw(fit_positions) + + circles = [p for p in ax.patches if isinstance(p, Circle)] + ellipses = [ + p for p in ax.patches if isinstance(p, Ellipse) and not isinstance(p, Circle) + ] + solid_ellipses = [e for e in ellipses if e.get_linestyle() in ("-", "solid")] + dashed_ellipses = [e for e in ellipses if e.get_linestyle() in ("--", "dashed")] + labels = _legend_labels(ax) + chi_squared_labels = [label for label in labels if label.startswith("χ² =")] + n = len(data.traced) + + assert ax.get_title() == "Fit Source Plane (point_0)" + # No subtitle or other free-floating fit-type text: only the image numbers are annotations. + assert all(text.get_text().isdigit() for text in ax.texts) + assert not any(label.startswith("χ²") for label in _texts(ax)) + assert data.presentation.centre_label in labels + assert "Caustic" in labels + assert len(ax.child_axes) == 1 # the caustic inset + + if plane == "image": + assert circles == [] and ellipses == [] + assert chi_squared_labels == [] + assert not any("noise" in label for label in labels) + elif weighting == "magnification": + assert len(circles) == n + assert len(dashed_ellipses) == n and solid_ellipses == [] + assert any(label.startswith("χ² noise σ/|μ|") for label in labels) + assert "Propagated noise A·σ (reference)" in labels + assert chi_squared_labels == [ + f"χ² = {spp._format_chi_squared(data.chi_squared)}" + ] + else: + assert circles == [] + assert len(solid_ellipses) == n and dashed_ellipses == [] + assert "χ² noise A·σ (jacobian)" in labels + assert "Propagated noise A·σ (reference)" not in labels + assert chi_squared_labels == [ + f"χ² = {spp._format_chi_squared(data.chi_squared)}" + ] + + # The zoom contains every traced position and the centre. + x_min, x_max, y_min, y_max = zoom + points = np.vstack([data.traced, data.centre[None, :]]) + assert np.all((points[:, 1] > x_min) & (points[:, 1] < x_max)) + assert np.all((points[:, 0] > y_min) & (points[:, 0] < y_max)) + + plt.close(fig) + + +def test__plot_source_plane_panel__no_inset_without_caustics(dataset, solver): + fig, ax, _, _ = _draw( + _fit(al.FitPositionsImagePair, dataset, solver).positions, caustics=False + ) + + assert ax.child_axes == [] + plt.close(fig) + + +def test__source_plane_data__matches_the_fits_own_chi_squared(dataset, solver): + magnification = _fit(al.FitPositionsSource, dataset, solver).positions + data = spp.source_plane_data_from(magnification) + + # Traced exactly as the fit traces them. + assert data.traced == pytest.approx(np.asarray(magnification.model_data.array)) + + # The σ/|μ| circle is the noise of the magnification chi-squared. + residuals = np.linalg.norm(data.traced - data.centre, axis=1) + assert data.chi_squared == pytest.approx( + np.sum(residuals**2 / data.magnification_radii**2) + ) + + # The A·σ ellipse is the covariance of the jacobian chi-squared. + jacobian = _fit(FitPositionsSourceJacobian, dataset, solver).positions + jacobian_data = spp.source_plane_data_from(jacobian) + assert spp.chi_squared_from( + jacobian_data.traced - jacobian_data.centre, jacobian_data.covariances + ) == pytest.approx(jacobian_data.chi_squared, rel=1e-6) + + # The solved centre is the fit's own β*. + solved = _fit(al.FitPositionsSourceSolved, dataset, solver).positions + assert spp.source_plane_data_from(solved).centre == pytest.approx( + np.asarray(solved.source_plane_coordinate) + ) + + +def test__image_plane_fit__no_chi_squared_computed(dataset, solver): + data = spp.source_plane_data_from( + _fit(al.FitPositionsImagePair, dataset, solver).positions + ) + + assert data.chi_squared is None + assert spp.chi_squared_label_from(data) is None + + +def test__ellipse_from__axes_and_angle(): + semi_major, semi_minor, angle = spp.ellipse_from(np.array([[1.0, 0.0], [0.0, 4.0]])) + + # (y,x) covariance: variance 4 along x. + assert semi_major == pytest.approx(2.0) + assert semi_minor == pytest.approx(1.0) + assert abs(np.cos(np.radians(angle))) == pytest.approx(1.0) + + assert spp.ellipse_from(np.full((2, 2), np.nan)) is None + + +def test__image_numbers_from__point_solver_order(): + positions = np.array([[-1.0, -1.0], [1.0, 0.1], [0.1, 1.0]]) + + order = image_order_from(positions) + numbers = image_numbers_from(positions) + + assert list(order) == [1, 2, 0] + assert list(numbers) == [3, 1, 2] + + +@pytest.fixture(name="figures") +def make_figures(monkeypatch): + figures = [] + monkeypatch.setattr( + fit_point_plots, "save_figure", lambda fig, **kwargs: figures.append(fig) + ) + return figures + + +def _number_texts(ax): + return sorted( + (t.get_text(), tuple(np.round(t.xy, 8))) + for t in ax.texts + if t.get_text().isdigit() + ) + + +@pytest.mark.parametrize("fluxes", [None, [1.0, 2.0, 3.0, 4.0]]) +def test__subplot_fit__layout_and_consistent_image_numbers( + figures, dataset, solver, fluxes +): + dataset = al.PointDataset( + name="point_0", + positions=dataset.positions, + positions_noise_map=dataset.positions_noise_map, + fluxes=None if fluxes is None else fluxes[: len(dataset.positions)], + fluxes_noise_map=None if fluxes is None else [0.1] * len(dataset.positions), + ) + fit = _fit(al.FitPositionsSource, dataset, solver) + + fit_point_plots.subplot_fit(fit=fit, image_plane_lines=[], source_plane_lines=[]) + + axes = figures[0].axes + titles = [ax.get_title() for ax in axes] + expected = ["Fit Positions (point_0)", "Fit Source Plane (point_0)"] + if fluxes is not None: + expected.append("Fit Fluxes (point_0)") + assert titles == expected + + # The source-plane fit's traced positions are not drawn on the image plane. + image_labels = _legend_labels(axes[0]) + assert "Model (source plane)" not in image_labels + + # Image k is the same observed position on both panels. + numbers = image_numbers_from(dataset.positions) + traced = spp.traced_positions_from(fit.positions) + observed = np.asarray(dataset.positions.array) + expected_image = sorted( + (str(n), tuple(np.round(p[::-1], 8))) for n, p in zip(numbers, observed) + ) + expected_source = sorted( + (str(n), tuple(np.round(p[::-1], 8))) for n, p in zip(numbers, traced) + ) + assert _number_texts(axes[0]) == expected_image + assert _number_texts(axes[1]) == expected_source + + +def _three_plane_fit(solver, dataset): + tracer = al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.4), + ), + al.Galaxy( + redshift=1.0, + mass=al.mp.Isothermal(centre=(0.3, 0.3), einstein_radius=0.2), + point_0=al.ps.Point(centre=SOURCE_CENTRE), + ), + al.Galaxy(redshift=2.0, mass=al.mp.Isothermal(einstein_radius=0.1)), + ] + ) + return al.FitPointDataset( + dataset=dataset, + tracer=tracer, + solver=solver, + fit_positions_cls=al.FitPositionsSource, + ) + + +def test__source_plane_caustics_from__point_plane_in_multi_plane_tracer( + monkeypatch, solver, dataset +): + fit = _three_plane_fit(solver, dataset) + assert fit.positions.plane_index == 1 + + calls = [] + + def fake_caustics(tracer, plane_index, grid): + calls.append(plane_index) + return [np.array([[0.0, 0.0], [0.1, 0.1]])], [] + + monkeypatch.setattr(spp, "point_plane_caustics_from", fake_caustics) + + # The pre-computed lines are the last plane's: not reused for a point in plane 1. + tangential, _ = fit_point_plots.source_plane_caustics_from( + fit=fit, + lines_given=True, + source_plane_lines=[np.array([[5.0, 5.0], [6.0, 6.0]])], + source_plane_line_colors=["white"], + ) + assert calls == [1] + assert tangential[0] == pytest.approx(np.array([[0.0, 0.0], [0.1, 0.1]])) + + # A two-plane point is in the last plane: the pre-computed lines are reused. + two_plane = _fit(al.FitPositionsSource, dataset, solver) + tangential, radial = fit_point_plots.source_plane_caustics_from( + fit=two_plane, + lines_given=True, + source_plane_lines=[ + np.array([[5.0, 5.0], [6.0, 6.0]]), + np.array([[1.0, 1.0], [2.0, 2.0]]), + ], + source_plane_line_colors=["white", "yellow"], + ) + assert calls == [1] + assert len(tangential) == 1 and len(radial) == 1 + + +def test__point_plane_caustics_from__own_plane_differs_from_last_plane(solver, dataset): + fit = _three_plane_fit(solver, dataset) + grid = al.Grid2D.uniform(shape_native=(80, 80), pixel_scales=0.05) + + own, _ = spp.point_plane_caustics_from(fit.tracer, plane_index=1, grid=grid) + last, _ = spp.point_plane_caustics_from(fit.tracer, plane_index=2, grid=grid) + + assert len(own) >= 1 and len(last) >= 1 + assert np.vstack(own).mean(axis=0) != pytest.approx(np.vstack(last).mean(axis=0)) + + +@pytest.mark.parametrize( + "cls, label", + [ + (al.FitPositionsImagePair, "Model"), + (al.FitPositionsImagePairAllSolved, "Model (from solved β*)"), + (al.FitPositionsImagePairRepeatSolved, "Model (from solved β*)"), + ], +) +def test__fit_positions_panel__model_label_names_a_solved_centre( + cls, label, dataset, solver +): + from autolens.point.plot.positions_panel import model_label_from + + assert model_label_from(_fit(cls, dataset, solver).positions) == label + + +@pytest.mark.parametrize( + "cls, linestyle, label", + [ + (al.FitPositionsImagePair, "-", "Residual"), + (al.FitPositionsImagePairRepeat, "-", "Residual"), + (al.FitPositionsImagePairRepeatSolved, "-", "Residual"), + (al.FitPositionsImagePairAll, "--", "Residual (nearest, soft pairing)"), + (al.FitPositionsImagePairAllSolved, "--", "Residual (nearest, soft pairing)"), + ], +) +def test__fit_positions_panel__residual_lines_styled_by_pairing( + cls, linestyle, label, dataset, solver +): + from autolens.point.plot.positions_panel import ( + plot_fit_model_positions, + residual_style_from, + ) + + fit_positions = _fit(cls, dataset, solver).positions + assert residual_style_from(fit_positions) == (linestyle, label) + + fig, ax = plt.subplots() + plot_fit_model_positions(ax, fit_positions) + residuals = [line for line in ax.lines if line.get_color() == "r"] + assert residuals + expected = {"-": ("-", "solid"), "--": ("--", "dashed")}[linestyle] + assert all(line.get_linestyle() in expected for line in residuals) + assert [line.get_label() for line in residuals][0] == label + plt.close(fig)