more updates

Author: Jammy2211

autoarray/__init__.py

@@ -9,9 +9,7 @@
 from . import util
 from . import fixtures
 from . import mock as m
-from .inversion.inversion.interferometer.inversion_interferometer_util import (
-    load_curvature_preload,
-)
+
 from .dataset import preprocess
 from .dataset.abstract.dataset import AbstractDataset
 from .dataset.grids import GridsInterface

autoarray/dataset/interferometer/dataset.py

@@ -201,14 +201,10 @@ def apply_sparse_linear_algebra(
                 "INTERFEROMETER - Computing W-Tilde; runtime scales with visibility count and mask resolution, CPU run times may exceed hours."
             )
 
-            curvature_preload = inversion_interferometer_util.w_tilde_curvature_preload_interferometer_from(
-                noise_map_real=self.noise_map.array.real,
-                uv_wavelengths=self.uv_wavelengths,
-                shape_masked_pixels_2d=self.transformer.grid.mask.shape_native_masked_pixels,
-                grid_radians_2d=self.transformer.grid.mask.derive_grid.all_false.in_radians.native.array,
+            curvature_preload = self.curvature_preload_from(
                 chunk_k=chunk_k,
-                show_memory=show_memory,
                 show_progress=show_progress,
+                show_memory=show_memory,
                 use_jax=use_jax,
             )
 
@@ -233,6 +229,24 @@ def apply_sparse_linear_algebra(
             sparse_linalg=sparse_linalg,
         )
 
+    def curvature_preload_from(
+        self,
+        chunk_k: int = 2048,
+        show_progress: bool = False,
+        show_memory: bool = False,
+        use_jax: bool = False,
+    ):
+        return inversion_interferometer_util.w_tilde_curvature_preload_interferometer_from(
+            noise_map_real=self.noise_map.array.real,
+            uv_wavelengths=self.uv_wavelengths,
+            shape_masked_pixels_2d=self.transformer.grid.mask.shape_native_masked_pixels,
+            grid_radians_2d=self.transformer.grid.mask.derive_grid.all_false.in_radians.native.array,
+            chunk_k=chunk_k,
+            show_memory=show_memory,
+            show_progress=show_progress,
+            use_jax=use_jax,
+        )
+
     @property
     def mask(self):
         return self.real_space_mask

autoarray/inversion/inversion/imaging/inversion_imaging_util.py

@@ -294,7 +294,7 @@ class ImagingSparseLinAlg:
 
     data_native: np.ndarray
     noise_map_native: np.ndarray
-    weight_map : np.ndarray
+    weight_map: np.ndarray
     inverse_variances_native: "jax.Array"  # (y, x) float64
     y_shape: int
     x_shape: int
@@ -341,10 +341,8 @@ def from_noise_map_and_psf(
         )
         inverse_variances_native = inverse_variances_native.native
 
-        weight_map = data.array / (noise_map.array ** 2)
-        weight_map = Array2D(
-            values=weight_map, mask=noise_map.mask
-        )
+        weight_map = data.array / (noise_map.array**2)
+        weight_map = Array2D(values=weight_map, mask=noise_map.mask)
 
         # If you *also* want to zero masked pixels explicitly:
         # mask_native = noise_map.mask  (depends on your API; might be bool native)

autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py

@@ -528,36 +528,6 @@ def w_tilde_via_preload_from(curvature_preload, native_index_for_slim_index):
     return w_tilde_via_preload
 
 
-def load_curvature_preload(
-    file: Union[str, Path],
-) -> Optional[np.ndarray]:
-    """
-    Load a saved curvature_preload if (and only if) it is compatible with the current mask geometry.
-
-    Parameters
-    ----------
-    file
-        Path to a previously saved NPZ.
-    require_mask_hash
-        If True, require the full mask content hash to match (safest).
-        If False, only bbox + shape + pixel scales are checked.
-
-    Returns
-    -------
-    np.ndarray
-        The loaded curvature_preload if compatible, otherwise raises ValueError.
-    """
-    file = Path(file)
-    if file.suffix.lower() != ".npz":
-        file = file.with_suffix(".npz")
-
-    if not file.exists():
-        raise FileNotFoundError(str(file))
-
-    with np.load(file, allow_pickle=False) as npz:
-        return np.asarray(npz["curvature_preload"])
-
-
 @dataclass(frozen=True)
 class InterferometerSparseLinAlg:
     """
@@ -718,44 +688,3 @@ def compute_block(start_col: int) -> jnp.ndarray:
             pix_weights_for_sub_slim_index,
             fft_index_for_masked_pixel,
         )
-
-    def save_curvature_preload(
-        self,
-        file: Union[str, Path],
-        *,
-        overwrite: bool = False,
-    ) -> Path:
-        """
-        Save curvature_preload plus enough metadata to ensure it is only reused when safe.
-
-        Uses NPZ so we can store:
-          - curvature_preload (array)
-          - meta_json (string)
-
-        Parameters
-        ----------
-        file
-            Path to save to. Recommended suffix: ".npz".
-            If you pass ".npy", we will still save an ".npz" next to it.
-        overwrite
-            If False and the file exists, raise FileExistsError.
-
-        Returns
-        -------
-        Path
-            The path actually written (will end with ".npz").
-        """
-        file = Path(file)
-
-        # Force .npz (storing metadata safely)
-        if file.suffix.lower() != ".npz":
-            file = file.with_suffix(".npz")
-
-        if file.exists() and not overwrite:
-            raise FileExistsError(f"File already exists: {file}")
-
-        np.savez_compressed(
-            file,
-            curvature_preload=np.asarray(self.curvature_preload),
-        )
-        return file

test_autoarray/dataset/interferometer/test_dataset.py

@@ -149,55 +149,3 @@ def test__different_interferometer_without_mock_objects__customize_constructor_i
     assert (dataset.data == 1.0 + 1.0j * np.ones((19,))).all()
     assert (dataset.noise_map == 2.0 + 2.0j * np.ones((19,))).all()
     assert (dataset.uv_wavelengths == 3.0 * np.ones((19, 2))).all()
-
-
-def test__curvature_preload_metadata_from(
-    visibilities_7,
-    visibilities_noise_map_7,
-    uv_wavelengths_7x2,
-    mask_2d_7x7,
-):
-
-    dataset = aa.Interferometer(
-        data=visibilities_7,
-        noise_map=visibilities_noise_map_7,
-        uv_wavelengths=uv_wavelengths_7x2,
-        real_space_mask=mask_2d_7x7,
-    )
-
-    dataset = dataset.apply_sparse_linear_algebra(use_jax=False)
-
-    file = f"{test_data_path}/curvature_preload_metadata"
-
-    dataset.sparse_linalg.save_curvature_preload(
-        file=file,
-        overwrite=True,
-    )
-
-    curvature_preload = aa.load_curvature_preload(
-        file=file, real_space_mask=dataset.real_space_mask
-    )
-
-    assert curvature_preload[0, 0] == pytest.approx(1.75, 1.0e-4)
-
-    real_space_mask_changed = np.array(
-        [
-            [True, True, True, True, True, True, True],
-            [True, True, True, True, True, True, True],
-            [True, True, False, False, False, True, True],
-            [True, True, False, True, False, True, True],
-            [True, True, False, False, False, True, True],
-            [True, True, True, True, True, True, True],
-            [True, True, True, True, True, True, True],
-        ]
-    )
-
-    real_space_mask_changed = aa.Mask2D(
-        mask=real_space_mask_changed, pixel_scales=(1.0, 1.0)
-    )
-
-    with pytest.raises(ValueError):
-
-        curvature_preload = aa.load_curvature_preload(
-            file=file, real_space_mask=real_space_mask_changed
-        )

test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py

@@ -37,35 +37,46 @@ def test__w_tilde_imaging_from():
 
 
 def test__weighted_data_imaging_from():
-    data = np.array(
-        [
+
+    mask = aa.Mask2D(
+        mask=[
+            [True, True, True, True],
+            [True, False, False, True],
+            [True, False, False, True],
+            [True, True, True, True],
+        ],
+        pixel_scales=(1.0, 1.0),
+    )
+
+    data = aa.Array2D(
+        values=[
             [0.0, 0.0, 0.0, 0.0],
             [0.0, 2.0, 1.0, 0.0],
             [0.0, 1.0, 2.0, 0.0],
             [0.0, 0.0, 0.0, 0.0],
-        ]
+        ],
+        mask=mask,
     )
 
-    noise_map = np.array(
-        [
+    noise_map = aa.Array2D(
+        values=[
             [0.0, 0.0, 0.0, 0.0],
             [0.0, 1.0, 1.0, 0.0],
             [0.0, 1.0, 2.0, 0.0],
             [0.0, 0.0, 0.0, 0.0],
-        ]
+        ],
+        mask=mask,
     )
 
     kernel = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [1.0, 2.0, 0.0]])
 
     native_index_for_slim_index = np.array([[1, 1], [1, 2], [2, 1], [2, 2]])
 
-    weight_map = data.array / (noise_map.array ** 2)
-    weight_map = aa.Array2D(
-        values=weight_map, mask=noise_map.mask
-    )
+    weight_map = data / (noise_map**2)
+    weight_map = aa.Array2D(values=weight_map, mask=mask)
 
     weighted_data = aa.util.inversion_imaging.weighted_data_imaging_from(
-        weight_map_native=weight_map.native,
+        weight_map_native=weight_map.native.array,
         kernel_native=kernel,
         native_index_for_slim_index=native_index_for_slim_index,
     )
@@ -230,9 +241,11 @@ def test__data_vector_via_weighted_data_two_methods_agree():
             sub_fraction_slim=mapper.over_sampler.sub_fraction.array,
         )
 
+        weight_map = image.array / (noise_map.array**2)
+        weight_map = aa.Array2D(values=weight_map, mask=noise_map.mask)
+
         weighted_data = aa.util.inversion_imaging.weighted_data_imaging_from(
-            image_native=image.native.array,
-            noise_map_native=noise_map.native.array,
+            weight_map_native=weight_map.native.array,
             kernel_native=kernel.native.array,
             native_index_for_slim_index=mask.derive_indexes.native_for_slim.astype(
                 "int"