add preloading in order to pass mapper indexes

Author: Jammy2211

autoarray/__init__.py

@@ -65,6 +65,7 @@
 from .operators.contour import Grid2DContour
 from .layout.layout import Layout1D
 from .layout.layout import Layout2D
+from .preloads import Preloads
 from .structures.arrays.uniform_1d import Array1D
 from .structures.arrays.uniform_2d import Array2D
 from .structures.arrays.rgb import Array2DRGB

autoarray/inversion/inversion/abstract.py

@@ -13,6 +13,7 @@
 from autoarray.inversion.pixelization.mappers.abstract import AbstractMapper
 from autoarray.inversion.regularization.abstract import AbstractRegularization
 from autoarray.inversion.inversion.settings import SettingsInversion
+from autoarray.preloads import Preloads
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.visibilities import Visibilities
 
@@ -27,6 +28,7 @@ def __init__(
         dataset: Union[Imaging, Interferometer, DatasetInterface],
         linear_obj_list: List[LinearObj],
         settings: SettingsInversion = SettingsInversion(),
+        preloads: Preloads = None,
     ):
         """
         An `Inversion` reconstructs an input dataset using a list of linear objects (e.g. a list of analytic functions
@@ -83,6 +85,8 @@ def __init__(
 
         self.settings = settings
 
+        self.preloads = preloads or Preloads()
+
     @property
     def data(self):
         return self.dataset.data
@@ -267,6 +271,22 @@ def no_regularization_index_list(self) -> List[int]:
 
         return no_regularization_index_list
 
+    @property
+    def mapper_index_list(self) -> List[int]:
+
+        if self.preloads.mapper_index_list is not None:
+            return self.preloads.mapper_index_list
+
+        mapper_index_list = []
+
+        param_range_list = self.param_range_list_from(cls=AbstractMapper)
+
+        for param_range in param_range_list:
+
+            mapper_index_list += range(param_range[0], param_range[1])
+
+        return mapper_index_list
+
     @property
     def mask(self) -> Array2D:
         return self.data.mask
@@ -358,14 +378,10 @@ def regularization_matrix_reduced(self) -> Optional[np.ndarray]:
             return self.regularization_matrix
 
         # ids of values which are on edge so zero-d and not solved for.
-        ids_to_not_solve_for = jnp.array(self.no_regularization_index_list, dtype=int)
-
-        # Create a boolean mask: True = keep, False = ignore
-        mask = jnp.ones(self.data_vector.shape[0], dtype=bool).at[ids_to_not_solve_for].set(False)
+        ids_to_keep = jnp.array(self.mapper_index_list, dtype=int)
 
         # Zero rows and columns in the matrix we want to ignore
-        mask_matrix = mask[:, None] * mask[None, :]
-        return self.regularization_matrix * mask_matrix
+        return self.regularization_matrix[ids_to_keep][:, ids_to_keep]
 
     @cached_property
     def curvature_reg_matrix(self) -> np.ndarray:
@@ -383,25 +399,28 @@ def curvature_reg_matrix(self) -> np.ndarray:
         return jnp.add(self.curvature_matrix, self.regularization_matrix)
 
     @cached_property
-    def curvature_reg_matrix_reduced(self) -> np.ndarray:
+    def curvature_reg_matrix_reduced(self) -> Optional[np.ndarray]:
         """
-        The linear system of equations solves for F + regularization_coefficient*H, which is computed below.
+        The regularization matrix H is used to impose smoothness on our inversion's reconstruction. This enters the
+        linear algebra system we solve for using D and F above and is given by
+        equation (12) in https://arxiv.org/pdf/astro-ph/0302587.pdf.
 
