Merge pull request #205 from Jammy2211/feature/linalg_mixed_precision

Feature/linalg mixed precision

Author: Jammy2211

autoarray/dataset/imaging/dataset.py

@@ -504,6 +504,10 @@ def apply_sparse_operator(
             Whether to use JAX to compute W-Tilde. This requires JAX to be installed.
         """
 
+        logger.info(
+            "IMAGING - Setting Up Sparse Operator For low Memory Pixelizations."
+        )
+
         sparse_operator = (
             inversion_imaging_util.ImagingSparseOperator.from_noise_map_and_psf(
                 data=self.data,

autoarray/dataset/interferometer/dataset.py

@@ -198,7 +198,7 @@ def apply_sparse_operator(
         if nufft_precision_operator is None:
 
             logger.info(
-                "INTERFEROMETER - Computing W-Tilde; runtime scales with visibility count and mask resolution, CPU run times may exceed hours."
+                "INTERFEROMETER - Computing NUFFT Precision Operator; runtime scales with visibility count and mask resolution, CPU run times may exceed hours."
             )
 
             nufft_precision_operator = self.psf_precision_operator_from(

autoarray/inversion/inversion/imaging/abstract.py

@@ -95,6 +95,7 @@ def operated_mapping_matrix_list(self) -> List[np.ndarray]:
                 self.psf.convolved_mapping_matrix_from(
                     mapping_matrix=linear_obj.mapping_matrix,
                     mask=self.mask,
+                    use_mixed_precision=self.settings.use_mixed_precision,
                     xp=self._xp,
                 )
                 if linear_obj.operated_mapping_matrix_override is None

autoarray/inversion/inversion/imaging/mapping.py

@@ -53,44 +53,6 @@ def __init__(
             xp=xp,
         )
 
-    @property
-    def _data_vector_mapper(self) -> np.ndarray:
-        """
-        Returns the `data_vector` of all mappers, a 1D vector whose values are solved for by the simultaneous
-        linear equations constructed by this object. The object is described in full in the method `data_vector`.
-
-        This method is used to compute part of the `data_vector` if there are also linear function list objects
-        in the inversion, and is separated into a separate method to enable preloading of the mapper `data_vector`.
-        """
-
-        if not self.has(cls=AbstractMapper):
-            return
-
-        data_vector = np.zeros(self.total_params)
-
-        mapper_list = self.cls_list_from(cls=AbstractMapper)
-        mapper_param_range_list = self.param_range_list_from(cls=AbstractMapper)
-
-        for i in range(len(mapper_list)):
-            mapper = mapper_list[i]
-            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
-            )
-
-            data_vector_mapper = (
-                inversion_imaging_util.data_vector_via_blurred_mapping_matrix_from(
-                    blurred_mapping_matrix=operated_mapping_matrix,
-                    image=self.data,
-                    noise_map=self.noise_map,
-                )
-            )
-
-            data_vector[param_range[0] : param_range[1],] = data_vector_mapper
-
-        return data_vector
-
     @cached_property
     def data_vector(self) -> np.ndarray:
         """
@@ -111,53 +73,6 @@ def data_vector(self) -> np.ndarray:
             noise_map=self.noise_map.array,
         )
 
