From 0edcf7d84e2be00486c3beec3752582eebf346df Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 12 Jun 2025 18:49:11 +0100 Subject: [PATCH 01/14] updated factory test on two mappers to use more stable solver --- autoarray/inversion/inversion/abstract.py | 13 +------ .../inversion/inversion/imaging/mapping.py | 1 - .../inversion/inversion/inversion_util.py | 34 ++++--------------- .../operators/over_sampling/over_sampler.py | 22 ++++++------ .../inversion/inversion/test_factory.py | 27 +++++++-------- 5 files changed, 33 insertions(+), 64 deletions(-) diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index 9398775dd..905aa5e16 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -73,17 +73,6 @@ def __init__( A dictionary which contains timing of certain functions calls which is used for profiling. """ - # try: - # import numba - # except ModuleNotFoundError: - # raise exc.InversionException( - # "Inversion functionality (linear light profiles, pixelized reconstructions) is " - # "disabled if numba is not installed.\n\n" - # "This is because the run-times without numba are too slow.\n\n" - # "Please install numba, which is described at the following web page:\n\n" - # "https://pyautolens.readthedocs.io/en/latest/installation/overview.html" - # ) - self.dataset = dataset self.linear_obj_list = linear_obj_list @@ -317,7 +306,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 np.hstack(self.operated_mapping_matrix_list) + return jnp.hstack(self.operated_mapping_matrix_list) @cached_property @profile_func diff --git a/autoarray/inversion/inversion/imaging/mapping.py b/autoarray/inversion/inversion/imaging/mapping.py index 12881a944..2ec0df160 100644 --- a/autoarray/inversion/inversion/imaging/mapping.py +++ b/autoarray/inversion/inversion/imaging/mapping.py @@ -109,7 +109,6 @@ def data_vector(self) -> np.ndarray: The calculation is described in more detail in `inversion_util.data_vector_via_blurred_mapping_matrix_from`. """ - return inversion_imaging_util.data_vector_via_blurred_mapping_matrix_from( blurred_mapping_matrix=self.operated_mapping_matrix, image=self.data.array, diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index 246e2cf56..94989842f 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/autoarray/inversion/inversion/inversion_util.py @@ -233,7 +233,9 @@ def reconstruction_positive_negative_from( The curvature_matrix plus regularization matrix, overwriting the curvature_matrix in memory. """ try: - reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector) + # print(curvature_reg_matrix) + # print(data_vector) + reconstruction = jnp.linalg.solve(curvature_reg_matrix, data_vector) except np.linalg.LinAlgError as e: raise exc.InversionException() from e @@ -299,32 +301,10 @@ 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 - - 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 - + try: + return jaxnnls.solve_nnls_primal(curvature_reg_matrix, data_vector) + except (RuntimeError, np.linalg.LinAlgError, ValueError) as e: + raise exc.InversionException() from e def preconditioner_matrix_via_mapping_matrix_from( mapping_matrix: np.ndarray, diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py index b78bb0829..499c6b37d 100644 --- a/autoarray/operators/over_sampling/over_sampler.py +++ b/autoarray/operators/over_sampling/over_sampler.py @@ -147,6 +147,17 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]): over_sample_size=sub_size, mask=mask ) + # Used for JAX based adaptive over sampling. + + # Define group sizes + group_sizes = np.array(self.sub_size.array ** 2) + + # Compute the cumulative sum of group sizes to get split points + self.split_indices = np.cumsum(group_sizes) + + # Ensure correct concatenation by making 0 a JAX array + self.start_indices = np.concatenate((np.array([0]), self.split_indices[:-1])) + @property def sub_is_uniform(self) -> bool: """ @@ -234,20 +245,11 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D": ).mean(axis=1) else: - # Define group sizes - group_sizes = jnp.array(self.sub_size.array**2) - - # Compute the cumulative sum of group sizes to get split points - split_indices = jnp.cumsum(group_sizes) - - # Ensure correct concatenation by making 0 a JAX array - start_indices = jnp.concatenate((jnp.array([0]), split_indices[:-1])) - # Compute the group means binned_array_2d = jnp.array( [ array[start:end].mean() - for start, end in zip(start_indices, split_indices) + for start, end in zip(self.start_indices, self.split_indices) ] ) diff --git a/test_autoarray/inversion/inversion/test_factory.py b/test_autoarray/inversion/inversion/test_factory.py index 743833b6a..a03be939e 100644 --- a/test_autoarray/inversion/inversion/test_factory.py +++ b/test_autoarray/inversion/inversion/test_factory.py @@ -43,6 +43,7 @@ def test__inversion_imaging__via_linear_obj_func_list(masked_imaging_7x7_no_blur linear_obj = aa.m.MockLinearObjFuncList( parameters=2, grid=grid, mapping_matrix=np.full(fill_value=0.5, shape=(9, 2)) ) + linear_obj.mapping_matrix[0, 0] = 1.0 inversion = aa.Inversion( dataset=masked_imaging_7x7_no_blur, @@ -52,8 +53,8 @@ def test__inversion_imaging__via_linear_obj_func_list(masked_imaging_7x7_no_blur assert isinstance(inversion.linear_obj_list[0], aa.m.MockLinearObjFuncList) assert isinstance(inversion, aa.InversionImagingMapping) + assert inversion.reconstruction == pytest.approx(np.array([0.0, 2.0]), abs=1.0e-4) assert inversion.mapped_reconstructed_image == pytest.approx(np.ones(9), 1.0e-4) - assert inversion.reconstruction == pytest.approx(np.array([1.0, 1.0]), 1.0e-4) def test__inversion_imaging__via_mapper( @@ -476,9 +477,11 @@ def test__inversion_matrices__x2_mappers( delaunay_mapper_9_3x3, regularization_constant, ): + inversion = aa.Inversion( dataset=masked_imaging_7x7_no_blur, linear_obj_list=[rectangular_mapper_7x7_3x3, delaunay_mapper_9_3x3], + settings=aa.SettingsInversion(use_positive_only_solver=True), ) assert ( @@ -524,27 +527,23 @@ def test__inversion_matrices__x2_mappers( assert (inversion.regularization_matrix[0:9, 9:18] == np.zeros((9, 9))).all() assert (inversion.regularization_matrix[9:18, 0:9] == np.zeros((9, 9))).all() - reconstruction_0 = 0.5 * np.ones(9) - reconstruction_1 = 0.5 * np.ones(9) - - assert inversion.reconstruction_dict[rectangular_mapper_7x7_3x3] == pytest.approx( - reconstruction_0, 1.0e-4 + assert inversion.reconstruction_dict[rectangular_mapper_7x7_3x3][4] == pytest.approx( + 0.05594123, 1.0e-4 ) - assert inversion.reconstruction_dict[delaunay_mapper_9_3x3] == pytest.approx( - reconstruction_1, 1.0e-4 + assert inversion.reconstruction_dict[delaunay_mapper_9_3x3][4] == pytest.approx( + 0.04686388, 1.0e-4 ) - assert inversion.reconstruction == pytest.approx( - np.concatenate([reconstruction_0, reconstruction_1]), 1.0e-4 + assert inversion.reconstruction[13] == pytest.approx( + 0.04686388, 1.0e-4 ) assert inversion.mapped_reconstructed_data_dict[ rectangular_mapper_7x7_3x3 - ] == pytest.approx(0.5 * np.ones(9), 1.0e-4) + ][4] == pytest.approx(0.05594123, 1.0e-4) assert inversion.mapped_reconstructed_data_dict[ delaunay_mapper_9_3x3 - ] == pytest.approx(0.5 * np.ones(9), 1.0e-4) - assert inversion.mapped_reconstructed_image == pytest.approx(np.ones(9), 1.0e-4) - + ][3] == pytest.approx(0.01521323, 1.0e-4) + assert inversion.mapped_reconstructed_image[4] == pytest.approx(0.10494037076075, 1.0e-4) def test__inversion_imaging__positive_only_solver(masked_imaging_7x7_no_blur): mask = masked_imaging_7x7_no_blur.mask From 0521b1b713aa1b352c86fbce6cca19fd583acbed Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 13 Jun 2025 10:36:07 +0100 Subject: [PATCH 02/14] fix test__identical_inversion_source_and_image_loops --- .../interferometer/test_inversion_interferometer_util.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index 28c45a9bd..ab9ce32b9 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -347,13 +347,13 @@ def test__identical_inversion_source_and_image_loops(): inversion_image_loop = aa.Inversion( dataset=dataset, linear_obj_list=[mapper], - settings=aa.SettingsInversion(use_w_tilde=True, use_source_loop=False), + settings=aa.SettingsInversion(use_w_tilde=True, use_source_loop=False, use_positive_only_solver=True), ) inversion_source_loop = aa.Inversion( dataset=dataset, linear_obj_list=[mapper], - settings=aa.SettingsInversion(use_w_tilde=True, use_source_loop=True), + settings=aa.SettingsInversion(use_w_tilde=True, use_source_loop=True, use_positive_only_solver=True), ) assert (inversion_image_loop.data == inversion_source_loop.data).all() From 506d51204867a0f00ef246e4a66ac0d61cd0fcc9 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 13 Jun 2025 10:43:40 +0100 Subject: [PATCH 03/14] check for nans when raising InversionException --- autoarray/inversion/inversion/inversion_util.py | 5 +++-- test_autoarray/inversion/inversion/test_abstract.py | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index 94989842f..dd54670db 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/autoarray/inversion/inversion/inversion_util.py @@ -233,12 +233,13 @@ def reconstruction_positive_negative_from( The curvature_matrix plus regularization matrix, overwriting the curvature_matrix in memory. """ try: - # print(curvature_reg_matrix) - # print(data_vector) reconstruction = jnp.linalg.solve(curvature_reg_matrix, data_vector) except np.linalg.LinAlgError as e: raise exc.InversionException() from e + if jnp.isnan(reconstruction).any(): + raise exc.InversionException + if ( conf.instance["general"]["inversion"]["check_reconstruction"] or force_check_reconstruction diff --git a/test_autoarray/inversion/inversion/test_abstract.py b/test_autoarray/inversion/inversion/test_abstract.py index 5ce2ad473..6ac97a4af 100644 --- a/test_autoarray/inversion/inversion/test_abstract.py +++ b/test_autoarray/inversion/inversion/test_abstract.py @@ -6,7 +6,6 @@ from autoarray import exc - directory = path.dirname(path.realpath(__file__)) @@ -455,7 +454,8 @@ def test__data_subtracted_dict(): def test__reconstruction_raises_exception_for_linalg_error(): # noinspection PyTypeChecker inversion = aa.m.MockInversion( - data_vector=np.ones(3), curvature_reg_matrix=np.ones((3, 3)) + data_vector=np.ones(3), + curvature_reg_matrix=np.ones((3, 3)) ) with pytest.raises(exc.InversionException): From e184330d6e5823130a090744aaa6758c0b842f4a Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 13 Jun 2025 10:45:52 +0100 Subject: [PATCH 04/14] test__inversion_imaging__via_linear_obj_func_and_mapper__force_edge_pixels_to_zero --- test_autoarray/inversion/inversion/test_factory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test_autoarray/inversion/inversion/test_factory.py b/test_autoarray/inversion/inversion/test_factory.py index a03be939e..8ffdc659e 100644 --- a/test_autoarray/inversion/inversion/test_factory.py +++ b/test_autoarray/inversion/inversion/test_factory.py @@ -265,7 +265,7 @@ def test__inversion_imaging__via_linear_obj_func_and_mapper__force_edge_pixels_t assert isinstance(inversion.linear_obj_list[1], aa.MapperDelaunay) assert isinstance(inversion, aa.InversionImagingMapping) assert inversion.reconstruction == pytest.approx( - np.array([2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), 1.0e-4 + np.array([2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), abs=1.0e-2 ) From 39fa6fc06d603375f3320a0ca13e62e735f20e64 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 13 Jun 2025 10:50:03 +0100 Subject: [PATCH 05/14] fix test__inversion_imaging__linear_obj_func_with_w_tilde --- test_autoarray/inversion/inversion/test_factory.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test_autoarray/inversion/inversion/test_factory.py b/test_autoarray/inversion/inversion/test_factory.py index 8ffdc659e..9516554f3 100644 --- a/test_autoarray/inversion/inversion/test_factory.py +++ b/test_autoarray/inversion/inversion/test_factory.py @@ -389,8 +389,8 @@ def test__inversion_imaging__linear_obj_func_with_w_tilde( assert inversion_mapping.reconstruction == pytest.approx( inversion_w_tilde.reconstruction, 1.0e-4 ) - assert inversion_mapping.mapped_reconstructed_image == pytest.approx( - inversion_w_tilde.mapped_reconstructed_image, 1.0e-4 + assert inversion_mapping.mapped_reconstructed_image.array == pytest.approx( + inversion_w_tilde.mapped_reconstructed_image.array, 1.0e-4 ) linear_obj_1 = aa.m.MockLinearObjFuncList( From 7eae31e8c446cf9adfa864b5664a37d09fd97c6f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 13 Jun 2025 10:54:36 +0100 Subject: [PATCH 06/14] fix test__inversion_imaging__linear_obj_func_and_non_func_give_same_terms --- test_autoarray/inversion/inversion/test_factory.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/test_autoarray/inversion/inversion/test_factory.py b/test_autoarray/inversion/inversion/test_factory.py index 9516554f3..0abe4d8d5 100644 --- a/test_autoarray/inversion/inversion/test_factory.py +++ b/test_autoarray/inversion/inversion/test_factory.py @@ -1,6 +1,7 @@ import copy import numpy as np import pytest +from dill import settings import autoarray as aa @@ -310,13 +311,13 @@ def test__inversion_imaging__linear_obj_func_and_non_func_give_same_terms( grid = aa.Grid2D.from_mask(mask=mask) linear_obj = aa.m.MockLinearObj( - parameters=2, grid=grid, mapping_matrix=np.full(fill_value=0.5, shape=(9, 2)) + parameters=2, grid=grid, mapping_matrix=np.full(fill_value=0.5, shape=(9, 2)), ) inversion = aa.Inversion( dataset=masked_imaging_7x7_no_blur, linear_obj_list=[linear_obj, rectangular_mapper_7x7_3x3], - settings=aa.SettingsInversion(use_w_tilde=False), + settings=aa.SettingsInversion(use_w_tilde=False, use_positive_only_solver=True), ) masked_imaging_7x7_no_blur = copy.copy(masked_imaging_7x7_no_blur) @@ -328,7 +329,7 @@ def test__inversion_imaging__linear_obj_func_and_non_func_give_same_terms( inversion_no_linear_func = aa.Inversion( dataset=masked_imaging_7x7_no_blur, linear_obj_list=[rectangular_mapper_7x7_3x3], - settings=aa.SettingsInversion(use_w_tilde=False), + settings=aa.SettingsInversion(use_w_tilde=False, use_positive_only_solver=True), ) assert inversion.regularization_term == pytest.approx( From c20897b1f1aac7f5361244dd8119ee32ac3c330c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 13 Jun 2025 10:58:50 +0100 Subject: [PATCH 07/14] fix test__inversion_imaging__linear_obj_func_with_w_tilde --- autoarray/inversion/inversion/inversion_util.py | 2 +- test_autoarray/inversion/inversion/test_factory.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index dd54670db..7169f7900 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/autoarray/inversion/inversion/inversion_util.py @@ -132,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 = np.dot(array.T, array) + curvature_matrix = jnp.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( diff --git a/test_autoarray/inversion/inversion/test_factory.py b/test_autoarray/inversion/inversion/test_factory.py index 0abe4d8d5..7bdf5646f 100644 --- a/test_autoarray/inversion/inversion/test_factory.py +++ b/test_autoarray/inversion/inversion/test_factory.py @@ -369,13 +369,13 @@ def test__inversion_imaging__linear_obj_func_with_w_tilde( inversion_mapping = aa.Inversion( dataset=masked_imaging_7x7, linear_obj_list=[linear_obj, rectangular_mapper_7x7_3x3], - settings=aa.SettingsInversion(use_w_tilde=False), + settings=aa.SettingsInversion(use_w_tilde=False, use_positive_only_solver=True), ) inversion_w_tilde = aa.Inversion( dataset=masked_imaging_7x7, linear_obj_list=[linear_obj, rectangular_mapper_7x7_3x3], - settings=aa.SettingsInversion(use_w_tilde=True), + settings=aa.SettingsInversion(use_w_tilde=True, use_positive_only_solver=True), ) assert inversion_mapping.data_vector == pytest.approx( From 949da4a37f44c46940ee494c206b02e598105d32 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 13 Jun 2025 10:59:42 +0100 Subject: [PATCH 08/14] test__identical_inversion_values_for_two_methods --- .../interferometer/test_inversion_interferometer_util.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index ab9ce32b9..256a67996 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -248,13 +248,13 @@ def test__identical_inversion_values_for_two_methods(): inversion_w_tilde = aa.Inversion( dataset=dataset, linear_obj_list=[mapper], - settings=aa.SettingsInversion(use_w_tilde=True), + settings=aa.SettingsInversion(use_w_tilde=True, use_positive_only_solver=True), ) inversion_mapping_matrices = aa.Inversion( dataset=dataset, linear_obj_list=[mapper], - settings=aa.SettingsInversion(use_w_tilde=False), + settings=aa.SettingsInversion(use_w_tilde=False, use_positive_only_solver=True), ) assert (inversion_w_tilde.data == inversion_mapping_matrices.data).all() From fd2e601e778e918f3fd8a6cf121b0427d0ad39a7 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 13 Jun 2025 15:28:08 +0100 Subject: [PATCH 09/14] midway through some refactoring --- autoarray/inversion/inversion/abstract.py | 122 ++++-------------- .../inversion/inversion/inversion_util.py | 55 ++------ autoarray/inversion/inversion/settings.py | 2 - .../operators/over_sampling/over_sampler.py | 2 +- .../test_inversion_interferometer_util.py | 8 +- .../inversion/inversion/test_abstract.py | 3 +- .../inversion/inversion/test_factory.py | 33 ++--- 7 files changed, 61 insertions(+), 164 deletions(-) diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index 905aa5e16..a048742a8 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -463,20 +463,14 @@ def reconstruction(self) -> np.ndarray: 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 + if self.has(cls=AbstractMapper) and self.settings.force_edge_pixels_to_zeros: + + ids_zeros = jnp.array(self.mapper_edge_pixel_list, dtype=int) - values_to_solve = np.ones( - np.shape(self.curvature_reg_matrix)[0], dtype=bool + values_to_solve = jnp.ones( + self.curvature_reg_matrix.shape[0], dtype=bool ) - values_to_solve[ids_zeros] = False + values_to_solve = values_to_solve.at[ids_zeros].set(False) data_vector_input = self.data_vector[values_to_solve] @@ -484,25 +478,32 @@ def reconstruction(self) -> np.ndarray: values_to_solve, : ][:, values_to_solve] - solutions = np.zeros(np.shape(self.curvature_reg_matrix)[0]) - - 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, - ) + # Get the values to assign (must be a JAX array) + reconstruction = inversion_util.reconstruction_positive_only_from( + data_vector=data_vector_input, + curvature_reg_matrix=curvature_reg_matrix_input, + settings=self.settings, ) + + # Allocate JAX array + solutions = jnp.zeros(self.curvature_reg_matrix.shape[0]) + + # Get indices where True + indices = jnp.where(values_to_solve)[0] + + # Set reconstruction values at those indices + solutions = solutions.at[indices].set(reconstruction) + return solutions + else: - solutions = inversion_util.reconstruction_positive_only_from( + + return 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( @@ -511,81 +512,6 @@ 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: diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index 7169f7900..b24f9ed2e 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/autoarray/inversion/inversion/inversion_util.py @@ -41,7 +41,6 @@ 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, @@ -61,55 +60,22 @@ 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() def curvature_matrix_mirrored_from( curvature_matrix: np.ndarray, ) -> np.ndarray: - curvature_matrix_mirrored = np.zeros( - (curvature_matrix.shape[0], curvature_matrix.shape[1]) - ) + # Copy the original matrix and its transpose + m1 = curvature_matrix + m2 = curvature_matrix.T - for i in range(curvature_matrix.shape[0]): - for j in range(curvature_matrix.shape[1]): - if curvature_matrix[i, j] != 0: - curvature_matrix_mirrored[i, j] = curvature_matrix[i, j] - curvature_matrix_mirrored[j, i] = curvature_matrix[i, j] - if curvature_matrix[j, i] != 0: - curvature_matrix_mirrored[i, j] = curvature_matrix[j, i] - curvature_matrix_mirrored[j, i] = curvature_matrix[j, i] + # For each entry, prefer the non-zero value from either the matrix or its transpose + mirrored = jnp.where(m1 != 0, m1, m2) - return curvature_matrix_mirrored + return mirrored def curvature_matrix_via_mapping_matrix_from( @@ -188,7 +154,7 @@ def mapped_reconstructed_data_via_mapping_matrix_from( The matrix representing the blurred mappings between sub-grid pixels and pixelization pixels. """ - return np.dot(mapping_matrix, reconstruction) + return jnp.dot(mapping_matrix, reconstruction) def reconstruction_positive_negative_from( @@ -307,6 +273,7 @@ def reconstruction_positive_only_from( except (RuntimeError, np.linalg.LinAlgError, ValueError) as e: raise exc.InversionException() from e + def preconditioner_matrix_via_mapping_matrix_from( mapping_matrix: np.ndarray, regularization_matrix: np.ndarray, diff --git a/autoarray/inversion/inversion/settings.py b/autoarray/inversion/inversion/settings.py index 2c0eba077..184e16977 100644 --- a/autoarray/inversion/inversion/settings.py +++ b/autoarray/inversion/inversion/settings.py @@ -15,7 +15,6 @@ def __init__( positive_only_uses_p_initial: Optional[bool] = None, use_border_relocator: Optional[bool] = None, force_edge_pixels_to_zeros: bool = True, - force_edge_image_pixels_to_zeros: bool = False, image_pixels_source_zero=None, no_regularization_add_to_curvature_diag_value: float = None, use_w_tilde_numpy: bool = False, @@ -84,7 +83,6 @@ def __init__( self._use_border_relocator = use_border_relocator self.use_linear_operators = use_linear_operators self.force_edge_pixels_to_zeros = force_edge_pixels_to_zeros - self.force_edge_image_pixels_to_zeros = force_edge_image_pixels_to_zeros self.image_pixels_source_zero = image_pixels_source_zero self._no_regularization_add_to_curvature_diag_value = ( no_regularization_add_to_curvature_diag_value diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py index 499c6b37d..5d9a99871 100644 --- a/autoarray/operators/over_sampling/over_sampler.py +++ b/autoarray/operators/over_sampling/over_sampler.py @@ -150,7 +150,7 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]): # Used for JAX based adaptive over sampling. # Define group sizes - group_sizes = np.array(self.sub_size.array ** 2) + group_sizes = np.array(self.sub_size.array**2) # Compute the cumulative sum of group sizes to get split points self.split_indices = np.cumsum(group_sizes) diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index 256a67996..cbb4744fc 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -347,13 +347,17 @@ def test__identical_inversion_source_and_image_loops(): inversion_image_loop = aa.Inversion( dataset=dataset, linear_obj_list=[mapper], - settings=aa.SettingsInversion(use_w_tilde=True, use_source_loop=False, use_positive_only_solver=True), + settings=aa.SettingsInversion( + use_w_tilde=True, use_source_loop=False, use_positive_only_solver=True + ), ) inversion_source_loop = aa.Inversion( dataset=dataset, linear_obj_list=[mapper], - settings=aa.SettingsInversion(use_w_tilde=True, use_source_loop=True, use_positive_only_solver=True), + settings=aa.SettingsInversion( + use_w_tilde=True, use_source_loop=True, use_positive_only_solver=True + ), ) assert (inversion_image_loop.data == inversion_source_loop.data).all() diff --git a/test_autoarray/inversion/inversion/test_abstract.py b/test_autoarray/inversion/inversion/test_abstract.py index 6ac97a4af..bf8f4a919 100644 --- a/test_autoarray/inversion/inversion/test_abstract.py +++ b/test_autoarray/inversion/inversion/test_abstract.py @@ -454,8 +454,7 @@ def test__data_subtracted_dict(): def test__reconstruction_raises_exception_for_linalg_error(): # noinspection PyTypeChecker inversion = aa.m.MockInversion( - data_vector=np.ones(3), - curvature_reg_matrix=np.ones((3, 3)) + data_vector=np.ones(3), curvature_reg_matrix=np.ones((3, 3)) ) with pytest.raises(exc.InversionException): diff --git a/test_autoarray/inversion/inversion/test_factory.py b/test_autoarray/inversion/inversion/test_factory.py index 7bdf5646f..984aab946 100644 --- a/test_autoarray/inversion/inversion/test_factory.py +++ b/test_autoarray/inversion/inversion/test_factory.py @@ -291,8 +291,8 @@ def test__inversion_imaging__compare_mapping_and_w_tilde_values( assert inversion_w_tilde.reconstruction == pytest.approx( inversion_mapping.reconstruction, 1.0e-4 ) - assert inversion_w_tilde.mapped_reconstructed_image == pytest.approx( - inversion_mapping.mapped_reconstructed_image, 1.0e-4 + assert inversion_w_tilde.mapped_reconstructed_image.array == pytest.approx( + inversion_mapping.mapped_reconstructed_image.array, 1.0e-4 ) assert inversion_w_tilde.log_det_curvature_reg_matrix_term == pytest.approx( inversion_mapping.log_det_curvature_reg_matrix_term @@ -311,7 +311,9 @@ def test__inversion_imaging__linear_obj_func_and_non_func_give_same_terms( grid = aa.Grid2D.from_mask(mask=mask) linear_obj = aa.m.MockLinearObj( - parameters=2, grid=grid, mapping_matrix=np.full(fill_value=0.5, shape=(9, 2)), + parameters=2, + grid=grid, + mapping_matrix=np.full(fill_value=0.5, shape=(9, 2)), ) inversion = aa.Inversion( @@ -528,23 +530,24 @@ def test__inversion_matrices__x2_mappers( assert (inversion.regularization_matrix[0:9, 9:18] == np.zeros((9, 9))).all() assert (inversion.regularization_matrix[9:18, 0:9] == np.zeros((9, 9))).all() - assert inversion.reconstruction_dict[rectangular_mapper_7x7_3x3][4] == pytest.approx( - 0.05594123, 1.0e-4 - ) + assert inversion.reconstruction_dict[rectangular_mapper_7x7_3x3][ + 4 + ] == pytest.approx(0.05594123, 1.0e-4) assert inversion.reconstruction_dict[delaunay_mapper_9_3x3][4] == pytest.approx( 0.04686388, 1.0e-4 ) - assert inversion.reconstruction[13] == pytest.approx( - 0.04686388, 1.0e-4 + assert inversion.reconstruction[13] == pytest.approx(0.04686388, 1.0e-4) + + assert inversion.mapped_reconstructed_data_dict[rectangular_mapper_7x7_3x3][ + 4 + ] == pytest.approx(0.05594123, 1.0e-4) + assert inversion.mapped_reconstructed_data_dict[delaunay_mapper_9_3x3][ + 3 + ] == pytest.approx(0.01521323, 1.0e-4) + assert inversion.mapped_reconstructed_image[4] == pytest.approx( + 0.10494037076075, 1.0e-4 ) - assert inversion.mapped_reconstructed_data_dict[ - rectangular_mapper_7x7_3x3 - ][4] == pytest.approx(0.05594123, 1.0e-4) - assert inversion.mapped_reconstructed_data_dict[ - delaunay_mapper_9_3x3 - ][3] == pytest.approx(0.01521323, 1.0e-4) - assert inversion.mapped_reconstructed_image[4] == pytest.approx(0.10494037076075, 1.0e-4) def test__inversion_imaging__positive_only_solver(masked_imaging_7x7_no_blur): mask = masked_imaging_7x7_no_blur.mask From c085e88d937791ba717a9a4a4523874a079c3d5c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sun, 15 Jun 2025 15:47:23 +0100 Subject: [PATCH 10/14] remove force_edge_image_pixels_to_zeros --- autoarray/inversion/inversion/abstract.py | 1 + autoarray/inversion/inversion/inversion_util.py | 16 ++++++++++++++-- .../inversion/inversion/test_settings_dict.py | 1 - 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index a048742a8..ff51a6e19 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -1,4 +1,5 @@ import copy +import jax import jax.numpy as jnp import numpy as np from scipy.linalg import block_diag diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index b24f9ed2e..cb9826e78 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/autoarray/inversion/inversion/inversion_util.py @@ -1,5 +1,7 @@ import jax.numpy as jnp import jaxnnls +import jax +import jax.lax as lax import numpy as np from typing import List, Optional, Tuple @@ -10,7 +12,6 @@ from autoarray import numba_util from autoarray import exc -from autoarray.util.fnnls import fnnls_cholesky def curvature_matrix_via_w_tilde_from( @@ -269,10 +270,21 @@ def reconstruction_positive_only_from( """ try: - return jaxnnls.solve_nnls_primal(curvature_reg_matrix, data_vector) + reconstruction = jaxnnls.solve_nnls_primal(curvature_reg_matrix, data_vector) except (RuntimeError, np.linalg.LinAlgError, ValueError) as e: raise exc.InversionException() from e + def handle_nan(reconstruction): + return jnp.zeros_like(reconstruction) + + def handle_valid(reconstruction): + return reconstruction + + has_nan = jnp.isnan(reconstruction).any() + reconstruction = lax.cond(has_nan, handle_nan, handle_valid, reconstruction) + + return reconstruction + def preconditioner_matrix_via_mapping_matrix_from( mapping_matrix: np.ndarray, diff --git a/test_autoarray/inversion/inversion/test_settings_dict.py b/test_autoarray/inversion/inversion/test_settings_dict.py index 6a6f8ca9f..d547014e1 100644 --- a/test_autoarray/inversion/inversion/test_settings_dict.py +++ b/test_autoarray/inversion/inversion/test_settings_dict.py @@ -17,7 +17,6 @@ def make_settings_dict(): "use_positive_only_solver": False, "positive_only_uses_p_initial": False, "force_edge_pixels_to_zeros": True, - "force_edge_image_pixels_to_zeros": False, "image_pixels_source_zero": None, "no_regularization_add_to_curvature_diag_value": 1e-08, "use_w_tilde_numpy": False, From 7d9115827878fb1aed9c8a55b3b0152f627af8c6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sun, 15 Jun 2025 15:51:22 +0100 Subject: [PATCH 11/14] convert JAX Array --- autoarray/inversion/pixelization/mappers/voronoi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/autoarray/inversion/pixelization/mappers/voronoi.py b/autoarray/inversion/pixelization/mappers/voronoi.py index c8e54cbf3..dcfdaa2eb 100644 --- a/autoarray/inversion/pixelization/mappers/voronoi.py +++ b/autoarray/inversion/pixelization/mappers/voronoi.py @@ -172,5 +172,5 @@ def interpolated_array_from( is input. """ return self.source_plane_mesh_grid.interpolated_array_from( - values=values, shape_native=shape_native, extent=extent, use_nn=True + values=np.array(values), shape_native=shape_native, extent=extent, use_nn=True ) From 9c7093238593c03893a5d81e355789abcba44dbd Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sun, 15 Jun 2025 15:53:58 +0100 Subject: [PATCH 12/14] ndarray conversions for JAX arrays --- autoarray/inversion/inversion/interferometer/mapping.py | 6 +++--- .../interferometer/test_inversion_interferometer_util.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/autoarray/inversion/inversion/interferometer/mapping.py b/autoarray/inversion/inversion/interferometer/mapping.py index 9cde492e9..33648f520 100644 --- a/autoarray/inversion/inversion/interferometer/mapping.py +++ b/autoarray/inversion/inversion/interferometer/mapping.py @@ -76,7 +76,7 @@ def data_vector(self) -> np.ndarray: """ return inversion_interferometer_util.data_vector_via_transformed_mapping_matrix_from( - transformed_mapping_matrix=self.operated_mapping_matrix, + transformed_mapping_matrix=np.array(self.operated_mapping_matrix), visibilities=np.array(self.data), noise_map=np.array(self.noise_map), ) @@ -152,8 +152,8 @@ def mapped_reconstructed_data_dict( visibilities = ( inversion_interferometer_util.mapped_reconstructed_visibilities_from( - transformed_mapping_matrix=operated_mapping_matrix_list[index], - reconstruction=reconstruction, + transformed_mapping_matrix=np.array(operated_mapping_matrix_list[index]), + reconstruction=np.array(reconstruction), ) ) diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index cbb4744fc..e02d38372 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -272,8 +272,8 @@ def test__identical_inversion_values_for_two_methods(): == inversion_mapping_matrices.regularization_matrix ).all() - assert inversion_w_tilde.data_vector == pytest.approx( - inversion_mapping_matrices.data_vector, 1.0e-8 + assert inversion_w_tilde.data_vector.array == pytest.approx( + inversion_mapping_matrices.data_vector.array, 1.0e-8 ) assert inversion_w_tilde.curvature_matrix == pytest.approx( inversion_mapping_matrices.curvature_matrix, 1.0e-8 @@ -384,8 +384,8 @@ def test__identical_inversion_source_and_image_loops(): assert inversion_image_loop.reconstruction == pytest.approx( inversion_source_loop.reconstruction, 1.0e-2 ) - assert inversion_image_loop.mapped_reconstructed_image == pytest.approx( - inversion_source_loop.mapped_reconstructed_image, 1.0e-2 + assert inversion_image_loop.mapped_reconstructed_image.array == pytest.approx( + inversion_source_loop.mapped_reconstructed_image.array, 1.0e-2 ) assert inversion_image_loop.mapped_reconstructed_data == pytest.approx( inversion_source_loop.mapped_reconstructed_data, 1.0e-2 From 97abd126bd71e5d8cfba1bfd993488d849657b32 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sun, 15 Jun 2025 15:59:38 +0100 Subject: [PATCH 13/14] fix all inversion unit tests --- .../inversion/inversion/inversion_util.py | 39 +++++++++++++++++-- .../test_inversion_interferometer_util.py | 8 ++-- 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index cb9826e78..41fdaf236 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/autoarray/inversion/inversion/inversion_util.py @@ -61,9 +61,42 @@ 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) + try: + return curvature_matrix.at[ + no_regularization_index_list, no_regularization_index_list + ].add(value) + except AttributeError: + return curvature_matrix_with_added_to_diag_from_numba( + curvature_matrix=curvature_matrix, + value=value, + no_regularization_index_list=no_regularization_index_list, + ) + +@numba_util.jit() +def curvature_matrix_with_added_to_diag_from_numba( + 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. + """ + + for i in no_regularization_index_list: + curvature_matrix[i, i] += value + + return curvature_matrix def curvature_matrix_mirrored_from( diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index e02d38372..c7dc03221 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -272,8 +272,8 @@ def test__identical_inversion_values_for_two_methods(): == inversion_mapping_matrices.regularization_matrix ).all() - assert inversion_w_tilde.data_vector.array == pytest.approx( - inversion_mapping_matrices.data_vector.array, 1.0e-8 + assert inversion_w_tilde.data_vector == pytest.approx( + inversion_mapping_matrices.data_vector, 1.0e-8 ) assert inversion_w_tilde.curvature_matrix == pytest.approx( inversion_mapping_matrices.curvature_matrix, 1.0e-8 @@ -285,8 +285,8 @@ def test__identical_inversion_values_for_two_methods(): assert inversion_w_tilde.reconstruction == pytest.approx( inversion_mapping_matrices.reconstruction, abs=1.0e-1 ) - assert inversion_w_tilde.mapped_reconstructed_image == pytest.approx( - inversion_mapping_matrices.mapped_reconstructed_image, abs=1.0e-1 + assert inversion_w_tilde.mapped_reconstructed_image.array == pytest.approx( + inversion_mapping_matrices.mapped_reconstructed_image.array, abs=1.0e-1 ) assert inversion_w_tilde.mapped_reconstructed_data == pytest.approx( inversion_mapping_matrices.mapped_reconstructed_data, abs=1.0e-1 From 8051e3515a0e1fdbadccd90663584470ecd08512 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sun, 15 Jun 2025 16:03:56 +0100 Subject: [PATCH 14/14] black --- autoarray/inversion/inversion/abstract.py | 5 ++++- autoarray/inversion/inversion/imaging/w_tilde.py | 2 +- autoarray/inversion/inversion/interferometer/mapping.py | 4 +++- autoarray/inversion/inversion/inversion_util.py | 7 ++++--- autoarray/inversion/pixelization/mappers/voronoi.py | 5 ++++- 5 files changed, 16 insertions(+), 7 deletions(-) diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index ff51a6e19..ebbb9927f 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -464,7 +464,10 @@ def reconstruction(self) -> np.ndarray: And the data_vector = ZTx, so the corresponding row is also taken out. """ - if self.has(cls=AbstractMapper) and self.settings.force_edge_pixels_to_zeros: + if ( + self.has(cls=AbstractMapper) + and self.settings.force_edge_pixels_to_zeros + ): ids_zeros = jnp.array(self.mapper_edge_pixel_list, dtype=int) diff --git a/autoarray/inversion/inversion/imaging/w_tilde.py b/autoarray/inversion/inversion/imaging/w_tilde.py index 1f6d08ac1..97fc9eb8f 100644 --- a/autoarray/inversion/inversion/imaging/w_tilde.py +++ b/autoarray/inversion/inversion/imaging/w_tilde.py @@ -440,7 +440,7 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: data_weights=mapper.unique_mappings.data_weights, pix_lengths=mapper.unique_mappings.pix_lengths, pix_pixels=mapper.params, - curvature_weights=curvature_weights, + curvature_weights=np.array(curvature_weights), image_frame_1d_lengths=self.convolver.image_frame_1d_lengths, image_frame_1d_indexes=self.convolver.image_frame_1d_indexes, image_frame_1d_kernels=self.convolver.image_frame_1d_kernels, diff --git a/autoarray/inversion/inversion/interferometer/mapping.py b/autoarray/inversion/inversion/interferometer/mapping.py index 33648f520..2b3219a2e 100644 --- a/autoarray/inversion/inversion/interferometer/mapping.py +++ b/autoarray/inversion/inversion/interferometer/mapping.py @@ -152,7 +152,9 @@ def mapped_reconstructed_data_dict( visibilities = ( inversion_interferometer_util.mapped_reconstructed_visibilities_from( - transformed_mapping_matrix=np.array(operated_mapping_matrix_list[index]), + transformed_mapping_matrix=np.array( + operated_mapping_matrix_list[index] + ), reconstruction=np.array(reconstruction), ) ) diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index 41fdaf236..29f1eaa07 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/autoarray/inversion/inversion/inversion_util.py @@ -72,11 +72,12 @@ def curvature_matrix_with_added_to_diag_from( no_regularization_index_list=no_regularization_index_list, ) + @numba_util.jit() def curvature_matrix_with_added_to_diag_from_numba( - curvature_matrix: np.ndarray, - value: float, - no_regularization_index_list: Optional[List] = None, + 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 diff --git a/autoarray/inversion/pixelization/mappers/voronoi.py b/autoarray/inversion/pixelization/mappers/voronoi.py index dcfdaa2eb..1ebae8ad9 100644 --- a/autoarray/inversion/pixelization/mappers/voronoi.py +++ b/autoarray/inversion/pixelization/mappers/voronoi.py @@ -172,5 +172,8 @@ def interpolated_array_from( is input. """ return self.source_plane_mesh_grid.interpolated_array_from( - values=np.array(values), shape_native=shape_native, extent=extent, use_nn=True + values=np.array(values), + shape_native=shape_native, + extent=extent, + use_nn=True, )