Feature/fft jax imaging

.github/workflows/main.yml

@@ -66,6 +66,7 @@ jobs:
     - name: Run tests
       run: |
         export ROOT_DIR=`pwd`
+        export JAX_ENABLE_X64=True
         export PYTHONPATH=$PYTHONPATH:$ROOT_DIR/PyAutoConf
         export PYTHONPATH=$PYTHONPATH:$ROOT_DIR/PyAutoArray
         pushd PyAutoArray

autoarray/__init__.py

@@ -9,17 +9,14 @@
 from . import util
 from . import fixtures
 from . import mock as m
-from .dataset.interferometer.w_tilde import load_curvature_preload_if_compatible
+
 from .dataset import preprocess
 from .dataset.abstract.dataset import AbstractDataset
-from .dataset.abstract.w_tilde import AbstractWTilde
 from .dataset.grids import GridsInterface
 from .dataset.imaging.dataset import Imaging
 from .dataset.imaging.simulator import SimulatorImaging
-from .dataset.imaging.w_tilde import WTildeImaging
 from .dataset.interferometer.dataset import Interferometer
 from .dataset.interferometer.simulator import SimulatorInterferometer
-from .dataset.interferometer.w_tilde import WTildeInterferometer
 from .dataset.dataset_model import DatasetModel
 from .fit.fit_dataset import AbstractFit
 from .fit.fit_dataset import FitDataset
@@ -46,9 +43,15 @@
 from .inversion.pixelization.image_mesh.abstract import AbstractImageMesh
 from .inversion.pixelization.mesh.abstract import AbstractMesh
 from .inversion.inversion.imaging.mapping import InversionImagingMapping
-from .inversion.inversion.imaging.w_tilde import InversionImagingWTilde
-from .inversion.inversion.interferometer.w_tilde import InversionInterferometerWTilde
+from .inversion.inversion.imaging.sparse import InversionImagingSparse
+from .inversion.inversion.imaging.inversion_imaging_util import ImagingSparseOperator
+from .inversion.inversion.interferometer.sparse import (
+    InversionInterferometerSparse,
+)
 from .inversion.inversion.interferometer.mapping import InversionInterferometerMapping
+from .inversion.inversion.interferometer.inversion_interferometer_util import (
+    InterferometerSparseOperator,
+)
 from .inversion.linear_obj.linear_obj import LinearObj
 from .inversion.linear_obj.func_list import AbstractLinearObjFuncList
 from .mask.derive.indexes_2d import DeriveIndexes2D

autoarray/dataset/grids.py

@@ -17,7 +17,7 @@ def __init__(
         over_sample_size_lp: Union[int, Array2D],
         over_sample_size_pixelization: Union[int, Array2D],
         psf: Optional[Kernel2D] = None,
-        use_w_tilde: bool = False,
+        use_sparse_operator: bool = False,
     ):
         """
         Contains grids of (y,x) Cartesian coordinates at the centre of every pixel in the dataset's image and
@@ -66,7 +66,7 @@ def __init__(
         self._blurring = None
         self._border_relocator = None
 
-        self.use_w_tilde = use_w_tilde
+        self.use_sparse_operator = use_sparse_operator
 
     @property
     def lp(self):
@@ -120,7 +120,7 @@ def border_relocator(self) -> BorderRelocator:
         self._border_relocator = BorderRelocator(
             mask=self.mask,
             sub_size=self.over_sample_size_pixelization,
-            use_w_tilde=self.use_w_tilde,
+            use_sparse_operator=self.use_sparse_operator,
         )
 
         return self._border_relocator

autoarray/dataset/imaging/dataset.py

@@ -3,19 +3,20 @@
 from pathlib import Path
 from typing import Optional, Union
 
-from autoconf import cached_property
-
 from autoarray.dataset.abstract.dataset import AbstractDataset
 from autoarray.dataset.grids import GridsDataset
-from autoarray.dataset.imaging.w_tilde import WTildeImaging
+from autoarray.inversion.inversion.imaging.inversion_imaging_util import (
+    ImagingSparseOperator,
+)
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.arrays.kernel_2d import Kernel2D
 from autoarray.mask.mask_2d import Mask2D
 from autoarray import type as ty
 
 from autoarray import exc
 from autoarray.operators.over_sampling import over_sample_util
-from autoarray.inversion.inversion.imaging import inversion_imaging_numba_util
+
+from autoarray.inversion.inversion.imaging import inversion_imaging_util
 
 logger = logging.getLogger(__name__)
 
@@ -32,7 +33,7 @@ def __init__(
         disable_fft_pad: bool = True,
         use_normalized_psf: Optional[bool] = True,
         check_noise_map: bool = True,
-        w_tilde: Optional[WTildeImaging] = None,
+        sparse_operator: Optional[ImagingSparseOperator] = None,
     ):
         """
         An imaging dataset, containing the image data, noise-map, PSF and associated quantities
