update extracted_array_2d_from to not use numba

Author: Jammy2211

autoarray/__init__.py

@@ -63,6 +63,7 @@
 from .layout.layout import Layout2D
 from .structures.arrays.uniform_1d import Array1D
 from .structures.arrays.uniform_2d import Array2D
+from .structures.arrays.rgb import Array2DRGB
 from .structures.arrays.irregular import ArrayIrregular
 from .structures.grids.uniform_1d import Grid1D
 from .structures.grids.uniform_2d import Grid2D

autoarray/abstract_ndarray.py

@@ -282,7 +282,7 @@ def output_to_fits(self, file_path: str, overwrite: bool = False):
             If a file already exists at the path, if overwrite=True it is overwritten else an error is raised.
         """
         output_to_fits(
-            values=self.native.array,
+            values=self.native.array.astype("float"),
             file_path=file_path,
             overwrite=overwrite,
             header_dict=self.mask.header_dict,

autoarray/fixtures.py

@@ -67,6 +67,9 @@ def make_array_1d_7():
 def make_array_2d_7x7():
     return aa.Array2D.ones(shape_native=(7, 7), pixel_scales=(1.0, 1.0))
 
+def make_array_2d_rgb_7x7():
+    return aa.Array2DRGB(values=np.ones((7, 7, 3)), mask=make_mask_2d_7x7())
+
 
 def make_layout_2d_7x7():
     return aa.Layout2D(

autoarray/inversion/inversion/abstract.py

@@ -317,7 +317,7 @@ def operated_mapping_matrix(self) -> np.ndarray:
         If there are multiple linear objects, the blurred mapping matrices are stacked such that their simultaneous
         linear equations are solved simultaneously.
         """
-        return jnp.hstack(self.operated_mapping_matrix_list)
+        return np.hstack(self.operated_mapping_matrix_list)
 
     @cached_property
     @profile_func
@@ -446,6 +446,7 @@ def mapper_zero_pixel_list(self) -> np.ndarray:
                 )
         return mapper_zero_pixel_list
 
+
     @cached_property
     @profile_func
     def reconstruction(self) -> np.ndarray:
@@ -495,15 +496,16 @@ def reconstruction(self) -> np.ndarray:
                     values_to_solve, :
                 ][:, values_to_solve]
 
-                solutions = inversion_util.reconstruction_positive_only_from(
-                    data_vector=data_vector_input,
-                    curvature_reg_matrix=curvature_reg_matrix_input,
-                    settings=self.settings,
-                )
-
-                mask = values_to_solve.astype(bool)
+                solutions = np.zeros(np.shape(self.curvature_reg_matrix)[0])
 
-                return solutions[mask]
+                solutions[values_to_solve] = (
+                    inversion_util.reconstruction_positive_only_from(
+                        data_vector=data_vector_input,
+                        curvature_reg_matrix=curvature_reg_matrix_input,
+                        settings=self.settings,
+                    )
+                )
+                return solutions
             else:
                 solutions = inversion_util.reconstruction_positive_only_from(
                     data_vector=self.data_vector,
@@ -521,6 +523,81 @@ def reconstruction(self) -> np.ndarray:
             mapper_param_range_list=mapper_param_range_list,
         )
 
