feature/vectorised triangles

autolens/mock.py

@@ -1,8 +1,19 @@
-from autofit.mock import *
-from autoarray.mock import *
-from autogalaxy.mock import *
-
-from autolens.imaging.mock.mock_fit_imaging import MockFitImaging
-from autolens.lens.mock.mock_tracer import MockTracer
-from autolens.lens.mock.mock_tracer import MockTracerPoint
-from autolens.point.mock.mock_point_solver import MockPointSolver
+import numpy as np
+
+from autofit.mock import *  # noqa
+from autoarray.mock import *  # noqa
+from autogalaxy.mock import *  # noqa
+from autolens import Tracer
+
+from autolens.imaging.mock.mock_fit_imaging import MockFitImaging  # noqa
+from autolens.lens.mock.mock_tracer import MockTracer  # noqa
+from autolens.lens.mock.mock_tracer import MockTracerPoint  # noqa
+from autolens.point.mock.mock_point_solver import MockPointSolver  # noqa
+
+
+class NullTracer(Tracer):
+    def __init__(self):
+        super().__init__([])
+
+    def deflections_yx_2d_from(self, grid):
+        return np.zeros_like(grid.array)

autolens/point/analysis.py

@@ -56,7 +56,7 @@ def log_likelihood_function(self, instance):
         """
         lens = instance.lens
 
-        solver = TriangleSolver(
+        solver = TriangleSolver.for_grid(
             lensing_obj=lens,
             grid=self.grid,
             pixel_scale_precision=self.pixel_scale_precision,
@@ -147,7 +147,7 @@ def _log_likelihood_for_coordinates(
                 "The number of predicted coordinates must be equal to the number of observed coordinates."
             )
 
-        predicted_coordinates = set(predicted_coordinates)
+        predicted_coordinates = set(map(tuple, predicted_coordinates))
         observed_coordinates = set(self.observed_coordinates)
 
         log_likelihood = 0.0

autolens/point/triangles/triangle_solver.py

@@ -1,24 +1,56 @@
 import logging
 import math
-from typing import Tuple, List
+from dataclasses import dataclass
+from typing import Tuple, List, Iterator, Type
 
 from autoarray import Grid2D, Grid2DIrregular
-from autoarray.structures.triangles.subsample_triangles import SubsampleTriangles
-from autoarray.structures.triangles.triangles import Triangles
+from autoarray.structures.triangles import array
+from autoarray.structures.triangles.abstract import AbstractTriangles
 from autoarray.type import Grid2DLike
 from autogalaxy import OperateDeflections
 
-
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class Step:
+    """
+    A step in the triangle solver algorithm.
+
+    Attributes
+    ----------
+    number
+        The number of the step.
+    initial_triangles
+        The triangles at the start of the step.
+    filtered_triangles
+        The triangles trace to triangles that contain the source plane coordinate.
+    neighbourhood
+        The neighbourhood of the filtered triangles.
+    up_sampled
+        The neighbourhood up-sampled to increase the resolution.
+    """
+
+    number: int
+    initial_triangles: AbstractTriangles
+    filtered_triangles: AbstractTriangles
+    neighbourhood: AbstractTriangles
+    up_sampled: AbstractTriangles
+
+
 class TriangleSolver:
+    # noinspection PyPep8Naming
     def __init__(
         self,
         lensing_obj: OperateDeflections,
-        grid: Grid2D,
+        scale: float,
+        y_min: float,
+        y_max: float,
+        x_min: float,
+        x_max: float,
         pixel_scale_precision: float,
         magnification_threshold=0.1,
+        ArrayTriangles: Type[AbstractTriangles] = array.ArrayTriangles,
     ):
         """
         Determine the image plane coordinates that are traced to be a source plane coordinate.