@@ -86,9 +87,9 @@ def __init__(
             the PSF kernel does not change the overall normalization of the image when it is convolved with it.
         check_noise_map
             If True, the noise-map is checked to ensure all values are above zero.
-        w_tilde
-            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`). Pass the `WTildeImaging` object here to
+        sparse_operator
+            The sparse linear algebra formalism of the linear algebra equations precomputes the convolution of every pair of masked
+            noise-map values given the PSF (see `inversion.inversion_util`). Pass the `ImagingSparseOperator` object here to
             enable this linear algebra formalism for pixelized reconstructions.
         """
 
@@ -191,17 +192,17 @@ def __init__(
             if psf.mask.shape[0] % 2 == 0 or psf.mask.shape[1] % 2 == 0:
                 raise exc.KernelException("Kernel2D Kernel2D must be odd")
 
-        use_w_tilde = True if w_tilde is not None else False
+        use_sparse_operator = True if sparse_operator is not None else False
 
         self.grids = GridsDataset(
             mask=self.data.mask,
             over_sample_size_lp=self.over_sample_size_lp,
             over_sample_size_pixelization=self.over_sample_size_pixelization,
             psf=self.psf,
-            use_w_tilde=use_w_tilde,
+            use_sparse_operator=use_sparse_operator,
         )
 
-        self.w_tilde = w_tilde
+        self.sparse_operator = sparse_operator
 
     @classmethod
     def from_fits(
@@ -474,9 +475,13 @@ def apply_over_sampling(
 
         return dataset
 
-    def apply_w_tilde(self, disable_fft_pad: bool = False):
+    def apply_sparse_operator(
+        self,
+        batch_size: int = 128,
+        disable_fft_pad: bool = False,
+    ):
         """
-        The w_tilde formalism of the linear algebra equations precomputes the convolution of every pair of masked
+        The sparse linear algebra formalism precomputes the convolution of every pair of masked
         noise-map values given the PSF (see `inversion.inversion_util`).
 
         The `WTilde` object stores these precomputed values in the imaging dataset ensuring they are only computed once
@@ -487,12 +492,66 @@ 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.")
+        sparse_operator = (
+            inversion_imaging_util.ImagingSparseOperator.from_noise_map_and_psf(
+                data=self.data,
+                noise_map=self.noise_map,
+                psf=self.psf.native,
+                batch_size=batch_size,
+            )
+        )
+
+        return Imaging(
+            data=self.data,
+            noise_map=self.noise_map,
+            psf=self.psf,
+            noise_covariance_matrix=self.noise_covariance_matrix,
+            over_sample_size_lp=self.over_sample_size_lp,
+            over_sample_size_pixelization=self.over_sample_size_pixelization,
+            disable_fft_pad=disable_fft_pad,
+            check_noise_map=False,
+            sparse_operator=sparse_operator,
+        )
 
+    def apply_sparse_operator_cpu(
+        self,
+        disable_fft_pad: bool = False,
+    ):
+        """
+        The sparse linear algebra formalism precomputes the convolution of every pair of masked
+        noise-map values given the PSF (see `inversion.inversion_util`).
+
+        The `WTilde` object stores these precomputed values in the imaging dataset ensuring they are only computed once
+        per analysis.
+
+        This uses lazy allocation such that the calculation is only performed when the wtilde matrices are used,
+        ensuring efficient set up of the `Imaging` class.
+
+        Returns
+        -------
+        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.
+        """
         try:
             import numba
         except ModuleNotFoundError:
@@ -504,20 +563,24 @@ def apply_w_tilde(self, disable_fft_pad: bool = False):
                 "https://pyautolens.readthedocs.io/en/latest/installation/overview.html"
             )
 
+        from autoarray.inversion.inversion.imaging_numba import (
+            inversion_imaging_numba_util,
+        )
+
         (
-            curvature_preload,
+            psf_precision_operator_sparse,
             indexes,
             lengths,
-        ) = inversion_imaging_numba_util.w_tilde_curvature_preload_imaging_from(
+        ) = inversion_imaging_numba_util.psf_precision_operator_sparse_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,
+        sparse_operator = inversion_imaging_numba_util.SparseLinAlgImagingNumba(
+            psf_precision_operator_sparse=psf_precision_operator_sparse,
             indexes=indexes.astype("int"),
             lengths=lengths.astype("int"),
             noise_map=self.noise_map,
@@ -534,7 +597,7 @@ def apply_w_tilde(self, disable_fft_pad: bool = False):
             over_sample_size_pixelization=self.over_sample_size_pixelization,
             disable_fft_pad=disable_fft_pad,
             check_noise_map=False,
-            w_tilde=w_tilde,
+            sparse_operator=sparse_operator,
         )
 
     def output_to_fits(