+    # @cached_property
+    # @profile_func
+    # def reconstruction(self) -> np.ndarray:
+    #     """
+    #     Solve the linear system [F + reg_coeff*H] S = D -> S = [F + reg_coeff*H]^-1 D given by equation (12)
+    #     of https://arxiv.org/pdf/astro-ph/0302587.pdf (Positive-Negative solution)
+    #
+    #     ============================================================================================
+    #
+    #     Solve the Eq.(2) of https://arxiv.org/pdf/astro-ph/0302587.pdf (Non-negative solution)
+    #     Find non-negative solution that minimizes |Z * S - x|^2.
+    #
+    #     We use fnnls (https://github.com/jvendrow/fnnls) to optimize the quadratic value. Two commonly used
+    #     variables in the code are defined as follows:
+    #         ZTZ := np.dot(Z.T, Z)
+    #         ZTx := np.dot(Z.T, x)
+    #     """
+    #     if self.settings.use_positive_only_solver:
+    #         """
+    #         For the new implementation, we now need to take out the cols and rows of
+    #         the curvature_reg_matrix that corresponds to the parameters we force to be 0.
+    #         Similar for the data vector.
+    #
+    #         What we actually doing is that we have set the correspoding cols of the Z to be 0.
+    #         As the curvature_reg_matrix = ZTZ, so the cols and rows are all taken out.
+    #         And the data_vector = ZTx, so the corresponding row is also taken out.
+    #         """
+    #
+    #         if self.settings.force_edge_pixels_to_zeros:
+    #             if self.settings.force_edge_image_pixels_to_zeros:
+    #                 ids_zeros = np.unique(
+    #                     np.append(
+    #                         self.mapper_edge_pixel_list, self.mapper_zero_pixel_list
+    #                     )
+    #                 )
+    #             else:
+    #                 ids_zeros = self.mapper_edge_pixel_list
+    #
+    #             values_to_solve = np.ones(
+    #                 np.shape(self.curvature_reg_matrix)[0], dtype=bool
+    #             )
+    #             values_to_solve[ids_zeros] = False
+    #
+    #             data_vector_input = self.data_vector[values_to_solve]
+    #
+    #             curvature_reg_matrix_input = self.curvature_reg_matrix[
+    #                 values_to_solve, :
+    #             ][:, values_to_solve]
+    #
+    #             solutions = inversion_util.reconstruction_positive_only_from(
+    #                 data_vector=data_vector_input,
+    #                 curvature_reg_matrix=curvature_reg_matrix_input,
+    #                 settings=self.settings,
+    #             )
+    #
+    #             mask = values_to_solve.astype(bool)
+    #
+    #             return solutions[mask]
+    #         else:
+    #             solutions = inversion_util.reconstruction_positive_only_from(
+    #                 data_vector=self.data_vector,
+    #                 curvature_reg_matrix=self.curvature_reg_matrix,
+    #                 settings=self.settings,
+    #             )
+    #
+    #             return solutions
+    #
+    #     mapper_param_range_list = self.param_range_list_from(cls=AbstractMapper)
+    #
+    #     return inversion_util.reconstruction_positive_negative_from(
+    #         data_vector=self.data_vector,
+    #         curvature_reg_matrix=self.curvature_reg_matrix,
+    #         mapper_param_range_list=mapper_param_range_list,
+    #     )
+
     @cached_property
     @profile_func
     def reconstruction_reduced(self) -> np.ndarray:

autoarray/inversion/inversion/inversion_util.py

