diff --git a/autoarray/__init__.py b/autoarray/__init__.py
index 1a6d79c53..5acddd535 100644
--- a/autoarray/__init__.py
+++ b/autoarray/__init__.py
@@ -20,6 +20,7 @@
 from .fit.fit_imaging import FitImaging
 from .fit.fit_interferometer import FitInterferometer
 from .geometry.geometry_2d import Geometry2D
+from .inversion.convolver import Convolver
 from .inversion.pixelization.mappers.abstract import AbstractMapper
 from .inversion.pixelization import mesh
 from .inversion.pixelization import image_mesh
@@ -44,7 +45,6 @@
 from .inversion.inversion.imaging.w_tilde import InversionImagingWTilde
 from .inversion.inversion.interferometer.w_tilde import InversionInterferometerWTilde
 from .inversion.inversion.interferometer.mapping import InversionInterferometerMapping
-from .inversion.inversion.interferometer.lop import InversionInterferometerMappingPyLops
 from .inversion.linear_obj.linear_obj import LinearObj
 from .inversion.linear_obj.func_list import AbstractLinearObjFuncList
 from .mask.derive.indexes_2d import DeriveIndexes2D
diff --git a/autoarray/abstract_ndarray.py b/autoarray/abstract_ndarray.py
index 8b5fbc00e..ded8c5452 100644
--- a/autoarray/abstract_ndarray.py
+++ b/autoarray/abstract_ndarray.py
@@ -4,10 +4,11 @@
 
 from abc import ABC
 from abc import abstractmethod
+import jax.numpy as jnp
 
 from autoconf.fitsable import output_to_fits
 
-from autoarray.numpy_wrapper import np, register_pytree_node, Array
+from autoarray.numpy_wrapper import register_pytree_node, Array
 
 from typing import TYPE_CHECKING
 
@@ -82,7 +83,7 @@ def __init__(self, array):
 
     def invert(self):
         new = self.copy()
-        new._array = np.invert(new._array)
+        new._array = jnp.invert(new._array)
         return new
 
     @classmethod
@@ -104,7 +105,7 @@ def instance_flatten(cls, instance):
     @staticmethod
     def flip_hdu_for_ds9(values):
         if conf.instance["general"]["fits"]["flip_for_ds9"]:
-            return np.flipud(values)
+            return jnp.flipud(values)
         return values
 
     @classmethod
@@ -117,7 +118,7 @@ def instance_unflatten(cls, aux_data, children):
             setattr(instance, key, value)
         return instance
 
-    def with_new_array(self, array: np.ndarray) -> "AbstractNDArray":
+    def with_new_array(self, array: jnp.ndarray) -> "AbstractNDArray":
         """
         Copy this object but give it a new array.
 
@@ -164,7 +165,7 @@ def __iter__(self):
 
     @to_new_array
     def sqrt(self):
-        return np.sqrt(self._array)
+        return jnp.sqrt(self._array)
 
     @property
     def array(self):
@@ -330,13 +331,13 @@ def __getitem__(self, item):
         result = self._array[item]
         if isinstance(item, slice):
             result = self.with_new_array(result)
-        if isinstance(result, np.ndarray):
+        if isinstance(result, jnp.ndarray):
             result = self.with_new_array(result)
         return result
 
     def __setitem__(self, key, value):
-        if isinstance(key, (np.ndarray, AbstractNDArray, Array)):
-            self._array = np.where(key, value, self._array)
+        if isinstance(key, (jnp.ndarray, AbstractNDArray, Array)):
+            self._array = jnp.where(key, value, self._array)
         else:
             self._array[key] = value
 
diff --git a/autoarray/config/grids.yaml b/autoarray/config/grids.yaml
deleted file mode 100644
index f82eaa5df..000000000
--- a/autoarray/config/grids.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-radial_minimum:
-  function_name:
-    class_name: 1.0e-08
diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py
index e3ec74d3b..3779a30d6 100644
--- a/autoarray/dataset/imaging/dataset.py
+++ b/autoarray/dataset/imaging/dataset.py
@@ -166,8 +166,9 @@ def __init__(
 
         self.psf = psf
 
-        if psf.mask.shape[0] % 2 == 0 or psf.mask.shape[1] % 2 == 0:
-            raise exc.KernelException("Kernel2D Kernel2D must be odd")
+        if psf is not None:
+            if psf.mask.shape[0] % 2 == 0 or psf.mask.shape[1] % 2 == 0:
+                raise exc.KernelException("Kernel2D Kernel2D must be odd")
 
     @cached_property
     def grids(self):
@@ -178,6 +179,27 @@ def grids(self):
             psf=self.psf,
         )
 
+    @cached_property
+    def convolver(self):
+        """
+        Returns a `Convolver` from a mask and 2D PSF kernel.
+
+        The `Convolver` stores in memory the array indexing between the mask and PSF, enabling efficient 2D PSF
+        convolution of images and matrices used for linear algebra calculations (see `operators.convolver`).
+
+        This uses lazy allocation such that the calculation is only performed when the convolver is used, ensuring
+        efficient set up of the `Imaging` class.
+
+        Returns
+        -------
+        Convolver
+            The convolver given the masked imaging data's mask and PSF.
+        """
+
+        from autoarray.inversion.convolver import Convolver
+
+        return Convolver(mask=self.mask, kernel=self.psf)
+
     @cached_property
     def w_tilde(self):
         """