-    @property
-    def _curvature_matrix_mapper_diag(self) -> Optional[np.ndarray]:
-        """
-        Returns the diagonal regions of the `curvature_matrix`, a 2D matrix which uses the mappings between the data
-        and the linear objects to construct the simultaneous linear equations. The object is described in full in
-        the method `curvature_matrix`.
-
-        This method computes the diagonal entries of all mapper objects in the `curvature_matrix`. It is separate from
-        other calculations to enable preloading of this calculation.
-        """
-
-        if not self.has(cls=AbstractMapper):
-            return None
-
-        curvature_matrix = np.zeros((self.total_params, self.total_params))
-
-        mapper_list = self.cls_list_from(cls=AbstractMapper)
-        mapper_param_range_list = self.param_range_list_from(cls=AbstractMapper)
-
-        for i in range(len(mapper_list)):
-            mapper_i = mapper_list[i]
-            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
-            )
-
-            diag = inversion_util.curvature_matrix_via_mapping_matrix_from(
-                mapping_matrix=operated_mapping_matrix,
-                noise_map=self.noise_map,
-                settings=self.settings,
-                add_to_curvature_diag=True,
-                no_regularization_index_list=self.no_regularization_index_list,
-                xp=self._xp,
-            )
-
-            curvature_matrix[
-                mapper_param_range_i[0] : mapper_param_range_i[1],
-                mapper_param_range_i[0] : mapper_param_range_i[1],
-            ] = diag
-
-        curvature_matrix = inversion_util.curvature_matrix_mirrored_from(
-            curvature_matrix=curvature_matrix, xp=self._xp
-        )
-
-        return curvature_matrix
-
     @cached_property
     def curvature_matrix(self):
         """

autoarray/inversion/inversion/interferometer/mapping.py

@@ -88,12 +88,14 @@ def curvature_matrix(self) -> np.ndarray:
         real_curvature_matrix = inversion_util.curvature_matrix_via_mapping_matrix_from(
             mapping_matrix=self.operated_mapping_matrix.real,
             noise_map=self.noise_map.real,
+            settings=self.settings,
             xp=self._xp,
         )
 
         imag_curvature_matrix = inversion_util.curvature_matrix_via_mapping_matrix_from(
             mapping_matrix=self.operated_mapping_matrix.imag,
             noise_map=self.noise_map.imag,
+            settings=self.settings,
             xp=self._xp,
         )
 

autoarray/inversion/inversion/inversion_util.py

@@ -78,11 +78,11 @@ def curvature_matrix_mirrored_from(curvature_matrix: np.ndarray, xp=np) -> np.nd
 
 
 def curvature_matrix_via_mapping_matrix_from(
-    mapping_matrix: np.ndarray,
-    noise_map: np.ndarray,
+    mapping_matrix: "np.ndarray",
+    noise_map: "np.ndarray",
     add_to_curvature_diag: bool = False,
     no_regularization_index_list: Optional[List] = None,
-    settings: SettingsInversion = SettingsInversion(),
+    settings: "SettingsInversion" = SettingsInversion(),
     xp=np,
 ) -> np.ndarray:
     """