@@ -41,6 +41,7 @@ def curvature_matrix_via_w_tilde_from(
     return np.dot(mapping_matrix.T, np.dot(w_tilde, mapping_matrix))
 
 
+@numba_util.jit()
 def curvature_matrix_with_added_to_diag_from(
     curvature_matrix: np.ndarray,
     value: float,
@@ -60,9 +61,35 @@ def curvature_matrix_with_added_to_diag_from(
     curvature_matrix
         The curvature matrix which is being constructed in order to solve a linear system of equations.
     """
-    return curvature_matrix.at[
-        no_regularization_index_list, no_regularization_index_list
-    ].add(value)
+
+    for i in no_regularization_index_list:
+        curvature_matrix[i, i] += value
+
+    return curvature_matrix
+
+
+# def curvature_matrix_with_added_to_diag_from(
+#     curvature_matrix: np.ndarray,
+#     value: float,
+#     no_regularization_index_list: Optional[List] = None,
+# ) -> np.ndarray:
+#     """
+#     It is common for the `curvature_matrix` computed to not be positive-definite, leading for the inversion
+#     via `np.linalg.solve` to fail and raise a `LinAlgError`.
+#
+#     In many circumstances, adding a small numerical value of `1.0e-8` to the diagonal of the `curvature_matrix`
+#     makes it positive definite, such that the inversion is performed without raising an error.
+#
+#     This function adds this numerical value to the diagonal of the curvature matrix.
+#
+#     Parameters
+#     ----------
+#     curvature_matrix
+#         The curvature matrix which is being constructed in order to solve a linear system of equations.
+#     """
+#     return curvature_matrix.at[
+#         no_regularization_index_list, no_regularization_index_list
+#     ].add(value)
 
 
 @numba_util.jit()
@@ -105,7 +132,7 @@ def curvature_matrix_via_mapping_matrix_from(
         Flattened 1D array of the noise-map used by the inversion during the fit.
     """
     array = mapping_matrix / noise_map[:, None]
-    curvature_matrix = jnp.dot(array.T, array)
+    curvature_matrix = np.dot(array.T, array)
 
     if add_to_curvature_diag and len(no_regularization_index_list) > 0:
         curvature_matrix = curvature_matrix_with_added_to_diag_from(
@@ -161,7 +188,7 @@ def mapped_reconstructed_data_via_mapping_matrix_from(
         The matrix representing the blurred mappings between sub-grid pixels and pixelization pixels.
 
     """
-    return jnp.dot(mapping_matrix, reconstruction)
+    return np.dot(mapping_matrix, reconstruction)
 
 
 def reconstruction_positive_negative_from(
@@ -272,31 +299,31 @@ def reconstruction_positive_only_from(
     Non-negative S that minimizes the Eq.(2) of https://arxiv.org/pdf/astro-ph/0302587.pdf.
     """
 
-    try:
-        return jaxnnls.solve_nnls_primal(curvature_reg_matrix, data_vector)
-    except (RuntimeError, np.linalg.LinAlgError, ValueError) as e:
-        raise exc.InversionException() from e
+    # try:
+    #     return jaxnnls.solve_nnls_primal(curvature_reg_matrix, data_vector)
+    # except (RuntimeError, np.linalg.LinAlgError, ValueError) as e:
+    #     raise exc.InversionException() from e
+
+    if len(data_vector):
+        try:
+            if settings.positive_only_uses_p_initial:
+                P_initial = np.linalg.solve(curvature_reg_matrix, data_vector) > 0
+            else:
+                P_initial = np.zeros(0, dtype=int)
+
+            reconstruction = fnnls_cholesky(
+                curvature_reg_matrix,
+                (data_vector).T,
+                P_initial=P_initial,
+            )
 
-    # if len(data_vector):
-    #     try:
-    #         if settings.positive_only_uses_p_initial:
-    #             P_initial = np.linalg.solve(curvature_reg_matrix, data_vector) > 0
-    #         else:
-    #             P_initial = np.zeros(0, dtype=int)
-    #
-    #         reconstruction = fnnls_cholesky(
-    #             curvature_reg_matrix,
-    #             (data_vector).T,
-    #             P_initial=P_initial,
-    #         )
-    #
-    #     except (RuntimeError, np.linalg.LinAlgError, ValueError) as e:
-    #         raise exc.InversionException() from e
-    #
-    # else:
-    #     raise exc.InversionException()
-    #
-    # return reconstruction
+        except (RuntimeError, np.linalg.LinAlgError, ValueError) as e:
+            raise exc.InversionException() from e
+
+    else:
+        raise exc.InversionException()
+
+    return reconstruction
 
 
 def preconditioner_matrix_via_mapping_matrix_from(

autoarray/mask/derive/zoom_2d.py

@@ -4,6 +4,7 @@
 
 if TYPE_CHECKING:
     from autoarray.structures.arrays.uniform_2d import Array2D
+    from autoarray.structures.arrays.rgb import Array2DRGB
 
 from autoarray.structures.arrays import array_2d_util
 from autoarray.structures.grids import grid_2d_util
@@ -229,7 +230,7 @@ def mask_2d_from(self, buffer: int = 1) -> "Mask2D":
             origin=self.mask.origin,
         )
 
-    def array_2d_from(self, array : Array2D, buffer: int = 1) -> Array2D:
+    def array_2d_from(self, array: Array2D, buffer: int = 1) -> Array2D:
         """
         Extract the 2D region of an array corresponding to the rectangle encompassing all unmasked values.
 
@@ -241,24 +242,82 @@ def array_2d_from(self, array : Array2D, buffer: int = 1) -> Array2D:
             The number pixels around the extracted array used as a buffer.
         """
         from autoarray.structures.arrays.uniform_2d import Array2D
+        from autoarray.structures.arrays.rgb import Array2DRGB
         from autoarray.mask.mask_2d import Mask2D
 
+        if isinstance(array, Array2DRGB):
+            return self.array_2d_rgb_from(array=array, buffer=buffer)
+
         extracted_array_2d = array_2d_util.extracted_array_2d_from(
-            array_2d=np.array(array.native),
+            array_2d=array.native.array,
             y0=self.region[0] - buffer,
             y1=self.region[1] + buffer,
             x0=self.region[2] - buffer,
             x1=self.region[3] + buffer,
         )
 
-        mask = Mask2D.all_false(
-            shape_native=extracted_array_2d.shape,
+        extracted_mask_2d = array_2d_util.extracted_array_2d_from(
+            array_2d=np.array(self.mask),
+            y0=self.region[0] - buffer,
+            y1=self.region[1] + buffer,
+            x0=self.region[2] - buffer,
+            x1=self.region[3] + buffer,
+        )
+
+        mask = Mask2D(
+            mask=extracted_mask_2d,
             pixel_scales=array.pixel_scales,
             origin=array.mask.mask_centre,
         )
 
-        arr = array_2d_util.convert_array_2d(
-            array_2d=extracted_array_2d, mask_2d=mask
+        arr = array_2d_util.convert_array_2d(array_2d=extracted_array_2d, mask_2d=mask)
+
+        return Array2D(values=arr, mask=mask, header=array.header)
+
+    def array_2d_rgb_from(self, array: Array2DRGB, buffer: int = 1) -> Array2DRGB:
+        """
+        Extract the 2D region of an RGB array corresponding to the rectangle encompassing all unmasked values.
+
+        This works the same as the `array_2d_from` method, but for RGB arrays, meaning that it iterates over the three
+        channels of the RGB array and extracts the region for each channel separately.
+
+        This is used to extract and visualize only the region of an RGB image that is used in an analysis.
+
+        Parameters
+        ----------
+        buffer
+            The number pixels around the extracted array used as a buffer.
+        """
+        from autoarray.structures.arrays.rgb import Array2DRGB
+        from autoarray.mask.mask_2d import Mask2D
+
+        for i in range(3):
+
+            extracted_array_2d = array_2d_util.extracted_array_2d_from(
+                array_2d=np.array(array.native[:, :, i]),
+                y0=self.region[0] - buffer,
+                y1=self.region[1] + buffer,
+                x0=self.region[2] - buffer,
+                x1=self.region[3] + buffer,
+            )
+
+            if i == 0:
+                array_2d_rgb = np.zeros((extracted_array_2d.shape[0], extracted_array_2d.shape[1], 3))
+
+            array_2d_rgb[:, :, i] = extracted_array_2d
+
+        extracted_mask_2d = array_2d_util.extracted_array_2d_from(
+            array_2d=np.array(self.mask),
+            y0=self.region[0] - buffer,
+            y1=self.region[1] + buffer,
+            x0=self.region[2] - buffer,
+            x1=self.region[3] + buffer,
+        )
+
+        mask = Mask2D(
+            mask=extracted_mask_2d,
+            pixel_scales=array.pixel_scales,
+            origin=array.mask.mask_centre,
         )
 
-        return Array2D(values=arr, mask=mask, header=array.header)
+        return Array2DRGB(values=array_2d_rgb.astype("int"), mask=mask)