Merge pull request #167 from Jammy2211/feature/unit_tests

Feature/unit tests

Author: Jammy2211

autoarray/__init__.py

@@ -4,7 +4,6 @@
 from . import fixtures
 from . import mock as m
 from .numba_util import profile_func
-from .preloads import Preloads
 from .dataset import preprocess
 from .dataset.abstract.dataset import AbstractDataset
 from .dataset.abstract.w_tilde import AbstractWTilde
@@ -55,8 +54,6 @@
 from .mask.derive.grid_2d import DeriveGrid2D
 from .mask.mask_1d import Mask1D
 from .mask.mask_2d import Mask2D
-from .operators.convolver import Convolver
-from .operators.convolver import Convolver
 from .operators.transformer import TransformerDFT
 from .operators.transformer import TransformerNUFFT
 from .operators.over_sampling.decorator import over_sample

autoarray/abstract_ndarray.py

@@ -4,11 +4,11 @@
 
 from abc import ABC
 from abc import abstractmethod
-import numpy as np
+import jax.numpy as jnp
 
 from autoconf.fitsable import output_to_fits
 
-from autoarray.numpy_wrapper import numpy as npw, register_pytree_node, Array
+from autoarray.numpy_wrapper import register_pytree_node, Array
 
 from typing import TYPE_CHECKING
 
@@ -83,7 +83,7 @@ def __init__(self, array):
 
     def invert(self):
         new = self.copy()
-        new._array = np.invert(new._array)
+        new._array = jnp.invert(new._array)
         return new
 
     @classmethod
@@ -105,7 +105,7 @@ def instance_flatten(cls, instance):
     @staticmethod
     def flip_hdu_for_ds9(values):
         if conf.instance["general"]["fits"]["flip_for_ds9"]:
-            return np.flipud(values)
+            return jnp.flipud(values)
         return values
 
     @classmethod
@@ -114,11 +114,11 @@ def instance_unflatten(cls, aux_data, children):
         Unflatten a tuple of attributes (i.e. a pytree) into an instance of an autoarray class
         """
         instance = cls.__new__(cls)
-        for key, value in zip(aux_data, children[1:]):
+        for key, value in zip(aux_data, children):
             setattr(instance, key, value)
         return instance
 
-    def with_new_array(self, array: np.ndarray) -> "AbstractNDArray":
+    def with_new_array(self, array: jnp.ndarray) -> "AbstractNDArray":
         """
         Copy this object but give it a new array.
 
@@ -165,7 +165,7 @@ def __iter__(self):
 
     @to_new_array
     def sqrt(self):
-        return np.sqrt(self._array)
+        return jnp.sqrt(self._array)
 
     @property
     def array(self):
@@ -331,13 +331,13 @@ def __getitem__(self, item):
         result = self._array[item]
         if isinstance(item, slice):
             result = self.with_new_array(result)
-        if isinstance(result, np.ndarray):
+        if isinstance(result, jnp.ndarray):
             result = self.with_new_array(result)
         return result
 
     def __setitem__(self, key, value):
-        if isinstance(key, (np.ndarray, AbstractNDArray, Array)):
-            self._array = npw.where(key, value, self._array)
+        if isinstance(key, (jnp.ndarray, AbstractNDArray, Array)):
+            self._array = jnp.where(key, value, self._array)
         else:
             self._array[key] = value
 

autoarray/config/general.yaml

@@ -16,6 +16,3 @@ pixelization:
   voronoi_nn_max_interpolation_neighbors: 300
 structures:
   native_binned_only: false           # If True, data structures are only stored in their native and binned format. This is used to reduce memory usage in autocti.
-test:
-  preloads_check_threshold: 1.0     # If the figure of merit of a fit with and without preloads is greater than this threshold, the check preload test fails and an exception raised for a model-fit.
-

autoarray/config/grids.yaml

@@ -9,7 +9,6 @@
 from autoarray.dataset.grids import GridsDataset
 from autoarray.dataset.imaging.w_tilde import WTildeImaging
 from autoarray.structures.arrays.uniform_2d import Array2D
-from autoarray.operators.convolver import Convolver
 from autoarray.structures.arrays.kernel_2d import Kernel2D
 from autoarray.mask.mask_2d import Mask2D
 from autoarray import type as ty
