feat: source-plane fit panel paired to each positions Fit class

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 (<name>)``.
       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 (<name>)``: 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,48 +71,60 @@ 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
         The solver whose grid sets the positions panel extent. Defaults to
         ``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,

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(