@@ -203,9 +225,11 @@ def w_tilde(self):
             indexes,
             lengths,
         ) = inversion_imaging_util.w_tilde_curvature_preload_imaging_from(
-            noise_map_native=np.array(self.noise_map.native),
-            kernel_native=np.array(self.psf.native),
-            native_index_for_slim_index=self.mask.derive_indexes.native_for_slim,
+            noise_map_native=np.array(self.noise_map.native.array).astype("float64"),
+            kernel_native=np.array(self.psf.native.array).astype("float64"),
+            native_index_for_slim_index=np.array(
+                self.mask.derive_indexes.native_for_slim
+            ).astype("int"),
         )
 
         return WTildeImaging(
@@ -408,20 +432,20 @@ def apply_noise_scaling(
         """
 
         if signal_to_noise_value is None:
-            noise_map = self.noise_map.native
-            noise_map[mask == False] = noise_value
+            noise_map = np.array(self.noise_map.native.array)
+            noise_map[mask.array == False] = noise_value
         else:
             noise_map = np.where(
                 mask == False,
-                np.median(self.data.native[mask.derive_mask.edge == False])
+                np.median(self.data.native.array[mask.derive_mask.edge == False])
                 / signal_to_noise_value,
-                self.noise_map.native,
+                self.noise_map.native.array,
             )
 
         if should_zero_data:
-            data = np.where(np.invert(mask), 0.0, self.data.native)
+            data = np.where(np.invert(mask.array), 0.0, self.data.native.array)
         else:
-            data = self.data.native
+            data = self.data.native.array
 
         data_unmasked = Array2D.no_mask(
             values=data,
diff --git a/autoarray/dataset/imaging/simulator.py b/autoarray/dataset/imaging/simulator.py
index 449a3d4b6..576dc6017 100644
--- a/autoarray/dataset/imaging/simulator.py
+++ b/autoarray/dataset/imaging/simulator.py
@@ -151,7 +151,7 @@ def via_image_from(
                 pixel_scales=image.pixel_scales,
             )
 
-        if np.isnan(noise_map).any():
+        if np.isnan(noise_map.array).any():
             raise exc.DatasetException(
                 "The noise-map has NaN values in it. This suggests your exposure time and / or"
                 "background sky levels are too low, creating signal counts at or close to 0.0."
@@ -161,7 +161,9 @@ def via_image_from(
             image = image - background_sky_map
 
         mask = Mask2D.all_false(
-            shape_native=image.shape_native, pixel_scales=image.pixel_scales
+            shape_native=image.shape_native,
+            pixel_scales=image.pixel_scales,
+            origin=image.origin,
         )
 
         image = Array2D(values=image, mask=mask)
diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py
index 06892ea68..01b5d84bd 100644
--- a/autoarray/dataset/interferometer/dataset.py
+++ b/autoarray/dataset/interferometer/dataset.py
@@ -193,7 +193,9 @@ def w_tilde(self):
 
         w_matrix = inversion_interferometer_util.w_tilde_via_preload_from(
             w_tilde_preload=curvature_preload,
-            native_index_for_slim_index=self.real_space_mask.derive_indexes.native_for_slim,
+            native_index_for_slim_index=np.array(
+                self.real_space_mask.derive_indexes.native_for_slim
+            ).astype("int"),
         )
 
         dirty_image = self.transformer.image_from(
@@ -205,7 +207,7 @@ def w_tilde(self):
         return WTildeInterferometer(
             w_matrix=w_matrix,
             curvature_preload=curvature_preload,
-            dirty_image=dirty_image,
+            dirty_image=np.array(dirty_image.array),
             real_space_mask=self.real_space_mask,
             noise_map_value=self.noise_map[0],
         )
diff --git a/autoarray/dataset/preprocess.py b/autoarray/dataset/preprocess.py
index 5c7338204..f4a7a1231 100644
--- a/autoarray/dataset/preprocess.py
+++ b/autoarray/dataset/preprocess.py
@@ -149,7 +149,8 @@ def noise_map_via_data_eps_and_exposure_time_map_from(data_eps, exposure_time_ma
         The exposure time at every data-point of the data.
     """
     return data_eps.with_new_array(
-        np.abs(data_eps * exposure_time_map) ** 0.5 / exposure_time_map
+        np.abs(data_eps.array * exposure_time_map.array) ** 0.5
+        / exposure_time_map.array
     )
 
 
@@ -263,15 +264,17 @@ def edges_from(image, no_edges):
     edges = []
 
     for edge_no in range(no_edges):
-        top_edge = image.native[edge_no, edge_no : image.shape_native[1] - edge_no]
-        bottom_edge = image.native[
+        top_edge = image.native.array[
+            edge_no, edge_no : image.shape_native[1] - edge_no
+        ]
+        bottom_edge = image.native.array[
             image.shape_native[0] - 1 - edge_no,
             edge_no : image.shape_native[1] - edge_no,
         ]
-        left_edge = image.native[
+        left_edge = image.native.array[
             edge_no + 1 : image.shape_native[0] - 1 - edge_no, edge_no
         ]
-        right_edge = image.native[
+        right_edge = image.native.array[
             edge_no + 1 : image.shape_native[0] - 1 - edge_no,
             image.shape_native[1] - 1 - edge_no,
         ]
@@ -406,9 +409,10 @@ def poisson_noise_via_data_eps_from(data_eps, exposure_time_map, seed=-1):
         An array describing simulated poisson noise_maps
     """
     setup_random_seed(seed)
-    image_counts = np.multiply(data_eps, exposure_time_map)
+
+    image_counts = np.multiply(data_eps.array, exposure_time_map.array)
     return data_eps - np.divide(
-        np.random.poisson(image_counts, data_eps.shape), exposure_time_map
+        np.random.poisson(image_counts, data_eps.shape), exposure_time_map.array
     )
 
 
@@ -506,8 +510,6 @@ def noise_map_with_signal_to_noise_limit_from(
     from autoarray.structures.arrays.uniform_1d import Array1D
     from autoarray.structures.arrays.uniform_2d import Array2D
 
-    # TODO : Refacotr into a util
-
     signal_to_noise_map = data / noise_map
     signal_to_noise_map[signal_to_noise_map < 0] = 0
 
@@ -517,12 +519,14 @@ def noise_map_with_signal_to_noise_limit_from(
     noise_map_limit = np.where(
         (signal_to_noise_map.native > signal_to_noise_limit)
         & (noise_limit_mask == False),
-        np.abs(data.native) / signal_to_noise_limit,
-        noise_map.native,
+        np.abs(data.native.array) / signal_to_noise_limit,
+        noise_map.native.array,
     )
 
     mask = Mask2D.all_false(
-        shape_native=data.shape_native, pixel_scales=data.pixel_scales
+        shape_native=data.shape_native,
+        pixel_scales=data.pixel_scales,
+        origin=data.origin,
     )
 
     if len(noise_map.native) == 1:
diff --git a/autoarray/fit/fit_dataset.py b/autoarray/fit/fit_dataset.py
index c2f498949..a21b98885 100644
--- a/autoarray/fit/fit_dataset.py
+++ b/autoarray/fit/fit_dataset.py
@@ -86,7 +86,7 @@ def chi_squared(self) -> float:
         """
         Returns the chi-squared terms of the model data's fit to an dataset, by summing the chi-squared-map.
         """
-        return fit_util.chi_squared_from(chi_squared_map=self.chi_squared_map)
+        return fit_util.chi_squared_from(chi_squared_map=self.chi_squared_map.array)
 
     @property
     def noise_normalization(self) -> float:
@@ -95,7 +95,7 @@ def noise_normalization(self) -> float:
 
         [Noise_Term] = sum(log(2*pi*[Noise]**2.0))
         """
-        return fit_util.noise_normalization_from(noise_map=self.noise_map)
+        return fit_util.noise_normalization_from(noise_map=self.noise_map.array)
 
     @property
     def log_likelihood(self) -> float:
diff --git a/autoarray/fit/fit_util.py b/autoarray/fit/fit_util.py
index d40f55d1c..c61262c1a 100644
--- a/autoarray/fit/fit_util.py
+++ b/autoarray/fit/fit_util.py
@@ -1,5 +1,6 @@
 from functools import wraps
-import jax.numpy as np
+import jax.numpy as jnp
+import numpy as np
 
 from autoarray.mask.abstract_mask import Mask
 
@@ -83,7 +84,7 @@ def chi_squared_from(*, chi_squared_map: ty.DataLike) -> float:
     chi_squared_map
         The chi-squared-map of values of the model-data fit to the dataset.
     """
-    return np.sum(chi_squared_map._array)
+    return jnp.sum(chi_squared_map)
 
 
 def noise_normalization_from(*, noise_map: ty.DataLike) -> float:
@@ -97,12 +98,12 @@ def noise_normalization_from(*, noise_map: ty.DataLike) -> float:
     noise_map
         The masked noise-map of the dataset.
     """
-    return np.sum(np.log(2 * np.pi * noise_map._array**2.0))
+    return jnp.sum(jnp.log(2 * jnp.pi * noise_map**2.0))
 
 
 def normalized_residual_map_complex_from(
-    *, residual_map: np.ndarray, noise_map: np.ndarray
-) -> np.ndarray:
+    *, residual_map: jnp.ndarray, noise_map: jnp.ndarray
+) -> jnp.ndarray:
     """
     Returns the normalized residual-map of the fit of complex model-data to a dataset, where:
 
@@ -126,8 +127,8 @@ def normalized_residual_map_complex_from(
 
 
 def chi_squared_map_complex_from(
-    *, residual_map: np.ndarray, noise_map: np.ndarray
-) -> np.ndarray:
+    *, residual_map: jnp.ndarray, noise_map: jnp.ndarray
+) -> jnp.ndarray:
     """
     Returnss the chi-squared-map of the fit of complex model-data to a dataset, where:
 
@@ -145,7 +146,7 @@ def chi_squared_map_complex_from(
     return chi_squared_map_real + 1j * chi_squared_map_imag
 
 
-def chi_squared_complex_from(*, chi_squared_map: np.ndarray) -> float:
+def chi_squared_complex_from(*, chi_squared_map: jnp.ndarray) -> float:
     """
     Returns the chi-squared terms of each complex model data's fit to a masked dataset, by summing the masked
     chi-squared-map of the fit.
@@ -157,12 +158,12 @@ def chi_squared_complex_from(*, chi_squared_map: np.ndarray) -> float:
     chi_squared_map
         The chi-squared-map of values of the model-data fit to the dataset.
     """
-    chi_squared_real = np.sum(chi_squared_map.real)
-    chi_squared_imag = np.sum(chi_squared_map.imag)
+    chi_squared_real = jnp.sum(np.array(chi_squared_map.real))
+    chi_squared_imag = jnp.sum(np.array(chi_squared_map.imag))
     return chi_squared_real + chi_squared_imag
 
 
-def noise_normalization_complex_from(*, noise_map: np.ndarray) -> float:
+def noise_normalization_complex_from(*, noise_map: jnp.ndarray) -> float:
     """
     Returns the noise-map normalization terms of a complex noise-map, summing the noise_map value in every pixel as:
 
@@ -173,8 +174,12 @@ def noise_normalization_complex_from(*, noise_map: np.ndarray) -> float:
     noise_map
         The masked noise-map of the dataset.
     """
-    noise_normalization_real = np.sum(np.log(2 * np.pi * noise_map.real**2.0))
-    noise_normalization_imag = np.sum(np.log(2 * np.pi * noise_map.imag**2.0))
+    noise_normalization_real = jnp.sum(
+        jnp.log(2 * jnp.pi * np.array(noise_map).real ** 2.0)
+    )
+    noise_normalization_imag = jnp.sum(
+        jnp.log(2 * jnp.pi * np.array(noise_map).imag ** 2.0)
+    )
     return noise_normalization_real + noise_normalization_imag
 
 
@@ -198,9 +203,7 @@ def residual_map_with_mask_from(
     model_data
         The model data used to fit the data.
     """
-    return np.subtract(
-        data, model_data, out=np.zeros_like(data), where=np.asarray(mask) == 0
-    )
+    return jnp.where(jnp.asarray(mask) == 0, jnp.subtract(data, model_data), 0)
 
 
 @to_new_array
@@ -223,12 +226,7 @@ def normalized_residual_map_with_mask_from(
     mask
         The mask applied to the residual-map, where `False` entries are included in the calculation.
     """
-    return np.divide(
-        residual_map,
-        noise_map,
-        out=np.zeros_like(residual_map),
-        where=np.asarray(mask) == 0,
-    )
+    return jnp.where(jnp.asarray(mask) == 0, jnp.divide(residual_map, noise_map), 0)
 
 
 @to_new_array
@@ -251,14 +249,7 @@ def chi_squared_map_with_mask_from(
     mask
         The mask applied to the residual-map, where `False` entries are included in the calculation.
     """
-    return np.square(
-        np.divide(
-            residual_map,
-            noise_map,
-            out=np.zeros_like(residual_map),
-            where=np.asarray(mask) == 0,
-        )
-    )
+    return jnp.where(jnp.asarray(mask) == 0, jnp.square(residual_map / noise_map), 0)
 
 
 def chi_squared_with_mask_from(*, chi_squared_map: ty.DataLike, mask: Mask) -> float:
@@ -275,7 +266,7 @@ def chi_squared_with_mask_from(*, chi_squared_map: ty.DataLike, mask: Mask) -> f
     mask
         The mask applied to the chi-squared-map, where `False` entries are included in the calculation.
     """
-    return float(np.sum(chi_squared_map[np.asarray(mask) == 0]))
+    return float(jnp.sum(chi_squared_map[jnp.asarray(mask) == 0]))
 
 
 def chi_squared_with_mask_fast_from(
@@ -302,14 +293,14 @@ def chi_squared_with_mask_fast_from(
         The mask applied to the chi-squared-map, where `False` entries are included in the calculation.
     """
     return float(
-        np.sum(
-            np.square(
-                np.divide(
-                    np.subtract(
+        jnp.sum(
+            jnp.square(
+                jnp.divide(
+                    jnp.subtract(
                         data,
                         model_data,
-                    )[np.asarray(mask) == 0],
-                    noise_map[np.asarray(mask) == 0],
+                    )[jnp.asarray(mask) == 0],
+                    noise_map[jnp.asarray(mask) == 0],
                 )
             )
         )
@@ -331,11 +322,13 @@ def noise_normalization_with_mask_from(*, noise_map: ty.DataLike, mask: Mask) ->
     mask
         The mask applied to the noise-map, where `False` entries are included in the calculation.
     """
-    return float(np.sum(np.log(2 * np.pi * noise_map[np.asarray(mask) == 0] ** 2.0)))
+    return float(
+        jnp.sum(jnp.log(2 * jnp.pi * noise_map[jnp.asarray(mask) == 0] ** 2.0))
+    )
 
 
 def chi_squared_with_noise_covariance_from(
-    *, residual_map: ty.DataLike, noise_covariance_matrix_inv: np.ndarray
+    *, residual_map: ty.DataLike, noise_covariance_matrix_inv: jnp.ndarray
 ) -> float:
     """
     Returns the chi-squared value of the fit of model-data to a masked dataset, where
@@ -351,7 +344,7 @@ def chi_squared_with_noise_covariance_from(
         The inverse of the noise covariance matrix.
     """
 
-    return residual_map @ noise_covariance_matrix_inv @ residual_map
+    return residual_map.array @ noise_covariance_matrix_inv @ residual_map.array
 
 
 def log_likelihood_from(*, chi_squared: float, noise_normalization: float) -> float:
@@ -431,8 +424,8 @@ def log_evidence_from(
 
 
 def residual_flux_fraction_map_from(
-    *, residual_map: np.ndarray, data: np.ndarray
-) -> np.ndarray:
+    *, residual_map: jnp.ndarray, data: jnp.ndarray
+) -> jnp.ndarray:
     """
     Returns the residual flux fraction map of the fit of model-data to a masked dataset, where:
 
@@ -445,12 +438,12 @@ def residual_flux_fraction_map_from(
     data
         The data of the dataset.
     """
-    return np.divide(residual_map, data, out=np.zeros_like(residual_map))
+    return jnp.where(data != 0, residual_map / data, 0)
 
 
 def residual_flux_fraction_map_with_mask_from(
-    *, residual_map: np.ndarray, data: np.ndarray, mask: Mask
-) -> np.ndarray:
+    *, residual_map: jnp.ndarray, data: jnp.ndarray, mask: Mask
+) -> jnp.ndarray:
     """
     Returnss the residual flux fraction map of the fit of model-data to a masked dataset, where:
 
@@ -467,9 +460,4 @@ def residual_flux_fraction_map_with_mask_from(
     mask
         The mask applied to the residual-map, where `False` entries are included in the calculation.
     """
-    return np.divide(
-        residual_map,
-        data,
-        out=np.zeros_like(residual_map),
-        where=np.asarray(mask) == 0,
-    )
+    return jnp.where(mask == 0, residual_map / data, 0)
diff --git a/autoarray/geometry/geometry_2d.py b/autoarray/geometry/geometry_2d.py
index e78f0f75a..9eea7e9f2 100644
--- a/autoarray/geometry/geometry_2d.py
+++ b/autoarray/geometry/geometry_2d.py
@@ -13,8 +13,6 @@
 from autoarray import type as ty
 from autoarray.geometry import geometry_util
 
-from autofit.jax_wrapper import use_jax
-
 logging.basicConfig()
 logger = logging.getLogger(__name__)
 
diff --git a/autoarray/geometry/geometry_util.py b/autoarray/geometry/geometry_util.py
index e2c6c8898..4cf32a082 100644
--- a/autoarray/geometry/geometry_util.py
+++ b/autoarray/geometry/geometry_util.py
@@ -57,7 +57,6 @@ def convert_pixel_scales_1d(pixel_scales: ty.PixelScales) -> Tuple[float]:
     return pixel_scales
 
 
-@numba_util.jit()
 def central_pixel_coordinates_1d_from(
     shape_slim: Tuple[int],
 ) -> Union[Tuple[float], Tuple[float]]:
@@ -85,7 +84,6 @@ def central_pixel_coordinates_1d_from(
     return (float(shape_slim[0] - 1) / 2,)
 
 
-@numba_util.jit()
 def central_scaled_coordinate_1d_from(
     shape_slim: Tuple[float],
     pixel_scales: ty.PixelScales,
@@ -121,7 +119,6 @@ def central_scaled_coordinate_1d_from(
     return (x_pixel,)
 
 
-@numba_util.jit()
 def pixel_coordinates_1d_from(
     scaled_coordinates_1d: Tuple[float],
     shape_slim: Tuple[int],
@@ -139,7 +136,6 @@ def pixel_coordinates_1d_from(
     return (x_pixel,)
 
 
-@numba_util.jit()
 def scaled_coordinates_1d_from(
     pixel_coordinates_1d: Tuple[float],
     shape_slim: Tuple[int],
@@ -445,7 +441,12 @@ def transform_grid_2d_to_reference_frame(
     grid
         The 2d grid of (y, x) coordinates which are transformed to a new reference frame.
     """
-    shifted_grid_2d = np.array(grid_2d) - jnp.array(centre)
+    try:
+        grid_2d = grid_2d.array
+    except AttributeError:
+        pass
+
+    shifted_grid_2d = grid_2d - jnp.array(centre)
 
     radius = jnp.sqrt(jnp.sum(shifted_grid_2d**2.0, axis=1))
     theta_coordinate_to_profile = jnp.arctan2(
@@ -477,23 +478,24 @@ def transform_grid_2d_from_reference_frame(
         The 2d grid of (y, x) coordinates which are transformed to a new reference frame.
     """
 
-    cos_angle = np.cos(np.radians(angle))
-    sin_angle = np.sin(np.radians(angle))
+    cos_angle = jnp.cos(jnp.radians(angle))
+    sin_angle = jnp.sin(jnp.radians(angle))
 
-    y = np.add(
-        np.add(
-            np.multiply(grid_2d[:, 1], sin_angle), np.multiply(grid_2d[:, 0], cos_angle)
+    y = jnp.add(
+        jnp.add(
+            jnp.multiply(grid_2d[:, 1], sin_angle),
+            jnp.multiply(grid_2d[:, 0], cos_angle),
         ),
         centre[0],
     )
-    x = np.add(
-        np.add(
-            np.multiply(grid_2d[:, 1], cos_angle),
-            -np.multiply(grid_2d[:, 0], sin_angle),
+    x = jnp.add(
+        jnp.add(
+            jnp.multiply(grid_2d[:, 1], cos_angle),
+            -jnp.multiply(grid_2d[:, 0], sin_angle),
         ),
         centre[1],
     )
-    return np.vstack((y, x)).T
+    return jnp.vstack((y, x)).T
 
 
 def grid_pixels_2d_slim_from(
diff --git a/autoarray/inversion/convolver.py b/autoarray/inversion/convolver.py
new file mode 100644
index 000000000..70eb6efdf
--- /dev/null
+++ b/autoarray/inversion/convolver.py
@@ -0,0 +1,544 @@
+from autoarray import numba_util
+import numpy as np
+
+from autoarray.structures.arrays.uniform_2d import Array2D
+
+from autoarray import exc
+from autoarray.mask import mask_2d_util
+
+
+class Convolver:
+    def __init__(self, mask, kernel):
+        """
+        Class to setup the 1D convolution of an / mapping matrix.
+
+        Take a simple 3x3 and masks:
+
+        [[2, 8, 2],
+        [5, 7, 5],
+        [3, 1, 4]]
+
+        [[True, False, True],   (True means that the value is masked)
+        [False, False, False],
+        [True, False, True]]
+
+        A set of values in a corresponding 1d array of this might be represented as:
+
+        [2, 8, 2, 5, 7, 5, 3, 1, 4]
+
+        and after masking as:
+
+        [8, 5, 7, 5, 1]
+
+        Setup is required to perform 2D real-space convolution on the masked array. This module finds the \
+        relationship between the unmasked 2D data, masked data and kernel, so that 2D real-space convolutions \
+        can be efficiently applied to reduced 1D masked structures.
+
+        This calculation also accounts for the blurring of light outside of the masked regions which blurs into \
+        the masked region.
+
+
+        **IMAGE FRAMES:**
+
+        For a masked in 2D, one can compute for every pixel all of the unmasked pixels it will blur light into for \
+        a given PSF kernel size, e.g.:
+
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI     This is an imaging.Mask2D, where:
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI     x = `True` (Pixel is masked and excluded from lens)
+        IxIxIxIoIoIoIxIxIxIxI     o = `False` (Pixel is not masked and included in lens)
+        IxIxIxIoIoIoIxIxIxIxI
+        IxIxIxIoIoIoIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+
+        Here, there are 9 unmasked pixels. Indexing of each unmasked pixel goes from the top-left corner right and \
+        downwards, therefore:
+
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxI0I1I2IxIxIxIxI
+        IxIxIxI3I4I5IxIxIxIxI
+        IxIxIxI6I7I8IxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+
+        For every unmasked pixel, the Convolver over-lays the PSF and computes three quantities;
+
+        image_frame_indexes - The indexes of all masked pixels it will blur light into.
+        image_frame_kernels - The kernel values that overlap each masked pixel it will blur light into.
+        image_frame_length - The number of masked pixels it will blur light into (unmasked pixels are excluded)
+
+        For example, if we had the following 3x3 kernel:
+
+        I0.1I0.2I0.3I
+        I0.4I0.5I0.6I
+        I0.7I0.8I0.9I
+
+        For pixel 0 above, when we overlap the kernel 4 unmasked pixels overlap this kernel, such that:
+
+        image_frame_indexes = [0, 1, 3, 4]
+        image_frame_kernels = [0.5, 0.6, 0.8, 0.9]
+        image_frame_length = 4
+
+        Noting that the other 5 kernel values (0.1, 0.2, 0.3, 0.4, 0.7) overlap masked pixels and are thus discarded.
+
+        For pixel 1, we get the following results:
+
+        image_frame_indexes = [0, 1, 2, 3, 4, 5]
+        image_frame_kernels = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+        image_frame_lengths = 6
+
+        In the majority of cases, the kernel will overlap only unmasked pixels. This is the case above for \
+        central pixel 4, where:
+
+        image_frame_indexes = [0, 1, 2, 3, 4, 5, 6, 7, 8]
+        image_frame_kernels = [0,1, 0.2, 0,3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+        image_frame_lengths = 9
+
+        Once we have set up all these quantities, the convolution routine simply uses them to convolve a 1D array of a
+        masked or the masked of a util in the inversion module.
+
+
+        **BLURRING FRAMES:**
+
+        Whilst the scheme above accounts for all blurred light within the masks, it does not account for the fact that \
+        pixels outside of the masks will also blur light into it. This effect is accounted for using blurring frames.
+
+        It is omitted for mapping matrix blurring, as an inversion does not fit data outside of the masks.
+
+        First, a blurring masks is computed from a masks, which describes all pixels which are close enough to the masks \
+        to blur light into it for a given kernel size. Following the example above, the following blurring masks is \
+        computed:
+
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI     This is an example grid.Mask2D, where:
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIoIoIoIoIoIxIxIxI     x = `True` (Pixel is masked and excluded from lens)
+        IxIxIoIxIxIxIoIxIxIxI     o = `False` (Pixel is not masked and included in lens)
+        IxIxIoIxIxIxIoIxIxIxI
+        IxIxIoIxIxIxIoIxIxIxI
+        IxIxIoIoIoIoIoIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+        IxIxIxIxIxIxIxIxIxIxI
+
+        Indexing again goes from the top-left corner right and downwards:
+
+        IxIxI xI xI xI xI xIxIxIxI
+        IxIxI xI xI xI xI xIxIxIxI
+        IxIxI xI xI xI xI xIxIxIxI
+        IxIxI 0I 1I 2I 3I 4IxIxIxI
+        IxIxI 5I xI xI xI 6IxIxIxI
+        IxIxI 7I xI xI xI 8IxIxIxI
+        IxIxI 9I xI xI xI10IxIxIxI
+        IxIxI11I12I13I14I15IxIxIxI
+        IxIxI xI xI xI xI xIxIxIxI
+        IxIxI xI xI xI xI xIxIxIxI
+
+        For every unmasked blurring-pixel, the Convolver over-lays the PSF kernel and computes three quantities;
+
+        blurring_frame_indexes - The indexes of all unmasked pixels (not unmasked blurring pixels) it will \
+        blur light into.
+        bluring_frame_kernels - The kernel values that overlap each pixel it will blur light into.
+        blurring_frame_length - The number of pixels it will blur light into.
+
+        The blurring frame therefore does not perform any blurring which blurs light into other blurring pixels. \
+        It only performs computations which add light inside of the masks.
+
+        For pixel 0 above, when we overlap the 3x3 kernel above only 1 unmasked pixels overlaps the kernel, such that:
+
+        blurring_frame_indexes = [0] (This 0 refers to pixel 0 within the masks, not blurring_frame_pixel 0)
+        blurring_frame_kernels = [0.9]
+        blurring_frame_length = 1
+
+        For pixel 1 above, when we overlap the 3x3 kernel above 2 unmasked pixels overlap the kernel, such that:
+
+        blurring_frame_indexes = [0, 1]  (This 0 and 1 refer to pixels 0 and 1 within the masks)
+        blurring_frame_kernels = [0.8, 0.9]
+        blurring_frame_length = 2
+
+        For pixel 3 above, when we overlap the 3x3 kernel above 3 unmasked pixels overlap the kernel, such that:
+
+        blurring_frame_indexes = [0, 1, 2]  (Again, these are pixels 0, 1 and 2)
+        blurring_frame_kernels = [0.7, 0.8, 0.9]
+        blurring_frame_length = 3
+
+        Parameters
+        ----------
+        mask
+            The mask within which the convolved signal is calculated.
+        blurring_mask
+            A masks of pixels outside the masks but whose light blurs into it after PSF convolution.
+        kernel : grid.PSF or ndarray
+            An array representing a PSF.
+        """
+        if kernel.shape_native[0] % 2 == 0 or kernel.shape_native[1] % 2 == 0:
+            raise exc.KernelException("PSF kernel must be odd")
+
+        self.mask = mask
+
+        self.mask_index_array = np.full(mask.shape, -1)
+        self.pixels_in_mask = int(np.size(mask) - np.sum(mask))
+
+        count = 0
+        for x in range(mask.shape[0]):
+            for y in range(mask.shape[1]):
+                if not mask[x, y]:
+                    self.mask_index_array[x, y] = count
+                    count += 1
+
+        self.kernel = kernel
+        self.kernel_max_size = self.kernel.shape_native[0] * self.kernel.shape_native[1]
+
+        mask_1d_index = 0
+        self.image_frame_1d_indexes = np.zeros(
+            (self.pixels_in_mask, self.kernel_max_size), dtype="int"
+        )
+        self.image_frame_1d_kernels = np.zeros(
+            (self.pixels_in_mask, self.kernel_max_size)
+        )
+        self.image_frame_1d_lengths = np.zeros((self.pixels_in_mask), dtype="int")
+        for x in range(self.mask_index_array.shape[0]):
+            for y in range(self.mask_index_array.shape[1]):
+                if not mask[x][y]:
+                    (
+                        image_frame_1d_indexes,
+                        image_frame_1d_kernels,
+                    ) = self.frame_at_coordinates_jit(
+                        coordinates=(x, y),
+                        mask=np.array(mask),
+                        mask_index_array=self.mask_index_array,
+                        kernel_2d=np.array(self.kernel.native.array),
+                    )
+                    self.image_frame_1d_indexes[mask_1d_index, :] = (
+                        image_frame_1d_indexes
+                    )
+                    self.image_frame_1d_kernels[mask_1d_index, :] = (
+                        image_frame_1d_kernels
+                    )
+                    self.image_frame_1d_lengths[mask_1d_index] = image_frame_1d_indexes[
+                        image_frame_1d_indexes >= 0
+                    ].shape[0]
+                    mask_1d_index += 1
+
+        self.blurring_mask = mask_2d_util.blurring_mask_2d_from(
+            mask_2d=np.array(mask),
+            kernel_shape_native=kernel.shape_native,
+        )
+
+        self.pixels_in_blurring_mask = int(
+            np.size(self.blurring_mask) - np.sum(self.blurring_mask)
+        )
+
+        mask_1d_index = 0
+        self.blurring_frame_1d_indexes = np.zeros(
+            (self.pixels_in_blurring_mask, self.kernel_max_size), dtype="int"
+        )
+        self.blurring_frame_1d_kernels = np.zeros(
+            (self.pixels_in_blurring_mask, self.kernel_max_size)
+        )
+        self.blurring_frame_1d_lengths = np.zeros(
+            (self.pixels_in_blurring_mask), dtype="int"
+        )
+        for x in range(mask.shape[0]):
+            for y in range(mask.shape[1]):
+                if mask[x][y] and not self.blurring_mask[x, y]:
+                    (
+                        image_frame_1d_indexes,
+                        image_frame_1d_kernels,
+                    ) = self.frame_at_coordinates_jit(
+                        coordinates=(x, y),
+                        mask=np.array(mask),
+                        mask_index_array=np.array(self.mask_index_array),
+                        kernel_2d=np.array(self.kernel.native.array),
+                    )
+                    self.blurring_frame_1d_indexes[mask_1d_index, :] = (
+                        image_frame_1d_indexes
+                    )
+                    self.blurring_frame_1d_kernels[mask_1d_index, :] = (
+                        image_frame_1d_kernels
+                    )
+                    self.blurring_frame_1d_lengths[mask_1d_index] = (
+                        image_frame_1d_indexes[image_frame_1d_indexes >= 0].shape[0]
+                    )
+                    mask_1d_index += 1
+
+    @staticmethod
+    @numba_util.jit()
+    def frame_at_coordinates_jit(coordinates, mask, mask_index_array, kernel_2d):
+        """
+        Returns the frame (indexes of pixels light is blurred into) and kernel_frame (kernel kernel values of those \
+        pixels) for a given coordinate in a masks and its PSF.
+
+        Parameters
+        ----------
+        coordinates: Tuple[int, int]
+            The coordinates of mask_index_array on which the frame should be centred
+        kernel_shape_native: Tuple[int, int]
+            The shape of the kernel for which this frame will be used
+        """
+
+        kernel_shape_native = kernel_2d.shape
+        kernel_max_size = kernel_shape_native[0] * kernel_shape_native[1]
+
+        half_x = int(kernel_shape_native[0] / 2)
+        half_y = int(kernel_shape_native[1] / 2)
+
+        frame = -1 * np.ones((kernel_max_size))
+        kernel_frame = -1.0 * np.ones((kernel_max_size))
+
+        count = 0
+        for i in range(kernel_shape_native[0]):
+            for j in range(kernel_shape_native[1]):
+                x = coordinates[0] - half_x + i
+                y = coordinates[1] - half_y + j
+                if (
+                    0 <= x < mask_index_array.shape[0]
+                    and 0 <= y < mask_index_array.shape[1]
+                ):
+                    value = mask_index_array[x, y]
+                    if value >= 0 and not mask[x, y]:
+                        frame[count] = value
+                        kernel_frame[count] = kernel_2d[i, j]
+                        count += 1
+
+        return frame, kernel_frame
+
+    def convolve_image(self, image, blurring_image):
+        """
+        For a given 1D array and blurring array, convolve the two using this convolver.
+
+        Parameters
+        ----------
+        image
+            1D array of the values which are to be blurred with the convolver's PSF.
+        blurring_image
+            1D array of the blurring values which blur into the array after PSF convolution.
+        """
+
+        if self.blurring_mask is None:
+            raise exc.KernelException(
+                "You cannot use the convolve_image function of a Convolver if the Convolver was"
+                "not created with a blurring_mask."
+            )
+
+        convolved_image = self.convolve_jit(
+            image_1d_array=np.array(image.slim.array),
+            image_frame_1d_indexes=self.image_frame_1d_indexes,
+            image_frame_1d_kernels=self.image_frame_1d_kernels,
+            image_frame_1d_lengths=self.image_frame_1d_lengths,
+            blurring_1d_array=np.array(blurring_image.slim.array),
+            blurring_frame_1d_indexes=self.blurring_frame_1d_indexes,
+            blurring_frame_1d_kernels=self.blurring_frame_1d_kernels,
+            blurring_frame_1d_lengths=self.blurring_frame_1d_lengths,
+        )
+
+        return Array2D(values=convolved_image, mask=self.mask)
+
+    @staticmethod
+    @numba_util.jit()
+    def convolve_jit(
+        image_1d_array,
+        image_frame_1d_indexes,
+        image_frame_1d_kernels,
+        image_frame_1d_lengths,
+        blurring_1d_array,
+        blurring_frame_1d_indexes,
+        blurring_frame_1d_kernels,
+        blurring_frame_1d_lengths,
+    ):
+        blurred_image_1d = np.zeros(image_1d_array.shape)
+
+        for image_1d_index in range(len(image_1d_array)):
+            frame_1d_indexes = image_frame_1d_indexes[image_1d_index]
+            frame_1d_kernel = image_frame_1d_kernels[image_1d_index]
+            frame_1d_length = image_frame_1d_lengths[image_1d_index]
+            image_value = image_1d_array[image_1d_index]
+
+            for kernel_1d_index in range(frame_1d_length):
+                vector_index = frame_1d_indexes[kernel_1d_index]
+                kernel_value = frame_1d_kernel[kernel_1d_index]
+                blurred_image_1d[vector_index] += image_value * kernel_value
+
+        for blurring_1d_index in range(len(blurring_1d_array)):
+            frame_1d_indexes = blurring_frame_1d_indexes[blurring_1d_index]
+            frame_1d_kernel = blurring_frame_1d_kernels[blurring_1d_index]
+            frame_1d_length = blurring_frame_1d_lengths[blurring_1d_index]
+            image_value = blurring_1d_array[blurring_1d_index]
+
+            for kernel_1d_index in range(frame_1d_length):
+                vector_index = frame_1d_indexes[kernel_1d_index]
+                kernel_value = frame_1d_kernel[kernel_1d_index]
+                blurred_image_1d[vector_index] += image_value * kernel_value
+
+        return blurred_image_1d
+
+    def convolve_image_no_blurring(self, image):
+        """For a given 1D array and blurring array, convolve the two using this convolver.
+
+        Parameters
+        ----------
+        image
+            1D array of the values which are to be blurred with the convolver's PSF.
+        """
+
+        convolved_image = self.convolve_no_blurring_jit(
+            image_1d_array=np.array(image.slim),
+            image_frame_1d_indexes=self.image_frame_1d_indexes,
+            image_frame_1d_kernels=self.image_frame_1d_kernels,
+            image_frame_1d_lengths=self.image_frame_1d_lengths,
+        )
+
+        return Array2D(values=convolved_image, mask=self.mask)
+
+    def convolve_image_no_blurring_interpolation(self, image):
+        """For a given 1D array and blurring array, convolve the two using this convolver.
+
+        Parameters
+        ----------
+        image
+            1D array of the values which are to be blurred with the convolver's PSF.
+        """
+
+        convolved_image = self.convolve_no_blurring_jit(
+            image_1d_array=image,
+            image_frame_1d_indexes=self.image_frame_1d_indexes,
+            image_frame_1d_kernels=self.image_frame_1d_kernels,
+            image_frame_1d_lengths=self.image_frame_1d_lengths,
+        )
+
+        return Array2D(values=convolved_image, mask=self.mask)
+
+    @staticmethod
+    @numba_util.jit()
+    def convolve_no_blurring_jit(
+        image_1d_array,
+        image_frame_1d_indexes,
+        image_frame_1d_kernels,
+        image_frame_1d_lengths,
+    ):
+        blurred_image_1d = np.zeros(image_1d_array.shape)
+
+        for image_1d_index in range(len(image_1d_array)):
+            frame_1d_indexes = image_frame_1d_indexes[image_1d_index]
+            frame_1d_kernel = image_frame_1d_kernels[image_1d_index]
+            frame_1d_length = image_frame_1d_lengths[image_1d_index]
+            image_value = image_1d_array[image_1d_index]
+
+            for kernel_1d_index in range(frame_1d_length):
+                vector_index = frame_1d_indexes[kernel_1d_index]
+                kernel_value = frame_1d_kernel[kernel_1d_index]
+                blurred_image_1d[vector_index] += image_value * kernel_value
+
+        return blurred_image_1d
+
+    def convolve_mapping_matrix(self, mapping_matrix):
+        """For a given inversion mapping matrix, convolve every pixel's mapped with the PSF kernel.
+
+        A mapping matrix provides non-zero entries in all elements which map two pixels to one another
+        (see *inversions.mappers*).
+
+        For example, lets take an which is masked using a 'cross' of 5 pixels:
+
+        [[ True, False,  True]],
+        [[False, False, False]],
+        [[ True, False,  True]]
+
+        As example mapping matrix of this cross is as follows (5 pixels x 3 source pixels):
+
+        [1, 0, 0] [0->0]
+        [1, 0, 0] [1->0]
+        [0, 1, 0] [2->1]
+        [0, 1, 0] [3->1]
+        [0, 0, 1] [4->2]
+
+        For each source-pixel, we can create an of its unit-surface brightnesses by util the non-zero
+        entries back to masks. For example, doing this for source pixel 1 gives:
+
+        [[0.0, 1.0, 0.0]],
+        [[1.0, 0.0, 0.0]]
+        [[0.0, 0.0, 0.0]]
+
+        And source pixel 2:
+
+        [[0.0, 0.0, 0.0]],
+        [[0.0, 1.0, 1.0]]
+        [[0.0, 0.0, 0.0]]
+
+        We then convolve each of these with our PSF kernel, in 2 dimensions, like we would a grid. For
+        example, using the kernel below:
+
+        kernel:
+
+        [[0.0, 0.1, 0.0]]
+        [[0.1, 0.6, 0.1]]
+        [[0.0, 0.1, 0.0]]
+
+        Blurred Source Pixel 1 (we don't need to perform the convolution into masked pixels):
+
+        [[0.0, 0.6, 0.0]],
+        [[0.6, 0.0, 0.0]],
+        [[0.0, 0.0, 0.0]]
+
+        Blurred Source pixel 2:
+
+        [[0.0, 0.0, 0.0]],
+        [[0.0, 0.7, 0.7]],
+        [[0.0, 0.0, 0.0]]
+
+        Finally, we map each of these blurred back to a blurred mapping matrix, which is analogous to the
+        mapping matrix.
+
+        [0.6, 0.0, 0.0] [0->0]
+        [0.6, 0.0, 0.0] [1->0]
+        [0.0, 0.7, 0.0] [2->1]
+        [0.0, 0.7, 0.0] [3->1]
+        [0.0, 0.0, 0.6] [4->2]
+
+        If the mapping matrix is sub-gridded, we perform the convolution on the fractional surface brightnesses in an
+        identical fashion to above.
+
+        Parameters
+        ----------
+        mapping_matrix
+            The 2D mapping matrix describing how every inversion pixel maps to a pixel on the data pixel.
+        """
+        return self.convolve_matrix_jit(
+            mapping_matrix=mapping_matrix,
+            image_frame_1d_indexes=self.image_frame_1d_indexes,
+            image_frame_1d_kernels=self.image_frame_1d_kernels,
+            image_frame_1d_lengths=self.image_frame_1d_lengths,
+        )
+
+    @staticmethod
+    @numba_util.jit()
+    def convolve_matrix_jit(
+        mapping_matrix,
+        image_frame_1d_indexes,
+        image_frame_1d_kernels,
+        image_frame_1d_lengths,
+    ):
+        blurred_mapping_matrix = np.zeros(mapping_matrix.shape)
+
+        for pixel_1d_index in range(mapping_matrix.shape[1]):
+            for image_1d_index in range(mapping_matrix.shape[0]):
+                value = mapping_matrix[image_1d_index, pixel_1d_index]
+
+                if value > 0:
+                    frame_1d_indexes = image_frame_1d_indexes[image_1d_index]
+                    frame_1d_kernel = image_frame_1d_kernels[image_1d_index]
+                    frame_1d_length = image_frame_1d_lengths[image_1d_index]
+
+                    for kernel_1d_index in range(frame_1d_length):
+                        vector_index = frame_1d_indexes[kernel_1d_index]
+                        kernel_value = frame_1d_kernel[kernel_1d_index]
+                        blurred_mapping_matrix[vector_index, pixel_1d_index] += (
+                            value * kernel_value
+                        )
+
+        return blurred_mapping_matrix
diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py
index 757a24ef0..d2546e9e3 100644
--- a/autoarray/inversion/inversion/abstract.py
+++ b/autoarray/inversion/inversion/abstract.py
@@ -100,6 +100,10 @@ def data(self):
     def noise_map(self):
         return self.dataset.noise_map
 
+    @property
+    def convolver(self):
+        return self.dataset.convolver
+
     def has(self, cls: Type) -> bool:
         """
         Does this `Inversion` have an attribute which is of type `cls`?
diff --git a/autoarray/inversion/inversion/dataset_interface.py b/autoarray/inversion/inversion/dataset_interface.py
index 0e417e71a..c1cf2db9a 100644
--- a/autoarray/inversion/inversion/dataset_interface.py
+++ b/autoarray/inversion/inversion/dataset_interface.py
@@ -5,6 +5,7 @@ def __init__(
         noise_map,
         grids=None,
         psf=None,
+        convolver=None,
         transformer=None,
         w_tilde=None,
         noise_covariance_matrix=None,
@@ -60,6 +61,7 @@ def __init__(
         self.noise_map = noise_map
         self.grids = grids
         self.psf = psf
+        self.convolver = convolver
         self.transformer = transformer
         self.w_tilde = w_tilde
         self.noise_covariance_matrix = noise_covariance_matrix
diff --git a/autoarray/inversion/inversion/factory.py b/autoarray/inversion/inversion/factory.py
index 327262786..350f19e65 100644
--- a/autoarray/inversion/inversion/factory.py
+++ b/autoarray/inversion/inversion/factory.py
@@ -9,9 +9,6 @@
 from autoarray.inversion.inversion.interferometer.w_tilde import (
     InversionInterferometerWTilde,
 )
-from autoarray.inversion.inversion.interferometer.lop import (
-    InversionInterferometerMappingPyLops,
-)
 from autoarray.inversion.inversion.dataset_interface import DatasetInterface
 from autoarray.inversion.linear_obj.linear_obj import LinearObj
 from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList
@@ -212,11 +209,3 @@ def inversion_interferometer_from(
                 settings=settings,
                 run_time_dict=run_time_dict,
             )
-
-    else:
-        return InversionInterferometerMappingPyLops(
-            dataset=dataset,
-            linear_obj_list=linear_obj_list,
-            settings=settings,
-            run_time_dict=run_time_dict,
-        )
diff --git a/autoarray/inversion/inversion/imaging/abstract.py b/autoarray/inversion/inversion/imaging/abstract.py
index 5efc4d0a9..f17e8fa4f 100644
--- a/autoarray/inversion/inversion/imaging/abstract.py
+++ b/autoarray/inversion/inversion/imaging/abstract.py
@@ -95,7 +95,7 @@ def operated_mapping_matrix_list(self) -> List[np.ndarray]:
 
         return [
             (
-                self.psf.convolve_mapping_matrix(
+                self.convolver.convolve_mapping_matrix(
                     mapping_matrix=linear_obj.mapping_matrix
                 )
                 if linear_obj.operated_mapping_matrix_override is None
@@ -138,7 +138,7 @@ def linear_func_operated_mapping_matrix_dict(self) -> Dict:
             if linear_func.operated_mapping_matrix_override is not None:
                 operated_mapping_matrix = linear_func.operated_mapping_matrix_override
             else:
-                operated_mapping_matrix = self.psf.convolve_mapping_matrix(
+                operated_mapping_matrix = self.convolver.convolve_mapping_matrix(
                     mapping_matrix=linear_func.mapping_matrix
                 )
 
@@ -220,7 +220,7 @@ def mapper_operated_mapping_matrix_dict(self) -> Dict:
         mapper_operated_mapping_matrix_dict = {}
 
         for mapper in self.cls_list_from(cls=AbstractMapper):
-            operated_mapping_matrix = self.psf.convolve_mapping_matrix(
+            operated_mapping_matrix = self.convolver.convolve_mapping_matrix(
                 mapping_matrix=mapper.mapping_matrix
             )
 
diff --git a/autoarray/inversion/inversion/imaging/mapping.py b/autoarray/inversion/inversion/imaging/mapping.py
index 03d73ff63..12881a944 100644
--- a/autoarray/inversion/inversion/imaging/mapping.py
+++ b/autoarray/inversion/inversion/imaging/mapping.py
@@ -78,7 +78,7 @@ def _data_vector_mapper(self) -> np.ndarray:
             mapper = mapper_list[i]
             param_range = mapper_param_range_list[i]
 
-            operated_mapping_matrix = self.psf.convolve_mapping_matrix(
+            operated_mapping_matrix = self.convolver.convolve_mapping_matrix(
                 mapping_matrix=mapper.mapping_matrix
             )
 
@@ -112,8 +112,8 @@ def data_vector(self) -> np.ndarray:
 
         return inversion_imaging_util.data_vector_via_blurred_mapping_matrix_from(
             blurred_mapping_matrix=self.operated_mapping_matrix,
-            image=np.array(self.data),
-            noise_map=np.array(self.noise_map),
+            image=self.data.array,
+            noise_map=self.noise_map.array,
         )
 
     @property
@@ -140,7 +140,7 @@ def _curvature_matrix_mapper_diag(self) -> Optional[np.ndarray]:
             mapper_i = mapper_list[i]
             mapper_param_range_i = mapper_param_range_list[i]
 
-            operated_mapping_matrix = self.psf.convolve_mapping_matrix(
+            operated_mapping_matrix = self.convolver.convolve_mapping_matrix(
                 mapping_matrix=mapper_i.mapping_matrix
             )
 
diff --git a/autoarray/inversion/inversion/imaging/w_tilde.py b/autoarray/inversion/inversion/imaging/w_tilde.py
index 199ba66b2..e5912452f 100644
--- a/autoarray/inversion/inversion/imaging/w_tilde.py
+++ b/autoarray/inversion/inversion/imaging/w_tilde.py
@@ -71,10 +71,12 @@ def __init__(
     @profile_func
     def w_tilde_data(self):
         return inversion_imaging_util.w_tilde_data_imaging_from(
-            image_native=np.array(self.data.native),
-            noise_map_native=np.array(self.noise_map.native),
-            kernel_native=np.array(self.psf.kernel.native),
-            native_index_for_slim_index=self.data.mask.derive_indexes.native_for_slim,
+            image_native=np.array(self.data.native.array).astype("float"),
+            noise_map_native=np.array(self.noise_map.native.array).astype("float"),
+            kernel_native=np.array(self.psf.native.array).astype("float"),
+            native_index_for_slim_index=np.array(
+                self.data.mask.derive_indexes.native_for_slim
+            ).astype("int"),
         )
 
     @property
@@ -207,8 +209,8 @@ def _data_vector_func_list_and_mapper(self) -> np.ndarray:
 
             diag = inversion_imaging_util.data_vector_via_blurred_mapping_matrix_from(
                 blurred_mapping_matrix=operated_mapping_matrix,
-                image=np.array(self.data),
-                noise_map=np.array(self.noise_map),
+                image=self.data.array,
+                noise_map=self.noise_map.array,
             )
 
             param_range = linear_func_param_range[linear_func_index]
@@ -521,12 +523,12 @@ def mapped_reconstructed_data_dict(self) -> Dict[LinearObj, Array2D]:
                     reconstruction=reconstruction,
                 )
 
-                mapped_reconstructed_image = Array2D(
-                    values=mapped_reconstructed_image, mask=self.mask
-                )
-
                 mapped_reconstructed_image = self.psf.convolve_image_no_blurring(
                     image=mapped_reconstructed_image, mask=self.mask
+                ).array
+
+                mapped_reconstructed_image = Array2D(
+                    values=np.array(mapped_reconstructed_image), mask=self.mask
                 )
 
             else:
diff --git a/autoarray/inversion/inversion/interferometer/lop.py b/autoarray/inversion/inversion/interferometer/lop.py
deleted file mode 100644
index fdd6b8adf..000000000
--- a/autoarray/inversion/inversion/interferometer/lop.py
+++ /dev/null
@@ -1,146 +0,0 @@
-from scipy import sparse
-
-import numpy as np
-from typing import Dict
-
-from autoconf import cached_property
-
-from autoarray.inversion.inversion.interferometer.abstract import (
-    AbstractInversionInterferometer,
-)
-from autoarray.inversion.linear_obj.linear_obj import LinearObj
-from autoarray.structures.visibilities import Visibilities
-
-from autoarray.numba_util import profile_func
-
-
-class InversionInterferometerMappingPyLops(AbstractInversionInterferometer):
-    """
-    Constructs linear equations (via vectors and matrices) which allow for sets of simultaneous linear equations
-    to be solved (see `inversion.inversion.abstract.AbstractInversion` for a full description).
-
-    A linear object describes the mappings between values in observed `data` and the linear object's model via its
-    `mapping_matrix`. This class constructs linear equations for `Interferometer` objects, where the data is an
-    an array of visibilities and the mappings include a non-uniform fast Fourier transform operation described by
-    the interferometer dataset's transformer.
-
-    This class uses the mapping formalism, which constructs the simultaneous linear equations using the
-    `mapping_matrix` of every linear object. This is performed using the library PyLops, which uses linear
-    operators to avoid these matrices being created explicitly in memory, making the calculation more efficient.
-    """
-
-    @cached_property
-    @profile_func
-    def reconstruction(self):
-        """
-        Solve the linear system [F + reg_coeff*H] S = D -> S = [F + reg_coeff*H]^-1 D given by equation (12)
-        of https://arxiv.org/pdf/astro-ph/0302587.pdf
-
-        S is the vector of reconstructed inversion values.
-        """
-
-        import pylops
-
-        Aop = pylops.MatrixMult(
-            sparse.bsr_matrix(self.linear_obj_list[0].mapping_matrix)
-        )
-
-        Fop = self.transformer
-
-        Op = Fop * Aop
-
-        MOp = pylops.MatrixMult(sparse.bsr_matrix(self.preconditioner_matrix_inverse))
-
-        try:
-            return pylops.NormalEquationsInversion(
-                Op=Op,
-                Regs=None,
-                epsNRs=[1.0],
-                data=self.data.ordered_1d,
-                Weight=pylops.Diagonal(diag=self.noise_map.weight_list_ordered_1d),
-                NRegs=[
-                    pylops.MatrixMult(sparse.bsr_matrix(self.regularization_matrix))
-                ],
-                M=MOp,
-                tol=self.settings.tolerance,
-                atol=self.settings.tolerance,
-                **dict(maxiter=self.settings.maxiter),
-            )
-        except AttributeError:
-            return pylops.normal_equations_inversion(
-                Op=Op,
-                Regs=None,
-                epsNRs=[1.0],
-                y=self.data.ordered_1d,
-                Weight=pylops.Diagonal(diag=self.noise_map.weight_list_ordered_1d),
-                NRegs=[
-                    pylops.MatrixMult(sparse.bsr_matrix(self.regularization_matrix))
-                ],
-                M=MOp,
-                tol=self.settings.tolerance,
-                atol=self.settings.tolerance,
-                **dict(maxiter=self.settings.maxiter),
-            )[0]
-
-    @property
-    @profile_func
-    def mapped_reconstructed_data_dict(
-        self,
-    ) -> Dict[LinearObj, Visibilities]:
-        """
-        When constructing the simultaneous linear equations (via vectors and matrices) the quantities of each individual
-        linear object (e.g. their `mapping_matrix`) are combined into single ndarrays. This does not track which
-        quantities belong to which linear objects, therefore the linear equation's solutions (which are returned as
-        ndarrays) do not contain information on which linear object(s) they correspond to.
-
-        For example, consider if two `Mapper` objects with 50 and 100 source pixels are used in an `Inversion`.
-        The `reconstruction` (which contains the solved for source pixels values) is an ndarray of shape [150], but
-        the ndarray itself does not track which values belong to which `Mapper`.
-
-        This function converts an ndarray of a `reconstruction` to a dictionary of ndarrays containing each linear
-        object's reconstructed images, where the keys are the instances of each mapper in the inversion.
-
-        The PyLops calculation bypasses the calculation of the `mapping_matrix` and it therefore cannot be used to map
-        the reconstruction's values to the image-plane. Instead, the unique data-to-pixelization mappings are used,
-        including the 2D non-uniform fast Fourier transform operation after mapping is complete.
-
-        Parameters
-        ----------
-        reconstruction
-            The reconstruction (in the source frame) whose values are mapped to a dictionary of values for each
-            individual mapper (in the image-plane).
-        """
-
-        mapped_reconstructed_image_dict = self.mapped_reconstructed_image_dict
-
-        return {
-            linear_obj: self.transformer.visibilities_from(image=image)
-            for linear_obj, image in mapped_reconstructed_image_dict.items()
-        }
-
-    @cached_property
-    @profile_func
-    def preconditioner_matrix(self):
-        curvature_matrix_approx = np.multiply(
-            np.sum(self.noise_map.weight_list_ordered_1d),
-            self.linear_obj_list[0].mapping_matrix.T
-            @ self.linear_obj_list[0].mapping_matrix,
-        )
-
-        return np.add(curvature_matrix_approx, self.regularization_matrix)
-
-    @cached_property
-    @profile_func
-    def preconditioner_matrix_inverse(self):
-        return np.linalg.inv(self.preconditioner_matrix)
-
-    @cached_property
-    @profile_func
-    def log_det_curvature_reg_matrix_term(self):
-        return 2.0 * np.sum(
-            np.log(np.diag(np.linalg.cholesky(self.preconditioner_matrix)))
-        )
-
-    @property
-    def reconstruction_noise_map(self):
-        return None
diff --git a/autoarray/inversion/inversion/interferometer/w_tilde.py b/autoarray/inversion/inversion/interferometer/w_tilde.py
index 18d7e1cec..83febb864 100644
--- a/autoarray/inversion/inversion/interferometer/w_tilde.py
+++ b/autoarray/inversion/inversion/interferometer/w_tilde.py
@@ -131,7 +131,9 @@ def curvature_matrix_diag(self) -> np.ndarray:
                 pix_indexes_for_sub_slim_index=mapper.pix_indexes_for_sub_slim_index,
                 pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index,
                 pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index,
-                native_index_for_slim_index=self.transformer.real_space_mask.derive_indexes.native_for_slim,
+                native_index_for_slim_index=np.array(
+                    self.transformer.real_space_mask.derive_indexes.native_for_slim
+                ).astype("int"),
                 pix_pixels=self.linear_obj_list[0].params,
             )
 
@@ -143,7 +145,9 @@ def curvature_matrix_diag(self) -> np.ndarray:
 
         return inversion_interferometer_util.curvature_matrix_via_w_tilde_curvature_preload_interferometer_from_2(
             curvature_preload=self.w_tilde.curvature_preload,
-            native_index_for_slim_index=self.transformer.real_space_mask.derive_indexes.native_for_slim,
+            native_index_for_slim_index=np.array(
+                self.transformer.real_space_mask.derive_indexes.native_for_slim
+            ).astype("int"),
             pix_pixels=self.linear_obj_list[0].params,
             sub_slim_indexes_for_pix_index=sub_slim_indexes_for_pix_index.astype("int"),
             sub_slim_sizes_for_pix_index=sub_slim_sizes_for_pix_index.astype("int"),
diff --git a/autoarray/inversion/mock/mock_inversion_imaging.py b/autoarray/inversion/mock/mock_inversion_imaging.py
index de7f2baa1..601c669b9 100644
--- a/autoarray/inversion/mock/mock_inversion_imaging.py
+++ b/autoarray/inversion/mock/mock_inversion_imaging.py
@@ -13,6 +13,7 @@ def __init__(
         data=None,
         noise_map=None,
         psf=None,
+        convolver=None,
         linear_obj_list=None,
         operated_mapping_matrix=None,
         linear_func_operated_mapping_matrix_dict=None,
@@ -23,6 +24,7 @@ def __init__(
             data=data,
             noise_map=noise_map,
             psf=psf,
+            convolver=convolver,
         )
 
         super().__init__(
diff --git a/autoarray/inversion/pixelization/border_relocator.py b/autoarray/inversion/pixelization/border_relocator.py
index 737444b53..88ec78091 100644
--- a/autoarray/inversion/pixelization/border_relocator.py
+++ b/autoarray/inversion/pixelization/border_relocator.py
@@ -269,13 +269,13 @@ def relocated_grid_from(self, grid: Grid2D) -> Grid2D:
             return grid
 
         values = grid_2d_util.relocated_grid_via_jit_from(
-            grid=np.array(grid),
-            border_grid=np.array(grid[self.border_slim]),
+            grid=np.array(grid.array),
+            border_grid=np.array(grid.array[self.border_slim]),
         )
 
         over_sampled = grid_2d_util.relocated_grid_via_jit_from(
-            grid=np.array(grid.over_sampled),
-            border_grid=np.array(grid.over_sampled[self.sub_border_slim]),
+            grid=np.array(grid.over_sampled.array),
+            border_grid=np.array(grid.over_sampled.array[self.sub_border_slim]),
         )
 
         return Grid2D(
@@ -303,7 +303,7 @@ def relocated_mesh_grid_from(
 
         return Grid2DIrregular(
             values=grid_2d_util.relocated_grid_via_jit_from(
-                grid=np.array(mesh_grid),
+                grid=np.array(mesh_grid.array),
                 border_grid=np.array(grid[self.sub_border_slim]),
             ),
         )
diff --git a/autoarray/inversion/pixelization/image_mesh/overlay.py b/autoarray/inversion/pixelization/image_mesh/overlay.py
index c5bc7eaef..654cf8ec7 100644
--- a/autoarray/inversion/pixelization/image_mesh/overlay.py
+++ b/autoarray/inversion/pixelization/image_mesh/overlay.py
@@ -220,10 +220,12 @@ def image_plane_mesh_grid_from(
             origin=origin,
         )
 
-        overlaid_centres = geometry_util.grid_pixel_centres_2d_slim_from(
-            grid_scaled_2d_slim=unmasked_overlay_grid,
-            shape_native=mask.shape_native,
-            pixel_scales=mask.pixel_scales,
+        overlaid_centres = np.array(
+            geometry_util.grid_pixel_centres_2d_slim_from(
+                grid_scaled_2d_slim=unmasked_overlay_grid,
+                shape_native=mask.shape_native,
+                pixel_scales=mask.pixel_scales,
+            )
         ).astype("int")
 
         total_pixels = total_pixels_2d_from(
diff --git a/autoarray/inversion/pixelization/mappers/mapper_grids.py b/autoarray/inversion/pixelization/mappers/mapper_grids.py
index 9a12e2f95..260266c5c 100644
--- a/autoarray/inversion/pixelization/mappers/mapper_grids.py
+++ b/autoarray/inversion/pixelization/mappers/mapper_grids.py
@@ -72,6 +72,7 @@ def image_plane_data_grid(self):
 
     @property
     def mesh_pixels_per_image_pixels(self):
+
         mesh_pixels_per_image_pixels = grid_2d_util.grid_pixels_in_mask_pixels_from(
             grid=np.array(self.image_plane_mesh_grid),
             shape_native=self.mask.shape_native,
diff --git a/autoarray/inversion/regularization/abstract.py b/autoarray/inversion/regularization/abstract.py
index 47cadc302..838eaf942 100644
--- a/autoarray/inversion/regularization/abstract.py
+++ b/autoarray/inversion/regularization/abstract.py
@@ -5,13 +5,6 @@
 if TYPE_CHECKING:
     from autoarray.inversion.linear_obj.linear_obj import LinearObj
 
-try:
-    import pylops
-
-    PyLopsOperator = pylops.LinearOperator
-except ModuleNotFoundError:
-    PyLopsOperator = object
-
 
 class AbstractRegularization:
     def __init__(self):
@@ -174,19 +167,3 @@ def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray:
         The regularization matrix.
         """
         raise NotImplementedError
-
-
-class RegularizationLop(PyLopsOperator):
-    def __init__(self, regularization_matrix):
-        self.regularization_matrix = regularization_matrix
-        self.pixels = regularization_matrix.shape[0]
-        self.dims = self.pixels
-        self.shape = (self.pixels, self.pixels)
-        self.dtype = dtype
-        self.explicit = False
-
-    def _matvec(self, x):
-        return np.dot(self.regularization_matrix, x)
-
-    def _rmatvec(self, x):
-        return np.dot(self.regularization_matrix.T, x)
diff --git a/autoarray/mask/derive/grid_2d.py b/autoarray/mask/derive/grid_2d.py
index 225c9a43b..068e40a1b 100644
--- a/autoarray/mask/derive/grid_2d.py
+++ b/autoarray/mask/derive/grid_2d.py
@@ -112,7 +112,7 @@ def all_false(self) -> Grid2D:
             origin=self.mask.origin,
         )
 
-        return Grid2D(values=grid_slim, mask=self.mask.derive_mask.all_false)
+        return Grid2D(values=np.array(grid_slim), mask=self.mask.derive_mask.all_false)
 
     @property
     def unmasked(self) -> Grid2D:
@@ -158,12 +158,12 @@ def unmasked(self) -> Grid2D:
         """
         from autoarray.structures.grids.uniform_2d import Grid2D
 
-        grid_1d = grid_2d_util.grid_2d_slim_via_mask_from(
-            mask_2d=np.array(self.mask),
+        grid_2d = grid_2d_util.grid_2d_slim_via_mask_from(
+            mask_2d=self.mask,
             pixel_scales=self.mask.pixel_scales,
             origin=self.mask.origin,
         )
-        return Grid2D(values=grid_1d, mask=self.mask)
+        return Grid2D(values=np.array(grid_2d), mask=self.mask)
 
     @property
     def edge(self) -> Grid2D:
diff --git a/autoarray/mask/derive/indexes_2d.py b/autoarray/mask/derive/indexes_2d.py
index 062c8e664..0d0a36b26 100644
--- a/autoarray/mask/derive/indexes_2d.py
+++ b/autoarray/mask/derive/indexes_2d.py
@@ -1,7 +1,8 @@
 from __future__ import annotations
 import logging
+import numpy as np
 
-from autoarray.numpy_wrapper import np, register_pytree_node_class
+from autoarray.numpy_wrapper import register_pytree_node_class
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
diff --git a/autoarray/mask/mask_1d.py b/autoarray/mask/mask_1d.py
index 8c36d8866..77f79c01e 100644
--- a/autoarray/mask/mask_1d.py
+++ b/autoarray/mask/mask_1d.py
@@ -59,7 +59,7 @@ def __init__(
             mask = np.asarray(mask).astype("bool")
 
         if invert:
-            mask = np.invert(mask)
+            mask = ~mask
 
         if type(pixel_scales) is float:
             pixel_scales = (pixel_scales,)
@@ -153,7 +153,9 @@ def from_fits(
         """
 
         return cls(
-            array_1d_util.numpy_array_1d_via_fits_from(file_path=file_path, hdu=hdu),
+            mask=np.array(
+                array_1d_util.numpy_array_1d_via_fits_from(file_path=file_path, hdu=hdu)
+            ),
             pixel_scales=pixel_scales,
             origin=origin,
         )
diff --git a/autoarray/mask/mask_1d_util.py b/autoarray/mask/mask_1d_util.py
index 3d9943c19..cfc9e7dff 100644
--- a/autoarray/mask/mask_1d_util.py
+++ b/autoarray/mask/mask_1d_util.py
@@ -1,43 +1,7 @@
+import jax.numpy as jnp
 import numpy as np
 
-from autoarray import numba_util
 
-
-@numba_util.jit()
-def total_pixels_1d_from(mask_1d: np.ndarray) -> int:
-    """
-    Returns the total number of unmasked pixels in a mask.
-
-    Parameters
-    ----------
-    mask_1d
-        A 2D array of bools, where `False` values are unmasked and included when counting pixels.
-
-    Returns
-    -------
-    int
-        The total number of pixels that are unmasked.
-
-    Examples
-    --------
-
-    mask = np.array([[True, False, True],
-                 [False, False, False]
-                 [True, False, True]])
-
-    total_regular_pixels = total_regular_pixels_from(mask=mask)
-    """
-
-    total_regular_pixels = 0
-
-    for x in range(mask_1d.shape[0]):
-        if not mask_1d[x]:
-            total_regular_pixels += 1
-
-    return total_regular_pixels
-
-
-@numba_util.jit()
 def native_index_for_slim_index_1d_from(
     mask_1d: np.ndarray,
 ) -> np.ndarray:
@@ -70,14 +34,5 @@ def native_index_for_slim_index_1d_from(
 
     """
 
-    total_pixels = total_pixels_1d_from(mask_1d=mask_1d)
-    native_index_for_slim_index_1d = np.zeros(shape=total_pixels)
-
-    slim_index = 0
-
-    for x in range(mask_1d.shape[0]):
-        if not mask_1d[x]:
-            native_index_for_slim_index_1d[slim_index] = x
-            slim_index += 1
-
+    native_index_for_slim_index_1d = jnp.flatnonzero(~mask_1d)
     return native_index_for_slim_index_1d
diff --git a/autoarray/mask/mask_2d.py b/autoarray/mask/mask_2d.py
index 05bd77852..5c65f8972 100644
--- a/autoarray/mask/mask_2d.py
+++ b/autoarray/mask/mask_2d.py
@@ -200,11 +200,8 @@ def __init__(
         if type(mask) is list:
             mask = np.asarray(mask).astype("bool")
 
-        if not isinstance(mask, np.ndarray):
-            mask = mask._array
-
         if invert:
-            mask = np.invert(mask)
+            mask = ~mask
 
         pixel_scales = geometry_util.convert_pixel_scales_2d(pixel_scales=pixel_scales)
 
@@ -327,7 +324,7 @@ def circular(
         )
 
         return cls(
-            mask=mask,
+            mask=np.array(mask),
             pixel_scales=pixel_scales,
             origin=origin,
             invert=invert,
@@ -381,7 +378,7 @@ def circular_annular(
         )
 
         return cls(
-            mask=mask,
+            mask=np.array(mask),
             pixel_scales=pixel_scales,
             origin=origin,
             invert=invert,
@@ -438,7 +435,7 @@ def elliptical(
         )
 
         return cls(
-            mask=mask,
+            mask=np.array(mask),
             pixel_scales=pixel_scales,
             origin=origin,
             invert=invert,
@@ -508,7 +505,7 @@ def elliptical_annular(
         )
 
         return cls(
-            mask=mask,
+            mask=np.array(mask),
             pixel_scales=pixel_scales,
             origin=origin,
             invert=invert,
@@ -556,7 +553,7 @@ def from_pixel_coordinates(
         )
 
         return cls(
-            mask=mask,
+            mask=np.array(mask),
             pixel_scales=pixel_scales,
             origin=origin,
             invert=invert,
@@ -594,7 +591,7 @@ def from_fits(
             mask = np.invert(mask.astype("bool"))
 
         mask = Mask2D(
-            mask=mask,
+            mask=np.array(mask),
             pixel_scales=pixel_scales,
             origin=origin,
         )
@@ -743,7 +740,7 @@ def rescaled_from(self, rescale_factor) -> Mask2D:
         from autoarray.mask.mask_2d import Mask2D
 
         rescaled_mask = mask_2d_util.rescaled_mask_2d_from(
-            mask_2d=np.array(self),
+            mask_2d=self.array,
             rescale_factor=rescale_factor,
         )
 
@@ -802,7 +799,7 @@ def resized_from(self, new_shape, pad_value: int = 0.0) -> Mask2D:
         """
 
         resized_mask = array_2d_util.resized_array_2d_from(
-            array_2d=np.array(self._array),
+            array_2d=self.array,
             resized_shape=new_shape,
             pad_value=pad_value,
         ).astype("bool")
diff --git a/autoarray/mask/mask_2d_util.py b/autoarray/mask/mask_2d_util.py
index 462448073..10a40b473 100644
--- a/autoarray/mask/mask_2d_util.py
+++ b/autoarray/mask/mask_2d_util.py
@@ -1,10 +1,10 @@
 import numpy as np
+import jax.numpy as jnp
 from scipy.ndimage import convolve
 from typing import Tuple
 import warnings
 
 from autoarray import exc
-from autoarray.numpy_wrapper import np as jnp
 
 
 def native_index_for_slim_index_2d_from(
diff --git a/autoarray/mock.py b/autoarray/mock.py
index 2261ba5e2..b3bc02c15 100644
--- a/autoarray/mock.py
+++ b/autoarray/mock.py
@@ -15,6 +15,8 @@
 from autoarray.fit.mock.mock_fit_imaging import MockFitImaging
 from autoarray.fit.mock.mock_fit_interferometer import MockFitInterferometer
 from autoarray.mask.mock.mock_mask import MockMask
+from autoarray.operators.mock.mock_psf import MockConvolver
+from autoarray.operators.mock.mock_psf import MockPSF
 from autoarray.structures.mock.mock_grid import MockGrid2DMesh
 from autoarray.structures.mock.mock_grid import MockMeshGrid
 from autoarray.structures.mock.mock_decorators import MockGridRadialMinimum
diff --git a/autoarray/numba_util.py b/autoarray/numba_util.py
index db34f3e1a..9e0298b73 100644
--- a/autoarray/numba_util.py
+++ b/autoarray/numba_util.py
@@ -33,7 +33,7 @@
 
 try:
     if os.environ.get("USE_JAX") == "1":
-        logger.warning("JAX and numba do not work together, so JAX is being used.")
+        1
     else:
         import numba
 
diff --git a/autoarray/operators/contour.py b/autoarray/operators/contour.py
index c7da5c7f1..abc9c7a25 100644
--- a/autoarray/operators/contour.py
+++ b/autoarray/operators/contour.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
-from autoarray.numpy_wrapper import np, use_jax
-import numpy
+import numpy as np
+import jax.numpy as jnp
 from skimage import measure
 from scipy.spatial import ConvexHull
 from scipy.spatial import QhullError
@@ -47,20 +47,23 @@ def contour_array(self):
             pixel_scales=self.pixel_scales,
         ).astype("int")
 
-        arr = np.zeros(self.shape_native)
-        if use_jax:
-            arr = arr.at[tuple(np.array(pixel_centres).T)].set(1)
-        else:
-            arr[tuple(np.array(pixel_centres).T)] = 1
+        arr = jnp.zeros(self.shape_native)
+        arr = arr.at[tuple(jnp.array(pixel_centres).T)].set(1)
 
         return arr
 
     @property
     def contour_list(self):
         # make sure to use base numpy to convert JAX array back to a normal array
-        contour_indices_list = measure.find_contours(
-            numpy.array(self.contour_array.array), 0
-        )
+
+        if isinstance(self.contour_array, jnp.ndarray):
+            contour_array = np.array(self.contour_array)
+        else:
+            contour_array = np.array(self.contour_array.array)
+
+        contour_indices_list = measure.find_contours(contour_array, 0)
+
+        print(contour_indices_list)
 
         if len(contour_indices_list) == 0:
             return []
@@ -89,10 +92,10 @@ def hull(
             return None
 
         # cast JAX arrays to base numpy arrays
-        grid_convex = numpy.zeros((len(self.grid), 2))
+        grid_convex = np.zeros((len(self.grid), 2))
 
-        grid_convex[:, 0] = numpy.array(self.grid[:, 1])
-        grid_convex[:, 1] = numpy.array(self.grid[:, 0])
+        grid_convex[:, 0] = np.array(self.grid[:, 1])
+        grid_convex[:, 1] = np.array(self.grid[:, 0])
 
         try:
             hull = ConvexHull(grid_convex)
@@ -104,9 +107,6 @@ def hull(
         hull_x = grid_convex[hull_vertices, 0]
         hull_y = grid_convex[hull_vertices, 1]
 
-        grid_hull = np.zeros((len(hull_vertices), 2))
-
-        grid_hull[:, 1] = hull_x
-        grid_hull[:, 0] = hull_y
+        grid_hull = jnp.stack((hull_y, hull_x), axis=-1)
 
         return grid_hull
diff --git a/autoarray/operators/mock/mock_convolver.py b/autoarray/operators/mock/mock_convolver.py
deleted file mode 100644
index 7dc5dfcbb..000000000
--- a/autoarray/operators/mock/mock_convolver.py
+++ /dev/null
@@ -1,6 +0,0 @@
-class MockPSF:
-    def __init__(self, operated_mapping_matrix=None):
-        self.operated_mapping_matrix = operated_mapping_matrix
-
-    def convolve_mapping_matrix(self, mapping_matrix):
-        return self.operated_mapping_matrix
diff --git a/autoarray/operators/mock/mock_psf.py b/autoarray/operators/mock/mock_psf.py
new file mode 100644
index 000000000..08864d236
--- /dev/null
+++ b/autoarray/operators/mock/mock_psf.py
@@ -0,0 +1,14 @@
+class MockPSF:
+    def __init__(self, operated_mapping_matrix=None):
+        self.operated_mapping_matrix = operated_mapping_matrix
+
+    def convolve_mapping_matrix(self, mapping_matrix):
+        return self.operated_mapping_matrix
+
+
+class MockConvolver:
+    def __init__(self, operated_mapping_matrix=None):
+        self.operated_mapping_matrix = operated_mapping_matrix
+
+    def convolve_mapping_matrix(self, mapping_matrix):
+        return self.operated_mapping_matrix
diff --git a/autoarray/operators/over_sampling/over_sample_util.py b/autoarray/operators/over_sampling/over_sample_util.py
index 084d0bd0d..d3488a457 100644
--- a/autoarray/operators/over_sampling/over_sample_util.py
+++ b/autoarray/operators/over_sampling/over_sample_util.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import numpy as np
-from typing import TYPE_CHECKING, Union, List, Tuple
+from typing import TYPE_CHECKING, Union
 from typing import List, Tuple
 
 from autoarray.structures.arrays.uniform_2d import Array2D
@@ -12,7 +12,6 @@
 from autoarray.mask.mask_2d import Mask2D
 
 from autoarray import numba_util
-from autoarray.mask import mask_2d_util
 
 from autoarray import type as ty
 
@@ -45,9 +44,9 @@ def over_sample_size_convert_to_array_2d_from(
     if isinstance(over_sample_size, int):
         over_sample_size = np.full(
             fill_value=over_sample_size, shape=mask.pixels_in_mask
-        ).astype("int")
+        )
 
-    return Array2D(values=over_sample_size, mask=mask)
+    return Array2D(values=np.array(over_sample_size).astype("int"), mask=mask)
 
 
 @numba_util.jit()
@@ -445,23 +444,6 @@ def grid_2d_slim_over_sampled_via_mask_from(
 
                 for y1 in range(sub):
                     for x1 in range(sub):
-                        # if use_jax:
-                        #     # while this makes it run, it is very, very slow
-                        #     grid_slim = grid_slim.at[sub_index, 0].set(
-                        #         -(
-                        #             y_scaled
-                        #             - y_sub_half
-                        #             + y1 * y_sub_step
-                        #             + (y_sub_step / 2.0)
-                        #         )
-                        #     )
-                        #     grid_slim = grid_slim.at[sub_index, 1].set(
-                        #         x_scaled
-                        #         - x_sub_half
-                        #         + x1 * x_sub_step
-                        #         + (x_sub_step / 2.0)
-                        #     )
-                        # else:
                         grid_slim[sub_index, 0] = -(
                             y_scaled - y_sub_half + y1 * y_sub_step + (y_sub_step / 2.0)
                         )
@@ -601,7 +583,7 @@ def over_sample_size_via_radial_bins_from(
         radial_grid = grid.distances_to_coordinate_from(coordinate=centre)
 
         sub_size_of_centre = sub_size_radial_bins_from(
-            radial_grid=np.array(radial_grid),
+            radial_grid=np.array(radial_grid.array),
             sub_size_list=np.array(sub_size_list),
             radial_list=np.array(radial_list),
         )
diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py
index 6f12f4b9f..b78bb0829 100644
--- a/autoarray/operators/over_sampling/over_sampler.py
+++ b/autoarray/operators/over_sampling/over_sampler.py
@@ -1,5 +1,6 @@
-from autoarray.numpy_wrapper import np
-from typing import List, Tuple, Union
+import numpy as np
+import jax.numpy as jnp
+from typing import Union
 
 from autoconf import conf
 from autoconf import cached_property
@@ -9,7 +10,7 @@
 
 from autoarray.operators.over_sampling import over_sample_util
 
-from autofit.jax_wrapper import register_pytree_node_class
+from autoarray.numpy_wrapper import register_pytree_node_class
 
 
 @register_pytree_node_class
@@ -146,6 +147,13 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]):
             over_sample_size=sub_size, mask=mask
         )
 
+    @property
+    def sub_is_uniform(self) -> bool:
+        """
+        Returns True if the sub_size is uniform across all pixels in the mask.
+        """
+        return np.all(np.isclose(self.sub_size, self.sub_size[0]))
+
     def tree_flatten(self):
         return (self.mask, self.sub_size), ()
 
@@ -216,19 +224,32 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D":
             return self
 
         try:
-            array = array.slim
+            array = array.slim.array
         except AttributeError:
             pass
 
-        # binned_array_2d = over_sample_util.binned_array_2d_from(
-        #     array_2d=np.array(array),
-        #     mask_2d=np.array(self.mask),
-        #     sub_size=np.array(self.sub_size).astype("int"),
-        # )
+        if self.sub_is_uniform:
+            binned_array_2d = array.reshape(
+                self.mask.shape_slim, self.sub_size[0] ** 2
+            ).mean(axis=1)
+        else:
 
-        binned_array_2d = array.reshape(
-            self.mask.shape_slim, self.sub_size[0] ** 2
-        ).mean(axis=1)
+            # Define group sizes
+            group_sizes = jnp.array(self.sub_size.array**2)
+
+            # Compute the cumulative sum of group sizes to get split points
+            split_indices = jnp.cumsum(group_sizes)
+
+            # Ensure correct concatenation by making 0 a JAX array
+            start_indices = jnp.concatenate((jnp.array([0]), split_indices[:-1]))
+
+            # Compute the group means
+            binned_array_2d = jnp.array(
+                [
+                    array[start:end].mean()
+                    for start, end in zip(start_indices, split_indices)
+                ]
+            )
 
         return Array2D(
             values=binned_array_2d,
@@ -290,7 +311,7 @@ def sub_mask_native_for_sub_mask_slim(self) -> np.ndarray:
             print(derive_indexes_2d.sub_mask_native_for_sub_mask_slim)
         """
         return over_sample_util.native_sub_index_for_slim_sub_index_2d_from(
-            mask_2d=self.mask.array, sub_size=np.array(self.sub_size)
+            mask_2d=self.mask.array, sub_size=self.sub_size.array
         ).astype("int")
 
     @cached_property
@@ -343,7 +364,7 @@ def slim_for_sub_slim(self) -> np.ndarray:
             print(derive_indexes_2d.slim_for_sub_slim)
         """
         return over_sample_util.slim_index_for_sub_slim_index_via_mask_2d_from(
-            mask_2d=np.array(self.mask), sub_size=np.array(self.sub_size)
+            mask_2d=np.array(self.mask), sub_size=self.sub_size.array
         ).astype("int")
 
     @property
@@ -366,7 +387,7 @@ def uniform_over_sampled(self):
         grid = over_sample_util.grid_2d_slim_over_sampled_via_mask_from(
             mask_2d=np.array(self.mask),
             pixel_scales=self.mask.pixel_scales,
-            sub_size=np.array(self.sub_size).astype("int"),
+            sub_size=self.sub_size.array.astype("int"),
             origin=self.mask.origin,
         )
 
diff --git a/autoarray/operators/transformer.py b/autoarray/operators/transformer.py
index acbe6bb73..c5b1d3a50 100644
--- a/autoarray/operators/transformer.py
+++ b/autoarray/operators/transformer.py
@@ -8,21 +8,11 @@ class NUFFTPlaceholder:
     pass
 
 
-class PyLopsPlaceholder:
-    pass
-
-
 try:
     from pynufft.linalg.nufft_cpu import NUFFT_cpu
 except ModuleNotFoundError:
     NUFFT_cpu = NUFFTPlaceholder
 
-try:
-    import pylops
-
-    PyLopsOperator = pylops.LinearOperator
-except ModuleNotFoundError:
-    PyLopsOperator = PyLopsPlaceholder
 
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.grids.uniform_2d import Grid2D
@@ -42,20 +32,8 @@ def pynufft_exception():
     )
 
 
-def pylops_exception():
-    raise ModuleNotFoundError(
-        "\n--------------------\n"
-        "You are attempting to perform interferometer analysis.\n\n"
-        "However, the optional library PyLops (https://github.com/PyLops/pylops) is not installed.\n\n"
-        "Install it via the command `pip install pylops==2.3.1`.\n\n"
-        "----------------------"
-    )
-
-
-class TransformerDFT(PyLopsOperator):
+class TransformerDFT:
     def __init__(self, uv_wavelengths, real_space_mask, preload_transform=True):
-        if isinstance(self, PyLopsPlaceholder):
-            pylops_exception()
 
         super().__init__()
 
@@ -70,12 +48,12 @@ def __init__(self, uv_wavelengths, real_space_mask, preload_transform=True):
 
         if preload_transform:
             self.preload_real_transforms = transformer_util.preload_real_transforms(
-                grid_radians=np.array(self.grid),
+                grid_radians=np.array(self.grid.array),
                 uv_wavelengths=self.uv_wavelengths,
             )
 
             self.preload_imag_transforms = transformer_util.preload_imag_transforms(
-                grid_radians=np.array(self.grid),
+                grid_radians=np.array(self.grid.array),
                 uv_wavelengths=self.uv_wavelengths,
             )
 
@@ -101,14 +79,14 @@ def __init__(self, uv_wavelengths, real_space_mask, preload_transform=True):
     def visibilities_from(self, image):
         if self.preload_transform:
             visibilities = transformer_util.visibilities_via_preload_jit_from(
-                image_1d=np.array(image),
+                image_1d=np.array(image.array),
                 preloaded_reals=self.preload_real_transforms,
                 preloaded_imags=self.preload_imag_transforms,
             )
 
         else:
             visibilities = transformer_util.visibilities_jit(
-                image_1d=np.array(image.slim),
+                image_1d=np.array(image.slim.array),
                 grid_radians=np.array(self.grid),
                 uv_wavelengths=self.uv_wavelengths,
             )
@@ -118,7 +96,7 @@ def visibilities_from(self, image):
     def image_from(self, visibilities, use_adjoint_scaling: bool = False):
         image_slim = transformer_util.image_via_jit_from(
             n_pixels=self.grid.shape[0],
-            grid_radians=np.array(self.grid),
+            grid_radians=np.array(self.grid.array),
             uv_wavelengths=self.uv_wavelengths,
             visibilities=visibilities.in_array,
         )
@@ -146,14 +124,11 @@ def transform_mapping_matrix(self, mapping_matrix):
             )
 
 
-class TransformerNUFFT(NUFFT_cpu, PyLopsOperator):
+class TransformerNUFFT(NUFFT_cpu):
     def __init__(self, uv_wavelengths, real_space_mask):
         if isinstance(self, NUFFTPlaceholder):
             pynufft_exception()
 
-        if isinstance(self, PyLopsPlaceholder):
-            pylops_exception()
-
         super(TransformerNUFFT, self).__init__()
 
         self.uv_wavelengths = uv_wavelengths
diff --git a/autoarray/plot/mat_plot/one_d.py b/autoarray/plot/mat_plot/one_d.py
index 733f2ff83..79b721bba 100644
--- a/autoarray/plot/mat_plot/one_d.py
+++ b/autoarray/plot/mat_plot/one_d.py
@@ -166,6 +166,17 @@ def plot_yx(
         text_manual_dict_y=None,
         bypass: bool = False,
     ):
+
+        try:
+            y = y.array
+        except AttributeError:
+            pass
+
+        try:
+            x = x.array
+        except AttributeError:
+            pass
+
         if (y is None) or np.count_nonzero(y) == 0 or np.isnan(y).all():
             return
 
@@ -185,7 +196,7 @@ def plot_yx(
             x = np.arange(len(y))
             use_integers = True
             pixel_scales = (x[1] - x[0],)
-            x = Array1D.no_mask(values=x, pixel_scales=pixel_scales)
+            x = Array1D.no_mask(values=x, pixel_scales=pixel_scales).array
 
         if self.yx_plot.plot_axis_type is None:
             plot_axis_type = "linear"
diff --git a/autoarray/plot/mat_plot/two_d.py b/autoarray/plot/mat_plot/two_d.py
index 9fcec3657..9b162d18a 100644
--- a/autoarray/plot/mat_plot/two_d.py
+++ b/autoarray/plot/mat_plot/two_d.py
@@ -313,7 +313,7 @@ def plot_array(
 
         aspect = self.figure.aspect_from(shape_native=array.shape_native)
 
-        norm = self.cmap.norm_from(array=array, use_log10=self.use_log10)
+        norm = self.cmap.norm_from(array=array.array, use_log10=self.use_log10)
 
         origin = conf.instance["visualize"]["general"]["general"]["imshow_origin"]
 
@@ -427,10 +427,10 @@ def plot_grid(
 
         if color_array is None:
             if y_errors is None and x_errors is None:
-                self.grid_scatter.scatter_grid(grid=grid_plot)
+                self.grid_scatter.scatter_grid(grid=grid_plot.array)
             else:
                 self.grid_errorbar.errorbar_grid(
-                    grid=grid_plot, y_errors=y_errors, x_errors=x_errors
+                    grid=grid_plot.array, y_errors=y_errors, x_errors=x_errors
                 )
 
         elif color_array is not None:
@@ -438,11 +438,11 @@ def plot_grid(
 
             if y_errors is None and x_errors is None:
                 self.grid_scatter.scatter_grid_colored(
-                    grid=grid, color_array=color_array, cmap=cmap
+                    grid=grid.array, color_array=color_array, cmap=cmap
                 )
             else:
                 self.grid_errorbar.errorbar_grid_colored(
-                    grid=grid,
+                    grid=grid.array,
                     cmap=cmap,
                     color_array=color_array,
                     y_errors=y_errors,
@@ -450,6 +450,7 @@ def plot_grid(
                 )
 
             if self.colorbar is not None:
+
                 colorbar = self.colorbar.set_with_color_values(
                     units=self.units,
                     cmap=self.cmap.cmap,
@@ -496,7 +497,7 @@ def plot_grid(
             self.contour.set(array=color_array, extent=extent, use_log10=self.use_log10)
 
         visuals_2d.plot_via_plotter(
-            plotter=self, grid_indexes=grid, geometry=grid.geometry
+            plotter=self, grid_indexes=grid.array, geometry=grid.geometry
         )
 
         if not self.is_for_subplot:
diff --git a/autoarray/plot/visuals/two_d.py b/autoarray/plot/visuals/two_d.py
index d1da534ff..e892a06a7 100644
--- a/autoarray/plot/visuals/two_d.py
+++ b/autoarray/plot/visuals/two_d.py
@@ -1,3 +1,4 @@
+import numpy as np
 from matplotlib import patches as ptch
 from typing import List, Optional, Union
 
@@ -54,16 +55,18 @@ def plot_via_plotter(self, plotter, grid_indexes=None, mapper=None, geometry=Non
             )
 
         if self.mask is not None:
-            plotter.mask_scatter.scatter_grid(grid=self.mask.derive_grid.edge)
+            plotter.mask_scatter.scatter_grid(
+                grid=np.array(self.mask.derive_grid.edge.array)
+            )
 
         if self.border is not None:
-            plotter.border_scatter.scatter_grid(grid=self.border)
+            plotter.border_scatter.scatter_grid(grid=self.border.array)
 
         if self.grid is not None:
-            plotter.grid_scatter.scatter_grid(grid=self.grid)
+            plotter.grid_scatter.scatter_grid(grid=self.grid.array)
 
         if self.mesh_grid is not None:
-            plotter.mesh_grid_scatter.scatter_grid(grid=self.mesh_grid)
+            plotter.mesh_grid_scatter.scatter_grid(grid=self.mesh_grid.array)
 
         if self.positions is not None:
             plotter.positions_scatter.scatter_grid(grid=self.positions)
diff --git a/autoarray/plot/wrap/base/colorbar.py b/autoarray/plot/wrap/base/colorbar.py
index b0650013a..1c9a8c5d8 100644
--- a/autoarray/plot/wrap/base/colorbar.py
+++ b/autoarray/plot/wrap/base/colorbar.py
@@ -194,6 +194,7 @@ def set_with_color_values(
         )
 
         if tick_values is None and tick_labels is None:
+
             cb = plt.colorbar(
                 mappable=mappable,
                 ax=ax,
diff --git a/autoarray/plot/wrap/two_d/array_overlay.py b/autoarray/plot/wrap/two_d/array_overlay.py
index 57652e8df..688866ade 100644
--- a/autoarray/plot/wrap/two_d/array_overlay.py
+++ b/autoarray/plot/wrap/two_d/array_overlay.py
@@ -19,4 +19,6 @@ def overlay_array(self, array, figure):
         aspect = figure.aspect_from(shape_native=array.shape_native)
         extent = array.extent_of_zoomed_array(buffer=0)
 
-        plt.imshow(X=array.native, aspect=aspect, extent=extent, **self.config_dict)
+        plt.imshow(
+            X=array.native._array, aspect=aspect, extent=extent, **self.config_dict
+        )
diff --git a/autoarray/plot/wrap/two_d/contour.py b/autoarray/plot/wrap/two_d/contour.py
index c80159813..164fe86be 100644
--- a/autoarray/plot/wrap/two_d/contour.py
+++ b/autoarray/plot/wrap/two_d/contour.py
@@ -93,10 +93,10 @@ def set(
         config_dict.pop("use_log10")
         config_dict.pop("include_values")
 
-        levels = self.levels_from(array)
+        levels = self.levels_from(array.array)
 
         ax = plt.contour(
-            array.native[::-1], levels=levels, extent=extent, **config_dict
+            array.native.array[::-1], levels=levels, extent=extent, **config_dict
         )
         if self.include_values:
             try:
diff --git a/autoarray/plot/wrap/two_d/grid_scatter.py b/autoarray/plot/wrap/two_d/grid_scatter.py
index 8399ae591..526aeffb9 100644
--- a/autoarray/plot/wrap/two_d/grid_scatter.py
+++ b/autoarray/plot/wrap/two_d/grid_scatter.py
@@ -1,4 +1,5 @@
 import matplotlib.pyplot as plt
+import jax.numpy as jnp
 import numpy as np
 import itertools
 from scipy.spatial import ConvexHull
@@ -8,6 +9,7 @@
 from autoarray.plot.wrap.two_d.abstract import AbstractMatWrap2D
 from autoarray.structures.grids.uniform_2d import Grid2D
 from autoarray.structures.grids.irregular_2d import Grid2DIrregular
+from autoarray.structures.mesh.abstract_2d import Abstract2DMesh
 
 
 class GridScatter(AbstractMatWrap2D):
diff --git a/autoarray/preloads.py b/autoarray/preloads.py
index 6808f0f6c..90c1deae9 100644
--- a/autoarray/preloads.py
+++ b/autoarray/preloads.py
@@ -198,13 +198,6 @@ def set_mapper_list(self, fit_0, fit_1):
         if fit_0.inversion.total(cls=AbstractMapper) == 0:
             return
 
-        from autoarray.inversion.inversion.interferometer.lop import (
-            InversionInterferometerMappingPyLops,
-        )
-
-        if isinstance(fit_0.inversion, InversionInterferometerMappingPyLops):
-            return
-
         inversion_0 = fit_0.inversion
         inversion_1 = fit_1.inversion
 
@@ -238,13 +231,6 @@ def set_operated_mapping_matrix_with_preloads(self, fit_0, fit_1):
 
         self.operated_mapping_matrix = None
 
-        from autoarray.inversion.inversion.interferometer.lop import (
-            InversionInterferometerMappingPyLops,
-        )
-
-        if isinstance(fit_0.inversion, InversionInterferometerMappingPyLops):
-            return
-
         inversion_0 = fit_0.inversion
         inversion_1 = fit_1.inversion
 
diff --git a/autoarray/structures/arrays/array_1d_util.py b/autoarray/structures/arrays/array_1d_util.py
index 5e6565326..915a2a4ad 100644
--- a/autoarray/structures/arrays/array_1d_util.py
+++ b/autoarray/structures/arrays/array_1d_util.py
@@ -1,11 +1,11 @@
 from __future__ import annotations
+import jax.numpy as jnp
 import numpy as np
 from typing import TYPE_CHECKING, List, Union
 
 if TYPE_CHECKING:
     from autoarray.mask.mask_1d import Mask1D
 
-from autoarray import numba_util
 from autoarray.mask import mask_1d_util
 from autoarray.structures.arrays import array_2d_util
 
@@ -38,26 +38,27 @@ def convert_array_1d(
         If True, the ndarray is stored in its native format [total_y_pixels, total_x_pixels]. This avoids
         mapping large data arrays to and from the slim / native formats, which can be a computational bottleneck.
     """
-
     array_1d = array_2d_util.convert_array(array=array_1d)
 
+    is_numpy = True if isinstance(array_1d, np.ndarray) else False
+
     is_native = array_1d.shape[0] == mask_1d.shape_native[0]
 
     if is_native == store_native:
-        return array_1d
+        array_1d = array_1d
     elif not store_native:
-        return array_1d_slim_from(
-            array_1d_native=np.array(array_1d),
-            mask_1d=np.array(mask_1d),
+        array_1d = array_1d_slim_from(
+            array_1d_native=array_1d,
+            mask_1d=mask_1d,
         )
-
-    return array_1d_native_from(
-        array_1d_slim=array_1d,
-        mask_1d=np.array(mask_1d),
-    )
+    else:
+        array_1d = array_1d_native_from(
+            array_1d_slim=array_1d,
+            mask_1d=mask_1d,
+        )
+    return np.array(array_1d) if is_numpy else jnp.array(array_1d)
 
 
-@numba_util.jit()
 def array_1d_slim_from(
     array_1d_native: np.ndarray,
     mask_1d: np.ndarray,
@@ -105,19 +106,8 @@ def array_1d_slim_from(
 
     array_1d_slim = array_1d_slim_from(array_1d_native, array_2d=array_2d)
     """
-
-    total_pixels = mask_1d_util.total_pixels_1d_from(
-        mask_1d=mask_1d,
-    )
-
-    line_1d_slim = np.zeros(shape=total_pixels)
-    index = 0
-
-    for x in range(mask_1d.shape[0]):
-        if not mask_1d[x]:
-            line_1d_slim[index] = array_1d_native[x]
-            index += 1
-
+    unmasked_indices = ~mask_1d
+    line_1d_slim = array_1d_native[unmasked_indices]
     return line_1d_slim
 
 
@@ -132,13 +122,12 @@ def array_1d_native_from(
     ).astype("int")
 
     return array_1d_via_indexes_1d_from(
-        array_1d_slim=np.array(array_1d_slim),
+        array_1d_slim=array_1d_slim,
         shape=shape,
         native_index_for_slim_index_1d=native_index_for_slim_index_1d,
     )
 
 
-@numba_util.jit()
 def array_1d_via_indexes_1d_from(
     array_1d_slim: np.ndarray,
     shape: int,
@@ -177,11 +166,5 @@ def array_1d_via_indexes_1d_from(
     ndarray
         The native 1D array of values mapped from the slimmed array with dimensions (total_x_pixels).
     """
-    array_1d_native = np.zeros(shape)
-
-    for slim_index in range(len(native_index_for_slim_index_1d)):
-        array_1d_native[native_index_for_slim_index_1d[slim_index]] = array_1d_slim[
-            slim_index
-        ]
-
-    return array_1d_native
+    array_1d_native = jnp.zeros(shape)
+    return array_1d_native.at[native_index_for_slim_index_1d].set(array_1d_slim)
diff --git a/autoarray/structures/arrays/array_2d_util.py b/autoarray/structures/arrays/array_2d_util.py
index d8a0f7e10..f8767e2f9 100644
--- a/autoarray/structures/arrays/array_2d_util.py
+++ b/autoarray/structures/arrays/array_2d_util.py
@@ -1,5 +1,4 @@
 from __future__ import annotations
-import jax
 import jax.numpy as jnp
 import numpy as np
 from typing import TYPE_CHECKING, List, Tuple, Union
@@ -11,7 +10,6 @@
 from autoarray.mask import mask_2d_util
 
 from autoarray import exc
-from functools import partial
 
 
 def convert_array(array: Union[np.ndarray, List]) -> np.ndarray:
@@ -23,12 +21,15 @@ def convert_array(array: Union[np.ndarray, List]) -> np.ndarray:
     array : list or ndarray
         The array which may be converted to an ndarray
     """
-    if isinstance(array, np.ndarray) or isinstance(array, list):
+
+    try:
+        array = array.array
+    except AttributeError:
+        pass
+
+    if isinstance(array, list):
         array = np.asarray(array)
-    elif isinstance(array, jnp.ndarray):
-        array = jax.lax.cond(
-            type(array) is list, lambda _: jnp.asarray(array), lambda _: array, None
-        )
+
     return array
 
 
@@ -120,25 +121,28 @@ def convert_array_2d(
     """
     array_2d = convert_array(array=array_2d).copy()
 
+    is_numpy = True if isinstance(array_2d, np.ndarray) else False
+
     check_array_2d_and_mask_2d(array_2d=array_2d, mask_2d=mask_2d)
 
     is_native = len(array_2d.shape) == 2
 
     if is_native and not skip_mask:
-        array_2d *= np.invert(mask_2d)
+        array_2d *= ~mask_2d
 
     if is_native == store_native:
-        return array_2d
+        array_2d = array_2d
     elif not store_native:
-        return array_2d_slim_from(
-            array_2d_native=np.array(array_2d),
-            mask_2d=np.array(mask_2d),
+        array_2d = array_2d_slim_from(
+            array_2d_native=array_2d,
+            mask_2d=mask_2d,
         )
-    array_2d = array_2d_native_from(
-        array_2d_slim=array_2d,
-        mask_2d=np.array(mask_2d),
-    )
-    return array_2d
+    else:
+        array_2d = array_2d_native_from(
+            array_2d_slim=array_2d,
+            mask_2d=mask_2d,
+        )
+    return np.array(array_2d) if is_numpy else jnp.array(array_2d)
 
 
 def convert_array_2d_to_slim(array_2d: np.ndarray, mask_2d: Mask2D) -> np.ndarray:
@@ -587,7 +591,6 @@ def array_2d_native_from(
     )
 
 
-@partial(jax.jit, static_argnums=(1,))
 def array_2d_via_indexes_from(
     array_2d_slim: np.ndarray,
     shape: Tuple[int, int],
diff --git a/autoarray/structures/arrays/irregular.py b/autoarray/structures/arrays/irregular.py
index 667ddfd43..f4b54ee11 100644
--- a/autoarray/structures/arrays/irregular.py
+++ b/autoarray/structures/arrays/irregular.py
@@ -38,12 +38,6 @@ def __init__(self, values: Union[List, np.ndarray]):
             A collection of values.
         """
 
-        # if len(values) == 0:
-        #     return []
-
-        # if isinstance(values, ArrayIrregular):
-        #     return values
-
         if type(values) is list:
             values = np.asarray(values)
 
diff --git a/autoarray/structures/arrays/kernel_2d.py b/autoarray/structures/arrays/kernel_2d.py
index bb62656a9..c2962dbbe 100644
--- a/autoarray/structures/arrays/kernel_2d.py
+++ b/autoarray/structures/arrays/kernel_2d.py
@@ -255,7 +255,7 @@ def from_gaussian(
         """
 
         grid = Grid2D.uniform(shape_native=shape_native, pixel_scales=pixel_scales)
-        grid_shifted = np.subtract(grid, centre)
+        grid_shifted = np.subtract(grid.array, centre)
         grid_radius = np.sqrt(np.sum(grid_shifted**2.0, 1))
         theta_coordinate_to_profile = np.arctan2(
             grid_shifted[:, 0], grid_shifted[:, 1]
@@ -475,7 +475,7 @@ def convolved_array_from(self, array: Array2D) -> Array2D:
         array_2d = array.native
 
         convolved_array_2d = scipy.signal.convolve2d(
-            array_2d._array, np.array(self.native._array), mode="same"
+            array_2d.array, np.array(self.native.array), mode="same"
         )
 
         convolved_array_1d = array_2d_util.array_2d_slim_from(
@@ -485,7 +485,7 @@ def convolved_array_from(self, array: Array2D) -> Array2D:
 
         return Array2D(values=convolved_array_1d, mask=array_2d.mask)
 
-    def convolve_image(self, image, blurring_image, jax_method="fft"):
+    def convolve_image(self, image, blurring_image, jax_method="direct"):
         """
         For a given 1D array and blurring array, convolve the two using this psf.
 
@@ -528,7 +528,7 @@ def convolve_image(self, image, blurring_image, jax_method="fft"):
 
         return Array2D(values=convolved_array_1d, mask=image.mask)
 
-    def convolve_image_no_blurring(self, image, mask, jax_method="fft"):
+    def convolve_image_no_blurring(self, image, mask, jax_method="direct"):
         """
         For a given 1D array and blurring array, convolve the two using this psf.
 
@@ -549,7 +549,12 @@ def convolve_image_no_blurring(self, image, mask, jax_method="fft"):
 
         expanded_array_native = jnp.zeros(mask.shape)
 
-        expanded_array_native = expanded_array_native.at[slim_to_native].set(image)
+        if isinstance(image, np.ndarray) or isinstance(image, jnp.ndarray):
+            expanded_array_native = expanded_array_native.at[slim_to_native].set(image)
+        else:
+            expanded_array_native = expanded_array_native.at[slim_to_native].set(
+                image.array
+            )
 
         kernel = np.array(self.native.array)
 
@@ -561,7 +566,7 @@ def convolve_image_no_blurring(self, image, mask, jax_method="fft"):
 
         return Array2D(values=convolved_array_1d, mask=mask)
 
-    def convolve_mapping_matrix(self, mapping_matrix, mask):
+    def convolve_mapping_matrix(self, mapping_matrix, mask, jax_method="direct"):
         """For a given 1D array and blurring array, convolve the two using this psf.
 
         Parameters
@@ -569,6 +574,6 @@ def convolve_mapping_matrix(self, mapping_matrix, mask):
         image
             1D array of the values which are to be blurred with the psf's PSF.
         """
-        return jax.vmap(self.convolve_image_no_blurring, in_axes=(1, None))(
-            mapping_matrix, mask
+        return jax.vmap(self.convolve_image_no_blurring, in_axes=(1, None, None))(
+            mapping_matrix, mask, jax_method
         ).T
diff --git a/autoarray/structures/arrays/uniform_1d.py b/autoarray/structures/arrays/uniform_1d.py
index 4dbf1692f..d708a68ad 100644
--- a/autoarray/structures/arrays/uniform_1d.py
+++ b/autoarray/structures/arrays/uniform_1d.py
@@ -24,6 +24,7 @@ def __init__(
         header: Optional[Header] = None,
         store_native: bool = False,
     ):
+
         values = array_1d_util.convert_array_1d(
             array_1d=values,
             mask_1d=mask,
@@ -87,7 +88,7 @@ def no_mask(
             origin=origin,
         )
 
-        return Array1D(values=values, mask=mask, header=header)
+        return Array1D(values=np.array(values), mask=mask, header=header)
 
     @classmethod
     def full(
diff --git a/autoarray/structures/arrays/uniform_2d.py b/autoarray/structures/arrays/uniform_2d.py
index 11c478ad5..66c881ff3 100644
--- a/autoarray/structures/arrays/uniform_2d.py
+++ b/autoarray/structures/arrays/uniform_2d.py
@@ -232,11 +232,6 @@ def __init__(
         if conf.instance["general"]["structures"]["native_binned_only"]:
             store_native = True
 
-        try:
-            values = values._array
-        except AttributeError:
-            values = values
-
         values = array_2d_util.convert_array_2d(
             array_2d=values,
             mask_2d=mask,
@@ -464,7 +459,7 @@ def zoomed_around_mask(self, buffer: int = 1) -> "Array2D":
         """
 
         extracted_array_2d = array_2d_util.extracted_array_2d_from(
-            array_2d=np.array(self.native._array),
+            array_2d=np.array(self.native.array),
             y0=self.mask.zoom_region[0] - buffer,
             y1=self.mask.zoom_region[1] + buffer,
             x0=self.mask.zoom_region[2] - buffer,
@@ -498,7 +493,7 @@ def extent_of_zoomed_array(self, buffer: int = 1) -> np.ndarray:
             The number pixels around the extracted array used as a buffer.
         """
         extracted_array_2d = array_2d_util.extracted_array_2d_from(
-            array_2d=np.array(self.native._array),
+            array_2d=np.array(self.native.array),
             y0=self.mask.zoom_region[0] - buffer,
             y1=self.mask.zoom_region[1] + buffer,
             x0=self.mask.zoom_region[2] - buffer,
@@ -532,7 +527,7 @@ def resized_from(
         """
 
         resized_array_2d = array_2d_util.resized_array_2d_from(
-            array_2d=np.array(self.native._array), resized_shape=new_shape
+            array_2d=np.array(self.native.array), resized_shape=new_shape
         )
 
         resized_mask = self.mask.resized_from(
@@ -592,14 +587,14 @@ def trimmed_after_convolution_from(
         psf_cut_x = int(np.ceil(kernel_shape[1] / 2)) - 1
         array_y = int(self.mask.shape[0])
         array_x = int(self.mask.shape[1])
-        trimmed_array_2d = self.native[
+        trimmed_array_2d = self.native.array[
             psf_cut_y : array_y - psf_cut_y, psf_cut_x : array_x - psf_cut_x
         ]
 
         resized_mask = self.mask.resized_from(new_shape=trimmed_array_2d.shape)
 
         array = array_2d_util.convert_array_2d(
-            array_2d=trimmed_array_2d._array, mask_2d=resized_mask
+            array_2d=trimmed_array_2d, mask_2d=resized_mask
         )
 
         return Array2D(
@@ -697,7 +692,7 @@ def no_mask(
             origin=origin,
         )
 
-        return Array2D(values=values, mask=mask, header=header)
+        return Array2D(values=np.array(values), mask=mask, header=header)
 
     @classmethod
     def full(
@@ -962,7 +957,7 @@ def from_yx_and_values(
         )
 
         grid_pixels = geometry_util.grid_pixel_indexes_2d_slim_from(
-            grid_scaled_2d_slim=np.array(grid.slim),
+            grid_scaled_2d_slim=grid.slim,
             shape_native=shape_native,
             pixel_scales=pixel_scales,
         )
diff --git a/autoarray/structures/decorators/__init__.py b/autoarray/structures/decorators/__init__.py
index d85cbe6ee..1efb9137e 100644
--- a/autoarray/structures/decorators/__init__.py
+++ b/autoarray/structures/decorators/__init__.py
@@ -4,4 +4,3 @@
 from .to_grid import to_grid
 from .to_vector_yx import to_vector_yx
 from .transform import transform
-from .relocate_radial import relocate_to_radial_minimum
diff --git a/autoarray/structures/decorators/abstract.py b/autoarray/structures/decorators/abstract.py
index c9e5fca87..e033815b4 100644
--- a/autoarray/structures/decorators/abstract.py
+++ b/autoarray/structures/decorators/abstract.py
@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import Union
 
 import numpy as np
 
diff --git a/autoarray/structures/decorators/relocate_radial.py b/autoarray/structures/decorators/relocate_radial.py
deleted file mode 100644
index 58411714f..000000000
--- a/autoarray/structures/decorators/relocate_radial.py
+++ /dev/null
@@ -1,106 +0,0 @@
-from autoarray.numpy_wrapper import np, use_jax
-from functools import wraps
-
-from typing import Union
-
-from autoconf.exc import ConfigException
-
-from autoarray.structures.grids.irregular_2d import Grid2DIrregular
-from autoarray.structures.grids.uniform_2d import Grid2D
-from autoconf import conf
-
-
-def relocate_to_radial_minimum(func):
-    """
-    Checks whether any coordinates in the grid are radially near (0.0, 0.0), which can lead to numerical faults in
-    the evaluation of a function (e.g. numerical integration reaching a singularity at (0.0, 0.0)).
-
-    If any coordinates are radially within the radial minimum threshold, their (y,x) coordinates are shifted to that
-    value to ensure they are evaluated at that coordinate.
-
-    The value the (y,x) coordinates are rounded to is set in the 'radial_minimum.yaml' config.
-
-    Parameters
-    ----------
-    func
-        A function that takes a grid of coordinates which may have a singularity as (0.0, 0.0)
-
-    Returns
-    -------
-        A function that has an input grid whose radial coordinates are relocated to the radial minimum.
-    """
-
-    @wraps(func)
-    def wrapper(
-        obj: object,
-        grid: Union[np.ndarray, Grid2D, Grid2DIrregular],
-        *args,
-        **kwargs,
-    ) -> Union[np.ndarray, Grid2D, Grid2DIrregular]:
-        """
-        Checks whether any coordinates in the grid are radially near (0.0, 0.0), which can lead to numerical faults in
-        the evaluation of a function (e.g. numerical integration reaching a singularity at (0.0, 0.0)).
-
-        If any coordinates are radially within the radial minimum threshold, their (y,x) coordinates are shifted to that
-        value to ensure they are evaluated at that coordinate.
-
-        The value the (y,x) coordinates are rounded to is set in the 'radial_minimum.yaml' config.
-
-        Parameters
-        ----------
-        obj
-            An object whose function uses grid_like inputs to compute quantities at every coordinate on the grid.
-        grid
-            The (y, x) coordinates which are to be radially moved from (0.0, 0.0).
-
-        Returns
-        -------
-            The grid_like object whose coordinates are radially moved from (0.0, 0.0).
-        """
-        if use_jax:
-            return func(obj, grid, *args, **kwargs)
-
-        try:
-            grid_radial_minimum = conf.instance["grids"]["radial_minimum"][
-                "radial_minimum"
-            ][obj.__class__.__name__]
-
-        except KeyError as e:
-            raise ConfigException(
-                rf"""
-                The {obj.__class__.__name__} profile you are using does not have a corresponding
-                entry in the `config/grid.yaml` config file.
-
-                When a profile is evaluated at (0.0, 0.0), they commonly break due to numericalinstabilities (e.g. 
-                division by zero). To prevent this, the code relocates the (y,x) coordinates of the grid to a 
-                minimum radial value, specified in the `config/grids.yaml` config file.
-
-                For example, if the value in `grid.yaml` is `radial_minimum: 1e-6`, then any (y,x) coordinates
-                with a radial distance less than 1e-6 to (0.0, 0.0) are relocated to 1e-6.
-
-                For a profile to be used it must have an entry in the `config/grids.yaml` config file. Go to this
-                file now and add your profile to the `radial_minimum` section. Adopting a value of 1e-6 is a good
-                default choice.
-
-                If you are going to make a pull request to add your profile to the source code, you should also
-                add an entry to the `config/grids.yaml` config file of the source code itself
-                (e.g. `PyAutoGalaxy/autogalaxy/config/grids.yaml`).
-                """
-            )
-
-        with np.errstate(all="ignore"):  # Division by zero fixed via isnan
-            grid_radii = obj.radial_grid_from(grid=grid)
-
-            grid_radial_scale = np.where(
-                grid_radii < grid_radial_minimum, grid_radial_minimum / grid_radii, 1.0
-            )
-            moved_grid = np.multiply(grid, grid_radial_scale[:, None])
-
-            if hasattr(grid, "with_new_array"):
-                moved_grid = grid.with_new_array(moved_grid)
-
-        moved_grid[np.isnan(np.array(moved_grid))] = grid_radial_minimum
-
-        return func(obj, moved_grid, *args, **kwargs)
-
-    return wrapper
diff --git a/autoarray/structures/decorators/to_grid.py b/autoarray/structures/decorators/to_grid.py
index 144137c69..4797c37ce 100644
--- a/autoarray/structures/decorators/to_grid.py
+++ b/autoarray/structures/decorators/to_grid.py
@@ -1,6 +1,5 @@
-from autoarray.numpy_wrapper import np
 from functools import wraps
-
+import numpy as np
 from typing import List, Union
 
 from autoarray.structures.decorators.abstract import AbstractMaker
diff --git a/autoarray/structures/decorators/to_vector_yx.py b/autoarray/structures/decorators/to_vector_yx.py
index 1cf23346d..90aea99ea 100644
--- a/autoarray/structures/decorators/to_vector_yx.py
+++ b/autoarray/structures/decorators/to_vector_yx.py
@@ -1,6 +1,5 @@
-from autoarray.numpy_wrapper import np
 from functools import wraps
-
+import numpy as np
 from typing import List, Union
 
 from autoarray.structures.decorators.abstract import AbstractMaker
diff --git a/autoarray/structures/decorators/transform.py b/autoarray/structures/decorators/transform.py
index bd837a399..eca0d883b 100644
--- a/autoarray/structures/decorators/transform.py
+++ b/autoarray/structures/decorators/transform.py
@@ -1,5 +1,5 @@
-from autoarray.numpy_wrapper import np
 from functools import wraps
+import numpy as np
 from typing import Union
 
 from autoarray.structures.grids.uniform_1d import Grid1D
diff --git a/autoarray/structures/grids/grid_1d_util.py b/autoarray/structures/grids/grid_1d_util.py
index aa0592c47..82aa4514e 100644
--- a/autoarray/structures/grids/grid_1d_util.py
+++ b/autoarray/structures/grids/grid_1d_util.py
@@ -1,15 +1,14 @@
 from __future__ import annotations
 import numpy as np
+import jax.numpy as jnp
 from typing import TYPE_CHECKING, List, Union, Tuple
 
 if TYPE_CHECKING:
     from autoarray.mask.mask_1d import Mask1D
 
 from autoarray.structures.arrays import array_1d_util
-from autoarray import numba_util
 from autoarray.geometry import geometry_util
 from autoarray.structures.grids import grid_2d_util
-from autoarray.mask import mask_1d_util
 from autoarray import type as ty
 
 
@@ -41,19 +40,23 @@ def convert_grid_1d(
 
     grid_1d = grid_2d_util.convert_grid(grid=grid_1d)
 
+    is_numpy = True if isinstance(grid_1d, np.ndarray) else False
+
     is_native = grid_1d.shape[0] == mask_1d.shape_native[0]
 
     if is_native == store_native:
-        return grid_1d
+        grid_1d = grid_1d
     elif not store_native:
-        return grid_1d_slim_from(
+        grid_1d = grid_1d_slim_from(
             grid_1d_native=grid_1d,
-            mask_1d=np.array(mask_1d),
+            mask_1d=mask_1d,
         )
-    return grid_1d_native_from(
-        grid_1d_slim=grid_1d,
-        mask_1d=np.array(mask_1d),
-    )
+    else:
+        grid_1d = grid_1d_native_from(
+            grid_1d_slim=grid_1d,
+            mask_1d=mask_1d,
+        )
+    return np.array(grid_1d) if is_numpy else jnp.array(grid_1d)
 
 
 def grid_1d_slim_via_shape_slim_from(
@@ -95,7 +98,6 @@ def grid_1d_slim_via_shape_slim_from(
     )
 
 
-@numba_util.jit()
 def grid_1d_slim_via_mask_from(
     mask_1d: np.ndarray,
     pixel_scales: ty.PixelScales,
@@ -131,23 +133,13 @@ def grid_1d_slim_via_mask_from(
     mask = np.array([True, False, True, False, False, False])
     grid_slim = grid_1d_via_mask_from(mask_1d=mask_1d, pixel_scales=(0.5, 0.5), origin=(0.0, 0.0))
     """
-
-    total_pixels = mask_1d_util.total_pixels_1d_from(mask_1d)
-
-    grid_1d = np.zeros(shape=(total_pixels,))
-
     centres_scaled = geometry_util.central_scaled_coordinate_1d_from(
         shape_slim=mask_1d.shape, pixel_scales=pixel_scales, origin=origin
     )
-
-    index = 0
-
-    for x in range(mask_1d.shape[0]):
-        if not mask_1d[x]:
-            grid_1d[index] = (x - centres_scaled[0]) * pixel_scales[0]
-            index += 1
-
-    return grid_1d
+    indices = jnp.arange(mask_1d.shape[0])
+    unmasked = jnp.logical_not(mask_1d)
+    coords = (indices - centres_scaled[0]) * pixel_scales[0]
+    return coords[unmasked]
 
 
 def grid_1d_slim_from(
@@ -179,7 +171,7 @@ def grid_1d_slim_from(
     """
 
     return array_1d_util.array_1d_slim_from(
-        array_1d_native=np.array(grid_1d_native),
+        array_1d_native=grid_1d_native,
         mask_1d=mask_1d,
     )
 
diff --git a/autoarray/structures/grids/grid_2d_util.py b/autoarray/structures/grids/grid_2d_util.py
index 6c72c00ef..b547fcb94 100644
--- a/autoarray/structures/grids/grid_2d_util.py
+++ b/autoarray/structures/grids/grid_2d_util.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 import numpy as np
 import jax.numpy as jnp
-import jax
 
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
@@ -15,6 +14,19 @@
 from autoarray import type as ty
 
 
+def convert_grid(grid: Union[np.ndarray, List]) -> np.ndarray:
+
+    try:
+        grid = grid.array
+    except AttributeError:
+        pass
+
+    if isinstance(grid, list):
+        grid = np.asarray(grid)
+
+    return grid
+
+
 def check_grid_slim(grid, shape_native):
     if shape_native is None:
         raise exc.GridException(
@@ -36,13 +48,6 @@ def check_grid_slim(grid, shape_native):
         )
 
 
-def convert_grid(grid: Union[np.ndarray, List]) -> np.ndarray:
-    if type(grid) is list:
-        grid = np.asarray(grid)
-
-    return grid
-
-
 def check_grid_2d(grid_2d: np.ndarray):
     if grid_2d.shape[-1] != 2:
         raise exc.GridException(
@@ -108,25 +113,33 @@ def convert_grid_2d(
 
     grid_2d = convert_grid(grid=grid_2d)
 
+    is_numpy = True if isinstance(grid_2d, np.ndarray) else False
+
     check_grid_2d_and_mask_2d(grid_2d=grid_2d, mask_2d=mask_2d)
 
     is_native = len(grid_2d.shape) == 3
 
     if is_native:
-        grid_2d[:, :, 0] *= np.invert(mask_2d)
-        grid_2d[:, :, 1] *= np.invert(mask_2d)
+        if not is_numpy:
+            grid_2d = grid_2d.at[:, :, 0].multiply(~mask_2d)
+            grid_2d = grid_2d.at[:, :, 1].multiply(~mask_2d)
+        else:
+            grid_2d[:, :, 0] *= ~mask_2d
+            grid_2d[:, :, 1] *= ~mask_2d
 
     if is_native == store_native:
-        return grid_2d
+        grid_2d = grid_2d
     elif not store_native:
-        return grid_2d_slim_from(
-            grid_2d_native=np.array(grid_2d),
-            mask=np.array(mask_2d),
+        grid_2d = grid_2d_slim_from(
+            grid_2d_native=grid_2d,
+            mask=mask_2d,
         )
-    return grid_2d_native_from(
-        grid_2d_slim=np.array(grid_2d),
-        mask_2d=np.array(mask_2d),
-    )
+    else:
+        grid_2d = grid_2d_native_from(
+            grid_2d_slim=grid_2d,
+            mask_2d=mask_2d,
+        )
+    return np.array(grid_2d) if is_numpy else jnp.array(grid_2d)
 
 
 def convert_grid_2d_to_slim(
@@ -553,7 +566,7 @@ def grid_scaled_2d_slim_radial_projected_from(
         grid_scaled_2d_slim_radii[slim_index, 1] = radii
         radii += pixel_scale
 
-    return grid_scaled_2d_slim_radii
+    return grid_scaled_2d_slim_radii + 1e-6
 
 
 @numba_util.jit()
@@ -671,16 +684,16 @@ def grid_2d_slim_from(
     """
 
     grid_1d_slim_y = array_2d_util.array_2d_slim_from(
-        array_2d_native=np.array(grid_2d_native[:, :, 0]),
-        mask_2d=np.array(mask),
+        array_2d_native=grid_2d_native[:, :, 0],
+        mask_2d=mask,
     )
 
     grid_1d_slim_x = array_2d_util.array_2d_slim_from(
-        array_2d_native=np.array(grid_2d_native[:, :, 1]),
-        mask_2d=np.array(mask),
+        array_2d_native=grid_2d_native[:, :, 1],
+        mask_2d=mask,
     )
 
-    return np.stack((grid_1d_slim_y, grid_1d_slim_x), axis=-1)
+    return jnp.stack((grid_1d_slim_y, grid_1d_slim_x), axis=-1)
 
 
 def grid_2d_native_from(
@@ -723,7 +736,7 @@ def grid_2d_native_from(
         mask_2d=mask_2d,
     )
 
-    return np.stack((grid_2d_native_y, grid_2d_native_x), axis=-1)
+    return jnp.stack((grid_2d_native_y, grid_2d_native_x), axis=-1)
 
 
 @numba_util.jit()
@@ -801,7 +814,6 @@ def compute_polygon_area(points):
     return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))
 
 
-@numba_util.jit()
 def grid_pixels_in_mask_pixels_from(
     grid, shape_native, pixel_scales, origin
 ) -> np.ndarray:
@@ -832,10 +844,11 @@ def grid_pixels_in_mask_pixels_from(
 
     mesh_pixels_per_image_pixel = np.zeros(shape=shape_native)
 
-    for i in range(grid_pixel_centres.shape[0]):
-        y = grid_pixel_centres[i, 0]
-        x = grid_pixel_centres[i, 1]
+    # Assuming grid_pixel_centres is a 2D array where each row contains (y, x) indices.
+    y_indices = grid_pixel_centres[:, 0]
+    x_indices = grid_pixel_centres[:, 1]
 
-        mesh_pixels_per_image_pixel[y, x] += 1
+    # Use np.add.at to increment the specific indices in a safe and efficient manner
+    np.add.at(mesh_pixels_per_image_pixel, (y_indices, x_indices), 1)
 
     return mesh_pixels_per_image_pixel
diff --git a/autoarray/structures/grids/irregular_2d.py b/autoarray/structures/grids/irregular_2d.py
index 57d4c4a6f..50fdf1a76 100644
--- a/autoarray/structures/grids/irregular_2d.py
+++ b/autoarray/structures/grids/irregular_2d.py
@@ -233,8 +233,8 @@ def furthest_distances_to_other_coordinates(self) -> ArrayIrregular:
         radial_distances_max = np.zeros((self.shape[0]))
 
         for i in range(self.shape[0]):
-            x_distances = np.square(np.subtract(self[i, 0], self[:, 0]))
-            y_distances = np.square(np.subtract(self[i, 1], self[:, 1]))
+            x_distances = np.square(np.subtract(self.array[i, 0], self.array[:, 0]))
+            y_distances = np.square(np.subtract(self.array[i, 1], self.array[:, 1]))
 
             radial_distances_max[i] = np.sqrt(np.max(np.add(x_distances, y_distances)))
 
diff --git a/autoarray/structures/grids/uniform_1d.py b/autoarray/structures/grids/uniform_1d.py
index 53a9ec756..7fb2e7688 100644
--- a/autoarray/structures/grids/uniform_1d.py
+++ b/autoarray/structures/grids/uniform_1d.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 import numpy as np
+import jax.numpy as jnp
 from typing import List, Union, Tuple
 
 from autoarray.structures.abstract_structure import Structure
@@ -178,7 +179,7 @@ def no_mask(
             origin=origin,
         )
 
-        return Grid1D(values=values, mask=mask)
+        return Grid1D(values=np.array(values), mask=mask)
 
     @classmethod
     def from_mask(cls, mask: Mask1D) -> "Grid1D":
@@ -195,12 +196,12 @@ def from_mask(cls, mask: Mask1D) -> "Grid1D":
         """
 
         grid_1d = grid_1d_util.grid_1d_slim_via_mask_from(
-            mask_1d=np.array(mask),
+            mask_1d=mask.array,
             pixel_scales=mask.pixel_scales,
             origin=mask.origin,
         )
 
-        return Grid1D(values=grid_1d, mask=mask)
+        return Grid1D(values=np.array(grid_1d), mask=mask)
 
     @classmethod
     def uniform(
@@ -311,11 +312,11 @@ def grid_2d_radial_projected_from(self, angle: float = 0.0) -> Grid2DIrregular:
             The projected and rotated 2D grid of (y,x) coordinates.
         """
 
-        grid = np.zeros((self.mask.pixels_in_mask, 2))
-        grid[:, 1] = self.slim
+        grid = jnp.zeros((self.mask.pixels_in_mask, 2))
+        grid = grid.at[:, 1].set(self.slim.array)
 
         grid = geometry_util.transform_grid_2d_to_reference_frame(
             grid_2d=grid, centre=(0.0, 0.0), angle=angle
         )
 
-        return Grid2DIrregular(values=grid)
+        return Grid2DIrregular(values=grid + 1e-6)
diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py
index 670cbfcbe..addfb2ec1 100644
--- a/autoarray/structures/grids/uniform_2d.py
+++ b/autoarray/structures/grids/uniform_2d.py
@@ -159,8 +159,9 @@ def __init__(
             not uniform (e.g. due to gravitational lensing) is cannot be computed internally in this function. If the
             over sampled grid is not passed in it is computed assuming uniformity.
         """
+
         values = grid_2d_util.convert_grid_2d(
-            grid_2d=np.array(values),
+            grid_2d=values,
             mask_2d=mask,
             store_native=store_native,
         )
@@ -180,15 +181,15 @@ def __init__(
         self.over_sampler = OverSampler(sub_size=over_sample_size, mask=mask)
 
         if over_sampled is None:
-            self.over_sampled = (
-                over_sample_util.grid_2d_slim_over_sampled_via_mask_from(
-                    mask_2d=np.array(self.mask),
-                    pixel_scales=self.mask.pixel_scales,
-                    sub_size=np.array(self.over_sampler.sub_size._array).astype("int"),
-                    origin=self.mask.origin,
-                )
+            over_sampled = over_sample_util.grid_2d_slim_over_sampled_via_mask_from(
+                mask_2d=np.array(self.mask),
+                pixel_scales=self.mask.pixel_scales,
+                sub_size=self.over_sampler.sub_size.array.astype("int"),
+                origin=self.mask.origin,
             )
 
+            self.over_sampled = Grid2DIrregular(values=over_sampled)
+
         else:
             self.over_sampled = over_sampled
 
@@ -244,7 +245,7 @@ def no_mask(
         )
 
         return Grid2D(
-            values=values,
+            values=np.array(values),
             mask=mask,
             over_sample_size=over_sample_size,
         )
@@ -544,14 +545,14 @@ def from_mask(
             The mask whose masked pixels are used to setup the grid.
         """
 
-        grid_1d = grid_2d_util.grid_2d_slim_via_mask_from(
-            mask_2d=mask._array,
+        grid_2d = grid_2d_util.grid_2d_slim_via_mask_from(
+            mask_2d=mask.array,
             pixel_scales=mask.pixel_scales,
             origin=mask.origin,
         )
 
         return Grid2D(
-            values=grid_1d,
+            values=np.array(grid_2d),
             mask=mask,
             over_sample_size=over_sample_size,
         )
@@ -839,14 +840,10 @@ def squared_distances_to_coordinate_from(
         coordinate
             The (y,x) coordinate from which the squared distance of every grid (y,x) coordinate is computed.
         """
-        if isinstance(self, jnp.ndarray):
-            squared_distances = jnp.square(
-                self.array[:, 0] - coordinate[0]
-            ) + jnp.square(self.array[:, 1] - coordinate[1])
-        else:
-            squared_distances = np.square(self[:, 0] - coordinate[0]) + np.square(
-                self[:, 1] - coordinate[1]
-            )
+        squared_distances = jnp.square(self.array[:, 0] - coordinate[0]) + jnp.square(
+            self.array[:, 1] - coordinate[1]
+        )
+
         return Array2D(values=squared_distances, mask=self.mask)
 
     def distances_to_coordinate_from(
@@ -863,7 +860,7 @@ def distances_to_coordinate_from(
         squared_distance = self.squared_distances_to_coordinate_from(
             coordinate=coordinate
         )
-        distances = np.sqrt(squared_distance.array)
+        distances = jnp.sqrt(squared_distance.array)
         return Array2D(values=distances, mask=self.mask)
 
     def grid_2d_radial_projected_shape_slim_from(
@@ -1099,6 +1096,7 @@ def padded_grid_from(self, kernel_shape_native: Tuple[int, int]) -> "Grid2D":
         padded_mask = Mask2D.all_false(
             shape_native=padded_shape,
             pixel_scales=self.mask.pixel_scales,
+            origin=self.origin,
         )
 
         pad_width = (
@@ -1107,7 +1105,7 @@ def padded_grid_from(self, kernel_shape_native: Tuple[int, int]) -> "Grid2D":
         )
 
         over_sample_size = np.pad(
-            self.over_sample_size.native._array,
+            self.over_sample_size.native.array,
             pad_width,
             mode="constant",
             constant_values=1,
diff --git a/autoarray/structures/mesh/delaunay_2d.py b/autoarray/structures/mesh/delaunay_2d.py
index a6b7f9012..e12d4898c 100644
--- a/autoarray/structures/mesh/delaunay_2d.py
+++ b/autoarray/structures/mesh/delaunay_2d.py
@@ -71,7 +71,7 @@ def interpolated_array_from(
 
         interpolated_array = mesh_util.delaunay_interpolated_array_from(
             shape_native=shape_native,
-            interpolation_grid_slim=interpolation_grid.slim,
+            interpolation_grid_slim=np.array(interpolation_grid.slim.array),
             delaunay=self.delaunay,
             pixel_values=values,
         )
diff --git a/autoarray/structures/mesh/rectangular_2d.py b/autoarray/structures/mesh/rectangular_2d.py
index 0e4eab108..c46fc258e 100644
--- a/autoarray/structures/mesh/rectangular_2d.py
+++ b/autoarray/structures/mesh/rectangular_2d.py
@@ -178,7 +178,9 @@ def interpolated_array_from(
             shape_native=shape_native, extent=extent
         )
 
-        interpolated_array = griddata(points=self, values=values, xi=interpolation_grid)
+        interpolated_array = griddata(
+            points=self.array, values=values, xi=interpolation_grid
+        )
 
         interpolated_array = interpolated_array.reshape(shape_native)
 
diff --git a/autoarray/structures/mesh/triangulation_2d.py b/autoarray/structures/mesh/triangulation_2d.py
index b5c9fedc9..635e02f62 100644
--- a/autoarray/structures/mesh/triangulation_2d.py
+++ b/autoarray/structures/mesh/triangulation_2d.py
@@ -96,7 +96,9 @@ def delaunay(self) -> scipy.spatial.Delaunay:
         `MeshException`, which helps exception handling in the `inversion` package.
         """
         try:
-            return scipy.spatial.Delaunay(np.asarray([self[:, 0], self[:, 1]]).T)
+            return scipy.spatial.Delaunay(
+                np.asarray([self.array[:, 0], self.array[:, 1]]).T
+            )
         except (ValueError, OverflowError, scipy.spatial.qhull.QhullError) as e:
             raise exc.MeshException() from e
 
diff --git a/autoarray/structures/mock/mock_decorators.py b/autoarray/structures/mock/mock_decorators.py
index 876b456d7..c02ebc0b8 100644
--- a/autoarray/structures/mock/mock_decorators.py
+++ b/autoarray/structures/mock/mock_decorators.py
@@ -116,7 +116,7 @@ def ndarray_2d_from(self, grid, *args, **kwargs):
 
         Such functions are common in **PyAutoGalaxy** for light and mass profile objects.
         """
-        return np.multiply(2.0, grid)
+        return np.multiply(2.0, grid.array)
 
     @decorators.to_vector_yx
     def ndarray_yx_2d_from(self, grid, *args, **kwargs):
@@ -146,7 +146,7 @@ def ndarray_2d_list_from(self, grid, *args, **kwargs):
 
         Such functions are common in **PyAutoGalaxy** for light and mass profile objects.
         """
-        return [np.multiply(1.0, grid), np.multiply(2.0, grid)]
+        return [np.multiply(1.0, grid.array), np.multiply(2.0, grid.array)]
 
     @decorators.to_vector_yx
     def ndarray_yx_2d_list_from(self, grid, *args, **kwargs):
@@ -156,7 +156,7 @@ def ndarray_yx_2d_list_from(self, grid, *args, **kwargs):
 
         Such functions are common in **PyAutoGalaxy** for light and mass profile objects.
         """
-        return [np.multiply(1.0, grid), np.multiply(2.0, grid)]
+        return [np.multiply(1.0, grid.array), np.multiply(2.0, grid.array)]
 
 
 class MockGridRadialMinimum:
@@ -165,7 +165,3 @@ def __init__(self):
 
     def radial_grid_from(self, grid):
         return np.sqrt(np.add(np.square(grid[:, 0]), np.square(grid[:, 1])))
-
-    @decorators.relocate_to_radial_minimum
-    def deflections_yx_2d_from(self, grid):
-        return grid
diff --git a/autoarray/structures/vectors/irregular.py b/autoarray/structures/vectors/irregular.py
index 5895149fb..b49ec3868 100644
--- a/autoarray/structures/vectors/irregular.py
+++ b/autoarray/structures/vectors/irregular.py
@@ -1,5 +1,6 @@
 import logging
 import numpy as np
+import jax.numpy as jnp
 from typing import List, Tuple, Union
 
 from autoarray.structures.vectors.abstract import AbstractVectorYX2D
@@ -43,7 +44,6 @@ def __init__(
         grid
             The irregular grid of (y,x) coordinates where each vector is located.
         """
-
         if type(values) is list:
             values = np.asarray(values)
 
@@ -120,13 +120,13 @@ def vectors_within_radius(
         squared_distances = self.grid.distances_to_coordinate_from(coordinate=centre)
         mask = squared_distances < radius
 
-        if np.all(mask == False):
+        if jnp.all(mask == False):
             raise exc.VectorYXException(
                 "The input radius removed all vectors / points on the grid."
             )
 
         return VectorYX2DIrregular(
-            values=self[mask], grid=Grid2DIrregular(self.grid[mask])
+            values=jnp.array(self.array)[mask], grid=Grid2DIrregular(self.grid[mask])
         )
 
     def vectors_within_annulus(
diff --git a/autoarray/structures/vectors/uniform.py b/autoarray/structures/vectors/uniform.py
index 89d589139..88dfe6a9c 100644
--- a/autoarray/structures/vectors/uniform.py
+++ b/autoarray/structures/vectors/uniform.py
@@ -1,7 +1,7 @@
 import logging
 
-# import numpy as np
-from autofit.jax_wrapper import numpy as np, use_jax
+import numpy as np
+import jax.numpy as jnp
 from typing import List, Optional, Tuple, Union
 
 from autoarray.structures.arrays.uniform_2d import Array2D
@@ -137,9 +137,6 @@ def __init__(
             mapping large data arrays to and from the slim / native formats, which can be a computational bottleneck.
         """
 
-        if type(values) is list:
-            values = np.asarray(values)
-
         values = grid_2d_util.convert_grid_2d(
             grid_2d=values, mask_2d=mask, store_native=store_native
         )
@@ -396,11 +393,10 @@ def magnitudes(self) -> Array2D:
         """
         Returns the magnitude of every vector which are computed as sqrt(y**2 + x**2).
         """
-        if use_jax:
-            s = self.array
-        else:
-            s = self
-        return Array2D(values=np.sqrt(s[:, 0] ** 2.0 + s[:, 1] ** 2.0), mask=self.mask)
+        return Array2D(
+            values=jnp.sqrt(self.array[:, 0] ** 2.0 + self.array[:, 1] ** 2.0),
+            mask=self.mask,
+        )
 
     @property
     def y(self) -> Array2D:
diff --git a/autoarray/structures/visibilities.py b/autoarray/structures/visibilities.py
index 2b4113c7c..8cc94dca5 100644
--- a/autoarray/structures/visibilities.py
+++ b/autoarray/structures/visibilities.py
@@ -213,9 +213,7 @@ def __init__(self, visibilities: Union[np.ndarray, List[complex]], *args, **kwar
         A collection of (real, imag) visibilities noise-map values which are used to represent the noise-map in
         an `Interferometer` dataset.
 
-        This data structure behaves the same as the `Visibilities` structure (see `AbstractVisibilities.__new__`). The
-        only difference is that it includes a `WeightOperator` used by `Inversion`'s which use `LinearOperators` and
-        the library `PyLops` to fit `Interferometer` data.
+        This data structure behaves the same as the `Visibilities` structure (see `AbstractVisibilities.__new__`).
 
         Parameters
         ----------
diff --git a/optional_requirements.txt b/optional_requirements.txt
index cbb64e17e..4ba479740 100644
--- a/optional_requirements.txt
+++ b/optional_requirements.txt
@@ -1,4 +1,3 @@
-pylops>=1.10.0,<=2.3.1
 pynufft
 #jax==0.4.3
 #jaxlib==0.4.3
diff --git a/test_autoarray/config/grids.yaml b/test_autoarray/config/grids.yaml
deleted file mode 100644
index 61c268a27..000000000
--- a/test_autoarray/config/grids.yaml
+++ /dev/null
@@ -1,12 +0,0 @@
-# Certain light and mass profile calculations become ill defined at (0.0, 0.0) or close to this value. This can lead
-# to numerical issues in the calculation of the profile, for example a np.nan may arise, crashing the code. 
-  
-# To avoid this, we set a minimum value for the radial coordinate of the profile. If the radial coordinate is below 
-# this value, it is rounded up to this value. This ensures that the profile cannot receive a radial coordinate of 0.0.
-
-# For example, if an input grid coordinate has a radial coordinate of 1e-12, for most profiles this will be rounded up
-# to radial_minimum=1e-08. This is a small enough value that it should not impact the results of the profile calculation.
-
-radial_minimum:
-  radial_minimum:
-    MockGridRadialMinimum: 2.5
\ No newline at end of file
diff --git a/test_autoarray/dataset/abstract/test_dataset.py b/test_autoarray/dataset/abstract/test_dataset.py
index dd99c1201..794af761c 100644
--- a/test_autoarray/dataset/abstract/test_dataset.py
+++ b/test_autoarray/dataset/abstract/test_dataset.py
@@ -32,9 +32,9 @@ def test__signal_to_noise_map():
 
     dataset = ds.AbstractDataset(data=array, noise_map=noise_map)
 
-    assert (
-        dataset.signal_to_noise_map.native == np.array([[0.1, 0.2], [0.1, 1.0]])
-    ).all()
+    assert dataset.signal_to_noise_map.native == pytest.approx(
+        np.array([[0.1, 0.2], [0.1, 1.0]]), 1.0e-4
+    )
     assert dataset.signal_to_noise_max == 1.0
 
     array = aa.Array2D.no_mask([[-1.0, 2.0], [3.0, -4.0]], pixel_scales=1.0)
@@ -43,9 +43,9 @@ def test__signal_to_noise_map():
 
     dataset = ds.AbstractDataset(data=array, noise_map=noise_map)
 
-    assert (
-        dataset.signal_to_noise_map.native == np.array([[0.0, 0.2], [0.1, 0.0]])
-    ).all()
+    assert dataset.signal_to_noise_map.native == pytest.approx(
+        np.array([[0.0, 0.2], [0.1, 0.0]]), 1.0e-4
+    )
     assert dataset.signal_to_noise_max == 0.2
 
 
@@ -115,15 +115,18 @@ def test__grid_settings__sub_size(image_7x7, noise_map_7x7):
 
 
 def test__new_imaging_with_arrays_trimmed_via_kernel_shape():
-    data = aa.Array2D.full(fill_value=20.0, shape_native=(3, 3), pixel_scales=1.0)
-    data[4] = 5.0
 
-    noise_map_array = aa.Array2D.full(
-        fill_value=5.0, shape_native=(3, 3), pixel_scales=1.0
+    data = aa.Array2D.no_mask(
+        values=[[20.0, 20.0, 20.0], [20.0, 5.0, 20.0], [20.0, 20.0, 20.0]],
+        pixel_scales=1.0,
+    )
+
+    noise_map = aa.Array2D.no_mask(
+        values=[[20.0, 20.0, 20.0], [20.0, 2.0, 20.0], [20.0, 20.0, 20.0]],
+        pixel_scales=1.0,
     )
-    noise_map_array[4] = 2.0
 
-    dataset = ds.AbstractDataset(data=data, noise_map=noise_map_array)
+    dataset = ds.AbstractDataset(data=data, noise_map=noise_map)
 
     dataset_trimmed = dataset.trimmed_after_convolution_from(kernel_shape=(3, 3))
 
diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py
index 45ef0a80f..ead9e51f5 100644
--- a/test_autoarray/dataset/imaging/test_dataset.py
+++ b/test_autoarray/dataset/imaging/test_dataset.py
@@ -78,7 +78,7 @@ def test__from_fits():
     )
 
     assert (dataset.data.native == np.ones((3, 3))).all()
-    assert (dataset.psf.native == (1.0 / 9.0) * np.ones((3, 3))).all()
+    assert dataset.psf.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4)
     assert (dataset.noise_map.native == 3.0 * np.ones((3, 3))).all()
 
     assert dataset.pixel_scales == (0.1, 0.1)
@@ -96,7 +96,7 @@ def test__from_fits():
     )
 
     assert (dataset.data.native == np.ones((3, 3))).all()
-    assert (dataset.psf.native == (1.0 / 9.0) * np.ones((3, 3))).all()
+    assert dataset.psf.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4)
     assert (dataset.noise_map.native == 3.0 * np.ones((3, 3))).all()
 
     assert dataset.pixel_scales == (0.1, 0.1)
@@ -105,6 +105,7 @@ def test__from_fits():
 
 
 def test__output_to_fits(imaging_7x7, test_data_path):
+
     imaging_7x7.output_to_fits(
         data_path=path.join(test_data_path, "data.fits"),
         psf_path=path.join(test_data_path, "psf.fits"),
@@ -139,7 +140,9 @@ def test__apply_mask(imaging_7x7, mask_2d_7x7, psf_3x3):
         == 2.0 * np.ones((7, 7)) * np.invert(mask_2d_7x7)
     ).all()
 
-    assert (masked_imaging_7x7.psf.slim == (1.0 / 3.0) * psf_3x3.slim).all()
+    assert masked_imaging_7x7.psf.slim == pytest.approx(
+        (1.0 / 3.0) * psf_3x3.slim.array, 1.0e-4
+    )
 
     assert type(masked_imaging_7x7.psf) == aa.Kernel2D
     assert masked_imaging_7x7.w_tilde.curvature_preload.shape == (35,)
@@ -156,10 +159,21 @@ def test__apply_noise_scaling(imaging_7x7, mask_2d_7x7):
     assert masked_imaging_7x7.noise_map.native[4, 4] == 1e5
 
 
-def test__apply_noise_scaling__use_signal_to_noise_value(imaging_7x7, mask_2d_7x7):
-    imaging_7x7 = copy.copy(imaging_7x7)
+def test__apply_noise_scaling__use_signal_to_noise_value(
+    image_7x7, psf_3x3, noise_map_7x7, mask_2d_7x7
+):
+
+    image_7x7 = np.array(image_7x7.native.array)
+    image_7x7[3, 3] = 2.0
 
-    imaging_7x7.data[24] = 2.0
+    image_7x7 = aa.Array2D(values=image_7x7, mask=mask_2d_7x7)
+
+    imaging_7x7 = aa.Imaging(
+        data=image_7x7,
+        psf=psf_3x3,
+        noise_map=noise_map_7x7,
+        over_sample_size_lp=1,
+    )
 
     masked_imaging_7x7 = imaging_7x7.apply_noise_scaling(
         mask=mask_2d_7x7, signal_to_noise_value=0.1, should_zero_data=False
@@ -244,4 +258,10 @@ def test__noise_map_unmasked_has_zeros_or_negative__raises_exception():
 def test__psf_not_odd_x_odd_kernel__raises_error():
 
     with pytest.raises(exc.KernelException):
-        aa.Kernel2D.no_mask(values=[[0.0, 1.0], [1.0, 2.0]], pixel_scales=1.0)
+        image = aa.Array2D.ones(shape_native=(3, 3), pixel_scales=1.0)
+        noise_map = aa.Array2D.ones(shape_native=(3, 3), pixel_scales=1.0)
+        psf = aa.Kernel2D.no_mask(values=[[0.0, 1.0], [1.0, 2.0]], pixel_scales=1.0)
+
+        dataset = aa.Imaging(
+            data=image, noise_map=noise_map, psf=psf, pad_for_psf=False
+        )
diff --git a/test_autoarray/dataset/imaging/test_simulator.py b/test_autoarray/dataset/imaging/test_simulator.py
index 3cc4182d3..54dc1f6ed 100644
--- a/test_autoarray/dataset/imaging/test_simulator.py
+++ b/test_autoarray/dataset/imaging/test_simulator.py
@@ -70,7 +70,7 @@ def test__via_image_from__psf_off__noise_off_value_is_noise_value(
         == np.array([[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]])
     ).all()
 
-    assert np.allclose(dataset.noise_map.native, 0.2 * np.ones((3, 3)))
+    assert np.allclose(dataset.noise_map.native.array, 0.2 * np.ones((3, 3)))
 
 
 def test__via_image_from__psf_off__background_sky_on(image_central_delta_3x3):
diff --git a/test_autoarray/dataset/test_preprocess.py b/test_autoarray/dataset/test_preprocess.py
index 74e8ef774..8e99a74e8 100644
--- a/test_autoarray/dataset/test_preprocess.py
+++ b/test_autoarray/dataset/test_preprocess.py
@@ -133,15 +133,15 @@ def test__noise_map_from_image_exposure_time_map():
         data_eps=image, exposure_time_map=exposure_time_map
     )
 
-    assert (
-        poisson_noise_map.native
-        == np.array(
+    assert poisson_noise_map.native == pytest.approx(
+        np.array(
             [
                 [np.sqrt(5.0), np.sqrt(6.0) / 2.0],
                 [np.sqrt(30.0) / 3.0, np.sqrt(80.0) / 4.0],
             ]
-        )
-    ).all()
+        ),
+        1.0e-4,
+    )
 
 
 def test__noise_map_from_image_exposure_time_map_and_background_noise_map():
@@ -462,7 +462,7 @@ def test__background_noise_map_via_edges_of_image_from_4():
     )
 
     assert np.allclose(
-        background_noise_map.native,
+        background_noise_map.native.array,
         np.full(fill_value=np.std(np.arange(28)), shape=image.shape_native),
     )
 
@@ -486,17 +486,17 @@ def test__background_noise_map_via_edges_of_image_from_5():
     )
 
     assert np.allclose(
-        background_noise_map.native,
+        background_noise_map.native.array,
         np.full(fill_value=np.std(np.arange(48)), shape=image.shape_native),
     )
 
 
 def test__exposure_time_map_from_exposure_time_and_inverse_noise_map():
     exposure_time = 6.0
-    background_noise_map = aa.Array2D.full(
-        fill_value=0.25, shape_native=(3, 3), pixel_scales=1.0
+
+    background_noise_map = aa.Array2D.no_mask(
+        [[0.5, 0.25, 0.25], [0.25, 0.25, 0.25], [0.25, 0.25, 0.25]], pixel_scales=1.0
     )
-    background_noise_map[0] = 0.5
 
     exposure_time_map = (
         aa.preprocess.exposure_time_map_via_exposure_time_and_background_noise_map_from(
@@ -556,6 +556,7 @@ def test__poisson_noise_from_data():
 def test__data_with_poisson_noised_added():
     data = aa.Array2D.zeros(shape_native=(2, 2), pixel_scales=1.0)
     exposure_time_map = aa.Array2D.ones(shape_native=(2, 2), pixel_scales=1.0)
+
     data_with_poisson_noise = aa.preprocess.data_eps_with_poisson_noise_added(
         data_eps=data, exposure_time_map=exposure_time_map, seed=1
     )
@@ -661,49 +662,32 @@ def test__data_with_complex_gaussian_noise_added():
 
 
 def test__noise_map_with_signal_to_noise_limit_from():
-    image = aa.Array2D.full(fill_value=20.0, shape_native=(2, 2), pixel_scales=1.0)
-    image[3] = 5.0
 
-    noise_map_array = aa.Array2D.full(
-        fill_value=5.0, shape_native=(2, 2), pixel_scales=1.0
-    )
-    noise_map_array[3] = 2.0
+    image = aa.Array2D.no_mask(values=[[20, 20], [20, 5]], pixel_scales=1.0)
+    noise_map = aa.Array2D.no_mask(values=[[5, 5], [5, 2]], pixel_scales=1.0)
 
     noise_map = aa.preprocess.noise_map_with_signal_to_noise_limit_from(
-        data=image, noise_map=noise_map_array, signal_to_noise_limit=100.0
+        data=image, noise_map=noise_map, signal_to_noise_limit=100.0
     )
 
     assert (noise_map.slim == np.array([5.0, 5.0, 5.0, 2.0])).all()
 
-    image = aa.Array2D.full(fill_value=20.0, shape_native=(2, 2), pixel_scales=1.0)
-    image[3] = 5.0
-
-    noise_map_array = aa.Array2D.full(
-        fill_value=5.0, shape_native=(2, 2), pixel_scales=1.0
-    )
-    noise_map_array[3] = 2.0
+    noise_map = aa.Array2D.no_mask(values=[[5, 5], [5, 2]], pixel_scales=1.0)
 
     noise_map = aa.preprocess.noise_map_with_signal_to_noise_limit_from(
-        data=image, noise_map=noise_map_array, signal_to_noise_limit=2.0
+        data=image, noise_map=noise_map, signal_to_noise_limit=2.0
     )
 
     assert (noise_map.native == np.array([[10.0, 10.0], [10.0, 2.5]])).all()
 
-    image = aa.Array2D.full(fill_value=20.0, shape_native=(2, 2), pixel_scales=1.0)
-    image[2] = 5.0
-    image[3] = 5.0
-
-    noise_map_array = aa.Array2D.full(
-        fill_value=5.0, shape_native=(2, 2), pixel_scales=1.0
-    )
-    noise_map_array[2] = 2.0
-    noise_map_array[3] = 2.0
+    image = aa.Array2D.no_mask(values=[[20, 20], [5, 5]], pixel_scales=1.0)
+    noise_map = aa.Array2D.no_mask(values=[[5, 5], [2, 2]], pixel_scales=1.0)
 
     mask = aa.Mask2D(mask=[[True, False], [False, True]], pixel_scales=1.0)
 
     noise_map = aa.preprocess.noise_map_with_signal_to_noise_limit_from(
         data=image,
-        noise_map=noise_map_array,
+        noise_map=noise_map,
         signal_to_noise_limit=2.0,
         noise_limit_mask=mask,
     )
diff --git a/test_autoarray/fit/test_fit_imaging.py b/test_autoarray/fit/test_fit_imaging.py
index c0c525e71..5b7e49446 100644
--- a/test_autoarray/fit/test_fit_imaging.py
+++ b/test_autoarray/fit/test_fit_imaging.py
@@ -30,7 +30,7 @@ def test__data_and_model_are_identical__no_masking__check_values_are_correct():
 
     assert fit.chi_squared == 0.0
     assert fit.reduced_chi_squared == 0.0
-    assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map**2.0))
+    assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map.array**2.0))
     assert fit.log_likelihood == -0.5 * (fit.chi_squared + fit.noise_normalization)
 
 
@@ -59,7 +59,7 @@ def test__data_and_model_are_different__include_masking__check_values_are_correc
 
     assert fit.chi_squared == 0.25
     assert fit.reduced_chi_squared == 0.25 / 3.0
-    assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map**2.0))
+    assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map.array**2.0))
     assert fit.log_likelihood == -0.5 * (fit.chi_squared + fit.noise_normalization)
 
 
@@ -92,7 +92,7 @@ def test__data_and_model_are_identical__inversion_included__changes_certain_prop
 
     assert fit.chi_squared == 0.0
     assert fit.reduced_chi_squared == 0.0
-    assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map**2.0))
+    assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map.array**2.0))
     assert fit.log_likelihood == -0.5 * (fit.chi_squared + fit.noise_normalization)
 
     assert fit.log_likelihood_with_regularization == -0.5 * (
diff --git a/test_autoarray/fit/test_fit_util.py b/test_autoarray/fit/test_fit_util.py
index 641dbf52e..0cbc811de 100644
--- a/test_autoarray/fit/test_fit_util.py
+++ b/test_autoarray/fit/test_fit_util.py
@@ -1,41 +1,41 @@
 import autoarray as aa
 
-import numpy as np
+import jax.numpy as jnp
 import pytest
 
 
 def test__residual_map_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    model_data = np.array([10.0, 10.0, 10.0, 10.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    model_data = jnp.array([10.0, 10.0, 10.0, 10.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
-    assert (residual_map == np.array([0.0, 0.0, 0.0, 0.0])).all()
+    assert (residual_map == jnp.array([0.0, 0.0, 0.0, 0.0])).all()
 
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
-    assert (residual_map == np.array([-1.0, 0.0, 1.0, 2.0])).all()
+    assert (residual_map == jnp.array([-1.0, 0.0, 1.0, 2.0])).all()
 
 
 def test__residual_map_with_mask_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    mask = np.array([True, False, False, True])
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    mask = jnp.array([True, False, False, True])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
     )
 
-    assert (residual_map == np.array([0.0, 0.0, 1.0, 0.0])).all()
+    assert (residual_map == jnp.array([0.0, 0.0, 1.0, 0.0])).all()
 
 
 def test__normalized_residual_map_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    noise_map = np.array([2.0, 2.0, 2.0, 2.0])
-    model_data = np.array([10.0, 10.0, 10.0, 10.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    noise_map = jnp.array([2.0, 2.0, 2.0, 2.0])
+    model_data = jnp.array([10.0, 10.0, 10.0, 10.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -43,9 +43,11 @@ def test__normalized_residual_map_from():
         residual_map=residual_map, noise_map=noise_map
     )
 
-    assert (normalized_residual_map == np.array([0.0, 0.0, 0.0, 0.0])).all()
+    assert normalized_residual_map == pytest.approx(
+        jnp.array([0.0, 0.0, 0.0, 0.0]), 1.0e-4
+    )
 
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -53,17 +55,16 @@ def test__normalized_residual_map_from():
         residual_map=residual_map, noise_map=noise_map
     )
 
-    assert (
-        normalized_residual_map
-        == np.array([-(1.0 / 2.0), 0.0, (1.0 / 2.0), (2.0 / 2.0)])
-    ).all()
+    assert normalized_residual_map == pytest.approx(
+        jnp.array([-(1.0 / 2.0), 0.0, (1.0 / 2.0), (2.0 / 2.0)]), 1.0e-4
+    )
 
 
 def test__normalized_residual_map_with_mask_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    mask = np.array([True, False, False, True])
-    noise_map = np.array([2.0, 2.0, 2.0, 2.0])
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    mask = jnp.array([True, False, False, True])
+    noise_map = jnp.array([2.0, 2.0, 2.0, 2.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
@@ -73,13 +74,15 @@ def test__normalized_residual_map_with_mask_from():
         residual_map=residual_map, mask=mask, noise_map=noise_map
     )
 
-    assert (normalized_residual_map == np.array([0.0, 0.0, (1.0 / 2.0), 0.0])).all()
+    assert normalized_residual_map == pytest.approx(
+        jnp.array([0.0, 0.0, (1.0 / 2.0), 0.0]), abs=1.0e-4
+    )
 
 
 def test__normalized_residual_map_complex_from():
-    data = np.array([10.0 + 10.0j, 10.0 + 10.0j])
-    noise_map = np.array([2.0 + 2.0j, 2.0 + 2.0j])
-    model_data = np.array([9.0 + 12.0j, 9.0 + 12.0j])
+    data = jnp.array([10.0 + 10.0j, 10.0 + 10.0j])
+    noise_map = jnp.array([2.0 + 2.0j, 2.0 + 2.0j])
+    model_data = jnp.array([9.0 + 12.0j, 9.0 + 12.0j])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -87,13 +90,13 @@ def test__normalized_residual_map_complex_from():
         residual_map=residual_map, noise_map=noise_map
     )
 
-    assert (normalized_residual_map == np.array([0.5 - 1.0j, 0.5 - 1.0j])).all()
+    assert (normalized_residual_map == jnp.array([0.5 - 1.0j, 0.5 - 1.0j])).all()
 
 
 def test__chi_squared_map_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    noise_map = np.array([2.0, 2.0, 2.0, 2.0])
-    model_data = np.array([10.0, 10.0, 10.0, 10.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    noise_map = jnp.array([2.0, 2.0, 2.0, 2.0])
+    model_data = jnp.array([10.0, 10.0, 10.0, 10.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -101,9 +104,9 @@ def test__chi_squared_map_from():
         residual_map=residual_map, noise_map=noise_map
     )
 
-    assert (chi_squared_map == np.array([0.0, 0.0, 0.0, 0.0])).all()
+    assert (chi_squared_map == jnp.array([0.0, 0.0, 0.0, 0.0])).all()
 
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -113,15 +116,15 @@ def test__chi_squared_map_from():
 
     assert (
         chi_squared_map
-        == np.array([(1.0 / 2.0) ** 2.0, 0.0, (1.0 / 2.0) ** 2.0, (2.0 / 2.0) ** 2.0])
+        == jnp.array([(1.0 / 2.0) ** 2.0, 0.0, (1.0 / 2.0) ** 2.0, (2.0 / 2.0) ** 2.0])
     ).all()
 
 
 def test__chi_squared_map_with_mask_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    mask = np.array([True, False, False, True])
-    noise_map = np.array([2.0, 2.0, 2.0, 2.0])
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    mask = jnp.array([True, False, False, True])
+    noise_map = jnp.array([2.0, 2.0, 2.0, 2.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
@@ -131,9 +134,9 @@ def test__chi_squared_map_with_mask_from():
         residual_map=residual_map, mask=mask, noise_map=noise_map
     )
 
-    assert (chi_squared_map == np.array([0.0, 0.0, (1.0 / 2.0) ** 2.0, 0.0])).all()
+    assert (chi_squared_map == jnp.array([0.0, 0.0, (1.0 / 2.0) ** 2.0, 0.0])).all()
 
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
@@ -143,13 +146,13 @@ def test__chi_squared_map_with_mask_from():
         residual_map=residual_map, mask=mask, noise_map=noise_map
     )
 
-    assert (chi_squared_map == np.array([0.0, 0.0, (1.0 / 2.0) ** 2.0, 0.0])).all()
+    assert (chi_squared_map == jnp.array([0.0, 0.0, (1.0 / 2.0) ** 2.0, 0.0])).all()
 
 
 def test__chi_squared_map_complex_from():
-    data = np.array([10.0 + 10.0j, 10.0 + 10.0j])
-    noise_map = np.array([2.0 + 2.0j, 2.0 + 2.0j])
-    model_data = np.array([9.0 + 12.0j, 9.0 + 12.0j])
+    data = jnp.array([10.0 + 10.0j, 10.0 + 10.0j])
+    noise_map = jnp.array([2.0 + 2.0j, 2.0 + 2.0j])
+    model_data = jnp.array([9.0 + 12.0j, 9.0 + 12.0j])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -157,13 +160,13 @@ def test__chi_squared_map_complex_from():
         residual_map=residual_map, noise_map=noise_map
     )
 
-    assert (chi_squared_map == np.array([0.25 + 1.0j, 0.25 + 1.0j])).all()
+    assert (chi_squared_map == jnp.array([0.25 + 1.0j, 0.25 + 1.0j])).all()
 
 
 def test__chi_squared_with_noise_covariance_from():
     resdiual_map = aa.Array2D.no_mask([[1.0, 1.0], [2.0, 2.0]], pixel_scales=1.0)
 
-    noise_covariance_matrix_inv = np.array(
+    noise_covariance_matrix_inv = jnp.array(
         [
             [1.0, 1.0, 4.0, 0.0],
             [0.0, 1.0, 9.0, 0.0],
@@ -181,10 +184,10 @@ def test__chi_squared_with_noise_covariance_from():
 
 
 def test__chi_squared_with_mask_fast_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    mask = np.array([True, False, False, True])
-    noise_map = np.array([1.0, 2.0, 3.0, 4.0])
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    mask = jnp.array([True, False, False, True])
+    noise_map = jnp.array([1.0, 2.0, 3.0, 4.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
@@ -207,10 +210,10 @@ def test__chi_squared_with_mask_fast_from():
 
     assert chi_squared == pytest.approx(chi_squared_fast, 1.0e-4)
 
-    data = np.array([[10.0, 10.0], [10.0, 10.0]])
-    mask = np.array([[True, False], [False, True]])
-    noise_map = np.array([[1.0, 2.0], [3.0, 4.0]])
-    model_data = np.array([[11.0, 10.0], [9.0, 8.0]])
+    data = jnp.array([[10.0, 10.0], [10.0, 10.0]])
+    mask = jnp.array([[True, False], [False, True]])
+    noise_map = jnp.array([[1.0, 2.0], [3.0, 4.0]])
+    model_data = jnp.array([[11.0, 10.0], [9.0, 8.0]])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
@@ -235,9 +238,9 @@ def test__chi_squared_with_mask_fast_from():
 
 
 def test__log_likelihood_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    noise_map = np.array([2.0, 2.0, 2.0, 2.0])
-    model_data = np.array([10.0, 10.0, 10.0, 10.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    noise_map = jnp.array([2.0, 2.0, 2.0, 2.0])
+    model_data = jnp.array([10.0, 10.0, 10.0, 10.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -255,17 +258,17 @@ def test__log_likelihood_from():
 
     chi_squared = 0.0
     noise_normalization = (
-        np.log(2.0 * np.pi * (2.0**2.0))
-        + np.log(2.0 * np.pi * (2.0**2.0))
-        + np.log(2.0 * np.pi * (2.0**2.0))
-        + np.log(2.0 * np.pi * (2.0**2.0))
+        jnp.log(2.0 * jnp.pi * (2.0**2.0))
+        + jnp.log(2.0 * jnp.pi * (2.0**2.0))
+        + jnp.log(2.0 * jnp.pi * (2.0**2.0))
+        + jnp.log(2.0 * jnp.pi * (2.0**2.0))
     )
 
     assert log_likelihood == pytest.approx(
         -0.5 * (chi_squared + noise_normalization), 1.0e-4
     )
 
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -288,17 +291,17 @@ def test__log_likelihood_from():
         ((1.0 / 2.0) ** 2.0) + 0.0 + ((1.0 / 2.0) ** 2.0) + ((2.0 / 2.0) ** 2.0)
     )
     noise_normalization = (
-        np.log(2.0 * np.pi * (2.0**2.0))
-        + np.log(2.0 * np.pi * (2.0**2.0))
-        + np.log(2.0 * np.pi * (2.0**2.0))
-        + np.log(2.0 * np.pi * (2.0**2.0))
+        jnp.log(2.0 * jnp.pi * (2.0**2.0))
+        + jnp.log(2.0 * jnp.pi * (2.0**2.0))
+        + jnp.log(2.0 * jnp.pi * (2.0**2.0))
+        + jnp.log(2.0 * jnp.pi * (2.0**2.0))
     )
 
     assert log_likelihood == pytest.approx(
         -0.5 * (chi_squared + noise_normalization), 1.0e-4
     )
 
-    noise_map = np.array([1.0, 2.0, 3.0, 4.0])
+    noise_map = jnp.array([1.0, 2.0, 3.0, 4.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -318,10 +321,10 @@ def test__log_likelihood_from():
 
     chi_squared = 1.0 + (1.0 / (3.0**2.0)) + 0.25
     noise_normalization = (
-        np.log(2 * np.pi * (1.0**2.0))
-        + np.log(2 * np.pi * (2.0**2.0))
-        + np.log(2 * np.pi * (3.0**2.0))
-        + np.log(2 * np.pi * (4.0**2.0))
+        jnp.log(2 * jnp.pi * (1.0**2.0))
+        + jnp.log(2 * jnp.pi * (2.0**2.0))
+        + jnp.log(2 * jnp.pi * (3.0**2.0))
+        + jnp.log(2 * jnp.pi * (4.0**2.0))
     )
 
     assert log_likelihood == pytest.approx(
@@ -330,10 +333,10 @@ def test__log_likelihood_from():
 
 
 def test__log_likelihood_from__with_mask():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    mask = np.array([True, False, False, True])
-    noise_map = np.array([1.0, 2.0, 3.0, 4.0])
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    mask = jnp.array([True, False, False, True])
+    noise_map = jnp.array([1.0, 2.0, 3.0, 4.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
@@ -358,18 +361,18 @@ def test__log_likelihood_from__with_mask():
     # chi squared = 0, 0.25, (0.25 and 1.0 are masked)
 
     chi_squared = 0.0 + (1.0 / 3.0) ** 2.0
-    noise_normalization = np.log(2 * np.pi * (2.0**2.0)) + np.log(
-        2 * np.pi * (3.0**2.0)
+    noise_normalization = jnp.log(2 * jnp.pi * (2.0**2.0)) + jnp.log(
+        2 * jnp.pi * (3.0**2.0)
     )
 
     assert log_likelihood == pytest.approx(
         -0.5 * (chi_squared + noise_normalization), 1e-4
     )
 
-    data = np.array([[10.0, 10.0], [10.0, 10.0]])
-    mask = np.array([[True, False], [False, True]])
-    noise_map = np.array([[1.0, 2.0], [3.0, 4.0]])
-    model_data = np.array([[11.0, 10.0], [9.0, 8.0]])
+    data = jnp.array([[10.0, 10.0], [10.0, 10.0]])
+    mask = jnp.array([[True, False], [False, True]])
+    noise_map = jnp.array([[1.0, 2.0], [3.0, 4.0]])
+    model_data = jnp.array([[11.0, 10.0], [9.0, 8.0]])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
@@ -394,8 +397,8 @@ def test__log_likelihood_from__with_mask():
     # chi squared = 0, 0.25, (0.25 and 1.0 are masked)
 
     chi_squared = 0.0 + (1.0 / 3.0) ** 2.0
-    noise_normalization = np.log(2 * np.pi * (2.0**2.0)) + np.log(
-        2 * np.pi * (3.0**2.0)
+    noise_normalization = jnp.log(2 * jnp.pi * (2.0**2.0)) + jnp.log(
+        2 * jnp.pi * (3.0**2.0)
     )
 
     assert log_likelihood == pytest.approx(
@@ -404,9 +407,9 @@ def test__log_likelihood_from__with_mask():
 
 
 def test__log_likelihood_from__complex_data():
-    data = np.array([10.0 + 10.0j, 10.0 + 10.0j])
-    noise_map = np.array([2.0 + 1.0j, 2.0 + 1.0j])
-    model_data = np.array([9.0 + 12.0j, 9.0 + 12.0j])
+    data = jnp.array([10.0 + 10.0j, 10.0 + 10.0j])
+    noise_map = jnp.array([2.0 + 1.0j, 2.0 + 1.0j])
+    model_data = jnp.array([9.0 + 12.0j, 9.0 + 12.0j])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -427,8 +430,8 @@ def test__log_likelihood_from__complex_data():
     # chi squared = 0.25 and 4.0
 
     chi_squared = 4.25
-    noise_normalization = np.log(2 * np.pi * (2.0**2.0)) + np.log(
-        2 * np.pi * (1.0**2.0)
+    noise_normalization = jnp.log(2 * jnp.pi * (2.0**2.0)) + jnp.log(
+        2 * jnp.pi * (1.0**2.0)
     )
 
     assert log_likelihood == pytest.approx(
@@ -457,8 +460,8 @@ def test__log_evidence_from():
 
 
 def test__residual_flux_fraction_map_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    model_data = np.array([10.0, 10.0, 10.0, 10.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    model_data = jnp.array([10.0, 10.0, 10.0, 10.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -466,9 +469,9 @@ def test__residual_flux_fraction_map_from():
         residual_map=residual_map, data=data
     )
 
-    assert (residual_flux_fraction_map == np.array([0.0, 0.0, 0.0, 0.0])).all()
+    assert (residual_flux_fraction_map == jnp.array([0.0, 0.0, 0.0, 0.0])).all()
 
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
 
@@ -476,13 +479,13 @@ def test__residual_flux_fraction_map_from():
         residual_map=residual_map, data=data
     )
 
-    assert (residual_flux_fraction_map == np.array([-0.1, 0.0, 0.1, 0.2])).all()
+    assert (residual_flux_fraction_map == jnp.array([-0.1, 0.0, 0.1, 0.2])).all()
 
 
 def test__residual_flux_fraction_map_with_mask_from():
-    data = np.array([10.0, 10.0, 10.0, 10.0])
-    mask = np.array([True, False, False, True])
-    model_data = np.array([11.0, 10.0, 9.0, 8.0])
+    data = jnp.array([10.0, 10.0, 10.0, 10.0])
+    mask = jnp.array([True, False, False, True])
+    model_data = jnp.array([11.0, 10.0, 9.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
@@ -492,9 +495,9 @@ def test__residual_flux_fraction_map_with_mask_from():
         residual_map=residual_map, mask=mask, data=data
     )
 
-    assert (residual_flux_fraction_map == np.array([0.0, 0.0, 0.1, 0.0])).all()
+    assert (residual_flux_fraction_map == jnp.array([0.0, 0.0, 0.1, 0.0])).all()
 
-    model_data = np.array([11.0, 9.0, 8.0, 8.0])
+    model_data = jnp.array([11.0, 9.0, 8.0, 8.0])
 
     residual_map = aa.util.fit.residual_map_with_mask_from(
         data=data, mask=mask, model_data=model_data
@@ -504,4 +507,4 @@ def test__residual_flux_fraction_map_with_mask_from():
         residual_map=residual_map, mask=mask, data=data
     )
 
-    assert (residual_flux_fraction_map == np.array([0.0, 0.1, 0.2, 0.0])).all()
+    assert (residual_flux_fraction_map == jnp.array([0.0, 0.1, 0.2, 0.0])).all()
diff --git a/test_autoarray/inversion/inversion/imaging/test_imaging.py b/test_autoarray/inversion/inversion/imaging/test_imaging.py
index df615ad85..0fad37e15 100644
--- a/test_autoarray/inversion/inversion/imaging/test_imaging.py
+++ b/test_autoarray/inversion/inversion/imaging/test_imaging.py
@@ -11,9 +11,13 @@
 directory = path.dirname(path.realpath(__file__))
 
 
-def test__operated_mapping_matrix_property(psf_7x7, rectangular_mapper_7x7_3x3):
+def test__operated_mapping_matrix_property(psf_3x3, rectangular_mapper_7x7_3x3):
     inversion = aa.m.MockInversionImaging(
-        psf=psf_7x7, linear_obj_list=[rectangular_mapper_7x7_3x3]
+        psf=psf_3x3,
+        linear_obj_list=[rectangular_mapper_7x7_3x3],
+        convolver=aa.Convolver(
+            kernel=psf_3x3, mask=rectangular_mapper_7x7_3x3.mapper_grids.mask
+        ),
     )
 
     assert inversion.operated_mapping_matrix_list[0][0, 0] == pytest.approx(1.0, 1e-4)
@@ -24,6 +28,7 @@ def test__operated_mapping_matrix_property(psf_7x7, rectangular_mapper_7x7_3x3):
     inversion = aa.m.MockInversionImaging(
         psf=psf,
         linear_obj_list=[rectangular_mapper_7x7_3x3, rectangular_mapper_7x7_3x3],
+        convolver=aa.m.MockConvolver(operated_mapping_matrix=np.ones((2, 2))),
     )
 
     operated_mapping_matrix_0 = np.array([[1.0, 1.0], [1.0, 1.0]])
@@ -42,7 +47,7 @@ def test__operated_mapping_matrix_property(psf_7x7, rectangular_mapper_7x7_3x3):
 
 
 def test__operated_mapping_matrix_property__with_operated_mapping_matrix_override(
-    psf_7x7, rectangular_mapper_7x7_3x3
+    psf_3x3, rectangular_mapper_7x7_3x3
 ):
     psf = aa.m.MockPSF(operated_mapping_matrix=np.ones((2, 2)))
 
@@ -54,7 +59,9 @@ def test__operated_mapping_matrix_property__with_operated_mapping_matrix_overrid
     )
 
     inversion = aa.m.MockInversionImaging(
-        psf=psf, linear_obj_list=[rectangular_mapper_7x7_3x3, linear_obj]
+        psf=psf,
+        linear_obj_list=[rectangular_mapper_7x7_3x3, linear_obj],
+        convolver=aa.m.MockConvolver(operated_mapping_matrix=np.ones((2, 2))),
     )
 
     operated_mapping_matrix_0 = np.array([[1.0, 1.0], [1.0, 1.0]])
@@ -88,6 +95,7 @@ def test__curvature_matrix(rectangular_mapper_7x7_3x3):
         data=np.ones(2),
         noise_map=noise_map,
         psf=psf,
+        convolver=aa.m.MockConvolver(operated_mapping_matrix=np.ones((2, 10))),
     )
 
     inversion = aa.InversionImagingMapping(
diff --git a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py
index e05e13350..570f79673 100644
--- a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py
+++ b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py
@@ -204,22 +204,24 @@ def test__data_vector_via_w_tilde_data_two_methods_agree():
         mapping_matrix = mapper.mapping_matrix
 
         blurred_mapping_matrix = psf.convolve_mapping_matrix(
-            mapping_matrix=mapping_matrix
+            mapping_matrix=mapping_matrix, mask=mask
         )
 
         data_vector = (
             aa.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(
-                blurred_mapping_matrix=blurred_mapping_matrix,
+                blurred_mapping_matrix=np.array(blurred_mapping_matrix),
                 image=np.array(image),
                 noise_map=np.array(noise_map),
             )
         )
 
         w_tilde_data = aa.util.inversion_imaging.w_tilde_data_imaging_from(
-            image_native=np.array(image.native),
-            noise_map_native=np.array(noise_map.native),
-            kernel_native=np.array(kernel.native),
-            native_index_for_slim_index=mask.derive_indexes.native_for_slim,
+            image_native=np.array(image.native.array),
+            noise_map_native=np.array(noise_map.native.array),
+            kernel_native=np.array(kernel.native.array),
+            native_index_for_slim_index=np.array(
+                mask.derive_indexes.native_for_slim
+            ).astype("int"),
         )
 
         (
@@ -273,16 +275,20 @@ def test__curvature_matrix_via_w_tilde_two_methods_agree():
     mapping_matrix = mapper.mapping_matrix
 
     w_tilde = aa.util.inversion_imaging.w_tilde_curvature_imaging_from(
-        noise_map_native=np.array(noise_map.native),
-        kernel_native=np.array(kernel.native),
-        native_index_for_slim_index=mask.derive_indexes.native_for_slim,
+        noise_map_native=np.array(noise_map.native.array),
+        kernel_native=np.array(kernel.native.array),
+        native_index_for_slim_index=np.array(
+            mask.derive_indexes.native_for_slim
+        ).astype("int"),
     )
 
     curvature_matrix_via_w_tilde = aa.util.inversion.curvature_matrix_via_w_tilde_from(
         w_tilde=w_tilde, mapping_matrix=mapping_matrix
     )
 
-    blurred_mapping_matrix = psf.convolve_mapping_matrix(mapping_matrix=mapping_matrix)
+    blurred_mapping_matrix = psf.convolve_mapping_matrix(
+        mapping_matrix=mapping_matrix, mask=mask
+    )
 
     curvature_matrix = aa.util.inversion.curvature_matrix_via_mapping_matrix_from(
         mapping_matrix=blurred_mapping_matrix,
@@ -326,9 +332,11 @@ def test__curvature_matrix_via_w_tilde_preload_two_methods_agree():
             w_tilde_indexes,
             w_tilde_lengths,
         ) = aa.util.inversion_imaging.w_tilde_curvature_preload_imaging_from(
-            noise_map_native=np.array(noise_map.native),
-            kernel_native=np.array(kernel.native),
-            native_index_for_slim_index=mask.derive_indexes.native_for_slim,
+            noise_map_native=np.array(noise_map.native.array),
+            kernel_native=np.array(kernel.native.array),
+            native_index_for_slim_index=np.array(
+                mask.derive_indexes.native_for_slim
+            ).astype("int"),
         )
 
         (
@@ -355,11 +363,12 @@ def test__curvature_matrix_via_w_tilde_preload_two_methods_agree():
         )
 
         blurred_mapping_matrix = psf.convolve_mapping_matrix(
-            mapping_matrix=mapping_matrix
+            mapping_matrix=mapping_matrix,
+            mask=mask,
         )
 
         curvature_matrix = aa.util.inversion.curvature_matrix_via_mapping_matrix_from(
-            mapping_matrix=blurred_mapping_matrix,
+            mapping_matrix=np.array(blurred_mapping_matrix),
             noise_map=np.array(noise_map),
         )
 
diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py
index 8875ea378..28c45a9bd 100644
--- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py
+++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py
@@ -68,23 +68,6 @@ def test__data_vector_via_transformed_mapping_matrix_from():
     assert (data_vector_complex_via_blurred == data_vector_via_transformed).all()
 
 
-def test__inversion_interferometer__via_mapper(
-    interferometer_7_no_fft,
-    rectangular_mapper_7x7_3x3,
-    delaunay_mapper_9_3x3,
-    voronoi_mapper_9_3x3,
-    regularization_constant,
-):
-    inversion = aa.Inversion(
-        dataset=interferometer_7_no_fft,
-        linear_obj_list=[rectangular_mapper_7x7_3x3],
-        settings=aa.SettingsInversion(use_linear_operators=True),
-    )
-
-    assert isinstance(inversion.linear_obj_list[0], aa.MapperRectangular)
-    assert isinstance(inversion, aa.InversionInterferometerMappingPyLops)
-
-
 def test__w_tilde_curvature_interferometer_from():
     noise_map = np.array([1.0, 2.0, 3.0])
     uv_wavelengths = np.array([[0.0001, 2.0, 3000.0], [3000.0, 2.0, 0.0001]])
@@ -289,20 +272,24 @@ def test__identical_inversion_values_for_two_methods():
         == inversion_mapping_matrices.regularization_matrix
     ).all()
 
+    assert inversion_w_tilde.data_vector == pytest.approx(
+        inversion_mapping_matrices.data_vector, 1.0e-8
+    )
     assert inversion_w_tilde.curvature_matrix == pytest.approx(
         inversion_mapping_matrices.curvature_matrix, 1.0e-8
     )
     assert inversion_w_tilde.curvature_reg_matrix == pytest.approx(
         inversion_mapping_matrices.curvature_reg_matrix, 1.0e-8
     )
+
     assert inversion_w_tilde.reconstruction == pytest.approx(
-        inversion_mapping_matrices.reconstruction, 1.0e-2
+        inversion_mapping_matrices.reconstruction, abs=1.0e-1
     )
     assert inversion_w_tilde.mapped_reconstructed_image == pytest.approx(
-        inversion_mapping_matrices.mapped_reconstructed_image, 1.0e-2
+        inversion_mapping_matrices.mapped_reconstructed_image, abs=1.0e-1
     )
     assert inversion_w_tilde.mapped_reconstructed_data == pytest.approx(
-        inversion_mapping_matrices.mapped_reconstructed_data, 1.0e-2
+        inversion_mapping_matrices.mapped_reconstructed_data, abs=1.0e-1
     )
 
 
diff --git a/test_autoarray/inversion/pixelization/image_mesh/test_overlay.py b/test_autoarray/inversion/pixelization/image_mesh/test_overlay.py
index 22536d5c8..54f242236 100644
--- a/test_autoarray/inversion/pixelization/image_mesh/test_overlay.py
+++ b/test_autoarray/inversion/pixelization/image_mesh/test_overlay.py
@@ -323,13 +323,13 @@ def test__image_plane_mesh_grid_from__simple():
 
     total_pixels = overlay_util.total_pixels_2d_from(
         mask_2d=mask.array,
-        overlaid_centres=overlaid_centres,
+        overlaid_centres=np.array(overlaid_centres),
     )
 
     overlay_for_mask_2d_util = overlay_util.overlay_for_mask_from(
         total_pixels=total_pixels,
         mask=mask.array,
-        overlaid_centres=overlaid_centres,
+        overlaid_centres=np.array(overlaid_centres),
     ).astype("int")
 
     image_mesh_util = overlay_util.overlay_via_unmasked_overlaid_from(
diff --git a/test_autoarray/mask/test_mask_1d_util.py b/test_autoarray/mask/test_mask_1d_util.py
index c8251471f..1b032c37f 100644
--- a/test_autoarray/mask/test_mask_1d_util.py
+++ b/test_autoarray/mask/test_mask_1d_util.py
@@ -1,14 +1,6 @@
-from autoarray import exc
 from autoarray import util
 
 import numpy as np
-import pytest
-
-
-def test__total_image_pixels_1d_from():
-    mask_1d = np.array([False, True, False, False, False, True])
-
-    assert util.mask_1d.total_pixels_1d_from(mask_1d=mask_1d) == 4
 
 
 def test__native_index_for_slim_index_1d_from():
diff --git a/test_autoarray/mask/test_mask_2d_util.py b/test_autoarray/mask/test_mask_2d_util.py
index f3db2938e..ee79a9543 100644
--- a/test_autoarray/mask/test_mask_2d_util.py
+++ b/test_autoarray/mask/test_mask_2d_util.py
@@ -738,8 +738,6 @@ def test__edge_1d_indexes_from():
 
     edge_pixels = util.mask_2d.edge_1d_indexes_from(mask_2d=mask)
 
-    print(edge_pixels)
-
     assert (edge_pixels == np.array([0])).all()
 
     mask = np.array(
diff --git a/test_autoarray/operators/over_sample/test_over_sampler.py b/test_autoarray/operators/over_sample/test_over_sampler.py
index 7da24d79a..a32b11e8f 100644
--- a/test_autoarray/operators/over_sample/test_over_sampler.py
+++ b/test_autoarray/operators/over_sample/test_over_sampler.py
@@ -70,6 +70,16 @@ def test__binned_array_2d_from():
         pixel_scales=1.0,
     )
 
+    over_sampling = aa.OverSampler(
+        mask=mask, sub_size=aa.Array2D(values=[2, 2], mask=mask)
+    )
+
+    arr = np.array([1.0, 5.0, 7.0, 10.0, 10.0, 10.0, 10.0, 10.0])
+
+    binned_array_2d = over_sampling.binned_array_2d_from(array=arr)
+
+    assert binned_array_2d.slim == pytest.approx(np.array([5.75, 10.0]), 1.0e-4)
+
     over_sampling = aa.OverSampler(
         mask=mask, sub_size=aa.Array2D(values=[1, 2], mask=mask)
     )
diff --git a/test_autoarray/plot/include/test_include.py b/test_autoarray/plot/include/test_include.py
index 6f4d29c77..b32616e9d 100644
--- a/test_autoarray/plot/include/test_include.py
+++ b/test_autoarray/plot/include/test_include.py
@@ -6,7 +6,7 @@ def test__loads_default_values_from_config_if_not_input():
 
     assert include.origin is True
     assert include.mask == True
-    assert include.border is True
+    assert include.border is False
     assert include.parallel_overscan is True
     assert include.serial_prescan is True
     assert include.serial_overscan is False
diff --git a/test_autoarray/structures/arrays/test_kernel_2d.py b/test_autoarray/structures/arrays/test_kernel_2d.py
index 4cf8b92d7..39ba25842 100644
--- a/test_autoarray/structures/arrays/test_kernel_2d.py
+++ b/test_autoarray/structures/arrays/test_kernel_2d.py
@@ -150,7 +150,7 @@ def test__rescaled_with_odd_dimensions_from__evens_to_odds():
         rescale_factor=0.5, normalize=True
     )
     assert kernel_2d.pixel_scales == (2.0, 2.0)
-    assert (kernel_2d.native == (1.0 / 9.0) * np.ones((3, 3))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4)
 
     array_2d = np.ones((9, 9))
     kernel_2d = aa.Kernel2D.no_mask(values=array_2d, pixel_scales=1.0, normalize=False)
@@ -158,7 +158,7 @@ def test__rescaled_with_odd_dimensions_from__evens_to_odds():
         rescale_factor=0.333333333333333, normalize=True
     )
     assert kernel_2d.pixel_scales == (3.0, 3.0)
-    assert (kernel_2d.native == (1.0 / 9.0) * np.ones((3, 3))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4)
 
     array_2d = np.ones((18, 6))
     kernel_2d = aa.Kernel2D.no_mask(values=array_2d, pixel_scales=1.0, normalize=False)
@@ -166,7 +166,7 @@ def test__rescaled_with_odd_dimensions_from__evens_to_odds():
         rescale_factor=0.5, normalize=True
     )
     assert kernel_2d.pixel_scales == (2.0, 2.0)
-    assert (kernel_2d.native == (1.0 / 27.0) * np.ones((9, 3))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 27.0) * np.ones((9, 3)), 1.0e-4)
 
     array_2d = np.ones((6, 18))
     kernel_2d = aa.Kernel2D.no_mask(values=array_2d, pixel_scales=1.0, normalize=False)
@@ -174,7 +174,7 @@ def test__rescaled_with_odd_dimensions_from__evens_to_odds():
         rescale_factor=0.5, normalize=True
     )
     assert kernel_2d.pixel_scales == (2.0, 2.0)
-    assert (kernel_2d.native == (1.0 / 27.0) * np.ones((3, 9))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 27.0) * np.ones((3, 9)), 1.0e-4)
 
 
 def test__rescaled_with_odd_dimensions_from__different_scalings():
@@ -183,7 +183,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings():
         rescale_factor=2.0, normalize=True
     )
     assert kernel_2d.pixel_scales == (0.4, 0.4)
-    assert (kernel_2d.native == (1.0 / 25.0) * np.ones((5, 5))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 25.0) * np.ones((5, 5)), 1.0e-4)
 
     kernel_2d = aa.Kernel2D.ones(
         shape_native=(40, 40), pixel_scales=1.0, normalize=False
@@ -192,7 +192,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings():
         rescale_factor=0.1, normalize=True
     )
     assert kernel_2d.pixel_scales == (8.0, 8.0)
-    assert (kernel_2d.native == (1.0 / 25.0) * np.ones((5, 5))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 25.0) * np.ones((5, 5)), 1.0e-4)
 
     kernel_2d = aa.Kernel2D.ones(shape_native=(2, 4), pixel_scales=1.0, normalize=False)
     kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from(
@@ -201,7 +201,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings():
 
     assert kernel_2d.pixel_scales[0] == pytest.approx(0.4, 1.0e-4)
     assert kernel_2d.pixel_scales[1] == pytest.approx(0.4444444, 1.0e-4)
-    assert (kernel_2d.native == (1.0 / 45.0) * np.ones((5, 9))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 45.0) * np.ones((5, 9)), 1.0e-4)
 
     kernel_2d = aa.Kernel2D.ones(shape_native=(4, 2), pixel_scales=1.0, normalize=False)
     kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from(
@@ -209,7 +209,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings():
     )
     assert kernel_2d.pixel_scales[0] == pytest.approx(0.4444444, 1.0e-4)
     assert kernel_2d.pixel_scales[1] == pytest.approx(0.4, 1.0e-4)
-    assert (kernel_2d.native == (1.0 / 45.0) * np.ones((9, 5))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 45.0) * np.ones((9, 5)), 1.0e-4)
 
     kernel_2d = aa.Kernel2D.ones(shape_native=(6, 4), pixel_scales=1.0, normalize=False)
     kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from(
@@ -217,7 +217,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings():
     )
 
     assert kernel_2d.pixel_scales == pytest.approx((2.0, 1.3333333333), 1.0e-4)
-    assert (kernel_2d.native == (1.0 / 9.0) * np.ones((3, 3))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4)
 
     kernel_2d = aa.Kernel2D.ones(
         shape_native=(9, 12), pixel_scales=1.0, normalize=False
@@ -227,7 +227,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings():
     )
 
     assert kernel_2d.pixel_scales == pytest.approx((3.0, 2.4), 1.0e-4)
-    assert (kernel_2d.native == (1.0 / 15.0) * np.ones((3, 5))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 15.0) * np.ones((3, 5)), 1.0e-4)
 
     kernel_2d = aa.Kernel2D.ones(shape_native=(4, 6), pixel_scales=1.0, normalize=False)
     kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from(
@@ -235,7 +235,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings():
     )
 
     assert kernel_2d.pixel_scales == pytest.approx((1.33333333333, 2.0), 1.0e-4)
-    assert (kernel_2d.native == (1.0 / 9.0) * np.ones((3, 3))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4)
 
     kernel_2d = aa.Kernel2D.ones(
         shape_native=(12, 9), pixel_scales=1.0, normalize=False
@@ -244,7 +244,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings():
         rescale_factor=0.33333333333, normalize=True
     )
     assert kernel_2d.pixel_scales == pytest.approx((2.4, 3.0), 1.0e-4)
-    assert (kernel_2d.native == (1.0 / 15.0) * np.ones((5, 3))).all()
+    assert kernel_2d.native == pytest.approx((1.0 / 15.0) * np.ones((5, 3)), 1.0e-4)
 
 
 def test__from_as_gaussian_via_alma_fits_header_parameters__identical_to_astropy_gaussian_model():
diff --git a/test_autoarray/structures/decorators/test_to_grid.py b/test_autoarray/structures/decorators/test_to_grid.py
index 60c70d71b..2e8b1be2f 100644
--- a/test_autoarray/structures/decorators/test_to_grid.py
+++ b/test_autoarray/structures/decorators/test_to_grid.py
@@ -15,11 +15,11 @@ def test__in_grid_1d__out_ndarray_2d():
 
     assert isinstance(ndarray_2d, aa.Grid2D)
     assert ndarray_2d.native == pytest.approx(
-        np.array([[[0.0, 0.0], [0.0, -1.0], [0.0, 1.0], [0.0, 0.0]]]), 1.0e-4
+        np.array([[[0.0, 0.0], [0.0, -1.0], [0.0, 1.0], [0.0, 0.0]]]), abs=1.0e-4
     )
 
 
-def test__in_grid_1d__out_ndarray_2d_list():
+def test__in_dgrid_1d__out_ndarray_2d_list():
     mask = aa.Mask1D(mask=[True, False, False, True], pixel_scales=(1.0,))
 
     grid_1d = aa.Grid1D.from_mask(mask=mask)
@@ -30,12 +30,12 @@ def test__in_grid_1d__out_ndarray_2d_list():
 
     assert isinstance(ndarray_2d_list[0], aa.Grid2D)
     assert ndarray_2d_list[0].native == pytest.approx(
-        np.array([[[0.0, 0.0], [0.0, -0.5], [0.0, 0.5], [0.0, 0.0]]]), 1.0e-4
+        np.array([[[0.0, 0.0], [0.0, -0.5], [0.0, 0.5], [0.0, 0.0]]]), abs=1.0e-4
     )
 
     assert isinstance(ndarray_2d_list[1], aa.Grid2D)
     assert ndarray_2d_list[1].native == pytest.approx(
-        np.array([[[0.0, 0.0], [0.0, -1.0], [0.0, 1.0], [0.0, 0.0]]]), 1.0e-4
+        np.array([[[0.0, 0.0], [0.0, -1.0], [0.0, 1.0], [0.0, 0.0]]]), abs=1.0e-4
     )
 
 
diff --git a/test_autoarray/structures/grids/test_grid_2d_util.py b/test_autoarray/structures/grids/test_grid_2d_util.py
index edf1d2497..463aa4e7b 100644
--- a/test_autoarray/structures/grids/test_grid_2d_util.py
+++ b/test_autoarray/structures/grids/test_grid_2d_util.py
@@ -236,7 +236,7 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(1.0, 1.0),
     )
 
-    assert (grid_radii == np.array([[0.0, 0.0], [0.0, 1.0]])).all()
+    assert grid_radii == pytest.approx(np.array([[0.0, 0.0], [0.0, 1.0]]), abs=1.0e-4)
 
     grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from(
         extent=np.array([-1.0, 3.0, -1.0, 1.0]),
@@ -244,9 +244,9 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(1.0, 1.0),
     )
 
-    assert (
-        grid_radii == np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]])
-    ).all()
+    assert grid_radii == pytest.approx(
+        np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4
+    )
 
     grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from(
         extent=np.array([-1.0, 3.0, -1.0, 1.0]),
@@ -254,7 +254,9 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(1.0, 1.0),
     )
 
-    assert (grid_radii == np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]])).all()
+    assert grid_radii == pytest.approx(
+        np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4
+    )
 
     grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from(
         extent=np.array([-2.0, 1.0, -1.0, 1.0]),
@@ -262,9 +264,9 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(1.0, 1.0),
     )
 
-    assert (
-        grid_radii == np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]])
-    ).all()
+    assert grid_radii == pytest.approx(
+        np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]]), abs=1.0e-4
+    )
 
     grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from(
         extent=np.array([-1.0, 1.0, -1.0, 1.0]),
@@ -272,10 +274,10 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(0.1, 0.5),
     )
 
-    assert (
-        grid_radii
-        == np.array([[0.0, 1.0], [0.0, 1.5], [0.0, 2.0], [0.0, 2.5], [0.0, 3.0]])
-    ).all()
+    assert grid_radii == pytest.approx(
+        np.array([[0.0, 1.0], [0.0, 1.5], [0.0, 2.0], [0.0, 2.5], [0.0, 3.0]]),
+        abs=1.0e-4,
+    )
 
     grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from(
         extent=np.array([5.0, 8.0, 99.9, 100.1]),
@@ -283,9 +285,8 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(10.0, 0.25),
     )
 
-    assert (
-        grid_radii
-        == np.array(
+    assert grid_radii == pytest.approx(
+        np.array(
             [
                 [100.0, 7.0],
                 [100.0, 7.25],
@@ -297,8 +298,9 @@ def test__grid_scaled_2d_slim_radial_projected_from():
                 [100.0, 8.75],
                 [100.0, 9.0],
             ]
-        )
-    ).all()
+        ),
+        abs=1.0e-4,
+    )
 
     grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from(
         extent=np.array([-1.0, 1.0, -1.0, 3.0]),
@@ -306,9 +308,9 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(1.0, 1.0),
     )
 
-    assert (
-        grid_radii == np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]])
-    ).all()
+    assert grid_radii == pytest.approx(
+        np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4
+    )
 
     grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from(
         extent=np.array([-1.0, 1.0, -2.0, 1.0]),
@@ -316,9 +318,9 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(1.0, 1.0),
     )
 
-    assert (
-        grid_radii == np.array([[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0]])
-    ).all()
+    assert grid_radii == pytest.approx(
+        np.array([[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]), abs=1.0e-4
+    )
 
     grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from(
         extent=np.array([-1.0, 1.0, -1.0, 1.0]),
@@ -326,10 +328,10 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(0.5, 0.1),
     )
 
-    assert (
-        grid_radii
-        == np.array([[1.0, 0.0], [1.0, 0.5], [1.0, 1.0], [1.0, 1.5], [1.0, 2.0]])
-    ).all()
+    assert grid_radii == pytest.approx(
+        np.array([[1.0, 0.0], [1.0, 0.5], [1.0, 1.0], [1.0, 1.5], [1.0, 2.0]]),
+        abs=1.0e-4,
+    )
 
     grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from(
         extent=np.array([99.9, 100.1, -1.0, 3.0]),
@@ -337,7 +339,9 @@ def test__grid_scaled_2d_slim_radial_projected_from():
         pixel_scales=(1.5, 10.0),
     )
 
-    assert (grid_radii == np.array([[-1.0, 100.0], [-1.0, 101.5], [-1.0, 103.0]])).all()
+    assert grid_radii == pytest.approx(
+        np.array([[-1.0, 100.0], [-1.0, 101.5], [-1.0, 103.0]]), abs=1.0e-4
+    )
 
 
 def test__grid_2d_slim_from():
diff --git a/test_autoarray/structures/grids/test_uniform_1d.py b/test_autoarray/structures/grids/test_uniform_1d.py
index b99da3608..e641c4015 100644
--- a/test_autoarray/structures/grids/test_uniform_1d.py
+++ b/test_autoarray/structures/grids/test_uniform_1d.py
@@ -123,7 +123,7 @@ def test__grid_2d_radial_projected_from():
 
     assert type(grid_2d) == aa.Grid2DIrregular
     assert grid_2d.slim == pytest.approx(
-        np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]]), 1.0e-4
+        np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]]), abs=1.0e-4
     )
 
     grid_2d = grid_1d.grid_2d_radial_projected_from(angle=90.0)
diff --git a/test_autoarray/structures/grids/test_uniform_2d.py b/test_autoarray/structures/grids/test_uniform_2d.py
index 78813329e..9352dfbb1 100644
--- a/test_autoarray/structures/grids/test_uniform_2d.py
+++ b/test_autoarray/structures/grids/test_uniform_2d.py
@@ -441,7 +441,7 @@ def test__from_mask():
     mask = aa.Mask2D(mask=mask, pixel_scales=(2.0, 2.0))
 
     grid_via_util = aa.util.grid_2d.grid_2d_slim_via_mask_from(
-        mask_2d=np.array(mask), pixel_scales=(2.0, 2.0)
+        mask_2d=mask, pixel_scales=(2.0, 2.0)
     )
 
     grid_2d = aa.Grid2D.from_mask(mask=mask)
@@ -451,8 +451,8 @@ def test__from_mask():
     assert grid_2d.pixel_scales == (2.0, 2.0)
 
     grid_2d_native = aa.util.grid_2d.grid_2d_native_from(
-        grid_2d_slim=np.array(grid_2d),
-        mask_2d=np.array(mask),
+        grid_2d_slim=grid_2d.array,
+        mask_2d=mask,
     )
 
     assert (grid_2d_native == grid_2d.native).all()
@@ -559,7 +559,7 @@ def test__grid_2d_radial_projected_shape_slim_from():
         pixel_scales=grid_2d.pixel_scales,
     )
 
-    assert (grid_radii == grid_radii_util).all()
+    assert grid_radii == pytest.approx(grid_radii_util, 1.0e-4)
     assert grid_radial_shape_slim == grid_radii_util.shape[0]
 
     grid_2d = aa.Grid2D.uniform(shape_native=(3, 4), pixel_scales=(3.0, 2.0))
diff --git a/test_autoarray/test_jax_changes.py b/test_autoarray/test_jax_changes.py
index f5104a942..2b6289317 100644
--- a/test_autoarray/test_jax_changes.py
+++ b/test_autoarray/test_jax_changes.py
@@ -1,8 +1,10 @@
-import autoarray as aa
+import jax.numpy as jnp
 import pytest
 
+
+import autoarray as aa
+
 from autoarray import Grid2D, Mask2D
-from autofit.jax_wrapper import numpy as np
 
 
 @pytest.fixture(name="array")
@@ -33,4 +35,4 @@ def test_boolean_issue():
         mask=Mask2D.all_false((10, 10), pixel_scales=1.0),
     )
     values, keys = Grid2D.instance_flatten(grid)
-    np.array(Grid2D.instance_unflatten(keys, values))
+    jnp.array(Grid2D.instance_unflatten(keys, values))