switch to inversion_imaging_util.w_tilde_data_imaging_from

Author: Jammy2211

autoarray/dataset/imaging/dataset.py

@@ -474,7 +474,12 @@ def apply_over_sampling(
 
         return dataset
 
-    def apply_w_tilde(self, disable_fft_pad: bool = False):
+    def apply_w_tilde(
+            self,
+            batch_size: int = 128,
+            disable_fft_pad: bool = False,
+            use_jax: bool = False,
+    ):
         """
         The w_tilde formalism of the linear algebra equations precomputes the convolution of every pair of masked
         noise-map values given the PSF (see `inversion.inversion_util`).
@@ -487,39 +492,19 @@ def apply_w_tilde(self, disable_fft_pad: bool = False):
 
         Returns
         -------
-        WTildeImaging
-            Precomputed values used for the w tilde formalism of linear algebra calculations.
+        batch_size
+            The size of batches used to compute the w-tilde curvature matrix via FFT-based convolution,
+            which can be reduced to produce lower memory usage at the cost of speed
+        disable_fft_pad
+            The FFT PSF convolution is optimal for a certain 2D FFT padding or trimming,
+            which places the fewest zeros around the image. If this is set to `True`, this optimal padding is not
+            performed and the image is used as-is. This is normally used to avoid repadding data that has already been
+            padded.
+        use_jax
+            Whether to use JAX to compute W-Tilde. This requires JAX to be installed.
         """
 
-        logger.info("IMAGING - Computing W-Tilde... May take a moment.")
-
-        try:
-            import numba
-        except ModuleNotFoundError:
-            raise exc.InversionException(
-                "Inversion w-tilde functionality (pixelized reconstructions) is "
-                "disabled if numba is not installed.\n\n"
-                "This is because the run-times without numba are too slow.\n\n"
-                "Please install numba, which is described at the following web page:\n\n"
-                "https://pyautolens.readthedocs.io/en/latest/installation/overview.html"
-            )
-
-        (
-            curvature_preload,
-            indexes,
-            lengths,
-        ) = inversion_imaging_numba_util.w_tilde_curvature_preload_imaging_from(
-            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"),
-        )
-
         w_tilde = WTildeImaging(
-            curvature_preload=curvature_preload,
-            indexes=indexes.astype("int"),
-            lengths=lengths.astype("int"),
             noise_map=self.noise_map,
             psf=self.psf,
             fft_mask=self.mask,

autoarray/dataset/imaging/w_tilde.py

@@ -12,9 +12,6 @@
 class WTildeImaging(AbstractWTilde):
     def __init__(
         self,
-        curvature_preload: np.ndarray,
-        indexes: np.ndim,
-        lengths: np.ndarray,
         noise_map: np.ndarray,
         psf: np.ndarray,
         fft_mask: np.ndarray,
@@ -40,52 +37,13 @@ def __init__(
             matrix efficienctly.
         """
         super().__init__(
-            curvature_preload=curvature_preload, fft_mask=fft_mask
+            curvature_preload=None,
+            fft_mask=fft_mask
         )
 
-        self.indexes = indexes
-        self.lengths = lengths
         self.noise_map = noise_map
         self.psf = psf
 
-    @property
-    def w_matrix(self):
-        """
-        The matrix `w_tilde_curvature` is a matrix of dimensions [image_pixels, image_pixels] that encodes the PSF
-        convolution of every pair of image pixels given the noise map. This can be used to efficiently compute the
-        curvature matrix via the mappings between image and source pixels, in a way that omits having to perform the
-        PSF convolution on every individual source pixel. This provides a significant speed up for inversions of imaging
-        datasets.
-
-        The limitation of this matrix is that the dimensions of [image_pixels, image_pixels] can exceed many 10s of GB's,
-        making it impossible to store in memory and its use in linear algebra calculations extremely. The method
-        `w_tilde_curvature_preload_imaging_from` describes a compressed representation that overcomes this hurdles. It is
-        advised `w_tilde` and this method are only used for testing.
-
-        Parameters
-        ----------
-        noise_map_native
-            The two dimensional masked noise-map of values which w_tilde is computed from.
-        kernel_native
-            The two dimensional PSF kernel that w_tilde encodes the convolution of.
-        native_index_for_slim_index
-            An array of shape [total_x_pixels*sub_size] that maps pixels from the slimmed array to the native array.
-
-        Returns
-        -------
-        ndarray
-            A matrix that encodes the PSF convolution values between the noise map that enables efficient calculation of
-            the curvature matrix.
-        """
-
-        return inversion_imaging_numba_util.w_tilde_curvature_imaging_from(
-            noise_map_native=self.noise_map.native.array,
-            kernel_native=self.psf.native.array,
-            native_index_for_slim_index=np.array(
-                self.mask.derive_indexes.native_for_slim
-            ).astype("int"),
-        )
-
     @property
     def psf_operator_matrix_dense(self):
 

autoarray/dataset/interferometer/dataset.py

@@ -195,7 +195,7 @@ def apply_w_tilde(
         if curvature_preload is None:
 
             logger.info(
-                "INTERFEROMETER - Computing W-Tilde; runtime scales with visibility count and mask resolution, extreme inputs may exceed hours."
+                "INTERFEROMETER - Computing W-Tilde; runtime scales with visibility count and mask resolution, CPU run times may exceed hours."
             )
 
             curvature_preload = inversion_interferometer_util.w_tilde_curvature_preload_interferometer_from(

autoarray/inversion/inversion/imaging/w_tilde.py

@@ -14,7 +14,7 @@
 from autoarray.structures.arrays.uniform_2d import Array2D
 
 from autoarray import exc
-from autoarray.inversion.inversion.imaging import inversion_imaging_numba_util
+from autoarray.inversion.inversion.imaging import inversion_imaging_util
 
 
 class InversionImagingWTilde(AbstractInversionImaging):
@@ -49,17 +49,6 @@ def __init__(
             the simultaneous linear equations are combined and solved simultaneously.
         """
 
-        try:
-            import numba
-        except ModuleNotFoundError:
-            raise exc.InversionException(
-                "Inversion w-tilde functionality (pixelized reconstructions) is "
-                "disabled if numba is not installed.\n\n"
-                "This is because the run-times without numba are too slow.\n\n"
-                "Please install numba, which is described at the following web page:\n\n"
-                "https://pyautolens.readthedocs.io/en/latest/installation/overview.html"
-            )
-
         super().__init__(
             dataset=dataset, linear_obj_list=linear_obj_list, settings=settings, xp=xp
         )
@@ -68,8 +57,8 @@ def __init__(
 
     @cached_property
     def w_tilde_data(self):
-        return inversion_imaging_numba_util.w_tilde_data_imaging_from(
-            image_native=np.array(self.data.native.array),
+        return inversion_imaging_util.w_tilde_data_imaging_from(
+            image_native=self.data.native.array,
             noise_map_native=self.noise_map.native.array,
             kernel_native=self.psf.native.array,
             native_index_for_slim_index=self.data.mask.derive_indexes.native_for_slim,