fix: point-source and weak-lensing critical-curve grids (clipped / transposed)

autolens/lens/plot/critical_curves.py

@@ -8,6 +8,7 @@
 
 import logging
 
+import autogalaxy as ag
 from autogalaxy.util.plot_utils import _critical_curves_from, _caustics_from
 from autoarray.plot.utils import numpy_lines as _to_lines
 
@@ -64,3 +65,47 @@ def _compute_critical_curve_lines(tracer, grid):
             exc_info=True,
         )
         return None, None, None, None
+
+
+def critical_curve_grid_from(
+    dataset, shape_native=(100, 100), scale: float = 2.0, buffer: float = 0.1
+) -> "ag.Grid2D":
+    """
+    Returns the grid on which the critical curves and caustics of a dataset's lens model are computed.
+
+    The grid is a square centred on the bounding box of the dataset's positions (``dataset.extent_from``), whose
+    half-width is ``scale`` times the half-span of that box (plus ``buffer``). For a point dataset critical curves
+    often extend beyond the multiple images, so a grid covering only their bounding box clips them into
+    disconnected arcs and ``scale=2.0`` is used. For a weak lensing dataset the source galaxies already span the
+    field, so ``scale=1.0`` covers it.
+
+    ``dataset.extent_from`` returns ``[y_min, y_max, x_min, x_max]`` whereas ``Grid2D.from_extent`` takes the
+    extent x first, which this function handles.
+
+    Parameters
+    ----------
+    dataset
+        A dataset with an ``extent_from`` method (e.g. a ``PointDataset`` or ``WeakDataset``) whose positions set
+        the centre and size of the grid.
+    shape_native
+        The 2D shape of the grid.
+    scale
+        The factor by which the grid's half-width exceeds the half-span of the positions.
+    buffer
+        A buffer (arc-seconds) added to the half-span of the positions.
+    """
+    y_min, y_max, x_min, x_max = dataset.extent_from(buffer=0.0)
+
+    y_centre, x_centre = (y_min + y_max) / 2.0, (x_min + x_max) / 2.0
+    half_width = scale * (max(y_max - y_min, x_max - x_min) / 2.0 + buffer)
+
+    # `Grid2D.from_extent` takes the extent x first: (x0, x1, y0, y1).
+    return ag.Grid2D.from_extent(
+        extent=(
+            x_centre - half_width,
+            x_centre + half_width,
+            y_centre - half_width,
+            y_centre + half_width,
+        ),
+        shape_native=shape_native,
+    )

autolens/point/model/plotter.py

@@ -1,7 +1,7 @@
-import autogalaxy as ag
 
 from autolens.analysis.plotter import Plotter
 from autolens.imaging.plot.fit_imaging_plots import _compute_critical_curve_lines
+from autolens.lens.plot.critical_curves import critical_curve_grid_from
 
 from autolens.point.fit.dataset import FitPointDataset
 from autolens.point.plot.fit_point_plots import subplot_fit as subplot_fit_point
@@ -115,9 +115,7 @@ def should_plot(name):
 
         # Use pre-computed critical curves if provided, otherwise compute once here.
         if image_plane_lines is None and source_plane_lines is None:
-            grid = ag.Grid2D.from_extent(
-                extent=fit.dataset.extent_from(), shape_native=(100, 100)
-            )
+            grid = critical_curve_grid_from(dataset=fit.dataset)
             ip_lines, ip_colors, sp_lines, sp_colors = _compute_critical_curve_lines(
                 fit.tracer, grid
             )

autolens/point/model/visualizer.py

@@ -1,7 +1,7 @@
 import autofit as af
-import autogalaxy as ag
 
 from autolens.point.model.plotter import PlotterPoint
+from autolens.lens.plot.critical_curves import critical_curve_grid_from
 from autolens.imaging.plot.fit_imaging_plots import _compute_critical_curve_lines
 
 