-        This is the curvature reg matrix for only the mappers, which is necessary for computing the log det
-        term without the linear light profiles included.
+        A complete description of regularization is given in the `regularization.py` and `regularization_util.py`
+        modules.
+
+        For multiple mappers, the regularization matrix is computed as the block diagonal of each individual mapper.
+        The scipy function `block_diag` has an overhead associated with it and if there is only one mapper and
+        regularization it is bypassed.
         """
+
         if self.all_linear_obj_have_regularization:
             return self.curvature_reg_matrix
 
         # ids of values which are on edge so zero-d and not solved for.
-        ids_to_not_solve_for = jnp.array(self.no_regularization_index_list, dtype=int)
-
-        # Create a boolean mask: True = keep, False = ignore
-        mask = jnp.ones(self.data_vector.shape[0], dtype=bool).at[ids_to_not_solve_for].set(False)
+        ids_to_keep = jnp.array(self.mapper_index_list, dtype=int)
 
         # Zero rows and columns in the matrix we want to ignore
-        mask_matrix = mask[:, None] * mask[None, :]
-        return self.curvature_reg_matrix * mask_matrix
+        return self.regularization_matrix[ids_to_keep][:, ids_to_keep]
 
     @property
     def mapper_zero_pixel_list(self) -> np.ndarray:
@@ -454,10 +473,14 @@ def reconstruction(self) -> np.ndarray:
             ):
 
                 # ids of values which are on edge so zero-d and not solved for.
-                ids_to_not_solve_for = jnp.array(self.mapper_edge_pixel_list, dtype=int)
+                ids_to_remove = jnp.array(self.mapper_edge_pixel_list, dtype=int)
 
                 # Create a boolean mask: True = keep, False = ignore
-                mask = jnp.ones(self.data_vector.shape[0], dtype=bool).at[ids_to_not_solve_for].set(False)
+                mask = (
+                    jnp.ones(self.data_vector.shape[0], dtype=bool)
+                    .at[ids_to_remove]
+                    .set(False)
+                )
 
                 # Zero out entries we don't want to solve for
                 data_vector_masked = self.data_vector * mask
@@ -502,13 +525,10 @@ def reconstruction_reduced(self) -> np.ndarray:
             return self.reconstruction
 
         # ids of values which are on edge so zero-d and not solved for.
-        ids_to_not_solve_for = jnp.array(self.no_regularization_index_list, dtype=int)
-
-        # Create a boolean mask: True = keep, False = ignore
-        mask = jnp.ones(self.reconstruction.shape[0], dtype=bool).at[ids_to_not_solve_for].set(False)
+        ids_to_keep = jnp.array(self.mapper_index_list, dtype=int)
 
-        # Zero out entries we don't want to solve for
-        return self.reconstruction * mask
+        # Zero rows and columns in the matrix we want to ignore
+        return self.reconstruction[ids_to_keep]
 
     @property
     def reconstruction_dict(self) -> Dict[LinearObj, np.ndarray]:
@@ -651,12 +671,6 @@ def regularization_term(self) -> float:
         if not self.has(cls=AbstractRegularization):
             return 0.0
 
-        print(self.reconstruction_reduced)
-        print(self.regularization_matrix_reduced)
-
-        print(self.reconstruction_reduced.shape)
-        print(self.regularization_matrix_reduced.shape)
-
         return np.matmul(
             self.reconstruction_reduced.T,
             np.matmul(self.regularization_matrix_reduced, self.reconstruction_reduced),
@@ -674,7 +688,7 @@ def log_det_curvature_reg_matrix_term(self) -> float:
 
         try:
             return 2.0 * np.sum(
-                np.log(np.diag(np.linalg.cholesky(self.curvature_reg_matrix_reduced)))
+                jnp.log(jnp.diag(jnp.linalg.cholesky(self.curvature_reg_matrix_reduced)))
             )
         except np.linalg.LinAlgError as e:
             raise exc.InversionException() from e

autoarray/inversion/inversion/factory.py

@@ -14,13 +14,15 @@
 from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList
 from autoarray.inversion.inversion.imaging.w_tilde import InversionImagingWTilde
 from autoarray.inversion.inversion.settings import SettingsInversion
+from autoarray.preloads import Preloads
 from autoarray.structures.arrays.uniform_2d import Array2D
 
 
 def inversion_from(
     dataset: Union[Imaging, Interferometer, DatasetInterface],
     linear_obj_list: List[LinearObj],
     settings: SettingsInversion = SettingsInversion(),
+    preloads :Preloads = None,
 ):
     """
     Factory which given an input dataset and list of linear objects, creates an `Inversion`.
