Merge pull request #83 from Jammy2211/feature/image_mesh_updates

Feature/image mesh updates

autoarray/inversion/pixelization/image_mesh/abstract.py

@@ -14,14 +14,36 @@ def __init__(self):
         """
         pass
 
-    @property
-    def is_stochastic(self):
-        return False
-
     @property
     def uses_adapt_images(self) -> bool:
         raise NotImplementedError
 
+    def weight_map_from(self, adapt_data: np.ndarray):
+        """
+        Returns the weight-map used by the image-mesh to compute the mesh pixel centres.
+
+        This is computed from an input adapt data, which is an image representing the data which the KMeans
+        clustering algorithm is applied too. This could be the image data itself, or a model fit which
+        only has certain features.
+
+        The ``weight_floor`` and ``weight_power`` attributes of the class are used to scale the weight map, which
+        gives the model flexibility in how it adapts the pixelization to the image data.
+
+        Parameters
+        ----------
+        adapt_data
+            A image which represents one or more components in the masked 2D data in the image-plane.
+
+        Returns
+        -------
+        The weight map which is used to adapt the Delaunay pixels in the image-plane to components in the data.
+        """
+
+        weight_map = (np.abs(adapt_data) + self.weight_floor) ** self.weight_power
+        weight_map /= np.sum(weight_map)
+
+        return weight_map
+
     def image_plane_mesh_grid_from(
         self, grid: Grid2D, adapt_data: Optional[np.ndarray] = None
     ) -> Grid2DIrregular:

autoarray/inversion/pixelization/image_mesh/abstract_weighted.py

@@ -0,0 +1,75 @@
+from typing import Optional
+
+import numpy as np
+
+from autoarray.inversion.pixelization.image_mesh.abstract import AbstractImageMesh
+from autoarray.structures.grids.uniform_2d import Grid2D
+from autoarray.structures.grids.irregular_2d import Grid2DIrregular
+
+
+class AbstractImageMeshWeighted(AbstractImageMesh):
+    def __init__(
+        self,
+        pixels=10.0,
+        weight_floor=0.0,
+        weight_power=0.0,
+    ):
+        """
+        An abstract image mesh, which is used by pixelizations to determine the (y,x) mesh coordinates from an adapt
+        image.
+
+        Parameters
+        ----------
+        pixels
+            The total number of pixels in the image mesh and drawn from the Hilbert curve.
+        weight_floor
+            The minimum weight value in the weight map, which allows more pixels to be drawn from the lower weight
+            regions of the adapt image.
+        weight_power
+            The power the weight values are raised too, which allows more pixels to be drawn from the higher weight
+            regions of the adapt image.
+        """
+
+        super().__init__()
+
+        self.pixels = pixels
+        self.weight_floor = weight_floor
+        self.weight_power = weight_power
+
+    @property
+    def uses_adapt_images(self) -> bool:
+        return True
+
+    def weight_map_from(self, adapt_data: np.ndarray):
+        """
+        Returns the weight-map used by the image-mesh to compute the mesh pixel centres.
+
+        This is computed from an input adapt data, which is an image representing the data which the KMeans
+        clustering algorithm is applied too. This could be the image data itself, or a model fit which
+        only has certain features.
+
+        The ``weight_floor`` and ``weight_power`` attributes of the class are used to scale the weight map, which
+        gives the model flexibility in how it adapts the pixelization to the image data.
+
+        Parameters
+        ----------
+        adapt_data
+            A image which represents one or more components in the masked 2D data in the image-plane.
+
+        Returns
+        -------
+        The weight map which is used to adapt the Delaunay pixels in the image-plane to components in the data.
+        """
+
+        weight_map = np.abs(adapt_data) / np.max(adapt_data)
+        weight_map += self.weight_floor
+        weight_map[weight_map > 1.0] = 1.0
+
+        weight_map = weight_map ** self.weight_power
+
+        return weight_map
+
+    def image_plane_mesh_grid_from(
+        self, grid: Grid2D, adapt_data: Optional[np.ndarray] = None
+    ) -> Grid2DIrregular:
+        raise NotImplementedError

autoarray/inversion/pixelization/image_mesh/hilbert.py

@@ -5,7 +5,9 @@
 
 from autoarray.structures.grids.uniform_2d import Grid2D
 from autoarray.mask.mask_2d import Mask2D
-from autoarray.inversion.pixelization.image_mesh.abstract import AbstractImageMesh
+from autoarray.inversion.pixelization.image_mesh.abstract_weighted import (
+    AbstractImageMeshWeighted,
+)
 from autoarray.structures.grids.irregular_2d import Grid2DIrregular
 
 from autoarray import exc
@@ -203,7 +205,7 @@ def inverse_transform_sampling_interpolated(probabilities, n_samples, gridx, gri
     return output_ids, output_x, output_y
 
 
-class Hilbert(AbstractImageMesh):
+class Hilbert(AbstractImageMeshWeighted):
     def __init__(
         self,
         pixels=10.0,
@@ -226,49 +228,21 @@ def __init__(
 
         Parameters
         ----------
-        total_pixels
+        pixels
             The total number of pixels in the image mesh and drawn from the Hilbert curve.
+        weight_floor
+            The minimum weight value in the weight map, which allows more pixels to be drawn from the lower weight
+            regions of the adapt image.
+        weight_power
+            The power the weight values are raised too, which allows more pixels to be drawn from the higher weight
+            regions of the adapt image.
         """
 
