fix now pointless unit ests

Author: Jammy2211

autoarray/dataset/imaging/w_tilde.py

@@ -71,4 +71,11 @@ def __init__(
             x_shape=data.shape_native[1],
         ))
 
+        self.curvature_matrix_off_diag_light_profiles_func = (inversion_imaging_util.build_curvature_matrix_off_diag_with_light_profiles_via_w_tilde_from_func(
+            psf=self.psf.native.array,
+            y_shape=data.shape_native[0],
+            x_shape=data.shape_native[1],
+        ))
+
+
         self.batch_size = batch_size

autoarray/inversion/inversion/imaging/inversion_imaging_util.py

@@ -521,6 +521,79 @@ def build_curvature_matrix_off_diag_via_w_tilde_from_func(psf: np.ndarray, y_sha
     return offdiag_jit
 
 
+def curvature_matrix_off_diag_with_light_profiles_via_w_tilde_from(
+    curvature_weights,      # (M_pix, n_funcs) = (H B) / noise^2  on slim grid
+    fft_index_for_masked_pixel,  # (M_pix,) slim -> rect(flat) indices
+    rows, cols, vals,            # triplets for sparse mapper A
+    y_shape: int,
+    x_shape: int,
+    S: int,
+    Khat_flip_r,                 # precomputed rfft2(flipped PSF padded)
+    Ky: int,
+    Kx: int,
+):
+    """
+    Computes: off_diag = A^T [ H^T(curvature_weights_native) ]
+    where curvature_weights = (H B) / noise^2 already.
+    """
+
+    import jax
+    import jax.numpy as jnp
+    from jax.ops import segment_sum
+
+    curvature_weights = jnp.asarray(curvature_weights, dtype=jnp.float64)
+    fft_index_for_masked_pixel = jnp.asarray(fft_index_for_masked_pixel, dtype=jnp.int32)
+
+    rows = jnp.asarray(rows, dtype=jnp.int32)
+    cols = jnp.asarray(cols, dtype=jnp.int32)
+    vals = jnp.asarray(vals, dtype=jnp.float64)
+
+    M_pix, n_funcs = curvature_weights.shape
+    M_rect = y_shape * x_shape
+    fft_shape = (y_shape + Ky - 1, x_shape + Kx - 1)
+
+    # 1) scatter slim weights onto rectangular grid (flat)
+    grid_flat = jnp.zeros((M_rect, n_funcs), dtype=jnp.float64)
+    grid_flat = grid_flat.at[fft_index_for_masked_pixel, :].set(curvature_weights)
+
+    # 2) apply H^T = convolution with flipped PSF (one convolution)
+    images = grid_flat.T.reshape((n_funcs, y_shape, x_shape))  # (B=n_funcs, Hy, Hx)
+    back_native = rfft_convolve2d_same(images, Khat_flip_r, Ky, Kx, fft_shape)
+
+    # 3) gather at mapper rows
+    back_flat = back_native.reshape((n_funcs, M_rect)).T       # (M_rect, n_funcs)
+    back_at_rows = back_flat[rows, :]                          # (nnz, n_funcs)
+
+    # 4) accumulate into sparse pixels
+    contrib = vals[:, None] * back_at_rows
+    off_diag = segment_sum(contrib, cols, num_segments=S)      # (S, n_funcs)
+    return off_diag
+
+
+def build_curvature_matrix_off_diag_with_light_profiles_via_w_tilde_from_func(psf: np.ndarray, y_shape: int, x_shape: int):
+
+    import jax
+    import jax.numpy as jnp
+
+    psf = jnp.asarray(psf, dtype=jnp.float64)
+    Ky, Kx = psf.shape
+    fft_shape = (y_shape + Ky - 1, x_shape + Kx - 1)
+
+    psf_flip = jnp.flip(psf, axis=(0, 1))
+    Khat_flip_r = precompute_Khat_rfft(psf_flip, fft_shape)
+
+    fn_jit = jax.jit(
+        partial(
+            curvature_matrix_off_diag_with_light_profiles_via_w_tilde_from,
+            Khat_flip_r=Khat_flip_r,
+            Ky=Ky,
+            Kx=Kx,
+        ),
+        static_argnames=("y_shape", "x_shape", "S"),
+    )
+    return fn_jit
+
+
 def mapped_image_rect_from_triplets(
     reconstruction,   # (S,)
     rows,

autoarray/inversion/inversion/imaging/w_tilde.py

@@ -88,12 +88,12 @@ def _data_vector_mapper(self) -> np.ndarray:
             rows, cols, vals = mapper.pixel_triplets_data
 
             data_vector_mapper = (
-                inversion_imaging_util.data_vector_via_w_tilde_data_imaging_from(
+                inversion_imaging_util.data_vector_via_w_tilde_from(
                     w_tilde_data=self.w_tilde_data,
                     rows=rows,
                     cols=cols,
                     vals=vals,
-                    S=mapper.total_params,
+                    S=mapper.params,
                 )
             )
             param_range = mapper_param_range[mapper_index]
@@ -199,15 +199,21 @@ def _data_vector_func_list_and_mapper(self) -> np.ndarray:
                 linear_func
             ]
 
-            diag = inversion_imaging_numba_util.data_vector_via_blurred_mapping_matrix_from(
-                blurred_mapping_matrix=np.array(operated_mapping_matrix),
+            diag = inversion_imaging_util.data_vector_via_blurred_mapping_matrix_from(
+                blurred_mapping_matrix=operated_mapping_matrix,
                 image=self.data.array,
                 noise_map=self.noise_map.array,
             )
 
             param_range = linear_func_param_range[linear_func_index]
 
-            data_vector[param_range[0] : param_range[1],] = diag
+            start = param_range[0]
+            end = param_range[1]
+
+            if self._xp is np:
+                data_vector[start:end] = diag
+            else:
+                data_vector = data_vector.at[start:end].set(diag)
 
         return data_vector
 
@@ -421,22 +427,35 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray:
                     / self.noise_map[:, None] ** 2
                 )
 
-                off_diag = inversion_imaging_numba_util.curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from(
-                    data_to_pix_unique=mapper.unique_mappings.data_to_pix_unique,
-                    data_weights=mapper.unique_mappings.data_weights,
-                    pix_lengths=mapper.unique_mappings.pix_lengths,
-                    pix_pixels=mapper.params,
-                    curvature_weights=np.array(curvature_weights),
-                    mask=self.mask.array,
-                    psf_kernel=self.psf.native.array,
+                rows, cols, vals = mapper.pixel_triplets_curvature
+
+                off_diag = self.w_tilde.curvature_matrix_off_diag_light_profiles_func(
+                    curvature_weights=curvature_weights,
+                    fft_index_for_masked_pixel=self.mask.fft_index_for_masked_pixel,
+                    rows=rows,
+                    cols=cols,
+                    vals=vals,
+                    y_shape=self.mask.shape_native[0],
+                    x_shape=self.mask.shape_native[1],
+                    S=mapper.params,
                 )
 
-                curvature_matrix[
-                    mapper_param_range[0] : mapper_param_range[1],
+                if self._xp is np:
+
+                    curvature_matrix[
+                        mapper_param_range[0] : mapper_param_range[1],
                     linear_func_param_range[0] : linear_func_param_range[1],
-                ] = off_diag
+                    ] = off_diag
+                else:
+
+                    curvature_matrix = curvature_matrix.at[
+                        mapper_param_range[0] : mapper_param_range[1],
+                        linear_func_param_range[0] : linear_func_param_range[1],
+                    ].set(off_diag)
+
 
         for index_0, linear_func_0 in enumerate(linear_func_list):
+
             linear_func_param_range_0 = linear_func_param_range_list[index_0]
 
             weighted_vector_0 = (
@@ -452,15 +471,24 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray:
                     / self.noise_map[:, None]
                 )
 
-                diag = np.dot(
+                diag = self._xp.dot(
                     weighted_vector_0.T,
                     weighted_vector_1,
                 )
 
-                curvature_matrix[
-                    linear_func_param_range_0[0] : linear_func_param_range_0[1],
-                    linear_func_param_range_1[0] : linear_func_param_range_1[1],
-                ] = diag
+                if self._xp is np:
+
+                    curvature_matrix[
+                        linear_func_param_range_0[0] : linear_func_param_range_0[1],
+                        linear_func_param_range_1[0] : linear_func_param_range_1[1],
+                    ] = diag
+
+                else:
+
+                    curvature_matrix = curvature_matrix.at[
+                        linear_func_param_range_0[0] : linear_func_param_range_0[1],
+                        linear_func_param_range_1[0] : linear_func_param_range_1[1],
+                    ].set(diag)
 
         return curvature_matrix
 

autoarray/inversion/inversion/interferometer/w_tilde.py

@@ -114,7 +114,7 @@ def curvature_matrix_diag(self) -> np.ndarray:
             pix_indexes_for_sub_slim_index=mapper.pix_indexes_for_sub_slim_index,
             pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index,
             pix_pixels=self.linear_obj_list[0].params,
-            fft_index_for_masked_pixel=self.w_tilde.fft_index_for_masked_pixel,
+            fft_index_for_masked_pixel=self.mask.fft_index_for_masked_pixel,
         )
 
     @property

test_autoarray/dataset/imaging/test_dataset.py

@@ -193,12 +193,6 @@ def test__apply_mask(imaging_7x7, mask_2d_7x7, psf_3x3):
 
     assert type(masked_imaging_7x7.psf) == aa.Kernel2D
 
-    masked_imaging_7x7 = masked_imaging_7x7.apply_w_tilde()
-
-    assert masked_imaging_7x7.w_tilde.curvature_preload.shape == (35,)
-    assert masked_imaging_7x7.w_tilde.indexes.shape == (35,)
-    assert masked_imaging_7x7.w_tilde.lengths.shape == (9,)
-
 
 def test__apply_noise_scaling(imaging_7x7, mask_2d_7x7):
     masked_imaging_7x7 = imaging_7x7.apply_noise_scaling(

test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py

@@ -128,7 +128,7 @@ def test__curvature_matrix_via_curvature_preload_from():
         fft_state=w_tilde.fft_state,
         pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,
         pix_weights_for_sub_slim_index=pix_weights_for_sub_slim_index,
-        fft_index_for_masked_pixel=w_tilde.fft_index_for_masked_pixel,
+        fft_index_for_masked_pixel=grid.mask.fft_index_for_masked_pixel,
         pix_pixels=3,
     )