@@ -55,6 +57,7 @@ def inversion_from(
             dataset=dataset,
             linear_obj_list=linear_obj_list,
             settings=settings,
+            preloads=preloads,
         )
 
     return inversion_interferometer_from(
@@ -68,6 +71,7 @@ def inversion_imaging_from(
     dataset,
     linear_obj_list: List[LinearObj],
     settings: SettingsInversion = SettingsInversion(),
+    preloads : Preloads = None,
 ):
     """
     Factory which given an input `Imaging` dataset and list of linear objects, creates an `InversionImaging`.
@@ -126,6 +130,7 @@ def inversion_imaging_from(
         dataset=dataset,
         linear_obj_list=linear_obj_list,
         settings=settings,
+        preloads=preloads,
     )
 
 

autoarray/inversion/inversion/imaging/abstract.py

@@ -10,6 +10,7 @@
 from autoarray.inversion.inversion.abstract import AbstractInversion
 from autoarray.inversion.linear_obj.linear_obj import LinearObj
 from autoarray.inversion.inversion.settings import SettingsInversion
+from autoarray.preloads import Preloads
 
 from autoarray.inversion.inversion.imaging import inversion_imaging_util
 
@@ -20,6 +21,7 @@ def __init__(
         dataset: Union[Imaging, DatasetInterface],
         linear_obj_list: List[LinearObj],
         settings: SettingsInversion = SettingsInversion(),
+        preloads: Preloads = None,
     ):
         """
         An `Inversion` reconstructs an input dataset using a list of linear objects (e.g. a list of analytic functions
@@ -66,6 +68,7 @@ def __init__(
             dataset=dataset,
             linear_obj_list=linear_obj_list,
             settings=settings,
+            preloads=preloads,
         )
 
     @property

autoarray/inversion/inversion/imaging/mapping.py

@@ -9,6 +9,7 @@
 from autoarray.inversion.linear_obj.linear_obj import LinearObj
 from autoarray.inversion.pixelization.mappers.abstract import AbstractMapper
 from autoarray.inversion.inversion.settings import SettingsInversion
+from autoarray.preloads import Preloads
 from autoarray.structures.arrays.uniform_2d import Array2D
 
 from autoarray.inversion.inversion import inversion_util
@@ -21,6 +22,7 @@ def __init__(
         dataset: Union[Imaging, DatasetInterface],
         linear_obj_list: List[LinearObj],
         settings: SettingsInversion = SettingsInversion(),
+        preloads: Preloads = None,
     ):
         """
         Constructs linear equations (via vectors and matrices) which allow for sets of simultaneous linear equations
@@ -46,6 +48,7 @@ def __init__(
             dataset=dataset,
             linear_obj_list=linear_obj_list,
             settings=settings,
+            preloads=preloads,
         )
 
     @property

autoarray/inversion/inversion/imaging/w_tilde.py

@@ -94,7 +94,9 @@ def _data_vector_mapper(self) -> np.ndarray:
             data_vector_mapper = (
                 inversion_imaging_util.data_vector_via_w_tilde_data_imaging_from(
                     w_tilde_data=self.w_tilde_data,
-                    data_to_pix_unique=np.array(mapper.unique_mappings.data_to_pix_unique),
+                    data_to_pix_unique=np.array(
+                        mapper.unique_mappings.data_to_pix_unique
+                    ),
                     data_weights=np.array(mapper.unique_mappings.data_weights),
                     pix_lengths=np.array(mapper.unique_mappings.pix_lengths),
                     pix_pixels=mapper.params,
@@ -276,7 +278,9 @@ def _curvature_matrix_mapper_diag(self) -> Optional[np.ndarray]:
                 curvature_preload=self.w_tilde.curvature_preload,
                 curvature_indexes=self.w_tilde.indexes,
                 curvature_lengths=self.w_tilde.lengths,
-                data_to_pix_unique=np.array(mapper_i.unique_mappings.data_to_pix_unique),
+                data_to_pix_unique=np.array(
+                    mapper_i.unique_mappings.data_to_pix_unique
+                ),
                 data_weights=np.array(mapper_i.unique_mappings.data_weights),
                 pix_lengths=np.array(mapper_i.unique_mappings.pix_lengths),
                 pix_pixels=mapper_i.params,

autoarray/inversion/pixelization/border_relocator.py

@@ -64,7 +64,7 @@ def sub_slim_indexes_for_slim_index_via_mask_2d_from(
 
 
 def furthest_grid_2d_slim_index_from(
-        grid_2d_slim: np.ndarray, slim_indexes: np.ndarray, coordinate: Tuple[float, float]
+    grid_2d_slim: np.ndarray, slim_indexes: np.ndarray, coordinate: Tuple[float, float]
 ) -> int:
     """
     Returns the index in `slim_indexes` corresponding to the 2D point in `grid_2d_slim`
@@ -87,7 +87,7 @@ def furthest_grid_2d_slim_index_from(
     subgrid = grid_2d_slim[slim_indexes]
     dy = subgrid[:, 0] - coordinate[0]
     dx = subgrid[:, 1] - coordinate[1]
-    squared_distances = dx ** 2 + dy ** 2
+    squared_distances = dx**2 + dy**2
 
     max_dist = np.max(squared_distances)
 
@@ -100,7 +100,6 @@ def furthest_grid_2d_slim_index_from(
     return slim_indexes[max_index]
 
 
-
 def sub_border_pixel_slim_indexes_from(
     mask_2d: np.ndarray, sub_size: Array2D
 ) -> np.ndarray:
@@ -154,12 +153,10 @@ def sub_border_pixel_slim_indexes_from(
             int(border_pixel)
         ]
 
-        sub_border_pixels[border_1d_index] = (
-            furthest_grid_2d_slim_index_from(
-                grid_2d_slim=sub_grid_2d_slim,
-                slim_indexes=sub_border_pixels_of_border_pixel,
-                coordinate=mask_centre,
-            )
+        sub_border_pixels[border_1d_index] = furthest_grid_2d_slim_index_from(
+            grid_2d_slim=sub_grid_2d_slim,
+            slim_indexes=sub_border_pixels_of_border_pixel,
+            coordinate=mask_centre,
         )
 
     return sub_border_pixels
@@ -235,32 +232,32 @@ def relocated_grid_from(grid, border_grid):
     border_origin = jnp.mean(border_grid, axis=0)
 
     # Radii from origin
-    grid_radii = jnp.linalg.norm(grid - border_origin, axis=1)                   # (N,)
-    border_radii = jnp.linalg.norm(border_grid - border_origin, axis=1)          # (M,)
+    grid_radii = jnp.linalg.norm(grid - border_origin, axis=1)  # (N,)
+    border_radii = jnp.linalg.norm(border_grid - border_origin, axis=1)  # (M,)
     border_min_radius = jnp.min(border_radii)
 
     # Determine which points are outside
     outside_mask = grid_radii > border_min_radius  # (N,)
 
     # To compute nearest border point for each grid point, we must do it for all and then mask later
     # Compute all distances: (N, M)
-    diffs = grid[:, None, :] - border_grid[None, :, :]     # (N, M, 2)
-    dists_squared = jnp.sum(diffs**2, axis=2)              # (N, M)
-    closest_indices = jnp.argmin(dists_squared, axis=1)    # (N,)
+    diffs = grid[:, None, :] - border_grid[None, :, :]  # (N, M, 2)
+    dists_squared = jnp.sum(diffs**2, axis=2)  # (N, M)
+    closest_indices = jnp.argmin(dists_squared, axis=1)  # (N,)
 
     # Get border radius for closest border point to each grid point
-    matched_border_radii = border_radii[closest_indices]   # (N,)
+    matched_border_radii = border_radii[closest_indices]  # (N,)
 
     # Ratio of border to grid radius
-    move_factors = matched_border_radii / grid_radii       # (N,)
+    move_factors = matched_border_radii / grid_radii  # (N,)
 
     # Only move if:
     #   - the point is outside the border
     #   - the matched border point is closer to the origin (i.e. move_factor < 1)
     apply_move = jnp.logical_and(outside_mask, move_factors < 1.0)  # (N,)
 
     # Compute moved positions (for all points, but will select with mask)
-    direction_vectors = grid - border_origin                # (N, 2)
+    direction_vectors = grid - border_origin  # (N, 2)
     moved_grid = move_factors[:, None] * direction_vectors + border_origin  # (N, 2)
 
     # Select which grid points to move

autoarray/inversion/pixelization/mappers/abstract.py

@@ -249,9 +249,13 @@ def unique_mappings(self) -> UniqueMappings:
             pix_lengths,
         ) = mapper_util.data_slim_to_pixelization_unique_from(
             data_pixels=self.over_sampler.mask.pixels_in_mask,
-            pix_indexes_for_sub_slim_index=np.array(self.pix_indexes_for_sub_slim_index),
+            pix_indexes_for_sub_slim_index=np.array(
+                self.pix_indexes_for_sub_slim_index
+            ),
             pix_sizes_for_sub_slim_index=np.array(self.pix_sizes_for_sub_slim_index),
-            pix_weights_for_sub_slim_index=np.array(self.pix_weights_for_sub_slim_index),
+            pix_weights_for_sub_slim_index=np.array(
+                self.pix_weights_for_sub_slim_index
+            ),
             pix_pixels=self.params,
             sub_size=np.array(self.over_sampler.sub_size).astype("int"),
         )
@@ -356,8 +360,12 @@ def data_weight_total_for_pix_from(self) -> np.ndarray:
         """
 
         return mapper_util.data_weight_total_for_pix_from(
-            pix_indexes_for_sub_slim_index=np.array(self.pix_indexes_for_sub_slim_index),
-            pix_weights_for_sub_slim_index=np.array(self.pix_weights_for_sub_slim_index),
+            pix_indexes_for_sub_slim_index=np.array(
+                self.pix_indexes_for_sub_slim_index
+            ),
+            pix_weights_for_sub_slim_index=np.array(
+                self.pix_weights_for_sub_slim_index
+            ),
             pixels=self.pixels,
         )