-        super().__init__()
-
-        self.pixels = pixels
-        self.weight_floor = weight_floor
-        self.weight_power = weight_power
-
-    def weight_map_from(self, adapt_data: np.ndarray):
-        """
-        Returns the weight-map used by the Hilbert curve to compute the mesh pixel centres.
-
-        This is computed from an input adapt data, which is an image representing the data which the KMeans
-        clustering algorithm is applied too. This could be the image data itself, or a model fit which
-        only has certain features.
-
-        The ``weight_floor`` and ``weight_power`` attributes of the class are used to scale the weight map, which
-        gives the model flexibility in how it adapts the pixelization to the image data.
-
-        Parameters
-        ----------
-        adapt_data
-            A image which represents one or more components in the masked 2D data in the image-plane.
-
-        Returns
-        -------
-        The weight map which is used to adapt the Delaunay pixels in the image-plane to components in the data.
-        """
-
-        # Do tests with this weight map, but use Qiuhans first
-
-        # weight_map = (adapt_data - np.min(adapt_data)) / (
-        #     np.max(adapt_data) - np.min(adapt_data)
-        # ) + self.weight_floor * np.max(adapt_data)
-
-        #        return np.power(weight_map, self.weight_power)
-
-        weight_map = (np.abs(adapt_data) + self.weight_floor) ** self.weight_power
-        weight_map /= np.sum(weight_map)
-
-        return weight_map
+        super().__init__(
+            pixels=pixels,
+            weight_floor=weight_floor,
+            weight_power=weight_power,
+        )
 
     def image_plane_mesh_grid_from(
         self, grid: Grid2D, adapt_data: Optional[np.ndarray]
@@ -309,6 +283,8 @@ def image_plane_mesh_grid_from(
 
         weight_map = self.weight_map_from(adapt_data=adapt_data_hb)
 
+        weight_map /= np.sum(weight_map)
+
         (
             drawn_id,
             drawn_x,
@@ -321,11 +297,3 @@ def image_plane_mesh_grid_from(
         )
 
         return Grid2DIrregular(values=np.stack((drawn_y, drawn_x), axis=-1))
-
-    @property
-    def uses_adapt_images(self) -> bool:
-        return True
-
-    @property
-    def is_stochastic(self):
-        return True

autoarray/inversion/pixelization/image_mesh/kmeans.py

@@ -7,13 +7,15 @@
 if TYPE_CHECKING:
     from autoarray.structures.grids.uniform_2d import Grid2D
 
-from autoarray.inversion.pixelization.image_mesh.abstract import AbstractImageMesh
+from autoarray.inversion.pixelization.image_mesh.abstract_weighted import (
+    AbstractImageMeshWeighted,
+)
 from autoarray.structures.grids.irregular_2d import Grid2DIrregular
 
 from autoarray import exc
 
 
-class KMeans(AbstractImageMesh):
+class KMeans(AbstractImageMeshWeighted):
     def __init__(
         self,
         pixels=10.0,
@@ -38,43 +40,14 @@ def __init__(
         ----------
         total_pixels
             The total number of pixels in the image mesh and input into the KMeans algortihm.
-        n_iter
-            The number of times the KMeans algorithm is repeated.
-        max_iter
-            The maximum number of iterations in one run of the KMeans algorithm.
+        weight_power
         """
 
-        super().__init__()
-
-        self.pixels = pixels
-        self.weight_floor = weight_floor
-        self.weight_power = weight_power
-
-    def weight_map_from(self, adapt_data: np.ndarray):
-        """
-        Returns the weight-map used by the KMeans clustering algorithm to compute the mesh pixel centres.
-
-        This is computed from an input adapt data, which is an image representing the data which the KMeans
-        clustering algorithm is applied too. This could be the image data itself, or a model fit which
-        only has certain features.
-
-        The ``weight_floor`` and ``weight_power`` attributes of the class are used to scale the weight map, which
-        gives the model flexibility in how it adapts the pixelization to the image data.
-
-        Parameters
-        ----------
-        adapt_data
-            A image which represents one or more components in the masked 2D data in the image-plane.
-
-        Returns
-        -------
-        The weight map which is used to adapt the Delaunay pixels in the image-plane to components in the data.
-        """
-        weight_map = (adapt_data - np.min(adapt_data)) / (
-            np.max(adapt_data) - np.min(adapt_data)
-        ) + self.weight_floor * np.max(adapt_data)
-
-        return np.power(weight_map, self.weight_power)
+        super().__init__(
+            pixels=pixels,
+            weight_floor=weight_floor,
+            weight_power=weight_power,
+        )
 
     def image_plane_mesh_grid_from(
         self, grid: Grid2D, adapt_data: Optional[np.ndarray]
@@ -119,11 +92,3 @@ def image_plane_mesh_grid_from(
         return Grid2DIrregular(
             values=kmeans.cluster_centers_,
         )
-
-    @property
-    def uses_adapt_images(self) -> bool:
-        return True
-
-    @property
-    def is_stochastic(self):
-        return True

test_autoarray/inversion/pixelization/image_mesh/test_abstract_weighted.py

@@ -0,0 +1,26 @@
+import numpy as np
+import pytest
+
+import autoarray as aa
+
+
+def test__weight_map_from():
+    adapt_data = np.array([-1.0, 1.0, 3.0])
+
+    pixelization = aa.image_mesh.KMeans(pixels=5, weight_floor=0.0, weight_power=1.0)
+
+    weight_map = pixelization.weight_map_from(adapt_data=adapt_data)
+
+    assert weight_map == pytest.approx([0.33333, 0.33333, 1.0], 1.0e-4)
+
+    pixelization = aa.image_mesh.KMeans(pixels=5, weight_floor=0.0, weight_power=2.0)
+
+    weight_map = pixelization.weight_map_from(adapt_data=adapt_data)
+
+    assert weight_map == pytest.approx([0.11111, 0.11111, 1.0], 1.0e-4)
+
+    pixelization = aa.image_mesh.KMeans(pixels=5, weight_floor=1.0, weight_power=1.0)
+
+    weight_map = pixelization.weight_map_from(adapt_data=adapt_data)
+
+    assert weight_map == pytest.approx([1.0, 1.0, 1.0], 1.0e-4)

test_autoarray/inversion/pixelization/image_mesh/test_hilbert.py

@@ -1,31 +1,8 @@
-import numpy as np
 import pytest
 
 import autoarray as aa
 
 
-# def test__weight_map_from():
-#     adapt_data = np.array([-1.0, 1.0, 3.0])
-#
-#     pixelization = aa.image_mesh.KMeans(pixels=5, weight_floor=0.0, weight_power=1.0)
-#
-#     weight_map = pixelization.weight_map_from(adapt_data=adapt_data)
-#
-#     assert (weight_map == np.array([0.0, 0.5, 1.0])).all()
-#
-#     pixelization = aa.image_mesh.KMeans(pixels=5, weight_floor=0.0, weight_power=2.0)
-#
-#     weight_map = pixelization.weight_map_from(adapt_data=adapt_data)
-#
-#     assert (weight_map == np.array([0.0, 0.25, 1.0])).all()
-#
-#     pixelization = aa.image_mesh.KMeans(pixels=5, weight_floor=1.0, weight_power=1.0)
-#
-#     weight_map = pixelization.weight_map_from(adapt_data=adapt_data)
-#
-#     assert (weight_map == np.array([3.0, 3.5, 4.0])).all()
-
-
 def test__image_plane_mesh_grid_from():
     mask = aa.Mask2D.circular(
         shape_native=(4, 4),

test_autoarray/inversion/pixelization/image_mesh/test_kmeans.py

@@ -4,28 +4,6 @@
 import autoarray as aa
 
 
-def test__weight_map_from():
-    adapt_data = np.array([-1.0, 1.0, 3.0])
-
-    pixelization = aa.image_mesh.KMeans(pixels=5, weight_floor=0.0, weight_power=1.0)
-
-    weight_map = pixelization.weight_map_from(adapt_data=adapt_data)
-
-    assert (weight_map == np.array([0.0, 0.5, 1.0])).all()
-
-    pixelization = aa.image_mesh.KMeans(pixels=5, weight_floor=0.0, weight_power=2.0)
-
-    weight_map = pixelization.weight_map_from(adapt_data=adapt_data)
-
-    assert (weight_map == np.array([0.0, 0.25, 1.0])).all()
-
-    pixelization = aa.image_mesh.KMeans(pixels=5, weight_floor=1.0, weight_power=1.0)
-
-    weight_map = pixelization.weight_map_from(adapt_data=adapt_data)
-
-    assert (weight_map == np.array([3.0, 3.5, 4.0])).all()
-
-
 def test__image_plane_mesh_grid_from():
     mask = aa.Mask2D(
         mask=np.array(