@@ -97,10 +97,26 @@ 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 add_to_curvature_diag and len(no_regularization_index_list) > 0:
+    # NumPy path: keep it simple + stable
+    if xp is np:
+        A = mapping_matrix / noise_map[:, None]
+        curvature_matrix = xp.dot(A.T, A)
+    else:
+        # Choose compute dtype
+
+        compute_dtype = xp.float32 if settings.use_mixed_precision else xp.float64
+        out_dtype = xp.float64  # always return float64 for downstream stability
+
+        A = mapping_matrix
+        w = (1.0 / noise_map).astype(compute_dtype)
+        A = A * w[:, None]
+        curvature_matrix = xp.dot(A.T, A).astype(out_dtype)
+
+    if (
+        add_to_curvature_diag
+        and no_regularization_index_list
+        and len(no_regularization_index_list) > 0
+    ):
         curvature_matrix = curvature_matrix_with_added_to_diag_from(
             curvature_matrix=curvature_matrix,
             value=settings.no_regularization_add_to_curvature_diag_value,

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
@@ -41,6 +48,7 @@ def __init__(
             For an interferometer inversion using the linear operators method, sets the maximum number of iterations
             of the solver (this input does nothing for dataset data and other interferometer methods).
         """
+        self.use_mixed_precision = use_mixed_precision
         self._use_positive_only_solver = use_positive_only_solver
         self._positive_only_uses_p_initial = positive_only_uses_p_initial
         self._use_border_relocator = use_border_relocator

autoarray/inversion/linear_obj/func_list.py

@@ -7,6 +7,7 @@
 from autoarray.inversion.linear_obj.neighbors import Neighbors
 from autoarray.inversion.linear_obj.unique_mappings import UniqueMappings
 from autoarray.inversion.regularization.abstract import AbstractRegularization
+from autoarray.inversion.inversion.settings import SettingsInversion
 from autoarray.type import Grid1D2DLike
 
 
@@ -15,6 +16,7 @@ def __init__(
         self,
         grid: Grid1D2DLike,
         regularization: Optional[AbstractRegularization],
+        settings=SettingsInversion(),
         xp=np,
     ):
         """
@@ -45,6 +47,7 @@ def __init__(
         super().__init__(regularization=regularization, xp=xp)
 
         self.grid = grid
+        self.settings = settings
 
     @cached_property
     def neighbors(self) -> Neighbors:

autoarray/inversion/pixelization/mappers/abstract.py

@@ -11,6 +11,7 @@
 from autoarray.inversion.pixelization.border_relocator import BorderRelocator
 from autoarray.inversion.pixelization.mappers.mapper_grids import MapperGrids
 from autoarray.inversion.regularization.abstract import AbstractRegularization
+from autoarray.inversion.inversion.settings import SettingsInversion
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.grids.uniform_2d import Grid2D
 from autoarray.structures.mesh.abstract_2d import Abstract2DMesh
@@ -25,6 +26,7 @@ def __init__(
         mapper_grids: MapperGrids,
         regularization: Optional[AbstractRegularization],
         border_relocator: BorderRelocator,
+        settings: SettingsInversion = SettingsInversion(),
         preloads=None,
         xp=np,
     ):
@@ -90,6 +92,7 @@ def __init__(
         self.border_relocator = border_relocator
         self.mapper_grids = mapper_grids
         self.preloads = preloads
+        self.settings = settings
 
     @property
     def params(self) -> int:
@@ -265,6 +268,7 @@ def mapping_matrix(self) -> np.ndarray:
             total_mask_pixels=self.over_sampler.mask.pixels_in_mask,
             slim_index_for_sub_slim_index=self.slim_index_for_sub_slim_index,
             sub_fraction=self.over_sampler.sub_fraction.array,
+            use_mixed_precision=self.settings.use_mixed_precision,
             xp=self._xp,
         )
 

autoarray/inversion/pixelization/mappers/factory.py

@@ -4,6 +4,7 @@
 from autoarray.inversion.pixelization.mappers.mapper_grids import MapperGrids
 from autoarray.inversion.pixelization.border_relocator import BorderRelocator
 from autoarray.inversion.regularization.abstract import AbstractRegularization
+from autoarray.inversion.inversion.settings import SettingsInversion
 from autoarray.structures.mesh.rectangular_2d import Mesh2DRectangular
 from autoarray.structures.mesh.rectangular_2d_uniform import Mesh2DRectangularUniform
 from autoarray.structures.mesh.delaunay_2d import Mesh2DDelaunay
@@ -13,6 +14,7 @@ def mapper_from(
     mapper_grids: MapperGrids,
     regularization: Optional[AbstractRegularization],
     border_relocator: Optional[BorderRelocator] = None,
+    settings=SettingsInversion(),
     preloads=None,
     xp=np,
 ):
@@ -53,20 +55,25 @@ def mapper_from(
             mapper_grids=mapper_grids,
             border_relocator=border_relocator,
             regularization=regularization,
+            settings=settings,
+            preloads=preloads,
             xp=xp,
         )
     elif isinstance(mapper_grids.source_plane_mesh_grid, Mesh2DRectangular):
         return MapperRectangular(
             mapper_grids=mapper_grids,
             border_relocator=border_relocator,
             regularization=regularization,
+            settings=settings,
+            preloads=preloads,
             xp=xp,
         )
     elif isinstance(mapper_grids.source_plane_mesh_grid, Mesh2DDelaunay):
         return MapperDelaunay(
             mapper_grids=mapper_grids,
             border_relocator=border_relocator,
             regularization=regularization,
+            settings=settings,
             preloads=preloads,
             xp=xp,
         )