@@ -31,22 +63,59 @@ def __init__(
         ----------
         lensing_obj
             A tracer describing the lensing system.
-        grid
-            The grid of image plane coordinates.
         pixel_scale_precision
             The target pixel scale of the image grid.
+        ArrayTriangles
+            The class to use for the triangles.
         """
         self.lensing_obj = lensing_obj
-        self.grid = grid
+        self.scale = scale
+        self.y_min = y_min
+        self.y_max = y_max
+        self.x_min = x_min
+        self.x_max = x_max
         self.pixel_scale_precision = pixel_scale_precision
         self.magnification_threshold = magnification_threshold
+        self.ArrayTriangles = ArrayTriangles
+
+    # noinspection PyPep8Naming
+    @classmethod
+    def for_grid(
+        cls,
+        lensing_obj: OperateDeflections,
+        grid: Grid2D,
+        pixel_scale_precision: float,
+        magnification_threshold=0.1,
+        ArrayTriangles: Type[AbstractTriangles] = array.ArrayTriangles,
+    ):
+        scale = grid.pixel_scale
+
+        y = grid[:, 0]
+        x = grid[:, 1]
+
+        y_min = y.min()
+        y_max = y.max()
+        x_min = x.min()
+        x_max = x.max()
+
+        return cls(
+            lensing_obj=lensing_obj,
+            scale=scale,
+            y_min=y_min,
+            y_max=y_max,
+            x_min=x_min,
+            x_max=x_max,
+            pixel_scale_precision=pixel_scale_precision,
+            magnification_threshold=magnification_threshold,
+            ArrayTriangles=ArrayTriangles,
+        )
 
     @property
     def n_steps(self) -> int:
         """
         How many times should triangles be subdivided?
         """
-        return math.ceil(math.log2(self.grid.pixel_scale / self.pixel_scale_precision))
+        return math.ceil(math.log2(self.scale / self.pixel_scale_precision))
 
     def _source_plane_grid(self, grid: Grid2DLike) -> Grid2DLike:
         """
@@ -87,31 +156,18 @@ def solve(
         -------
         A list of image plane coordinates that are traced to the source plane coordinate.
         """
-        triangles = Triangles.for_grid(grid=self.grid)
-
         if self.n_steps == 0:
             raise ValueError(
                 "The target pixel scale is too large to subdivide the triangles."
             )
 
-        kept_triangles = []
-
-        for _ in range(self.n_steps):
-            kept_triangles = self._filter_triangles(
-                triangles=triangles,
-                source_plane_coordinate=source_plane_coordinate,
-            )
-            with_neighbourhood = {
-                triangle
-                for kept_triangle in kept_triangles
-                for triangle in kept_triangle.neighbourhood
-            }
-            triangles = SubsampleTriangles(parent_triangles=list(with_neighbourhood))
+        steps = list(self.steps(source_plane_coordinate=source_plane_coordinate))
+        final_step = steps[-1]
+        kept_triangles = final_step.filtered_triangles
 
-        means = [triangle.mean for triangle in kept_triangles]
-        filtered_means = self._filter_low_magnification(points=means)
+        filtered_means = self._filter_low_magnification(points=kept_triangles.means)
 
-        difference = len(means) - len(filtered_means)
+        difference = len(kept_triangles.means) - len(filtered_means)
         if difference > 0:
             logger.debug(
                 f"Filtered one multiple-image with magnification below threshold."
@@ -144,15 +200,15 @@ def _filter_low_magnification(
                 points,
                 self.lensing_obj.magnification_2d_via_hessian_from(
                     grid=Grid2DIrregular(points),
-                    buffer=self.grid.pixel_scale,
+                    buffer=self.scale,
                 ),
             )
             if abs(magnification) > self.magnification_threshold
         ]
 
     def _filter_triangles(
         self,
-        triangles: Triangles,
+        triangles: AbstractTriangles,
         source_plane_coordinate: Tuple[float, float],
     ):
         """
@@ -169,16 +225,51 @@ def _filter_triangles(
         -------
         The triangles that contain the source plane coordinate.
         """
-        source_plane_grid = self._source_plane_grid(grid=triangles.grid_2d)
-
-        kept_triangles = []
-        for image_triangle, source_triangle in zip(
-            triangles.triangles,
-            triangles.with_updated_grid(source_plane_grid),
-        ):
-            if source_triangle.contains(
-                point=source_plane_coordinate,
-            ):
-                kept_triangles.append(image_triangle)
-
-        return kept_triangles
+        source_plane_grid = self._source_plane_grid(
+            grid=Grid2DIrregular(triangles.vertices)
+        )
+        source_triangles = triangles.with_vertices(source_plane_grid.array)
+        indexes = source_triangles.containing_indices(point=source_plane_coordinate)
+        return triangles.for_indexes(indexes=indexes)
+
+    def steps(
+        self,
+        source_plane_coordinate: Tuple[float, float],
+    ) -> Iterator[Step]:
+        """
+        Iterate over the steps of the triangle solver algorithm.
+
+        Parameters
+        ----------
+        source_plane_coordinate
+            The source plane coordinate to trace to the image plane.
+
+        Returns
+        -------
+        An iterator over the steps of the triangle solver algorithm.
+        """
+        initial_triangles = self.ArrayTriangles.for_limits_and_scale(
+            y_min=self.y_min,
+            y_max=self.y_max,
+            x_min=self.x_min,
+            x_max=self.x_max,
+            scale=self.scale,
+        )
+
+        for number in range(self.n_steps):
+            kept_triangles = self._filter_triangles(
+                initial_triangles,
+                source_plane_coordinate,
+            )
+            neighbourhood = kept_triangles.neighborhood()
+            up_sampled = neighbourhood.up_sample()
+
+            yield Step(
+                number=number,
+                initial_triangles=initial_triangles,
+                filtered_triangles=kept_triangles,
+                neighbourhood=neighbourhood,
+                up_sampled=up_sampled,
+            )
+
+            initial_triangles = up_sampled

autolens/point/triangles/visualise.py

@@ -0,0 +1,23 @@
+from .triangle_solver import Step
+from matplotlib import pyplot as plt
+import numpy as np
+
+
+def add_triangles(triangles, color):
+    for triangle in triangles:
+        triangle = np.append(triangle, [triangle[0]], axis=0)
+        plt.plot(triangle[:, 0], triangle[:, 1], "o-", color=color)
+
+
+def visualise(step: Step):
+    plt.figure(figsize=(8, 8))
+    add_triangles(step.initial_triangles, color="black")
+    add_triangles(step.up_sampled, color="green")
+    add_triangles(step.neighbourhood, color="red")
+    add_triangles(step.filtered_triangles, color="blue")
+
+    plt.xlabel("X")
+    plt.ylabel("Y")
+    plt.title(f"Step {step.number}")
+    plt.gca().set_aspect("equal", adjustable="box")
+    plt.show()

test_autolens/point/triangles/conftest.py

@@ -5,6 +5,26 @@
 @pytest.fixture
 def grid():
     return al.Grid2D.uniform(
-        shape_native=(100, 100),
-        pixel_scales=0.05,
+        shape_native=(10, 10),
+        pixel_scales=1.0,
     )
+
+
+@pytest.fixture
+def tracer():
+    isothermal_mass_profile = al.mp.Isothermal(
+        centre=(0.0, 0.0),
+        einstein_radius=1.6,
+        ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),
+    )
+
+    lens_galaxy = al.Galaxy(
+        redshift=0.5,
+        mass=isothermal_mass_profile,
+    )
+
+    point_source = al.ps.PointSourceChi(centre=(0.07, 0.07))
+
+    source_galaxy = al.Galaxy(redshift=1.0, point_0=point_source)
+
+    return al.Tracer(galaxies=[lens_galaxy, source_galaxy])

test_autolens/point/triangles/test_regressions.py

@@ -66,7 +66,7 @@ def test_missing_multiple_image(grid):
 
     tracer = al.Tracer(galaxies=[instance.lens_galaxy, instance.source_galaxy])
 
-    solver = TriangleSolver(
+    solver = TriangleSolver.for_grid(
         grid=grid,
         lensing_obj=tracer,
         pixel_scale_precision=0.001,