@@ -30,7 +29,7 @@ def __init__(
         noise_covariance_matrix: Optional[np.ndarray] = None,
         over_sample_size_lp: Union[int, Array2D] = 4,
         over_sample_size_pixelization: Union[int, Array2D] = 4,
-        pad_for_convolver: bool = False,
+        pad_for_psf: bool = False,
         use_normalized_psf: Optional[bool] = True,
         check_noise_map: bool = True,
     ):
@@ -77,7 +76,7 @@ def __init__(
         over_sample_size_pixelization
             How over sampling is performed for the grid which is associated with a pixelization, which is therefore
             passed into the calculations performed in the `inversion` module.
-        pad_for_convolver
+        pad_for_psf
             The PSF convolution may extend beyond the edges of the image mask, which can lead to edge effects in the
             convolved image. If `True`, the image and noise-map are padded to ensure the PSF convolution does not
             extend beyond the edge of the image.
@@ -90,9 +89,9 @@ def __init__(
 
         self.unmasked = None
 
-        self.pad_for_convolver = pad_for_convolver
+        self.pad_for_psf = pad_for_psf
 
-        if pad_for_convolver and psf is not None:
+        if pad_for_psf and psf is not None:
             try:
                 data.mask.derive_mask.blurring_from(
                     kernel_shape_native=psf.shape_native
@@ -162,11 +161,15 @@ def __init__(
 
         if psf is not None and use_normalized_psf:
             psf = Kernel2D.no_mask(
-                values=psf.native, pixel_scales=psf.pixel_scales, normalize=True
+                values=psf.native._array, pixel_scales=psf.pixel_scales, normalize=True
             )
 
         self.psf = psf
 
+        if psf is not None:
+            if psf.mask.shape[0] % 2 == 0 or psf.mask.shape[1] % 2 == 0:
+                raise exc.KernelException("Kernel2D Kernel2D must be odd")
+
     @cached_property
     def grids(self):
         return GridsDataset(
@@ -176,25 +179,6 @@ def grids(self):
             psf=self.psf,
         )
 
-    @cached_property
-    def convolver(self):
-        """
-        Returns a `Convolver` from a mask and 2D PSF kernel.
-
-        The `Convolver` stores in memory the array indexing between the mask and PSF, enabling efficient 2D PSF
-        convolution of images and matrices used for linear algebra calculations (see `operators.convolver`).
-
-        This uses lazy allocation such that the calculation is only performed when the convolver is used, ensuring
-        efficient set up of the `Imaging` class.
-
-        Returns
-        -------
-        Convolver
-            The convolver given the masked imaging data's mask and PSF.
-        """
-
-        return Convolver(mask=self.mask, kernel=self.psf)
-
     @cached_property
     def w_tilde(self):
         """
@@ -220,9 +204,9 @@ def w_tilde(self):
             indexes,
             lengths,
         ) = inversion_imaging_util.w_tilde_curvature_preload_imaging_from(
-            noise_map_native=np.array(self.noise_map.native),
-            kernel_native=np.array(self.psf.native),
-            native_index_for_slim_index=self.mask.derive_indexes.native_for_slim,
+            noise_map_native=np.array(self.noise_map.native.array).astype("float64"),
+            kernel_native=np.array(self.psf.native.array).astype("float64"),
+            native_index_for_slim_index=np.array(self.mask.derive_indexes.native_for_slim).astype("int"),
         )
 
         return WTildeImaging(
@@ -370,7 +354,7 @@ def apply_mask(self, mask: Mask2D) -> "Imaging":
             noise_covariance_matrix=noise_covariance_matrix,
             over_sample_size_lp=over_sample_size_lp,
             over_sample_size_pixelization=over_sample_size_pixelization,
-            pad_for_convolver=True,
+            pad_for_psf=True,
         )
 
         dataset.unmasked = unmasked_dataset
@@ -425,20 +409,20 @@ def apply_noise_scaling(
         """
 
         if signal_to_noise_value is None:
-            noise_map = self.noise_map.native
-            noise_map[mask == False] = noise_value
+            noise_map = np.array(self.noise_map.native.array)
+            noise_map[mask.array == False] = noise_value
         else:
             noise_map = np.where(
                 mask == False,
-                np.median(self.data.native[mask.derive_mask.edge == False])
+                np.median(self.data.native.array[mask.derive_mask.edge == False])
                 / signal_to_noise_value,
-                self.noise_map.native,
+                self.noise_map.native.array,
             )
 
         if should_zero_data:
-            data = np.where(np.invert(mask), 0.0, self.data.native)
+            data = np.where(np.invert(mask.array), 0.0, self.data.native.array)
         else:
-            data = self.data.native
+            data = self.data.native.array
 
         data_unmasked = Array2D.no_mask(
             values=data,
@@ -463,7 +447,7 @@ def apply_noise_scaling(
             noise_covariance_matrix=self.noise_covariance_matrix,
             over_sample_size_lp=self.over_sample_size_lp,
             over_sample_size_pixelization=self.over_sample_size_pixelization,
-            pad_for_convolver=False,
+            pad_for_psf=False,
             check_noise_map=False,
         )
 
@@ -511,7 +495,7 @@ def apply_over_sampling(
             over_sample_size_lp=over_sample_size_lp or self.over_sample_size_lp,
             over_sample_size_pixelization=over_sample_size_pixelization
             or self.over_sample_size_pixelization,
-            pad_for_convolver=False,
+            pad_for_psf=False,
             check_noise_map=False,
         )
 

autoarray/dataset/imaging/dataset.py

@@ -276,5 +276,5 @@ def output_to_fits(
             )
 
     @property
-    def convolver(self):
+    def psf(self):
         return None

autoarray/dataset/interferometer/dataset.py

@@ -263,15 +263,15 @@ def edges_from(image, no_edges):
     edges = []
 
     for edge_no in range(no_edges):
-        top_edge = image.native[edge_no, edge_no : image.shape_native[1] - edge_no]
-        bottom_edge = image.native[
+        top_edge = image.native.array[edge_no, edge_no : image.shape_native[1] - edge_no]
+        bottom_edge = image.native.array[
             image.shape_native[0] - 1 - edge_no,
             edge_no : image.shape_native[1] - edge_no,
         ]
-        left_edge = image.native[
+        left_edge = image.native.array[
             edge_no + 1 : image.shape_native[0] - 1 - edge_no, edge_no
         ]
-        right_edge = image.native[
+        right_edge = image.native.array[
             edge_no + 1 : image.shape_native[0] - 1 - edge_no,
             image.shape_native[1] - 1 - edge_no,
         ]
@@ -328,7 +328,7 @@ def background_noise_map_via_edges_from(image, no_edges):
 def psf_with_odd_dimensions_from(psf):
     """
     If the PSF kernel has one or two even-sized dimensions, return a PSF object where the kernel has odd-sized
-    dimensions (odd-sized dimensions are required by a *Convolver*).
+    dimensions (odd-sized dimensions are required for 2D convolution).
 
     Kernels are rescaled using the scikit-image routine rescale, which performs rescaling via an interpolation
     routine. This may lead to loss of accuracy in the PSF kernel and it is advised that users, where possible,
@@ -517,8 +517,8 @@ def noise_map_with_signal_to_noise_limit_from(
     noise_map_limit = np.where(
         (signal_to_noise_map.native > signal_to_noise_limit)
         & (noise_limit_mask == False),
-        np.abs(data.native) / signal_to_noise_limit,
-        noise_map.native,
+        np.abs(data.native.array) / signal_to_noise_limit,
+        noise_map.native.array,
     )
 
     mask = Mask2D.all_false(

autoarray/dataset/preprocess.py

@@ -106,17 +106,6 @@ class PlottingException(Exception):
     pass
 
 
-class PreloadsException(Exception):
-    """
-    Raises exceptions associated with the `preloads.py` module and `Preloads` class.
-
-    For example if the preloaded quantities lead to a change in figure of merit of a fit compared to a fit without
-    preloading.
-    """
-
-    pass
-
-
 class ProfilingException(Exception):
     """
     Raises exceptions associated with in-built profiling tools (e.g. the `profile_func` decorator).