Merge pull request #78 from Jammy2211/feature/image_cluster_improve

Feature/image cluster improve

autoarray/fixtures.py

@@ -399,9 +399,6 @@ def make_voronoi_mesh_grid_9():
 
     return aa.Mesh2DVoronoi(
         values=grid_9,
-        nearest_pixelization_index_for_slim_index=np.zeros(
-            shape=make_grid_2d_7x7().shape_slim, dtype="int"
-        ),
     )
 
 

autoarray/inversion/mock/mock_mesh.py

@@ -36,7 +36,11 @@ def mapper_grids_from(
         )
 
     def image_plane_mesh_grid_from(
-        self, image_plane_data_grid, adapt_data, settings=None
+        self,
+        image_plane_data_grid,
+        adapt_data,
+        settings=None,
+        noise_map: np.ndarray = None,
     ):
         if adapt_data is not None and self.image_plane_mesh_grid is not None:
             return adapt_data * self.image_plane_mesh_grid

autoarray/inversion/mock/mock_pixelization.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from autoarray.inversion.pixelization.pixelization import Pixelization
 
 
@@ -24,7 +26,11 @@ def mapper_grids_from(
         return self.mapper
 
     def image_plane_mesh_grid_from(
-        self, image_plane_data_grid, adapt_data, settings=None
+        self,
+        image_plane_data_grid,
+        adapt_data,
+        settings=None,
+        noise_map: np.ndarray = None,
     ):
         if adapt_data is not None and self.image_plane_mesh_grid is not None:
             return adapt_data * self.image_plane_mesh_grid

autoarray/inversion/pixelization/mappers/mapper_util.py

@@ -172,10 +172,47 @@ def pix_indexes_for_sub_slim_index_delaunay_from(
     return pix_indexes_for_sub_slim_index, pix_indexes_for_sub_slim_index_sizes
 
 
+@numba_util.jit()
+def nearest_pixelization_index_for_slim_index_from(grid, mesh_grid):
+    """
+    Uses a nearest neighbor search to determine for each data pixel its nearest pixelization pixel.
+
+    This is used to speed up the `pix_indexes_for_sub_slim_index_voronoi_from` function, which otherwise would
+    have to loop over every pixelization pixel to determine the nearest pixelization pixel to each data pixel.
+
+    This is only used for a regular `Voronoi` mesh, not a `Delaunay` or `VoronoiNN`, and therefore has limited
+    use in general given the `VoronoiNN` is a superior mesh because it uses natural neighbor interpolation.
+
+
+    Parameters
+    ----------
+    grid
+        The grid of (y,x) scaled coordinates at the centre of every unmasked pixel, which has been traced to
+        to an irgrid via lens.
+    mesh_grid
+        The (y,x) centre of every Voronoi pixel in arc-seconds.
+
+    Returns
+    -------
+    A 1D array of length (total_unmasked_pixels) where each entry corresponds to the index of the nearest
+    pixelization pixel to each data pixel.
+    """
+
+    nearest_pixelization_index_for_slim_index = np.zeros((grid.shape[0],))
+
+    for image_index in range(grid.shape[0]):
+        distances = (grid[image_index, 0] - mesh_grid[:, 0]) ** 2 + (
+            grid[image_index, 1] - mesh_grid[:, 1]
+        ) ** 2
+
+        nearest_pixelization_index_for_slim_index[image_index] = np.argmin(distances)
+
+    return nearest_pixelization_index_for_slim_index
+
+
 @numba_util.jit()
 def pix_indexes_for_sub_slim_index_voronoi_from(
     grid: np.ndarray,
-    nearest_pixelization_index_for_slim_index: np.ndarray,
     slim_index_for_sub_slim_index: np.ndarray,
     mesh_grid: np.ndarray,
     neighbors: np.ndarray,
@@ -195,8 +232,6 @@ def pix_indexes_for_sub_slim_index_voronoi_from(
     grid
         The grid of (y,x) scaled coordinates at the centre of every unmasked pixel, which has been traced to
         to an irgrid via lens.
-    nearest_pixelization_index_for_slim_index
-        A 1D array that maps every slimmed data-plane pixel to its nearest pixelization pixel.
     slim_index_for_sub_slim_index
         The mappings between the data slimmed sub-pixels and their regular pixels.
     mesh_grid
@@ -209,6 +244,12 @@ def pix_indexes_for_sub_slim_index_voronoi_from(
         Voronoi grid.
     """
 
+    nearest_pixelization_index_for_slim_index = (
+        nearest_pixelization_index_for_slim_index_from(
+            grid=grid, mesh_grid=mesh_grid
+        ).astype("int")
+    )
+
     pix_indexes_for_sub_slim_index = np.zeros(shape=(grid.shape[0], 1))
 
     for sub_slim_index in range(grid.shape[0]):
@@ -342,8 +383,6 @@ def pix_size_weights_voronoi_nn_from(
     grid
         The grid of (y,x) scaled coordinates at the centre of every unmasked pixel, which has been traced to
         to an irgrid via lens.
-    nearest_pixelization_index_for_slim_index
-        A 1D array that maps every slimmed data-plane pixel to its nearest pixelization pixel.
     slim_index_for_sub_slim_index
         The mappings between the data slimmed sub-pixels and their regular pixels.
     mesh_grid

autoarray/inversion/pixelization/mappers/voronoi.py

@@ -230,7 +230,6 @@ def pix_sub_weights(self) -> PixSubWeights:
         """
         mappings = mapper_util.pix_indexes_for_sub_slim_index_voronoi_from(
             grid=self.source_plane_data_grid,
-            nearest_pixelization_index_for_slim_index=self.source_plane_mesh_grid.nearest_pixelization_index_for_slim_index,
             slim_index_for_sub_slim_index=self.source_plane_data_grid.mask.derive_indexes.slim_for_sub_slim,
             mesh_grid=self.source_plane_mesh_grid,
             neighbors=self.source_plane_mesh_grid.neighbors,

autoarray/inversion/pixelization/mesh/__init__.py

@@ -2,7 +2,9 @@
 from .rectangular import Rectangular
 from .voronoi import VoronoiMagnification
 from .voronoi import VoronoiBrightnessImage
+from .voronoi import VoronoiSNRImage
 from .voronoi_nn import VoronoiNNMagnification
 from .voronoi_nn import VoronoiNNBrightnessImage
+from .voronoi_nn import VoronoiNNSNRImage
 from .delaunay import DelaunayMagnification
 from .delaunay import DelaunayBrightnessImage

autoarray/inversion/pixelization/mesh/abstract.py

@@ -106,7 +106,6 @@ def mesh_grid_from(
         self,
         source_plane_data_grid: Grid2D,
         source_plane_mesh_grid: Grid2DSparse,
-        sparse_index_for_slim_index: np.ndarray = None,
     ):
         raise NotImplementedError
 

autoarray/inversion/pixelization/mesh/delaunay.py

@@ -45,7 +45,6 @@ def mesh_grid_from(
         self,
         source_plane_data_grid=None,
         source_plane_mesh_grid=None,
-        sparse_index_for_slim_index=None,
     ):
         """
         Return the Delaunay ``source_plane_mesh_grid`` as a ``Mesh2DDelaunay`` object, which provides additional
@@ -114,6 +113,7 @@ def image_plane_mesh_grid_from(
         image_plane_data_grid: Grid2D,
         adapt_data: np.ndarray = None,
         settings=SettingsPixelization(),
+        noise_map: np.ndarray = None,
     ):
         """
         Computes the ``mesh_grid`` in the image-plane, by overlaying a uniform grid of coordinates over the
@@ -214,6 +214,7 @@ def image_plane_mesh_grid_from(
         image_plane_data_grid: Grid2D,
         adapt_data: np.ndarray,
         settings=SettingsPixelization(),
+        noise_map: np.ndarray = None,
     ):
         """
         Computes the ``mesh_grid`` in the image-plane, by overlaying a uniform grid of coordinates over the

autoarray/inversion/pixelization/mesh/rectangular.py

@@ -130,7 +130,6 @@ def mesh_grid_from(
         self,
         source_plane_data_grid: Optional[Grid2D] = None,
         source_plane_mesh_grid: Optional[Grid2D] = None,
-        sparse_index_for_slim_index: Optional[np.ndarray] = None,
     ) -> Mesh2DRectangular:
         """
         Return the rectangular `source_plane_mesh_grid` as a `Mesh2DRectangular` object, which provides additional
@@ -144,8 +143,6 @@ def mesh_grid_from(
         source_plane_mesh_grid
             Not used for a rectangular pixelization, because the pixelization grid in the `source` frame is computed
             by overlaying the `source_plane_data_grid` with the rectangular pixelization.
-        sparse_index_for_slim_index
-            Not used for a rectangular pixelization.
         """
         return Mesh2DRectangular.overlay_grid(
             shape_native=self.shape, grid=source_plane_data_grid
@@ -156,6 +153,7 @@ def image_plane_mesh_grid_from(
         image_plane_data_grid: Grid2D,
         adapt_data: np.ndarray = None,
         settings=SettingsPixelization(),
+        noise_map: np.ndarray = None,
     ):
         """
         Not used for rectangular pixelization.

autoarray/inversion/pixelization/mesh/triangulation.py

@@ -84,7 +84,6 @@ def mapper_grids_from(
             source_plane_mesh_grid = self.mesh_grid_from(
                 source_plane_data_grid=source_plane_data_grid,
                 source_plane_mesh_grid=relocated_source_plane_mesh_grid,
-                sparse_index_for_slim_index=source_plane_mesh_grid.sparse_index_for_slim_index,
             )
         except ValueError as e:
             raise e

autoarray/inversion/pixelization/mesh/voronoi.py

@@ -48,16 +48,11 @@ def mesh_grid_from(
         self,
         source_plane_data_grid=None,
         source_plane_mesh_grid=None,
-        sparse_index_for_slim_index=None,
     ) -> Mesh2DVoronoi:
         """
         Return the Voronoi `source_plane_mesh_grid` as a `Mesh2DVoronoi` object, which provides additional
         functionality for performing operations that exploit the geometry of a Voronoi mesh.
 
-        The array `sparse_index_for_slim_index` encodes the closest source pixel of every pixel on the
-        (full resolution) sub image-plane grid. This is used for efficiently pairing every image-plane pixel to its
-        corresponding source-plane pixel.
-
         Parameters
         ----------
         source_plane_data_grid
@@ -73,7 +68,6 @@ def mesh_grid_from(
 
         return Mesh2DVoronoi(
             values=source_plane_mesh_grid,
-            nearest_pixelization_index_for_slim_index=sparse_index_for_slim_index,
             uses_interpolation=self.uses_interpolation,
         )
 
@@ -125,6 +119,7 @@ def image_plane_mesh_grid_from(
         image_plane_data_grid: Grid2D,
         adapt_data: np.ndarray = None,
         settings=SettingsPixelization(),
+        noise_map: np.ndarray = None,
     ) -> Grid2DSparse:
         """
         Computes the `mesh_grid` in the `data` frame, by overlaying a uniform grid of coordinates over the
@@ -228,6 +223,7 @@ def image_plane_mesh_grid_from(
         image_plane_data_grid: Grid2D,
         adapt_data: np.ndarray,
         settings=SettingsPixelization(),
+        noise_map: np.ndarray = None,
     ):
         """
         Computes the `mesh_grid` in the `data` frame, by overlaying a uniform grid of coordinates over the
@@ -264,3 +260,120 @@ def image_plane_mesh_grid_from(
     @property
     def is_stochastic(self) -> bool:
         return True
+
+
+class VoronoiSNRImage(Voronoi):
+    def __init__(
+        self, pixels=10, fraction_high_snr=0.5, snr_cut: float = 3.0, **kwargs
+    ):
+        """
+        An irregular mesh of Voronoi pixels, which using no interpolation are paired with a 2D grid of (y,x)
+        coordinates. The Voronoi cell centers are derived in the image-plane by applying a KMeans
+        clustering algorithm to the image's weight map.
+
+        For a full description of how a mesh is paired with another grid,
+        see the :meth:`Pixelization API documentation <autoarray.inversion.pixelization.pixelization.Pixelization>`.
+
+        The Voronoi mesh represents pixels as an irregular 2D grid of Voronoi cells.
+
+        A ``Pixelization`` using a ``Voronoi`` mesh has four grids associated with it:
+
+        - ``image_plane_data_grid``: The observed data grid in the image-plane (which is paired with the mesh in
+          the source-plane).
+        - ``image_plane_mesh_grid``: The (y,x) mesh coordinates in the image-plane (which are the centres of Voronoi
+          cells in the source-plane).
+        - ``source_plane_data_grid``: The observed data grid mapped to the source-plane (e.g. after gravitational lensing).
+        - ``source_plane_mesh_grid``: The centre of each Voronoi cell in the source-plane
+          (the ``image_plane_mesh_grid`` maps to this after gravitational lensing).
+
+        Each (y,x) coordinate in the ``source_plane_data_grid`` is paired with all Voronoi cells it falls within,
+        without using an interpolation scheme.
+
+        The centers of the Voronoi cell pixels are derived in the image plane, by applying a KMeans clustering algorithm
+        to the masked image data's weight-map. The ``weight_floor`` and ``weight_power`` allow the KMeans algorithm to
+        adapt the image-plane coordinates to the image's brightest or faintest values. The computed valies are
+        mapped to the source-plane  via gravitational lensing, where they form the Voronoi cell pixel centers.
+
+        Parameters
+        ----------
+        pixels
+            The total number of pixels in the mesh, which is therefore also the number of (y,x) coordinates computed
+            via the KMeans clustering algorithm in image-plane.
+        weight_floor
+            A parameter which reweights the data values the KMeans algorithm is applied too; as the floor increases
+            more weight is applied to values with lower values thus allowing mesh pixels to be placed in these
+            regions of the data.
+        weight_power
+            A parameter which reweights the data values the KMeans algorithm is applied too; as the power increases
+            more weight is applied to values with higher values thus allowing mesh pixels to be placed in these
+            regions of the data.
+        """
+        super().__init__()
+
+        self.pixels = int(pixels)
+        self.fraction_high_snr = fraction_high_snr
+        self.snr_cut = snr_cut
+
+    def weight_map_from(
+        self, adapt_data: np.ndarray, noise_map: np.ndarray
+    ) -> np.ndarray:
+        """
+        Computes a `weight_map` from an input `adapt_data`, where this image represents components in the masked 2d
+        data in the `data` frame. This applies the `weight_floor` and `weight_power` attributes of the class, which
+        scale the weights to make different components upweighted relative to one another.
+
+        Parameters
+        ----------
+        adapt_data
+            A image which represents one or more components in the masked 2D data in the `data` frame.
+
+        Returns
+        -------
+        The weight map which is used to adapt the Voronoi pixels in the `data` frame to components in the data.
+        """
+        return adapt_data / noise_map
+
+    def image_plane_mesh_grid_from(
+        self,
+        image_plane_data_grid: Grid2D,
+        adapt_data: np.ndarray,
+        settings=SettingsPixelization(),
+        noise_map: np.ndarray = None,
+    ):
+        """
+        Computes the `mesh_grid` in the `data` frame, by overlaying a uniform grid of coordinates over the
+        masked 2D data (see `Grid2DSparse.from_grid_and_unmasked_2d_grid_shape()`).
+
+        The `data_pixelization_grid` is transformed to the `source_plane_mesh_grid`, and it is these (y,x) values
+        which then act the centres of the Voronoi mesh's pixels.
+
+        For a `VoronoiBrightnessImage` this grid is computed by applying a KMeans clustering algorithm to the masked
+        data's values, where these values are reweighted by the `adapt_data` so that the algorithm can adapt to
+        specific parts of the data.
+
+        Parameters
+        ----------
+        image_plane_mesh_grid
+            The sparse set of (y,x) coordinates computed from the unmasked data in the `data` frame. This has a
+            transformation applied to it to create the `source_plane_mesh_grid`.
+        adapt_data
+            An image which is used to determine the `image_plane_mesh_grid` and therefore adapt the distribution of
+            pixels of the Voronoi grid to the data it discretizes.
+        settings
+            Settings controlling the pixelization for example if a border is used to relocate its exterior coordinates.
+        """
+        snr_map = self.weight_map_from(adapt_data=adapt_data, noise_map=noise_map)
+
+        return Grid2DSparse.from_snr_split(
+            pixels=self.pixels,
+            fraction_high_snr=self.fraction_high_snr,
+            snr_cut=self.snr_cut,
+            grid=image_plane_data_grid,
+            snr_map=snr_map,
+            seed=settings.kmeans_seed,
+            stochastic=settings.is_stochastic,
+        )
+
+    @property
+    def is_stochastic(self) -> bool:
+        return True