fix or remove curvature preload unit tests

Author: Jammy2211

autoarray/dataset/interferometer/w_tilde.py

@@ -45,7 +45,7 @@ def __init__(
             inversion_interferometer_util,
         )
 
-        self.operator_state = inversion_interferometer_util.w_tilde_fft_state_from(
+        self.fft_state = inversion_interferometer_util.w_tilde_fft_state_from(
             curvature_preload=self.curvature_preload, batch_size=450
         )
 

autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py

@@ -67,12 +67,12 @@ def _report_memory(arr):
     """
     try:
         import resource
+
         rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
         arr_mb = arr.nbytes / 1024**2
         from tqdm import tqdm
-        tqdm.write(
-            f"    Memory: array={arr_mb:.1f} MB, RSS≈{rss_mb:.1f} MB"
-        )
+
+        tqdm.write(f"    Memory: array={arr_mb:.1f} MB, RSS≈{rss_mb:.1f} MB")
     except Exception:
         pass
 
@@ -141,26 +141,26 @@ def w_tilde_curvature_preload_interferometer_from(
     - The values of pixels paired with themselves are also computed repeatedly for the standard calculation (e.g. 9
     times using the mask above).
 
-    The `w_tilde_preload` method instead only computes each value once. To do this, it stores the preload values in a
+    The `curvature_preload` method instead only computes each value once. To do this, it stores the preload values in a
     matrix of dimensions [shape_masked_pixels_y, shape_masked_pixels_x, 2], where `shape_masked_pixels` is the (y,x)
     size of the vertical and horizontal extent of unmasked pixels, e.g. the spatial extent over which the real space
     grid extends.
 
-    Each entry in the matrix `w_tilde_preload[:,:,0]` provides the the precomputed NUFFT value mapping an image pixel
+    Each entry in the matrix `curvature_preload[:,:,0]` provides the the precomputed NUFFT value mapping an image pixel
     to a pixel offset by that much in the y and x directions, for example:
 
-    - w_tilde_preload[0,0,0] gives the precomputed values of image pixels that are offset in the y direction by 0 and
+    - curvature_preload[0,0,0] gives the precomputed values of image pixels that are offset in the y direction by 0 and
     in the x direction by 0 - the values of pixels paired with themselves.
-    - w_tilde_preload[1,0,0] gives the precomputed values of image pixels that are offset in the y direction by 1 and
+    - curvature_preload[1,0,0] gives the precomputed values of image pixels that are offset in the y direction by 1 and
     in the x direction by 0 - the values of pixel pairs [0,3], [1,4], [2,5], [3,6], [4,7] and [5,8]
-    - w_tilde_preload[0,1,0] gives the precomputed values of image pixels that are offset in the y direction by 0 and
+    - curvature_preload[0,1,0] gives the precomputed values of image pixels that are offset in the y direction by 0 and
     in the x direction by 1 - the values of pixel pairs [0,1], [1,2], [3,4], [4,5], [6,7] and [7,9].
 
     Flipped pairs:
 
     The above preloaded values pair all image pixel NUFFT values when a pixel is to the right and / or down of the
     first image pixel. However, one must also precompute pairs where the paired pixel is to the left of the host
-    pixels. These pairings are stored in `w_tilde_preload[:,:,1]`, and the ordering of these pairings is flipped in the
+    pixels. These pairings are stored in `curvature_preload[:,:,1]`, and the ordering of these pairings is flipped in the
     x direction to make it straight forward to use this matrix when computing w_tilde.
 
     Parameters
@@ -196,7 +196,7 @@ def w_tilde_curvature_preload_interferometer_from(
 
     K = uv_wavelengths.shape[0]
 
-    w = 1.0 / (noise_map_real ** 2)
+    w = 1.0 / (noise_map_real**2)
     ku = 2.0 * np.pi * uv_wavelengths[:, 0]
     kv = 2.0 * np.pi * uv_wavelengths[:, 1]
 
@@ -225,10 +225,7 @@ def accum_from_corner(y_ref, x_ref, gy_block, gx_block, label=""):
         for k0 in iterator:
             k1 = min(K, k0 + chunk_k)
 
-            phase = (
-                dx[..., None] * ku[k0:k1]
-                + dy[..., None] * kv[k0:k1]
-            )
+            phase = dx[..., None] * ku[k0:k1] + dy[..., None] * kv[k0:k1]
             acc += np.sum(
                 np.cos(phase) * w[k0:k1],
                 axis=2,
@@ -242,27 +239,21 @@ def accum_from_corner(y_ref, x_ref, gy_block, gx_block, label=""):
     # -----------------------------
     # Main quadrant (+,+)
     # -----------------------------
-    out[:y_shape, :x_shape] = accum_from_corner(
-        y00, x00, gy, gx, label="(+,+)"
-    )
+    out[:y_shape, :x_shape] = accum_from_corner(y00, x00, gy, gx, label="(+,+)")
 
     # -----------------------------
     # Flip in x (+,-)
     # -----------------------------
     if x_shape > 1:
-        block = accum_from_corner(
-            y0m, x0m, gy[:, ::-1], gx[:, ::-1], label="(+,-)"
-        )
-        out[:y_shape, -1:-(x_shape): -1] = block[:, 1:]
+        block = accum_from_corner(y0m, x0m, gy[:, ::-1], gx[:, ::-1], label="(+,-)")
+        out[:y_shape, -1:-(x_shape):-1] = block[:, 1:]
 
     # -----------------------------
     # Flip in y (-,+)
     # -----------------------------
     if y_shape > 1:
-        block = accum_from_corner(
-            ym0, xm0, gy[::-1, :], gx[::-1, :], label="(-,+)"
-        )
-        out[-1:-(y_shape): -1, :x_shape] = block[1:, :]
+        block = accum_from_corner(ym0, xm0, gy[::-1, :], gx[::-1, :], label="(-,+)")
+        out[-1:-(y_shape):-1, :x_shape] = block[1:, :]
 
     # -----------------------------
     # Flip in x and y (-,-)
@@ -271,21 +262,19 @@ def accum_from_corner(y_ref, x_ref, gy_block, gx_block, label=""):
         block = accum_from_corner(
             ymm, xmm, gy[::-1, ::-1], gx[::-1, ::-1], label="(-,-)"
         )
-        out[-1:-(y_shape): -1, -1:-(x_shape): -1] = block[1:, 1:]
+        out[-1:-(y_shape):-1, -1:-(x_shape):-1] = block[1:, 1:]
 
     return out
 
 
-
-
-def w_tilde_via_preload_from(w_tilde_preload, native_index_for_slim_index):
+def w_tilde_via_preload_from(curvature_preload, native_index_for_slim_index):
     """
-    Use the preloaded w_tilde matrix (see `w_tilde_preload_interferometer_from`) to compute
+    Use the preloaded w_tilde matrix (see `curvature_preload_interferometer_from`) to compute
     w_tilde (see `w_tilde_interferometer_from`) efficiently.
 
     Parameters
     ----------
-    w_tilde_preload
+    curvature_preload
         The preloaded values of the NUFFT that enable efficient computation of w_tilde.
     native_index_for_slim_index
         An array of shape [total_unmasked_pixels*sub_size] that maps every unmasked sub-pixel to its corresponding
@@ -311,7 +300,7 @@ def w_tilde_via_preload_from(w_tilde_preload, native_index_for_slim_index):
             y_diff = j_y - i_y
             x_diff = j_x - i_x
 
-            w_tilde_via_preload[i, j] = w_tilde_preload[y_diff, x_diff]
+            w_tilde_via_preload[i, j] = curvature_preload[y_diff, x_diff]
 
     for i in range(slim_size):
         for j in range(i, slim_size):

autoarray/inversion/inversion/interferometer/w_tilde.py

@@ -110,7 +110,7 @@ def curvature_matrix_diag(self) -> np.ndarray:
         mapper = self.cls_list_from(cls=AbstractMapper)[0]
 
         return inversion_interferometer_util.curvature_matrix_via_w_tilde_interferometer_from(
-            fft_state=self.w_tilde.operator_state,
+            fft_state=self.w_tilde.fft_state,
             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,

test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py

@@ -68,49 +68,14 @@ def test__data_vector_via_transformed_mapping_matrix_from():
     assert (data_vector_complex_via_blurred == data_vector_via_transformed).all()
 
 
-def test__w_tilde_curvature_interferometer_from():
-    noise_map = np.array([1.0, 2.0, 3.0])
-    uv_wavelengths = np.array([[0.0001, 2.0, 3000.0], [3000.0, 2.0, 0.0001]])
-
-    grid = aa.Grid2D.uniform(shape_native=(2, 2), pixel_scales=0.0005)
-
-    w_tilde = (
-        aa.util.inversion_interferometer_numba.w_tilde_curvature_interferometer_from(
-            noise_map_real=noise_map,
-            uv_wavelengths=uv_wavelengths,
-            grid_radians_slim=grid.array,
-        )
-    )
-
-    assert w_tilde == pytest.approx(
-        np.array(
-            [
-                [1.25, 0.75, 1.24997, 0.74998],
-                [0.75, 1.25, 0.74998, 1.24997],
-                [1.24994, 0.74998, 1.25, 0.75],
-                [0.74998, 1.24997, 0.75, 1.25],
-            ]
-        ),
-        1.0e-4,
-    )
-
-
-def test__curvature_matrix_via_w_tilde_preload_from():
+def test__curvature_matrix_via_curvature_preload_from():
     noise_map = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
     uv_wavelengths = np.array(
         [[0.0001, 2.0, 3000.0, 50000.0, 200000.0], [3000.0, 2.0, 0.0001, 10.0, 5000.0]]
     )
 
     grid = aa.Grid2D.uniform(shape_native=(3, 3), pixel_scales=0.0005)
 
-    w_tilde = (
-        aa.util.inversion_interferometer_numba.w_tilde_curvature_interferometer_from(
-            noise_map_real=noise_map,
-            uv_wavelengths=uv_wavelengths,
-            grid_radians_slim=grid.array,
-        )
-    )
-
     mapping_matrix = np.array(
         [
             [1.0, 0.0, 0.0],
@@ -125,34 +90,45 @@ def test__curvature_matrix_via_w_tilde_preload_from():
         ]
     )
 
-    curvature_matrix_via_w_tilde = aa.util.inversion.curvature_matrix_via_w_tilde_from(
-        w_tilde=w_tilde, mapping_matrix=mapping_matrix
+    curvature_preload = (
+        aa.util.inversion_interferometer.w_tilde_curvature_preload_interferometer_from(
+            noise_map_real=noise_map,
+            uv_wavelengths=uv_wavelengths,
+            shape_masked_pixels_2d=(3, 3),
+            grid_radians_2d=np.array(grid.native),
+        )
     )
 
-    w_tilde_preload = aa.util.inversion_interferometer_numba.w_tilde_curvature_preload_interferometer_from(
-        noise_map_real=noise_map,
-        uv_wavelengths=uv_wavelengths,
-        shape_masked_pixels_2d=(3, 3),
-        grid_radians_2d=np.array(grid.native),
+    native_index_for_slim_index = np.array(
+        [[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2], [2, 0], [2, 1], [2, 2]]
+    )
+
+    w_tilde = aa.util.inversion_interferometer.w_tilde_via_preload_from(
+        curvature_preload=curvature_preload,
+        native_index_for_slim_index=native_index_for_slim_index,
+    )
+
+    curvature_matrix_via_w_tilde = aa.util.inversion.curvature_matrix_via_w_tilde_from(
+        w_tilde=w_tilde, mapping_matrix=mapping_matrix
     )
 
     pix_indexes_for_sub_slim_index = np.array(
         [[0], [2], [1], [1], [2], [2], [0], [2], [0]]
     )
 
-    pix_size_for_sub_slim_index = np.ones(shape=(9,)).astype("int")
     pix_weights_for_sub_slim_index = np.ones(shape=(9, 1))
 
-    native_index_for_slim_index = np.array(
-        [[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2], [2, 0], [2, 1], [2, 2]]
+    w_tilde = aa.WTildeInterferometer(
+        curvature_preload=curvature_preload,
+        dirty_image=None,
+        real_space_mask=grid.mask,
     )
 
-    curvature_matrix_via_preload = aa.util.inversion_interferometer_numba.curvature_matrix_via_w_tilde_curvature_preload_interferometer_from(
-        curvature_preload=w_tilde_preload,
+    curvature_matrix_via_preload = aa.util.inversion_interferometer.curvature_matrix_via_w_tilde_interferometer_from(
+        fft_state=w_tilde.fft_state,
         pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,
-        pix_size_for_sub_slim_index=pix_size_for_sub_slim_index,
         pix_weights_for_sub_slim_index=pix_weights_for_sub_slim_index,
-        native_index_for_slim_index=native_index_for_slim_index,
+        rect_index_for_mask_index=w_tilde.rect_index_for_mask_index,
         pix_pixels=3,
     )
 
@@ -161,43 +137,6 @@ def test__curvature_matrix_via_w_tilde_preload_from():
     )
 
 
-def test__curvature_matrix_via_w_tilde_two_methods_agree():
-    noise_map = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
-    uv_wavelengths = np.array(
-        [[0.0001, 2.0, 3000.0, 50000.0, 200000.0], [3000.0, 2.0, 0.0001, 10.0, 5000.0]]
-    )
-
-    grid = aa.Grid2D.uniform(shape_native=(3, 3), pixel_scales=0.0005)
-
-    w_tilde = (
-        aa.util.inversion_interferometer_numba.w_tilde_curvature_interferometer_from(
-            noise_map_real=noise_map,
-            uv_wavelengths=uv_wavelengths,
-            grid_radians_slim=grid.array,
-        )
-    )
-
-    w_tilde_preload = aa.util.inversion_interferometer_numba.w_tilde_curvature_preload_interferometer_from(
-        noise_map_real=np.array(noise_map),
-        uv_wavelengths=np.array(uv_wavelengths),
-        shape_masked_pixels_2d=(3, 3),
-        grid_radians_2d=np.array(grid.native),
-    )
-
-    native_index_for_slim_index = np.array(
-        [[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2], [2, 0], [2, 1], [2, 2]]
-    )
-
-    w_tilde_via_preload = (
-        aa.util.inversion_interferometer_numba.w_tilde_via_preload_from(
-            w_tilde_preload=w_tilde_preload,
-            native_index_for_slim_index=native_index_for_slim_index,
-        )
-    )
-
-    assert (w_tilde == w_tilde_via_preload).all()
-
-
 def test__identical_inversion_values_for_two_methods():
     real_space_mask = aa.Mask2D.all_false(
         shape_native=(7, 7),