couple of tests

Author: Jammy2211

autoarray/inversion/inversion/imaging/abstract.py

@@ -1,6 +1,8 @@
 import numpy as np
 from typing import Dict, List, Union, Type
 
+from autoconf import cached_property
+
 from autoarray.dataset.imaging.dataset import Imaging
 from autoarray.inversion.inversion.dataset_interface import DatasetInterface
 from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList
@@ -136,6 +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:
+                vvv
                 operated_mapping_matrix = self.psf.convolved_mapping_matrix_from(
                     mapping_matrix=linear_func.mapping_matrix,
                     mask=self.mask,
@@ -200,7 +203,7 @@ def data_linear_func_matrix_dict(self):
 
         return data_linear_func_matrix_dict
 
-    @property
+    @cached_property
     def mapper_operated_mapping_matrix_dict(self) -> Dict:
         """
         The `operated_mapping_matrix` of a `Mapper` object describes the mappings between the observed data's values

autoarray/inversion/inversion/imaging/mapping.py

@@ -76,7 +76,10 @@ def _data_vector_mapper(self) -> np.ndarray:
             param_range = mapper_param_range_list[i]
 
             operated_mapping_matrix = self.psf.convolved_mapping_matrix_from(
-                mapping_matrix=mapper.mapping_matrix, mask=self.mask, xp=self._xp
+                mapping_matrix=mapper.mapping_matrix,
+                mask=self.mask,
+                use_mixed_precision=self.settings.use_mixed_precision,
+                xp=self._xp
             )
 
             data_vector_mapper = (
@@ -135,7 +138,10 @@ def _curvature_matrix_mapper_diag(self) -> Optional[np.ndarray]:
             mapper_param_range_i = mapper_param_range_list[i]
 
             operated_mapping_matrix = self.psf.convolved_mapping_matrix_from(
-                mapping_matrix=mapper_i.mapping_matrix, mask=self.mask, xp=self._xp
+                mapping_matrix=mapper_i.mapping_matrix,
+                mask=self.mask,
+                use_mixed_precision=self.settings.use_mixed_precision,
+                xp=self._xp
             )
 
             diag = inversion_util.curvature_matrix_via_mapping_matrix_from(

autoarray/inversion/inversion/inversion_util.py

@@ -84,6 +84,10 @@ def curvature_matrix_via_mapping_matrix_from(
     no_regularization_index_list: Optional[List] = None,
     settings: SettingsInversion = SettingsInversion(),
     xp=np,
+    *,
+    mp_gemm: bool = True,  # mixed precision matmul
+    gemm_dtype=None,  # e.g. xp.float32
+    out_dtype=None,  # e.g. xp.float64
 ) -> np.ndarray:
     """
     Returns the curvature matrix `F` from a blurred mapping matrix `f` and the 1D noise-map $\sigma$
@@ -97,8 +101,13 @@ def curvature_matrix_via_mapping_matrix_from(
     noise_map
         Flattened 1D array of the noise-map used by the inversion during the fit.
     """
-    array = mapping_matrix / noise_map[:, None]
-    curvature_matrix = xp.dot(array.T, array)
+    if gemm_dtype is None:
+        gemm_dtype = xp.float32 if (mp_gemm and xp is not np) else mapping_matrix.dtype
+
+    # form A in chosen dtype (usually float32 on device)
+    A = (mapping_matrix / noise_map[:, None]).astype(gemm_dtype)
+
+    curvature_matrix = xp.dot(A.T, A)  # float32 GEMM if A is float32
 
     if add_to_curvature_diag and len(no_regularization_index_list) > 0:
         curvature_matrix = curvature_matrix_with_added_to_diag_from(

autoarray/inversion/inversion/settings.py

@@ -10,6 +10,7 @@
 class SettingsInversion:
     def __init__(
         self,
+        use_mixed_precision : bool = False,
         use_positive_only_solver: Optional[bool] = None,
         positive_only_uses_p_initial: Optional[bool] = None,
         use_border_relocator: Optional[bool] = None,
@@ -24,6 +25,12 @@ def __init__(
 
         Parameters
         ----------
+        use_mixed_precision
+            If `True`, the linear algebra calculations of the inversion are performed using single precision on a
+            targeted subset of functions which provide significant speed up when using a GPU (x4), reduces VRAM
+            use and are expected to have minimal impact on the accuracy of the results. If `False`, all linear algebra
+            calculations are performed using double precision, which is the default and is more accurate but
+            slower on a GPU.
         use_positive_only_solver
             Whether to use a positive-only linear system solver, which requires that every reconstructed value is
             positive but is computationally much slower than the default solver (which allows for positive and

autoarray/structures/arrays/kernel_2d.py

@@ -533,6 +533,7 @@ def mapping_matrix_native_from(
         mask: "Mask2D",
         blurring_mapping_matrix: Optional[np.ndarray] = None,
         blurring_mask: Optional["Mask2D"] = None,
+        use_mixed_precision: bool = False,
         xp=np,
     ) -> np.ndarray:
         """
@@ -558,6 +559,10 @@ def mapping_matrix_native_from(
             Mask defining the blurring region pixels. Must be provided if
             `blurring_mapping_matrix` is given and `slim_to_native_blurring_tuple`
             is not already cached.
+        use_mixed_precision
+            If True, the mapping matrices are cast to single precision (float32) to
+            speed up GPU computations and reduce VRAM usage. If False, double precision
+            (float64) is used for maximum accuracy.
 
         Returns
         -------
@@ -566,33 +571,29 @@ def mapping_matrix_native_from(
             Contains contributions from both the main mapping matrix and, if provided,
             the blurring mapping matrix.
         """
+        dtype_native = xp.float32 if use_mixed_precision else xp.float64
+
         n_src = mapping_matrix.shape[1]
 
-        # Allocate full native grid (ny, nx, n_src)
-        mapping_matrix_native = xp.zeros(
-            mask.shape + (n_src,), dtype=mapping_matrix.dtype
-        )
+        mapping_matrix_native = xp.zeros(mask.shape + (n_src,), dtype=dtype_native)
+
+        # Cast inputs to the target dtype to avoid implicit up/downcasts inside scatter
+        mm = mapping_matrix if mapping_matrix.dtype == dtype_native else xp.asarray(mapping_matrix, dtype=dtype_native)
 
-        # Scatter main mapping matrix into native cube
         if xp.__name__.startswith("jax"):
-            mapping_matrix_native = mapping_matrix_native.at[
-                mask.slim_to_native_tuple
-            ].set(mapping_matrix)
+            mapping_matrix_native = mapping_matrix_native.at[mask.slim_to_native_tuple].set(mm)
         else:
-            mapping_matrix_native[mask.slim_to_native_tuple] = mapping_matrix
-
-        # Optionally scatter blurring mapping matrix
+            mapping_matrix_native[mask.slim_to_native_tuple] = np.asarray(mm)
 
         if blurring_mapping_matrix is not None:
+            bm = blurring_mapping_matrix
+            if getattr(bm, "dtype", None) != dtype_native:
+                bm = xp.asarray(bm, dtype=dtype_native)
 
             if xp.__name__.startswith("jax"):
-                mapping_matrix_native = mapping_matrix_native.at[
-                    blurring_mask.slim_to_native_tuple
-                ].set(blurring_mapping_matrix)
+                mapping_matrix_native = mapping_matrix_native.at[blurring_mask.slim_to_native_tuple].set(bm)
             else:
-                mapping_matrix_native[blurring_mask.slim_to_native_tuple] = (
-                    blurring_mapping_matrix
-                )
+                mapping_matrix_native[blurring_mask.slim_to_native_tuple] = np.asarray(bm)
 
         return mapping_matrix_native
 
@@ -730,6 +731,7 @@ def convolved_mapping_matrix_from(
         blurring_mapping_matrix=None,
         blurring_mask: Optional[Mask2D] = None,
         jax_method="direct",
+        use_mixed_precision: bool = False,
         xp=np,
     ):
         """
@@ -770,12 +772,19 @@ def convolved_mapping_matrix_from(
             Mapping matrix for the blurring region, outside the mask core.
         jax_method : str
             Backend passed to real-space convolution if ``use_fft=False``.
+        use_mixed_precision
+            If `True`, the FFT is performed using single precision, which provide significant speed up when using a
+            GPU (x4), reduces VRAM use and is expected to have minimal impact on the accuracy of the results. If `False`,
+            the FFT is performed using double precision, which is the default and is more accurate but slower on a GPU.
 
         Returns
         -------
         ndarray of shape (N_pix, N_src)
             Convolved mapping matrix in slim form.
         """
+        # -------------------------------------------------------------------------
+        # NumPy path unchanged
+        # -------------------------------------------------------------------------
         if xp is np:
             return self.convolved_mapping_matrix_via_real_space_np_from(
                 mapping_matrix=mapping_matrix,
@@ -785,6 +794,9 @@ def convolved_mapping_matrix_from(
                 xp=xp,
             )
 
+        # -------------------------------------------------------------------------
+        # Non-FFT JAX path unchanged
+        # -------------------------------------------------------------------------
         if not self.use_fft:
             return self.convolved_mapping_matrix_via_real_space_from(
                 mapping_matrix=mapping_matrix,
@@ -796,34 +808,50 @@ def convolved_mapping_matrix_from(
             )
 
         import jax
+        import jax.numpy as jnp
 
+        # -------------------------------------------------------------------------
+        # Validate cached FFT shapes / state
+        # -------------------------------------------------------------------------
         if self.fft_shape is None:
-
             full_shape, fft_shape, mask_shape = self.fft_shape_from(mask=mask)
-
             raise ValueError(
                 f"FFT convolution requires precomputed padded shapes, but `self.fft_shape` is None.\n"
                 f"Expected mapping matrix padded to match FFT shape of PSF.\n"
                 f"PSF fft_shape: {fft_shape}, mask shape: {mask.shape}, "
                 f"mapping_matrix shape: {getattr(mapping_matrix, 'shape', 'unknown')}."
             )
-
         else:
-
             fft_shape = self.fft_shape
             full_shape = self.full_shape
             mask_shape = self.mask_shape
             fft_psf_mapping = self.fft_psf_mapping
 
+        # -------------------------------------------------------------------------
+        # Mixed precision dtypes (JAX only)
+        # -------------------------------------------------------------------------
+        fft_real_dtype = jnp.float32 if use_mixed_precision else jnp.float64
+        fft_complex_dtype = jnp.complex64 if use_mixed_precision else jnp.complex128
+
+        # Ensure PSF FFT dtype matches the FFT path
+        fft_psf_mapping = jnp.asarray(fft_psf_mapping, dtype=fft_complex_dtype)
+
+        # -------------------------------------------------------------------------
+        # Build native cube in the FFT dtype (THIS IS THE KEY)
+        # Requires mapping_matrix_native_from to accept dtype_native kwarg.
+        # -------------------------------------------------------------------------
         mapping_matrix_native = self.mapping_matrix_native_from(
             mapping_matrix=mapping_matrix,
             mask=mask,
             blurring_mapping_matrix=blurring_mapping_matrix,
             blurring_mask=blurring_mask,
+            use_mixed_precision=use_mixed_precision,
             xp=xp,
         )
 
+        # -------------------------------------------------------------------------
         # FFT convolution
+        # -------------------------------------------------------------------------
         fft_mapping_matrix_native = xp.fft.rfft2(
             mapping_matrix_native, s=fft_shape, axes=(0, 1)
         )
@@ -833,7 +861,9 @@ def convolved_mapping_matrix_from(
             axes=(0, 1),
         )
 
-        # crop back
+        # -------------------------------------------------------------------------
+        # Crop back to mask-shape
+        # -------------------------------------------------------------------------
         start_indices = tuple(
             (full_size - out_size) // 2
             for full_size, out_size in zip(full_shape, mask_shape)
@@ -846,8 +876,10 @@ def convolved_mapping_matrix_from(
             out_shape_full,
         )
 
-        # return slim form
-        return blurred_mapping_matrix_native[mask.slim_to_native_tuple]
+        # Return slim form
+        blurred_slim = blurred_mapping_matrix_native[mask.slim_to_native_tuple]
+
+        return blurred_slim
 
     def rescaled_with_odd_dimensions_from(
         self, rescale_factor: float, normalize: bool = False