@@ -76,9 +76,7 @@ def visualize(
 
         tracer = fit.tracer
 
-        grid = ag.Grid2D.from_extent(
-            extent=fit.dataset.extent_from(), shape_native=(100, 100)
-        )
+        grid = critical_curve_grid_from(dataset=fit.dataset)
 
         # Compute critical curves and caustics once for all plot functions.
         ip_lines, ip_colors, sp_lines, sp_colors = _compute_critical_curve_lines(

autolens/weak/model/visualizer.py

@@ -1,6 +1,6 @@
 import autofit as af
-import autogalaxy as ag
 
+from autolens.lens.plot.critical_curves import critical_curve_grid_from
 from autolens.weak.model.plotter import PlotterWeak
 
 
@@ -77,9 +77,8 @@ def visualize(
 
         tracer = fit.tracer
 
-        grid = ag.Grid2D.from_extent(
-            extent=fit.dataset.extent_from(), shape_native=(100, 100)
-        )
+        # The source galaxies span the field, so the grid covers their (square) extent.
+        grid = critical_curve_grid_from(dataset=fit.dataset, scale=1.0)
 
         plotter.tracer(
             tracer=tracer,

test_autolens/point/model/test_plotter_point.py

@@ -147,3 +147,59 @@ def test__visualizer__no_image__no_image_with_points(
     assert str(plot_path / "fit.png") in plot_patch.paths
     assert str(plot_path / "image_with_points.png") not in plot_patch.paths
     assert str(plot_path / "fit_image_with_points.png") not in plot_patch.paths
+
+
+def test__critical_curve_grid_from__contains_curve_beyond_multiple_images():
+    import numpy as np
+
+    import autolens as al
+    from autolens.lens.plot.critical_curves import _compute_critical_curve_lines
+    from autolens.lens.plot.critical_curves import critical_curve_grid_from
+
+    # Multiple images inside a 1" box, an isothermal tangential critical curve of radius 1.6": a grid covering
+    # only the images' bounding box clips the curve into arcs.
+    dataset = al.PointDataset(
+        name="point_0",
+        positions=al.Grid2DIrregular([(1.0, 0.5), (0.5, 1.0), (-0.5, 0.0), (0.0, -0.5)]),
+        positions_noise_map=al.ArrayIrregular([0.05, 0.05, 0.05, 0.05]),
+    )
+
+    tracer = al.Tracer(
+        galaxies=[
+            al.Galaxy(redshift=0.5, mass=al.mp.Isothermal(einstein_radius=1.6)),
+            al.Galaxy(redshift=1.0),
+        ]
+    )
+
+    grid = critical_curve_grid_from(dataset=dataset)
+
+    image_plane_lines, _, _, _ = _compute_critical_curve_lines(tracer, grid)
+
+    assert len(image_plane_lines) == 1
+
+    line = np.asarray(image_plane_lines[0])
+
+    assert np.allclose(line[0], line[-1], atol=0.05)
+    assert np.max(np.abs(line)) == pytest.approx(1.6, abs=0.05)
+
+
+def test__critical_curve_grid_from__extent_is_x_first():
+    import numpy as np
+
+    import autolens as al
+    from autolens.lens.plot.critical_curves import critical_curve_grid_from
+
+    # Images spread far more in y than x: the grid must be centred on them in (y,x).
+    dataset = al.PointDataset(
+        name="point_0",
+        positions=al.Grid2DIrregular([(3.0, 1.0), (5.0, 1.0)]),
+        positions_noise_map=al.ArrayIrregular([0.05, 0.05]),
+    )
+
+    grid = np.asarray(critical_curve_grid_from(dataset=dataset).array)
+
+    assert np.mean(grid[:, 0]) == pytest.approx(4.0, abs=0.05)
+    assert np.mean(grid[:, 1]) == pytest.approx(1.0, abs=0.05)
+    assert grid[:, 0].max() - grid[:, 0].min() == pytest.approx(
+        grid[:, 1].max() - grid[:, 1].min(), abs=0.05
+    )

test_autolens/weak/test_dataset.py

@@ -99,3 +99,21 @@ def test__weak_dataset__json_round_trip(tmp_path):
     assert np.allclose(loaded.shear_yx.array, dataset.shear_yx.array)
     assert np.allclose(loaded.shear_yx.grid.array, dataset.shear_yx.grid.array)
     assert np.allclose(loaded.noise_map.array, dataset.noise_map.array)
+
+
+def test__critical_curve_grid_from__covers_field_in_yx_order():
+    from autolens.lens.plot.critical_curves import critical_curve_grid_from
+
+    # A field far longer in y than x, offset from the origin: the old `from_extent(extent_from())` transposed it.
+    shear_yx = _make_shear_yx(
+        positions=[(-1.0, 1.0), (9.0, 1.0), (4.0, 3.0)],
+        values=[(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
+    )
+    dataset = al.WeakDataset(shear_yx=shear_yx, noise_map=0.3)
+
+    grid = np.asarray(critical_curve_grid_from(dataset=dataset, scale=1.0).array)
+
+    assert np.mean(grid[:, 0]) == pytest.approx(4.0, abs=0.1)
+    assert np.mean(grid[:, 1]) == pytest.approx(2.0, abs=0.1)
+    assert grid[:, 0].min() < -1.0 and grid[:, 0].max() > 9.0
+    assert grid[:, 1].min() < 1.0 and grid[:, 1].max() > 3.0