Merge pull request #170 from Jammy2211/feature/jax_ndarray_casting_rules

Feature/jax ndarray casting rules

Author: Jammy2211

autoarray/__init__.py

@@ -20,6 +20,7 @@
 from .fit.fit_imaging import FitImaging
 from .fit.fit_interferometer import FitInterferometer
 from .geometry.geometry_2d import Geometry2D
+from .inversion.convolver import Convolver
 from .inversion.pixelization.mappers.abstract import AbstractMapper
 from .inversion.pixelization import mesh
 from .inversion.pixelization import image_mesh
@@ -44,7 +45,6 @@
 from .inversion.inversion.imaging.w_tilde import InversionImagingWTilde
 from .inversion.inversion.interferometer.w_tilde import InversionInterferometerWTilde
 from .inversion.inversion.interferometer.mapping import InversionInterferometerMapping
-from .inversion.inversion.interferometer.lop import InversionInterferometerMappingPyLops
 from .inversion.linear_obj.linear_obj import LinearObj
 from .inversion.linear_obj.func_list import AbstractLinearObjFuncList
 from .mask.derive.indexes_2d import DeriveIndexes2D

autoarray/abstract_ndarray.py

@@ -4,10 +4,11 @@
 
 from abc import ABC
 from abc import abstractmethod
+import jax.numpy as jnp
 
 from autoconf.fitsable import output_to_fits
 
-from autoarray.numpy_wrapper import np, register_pytree_node, Array
+from autoarray.numpy_wrapper import register_pytree_node, Array
 
 from typing import TYPE_CHECKING
 
@@ -82,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
@@ -104,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
@@ -117,7 +118,7 @@ def instance_unflatten(cls, 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.
 
@@ -164,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):
@@ -330,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 = np.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/grids.yaml

@@ -166,8 +166,9 @@ def __init__(
 
         self.psf = psf
 
-        if psf.mask.shape[0] % 2 == 0 or psf.mask.shape[1] % 2 == 0:
-            raise exc.KernelException("Kernel2D Kernel2D must be odd")
+        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):
@@ -178,6 +179,27 @@ 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.
+        """
+
+        from autoarray.inversion.convolver import Convolver
+
+        return Convolver(mask=self.mask, kernel=self.psf)
+
     @cached_property
     def w_tilde(self):
         """
@@ -203,9 +225,11 @@ 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(
@@ -408,20 +432,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,

autoarray/dataset/imaging/dataset.py

@@ -151,7 +151,7 @@ def via_image_from(
                 pixel_scales=image.pixel_scales,
             )
 
-        if np.isnan(noise_map).any():
+        if np.isnan(noise_map.array).any():
             raise exc.DatasetException(
                 "The noise-map has NaN values in it. This suggests your exposure time and / or"
                 "background sky levels are too low, creating signal counts at or close to 0.0."
@@ -161,7 +161,9 @@ def via_image_from(
             image = image - background_sky_map
 
         mask = Mask2D.all_false(
-            shape_native=image.shape_native, pixel_scales=image.pixel_scales
+            shape_native=image.shape_native,
+            pixel_scales=image.pixel_scales,
+            origin=image.origin,
         )
 
         image = Array2D(values=image, mask=mask)

autoarray/dataset/imaging/simulator.py

@@ -193,7 +193,9 @@ def w_tilde(self):
 
         w_matrix = inversion_interferometer_util.w_tilde_via_preload_from(
             w_tilde_preload=curvature_preload,
-            native_index_for_slim_index=self.real_space_mask.derive_indexes.native_for_slim,
+            native_index_for_slim_index=np.array(
+                self.real_space_mask.derive_indexes.native_for_slim
+            ).astype("int"),
         )
 
         dirty_image = self.transformer.image_from(
@@ -205,7 +207,7 @@ def w_tilde(self):
         return WTildeInterferometer(
             w_matrix=w_matrix,
             curvature_preload=curvature_preload,
-            dirty_image=dirty_image,
+            dirty_image=np.array(dirty_image.array),
             real_space_mask=self.real_space_mask,
             noise_map_value=self.noise_map[0],
         )

autoarray/dataset/interferometer/dataset.py

@@ -149,7 +149,8 @@ def noise_map_via_data_eps_and_exposure_time_map_from(data_eps, exposure_time_ma
         The exposure time at every data-point of the data.
     """
     return data_eps.with_new_array(
-        np.abs(data_eps * exposure_time_map) ** 0.5 / exposure_time_map
+        np.abs(data_eps.array * exposure_time_map.array) ** 0.5
+        / exposure_time_map.array
     )
 
 
@@ -263,15 +264,17 @@ 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,
         ]
@@ -406,9 +409,10 @@ def poisson_noise_via_data_eps_from(data_eps, exposure_time_map, seed=-1):
         An array describing simulated poisson noise_maps
     """
     setup_random_seed(seed)
-    image_counts = np.multiply(data_eps, exposure_time_map)
+
+    image_counts = np.multiply(data_eps.array, exposure_time_map.array)
     return data_eps - np.divide(
-        np.random.poisson(image_counts, data_eps.shape), exposure_time_map
+        np.random.poisson(image_counts, data_eps.shape), exposure_time_map.array
     )
 
 
@@ -506,8 +510,6 @@ def noise_map_with_signal_to_noise_limit_from(
     from autoarray.structures.arrays.uniform_1d import Array1D
     from autoarray.structures.arrays.uniform_2d import Array2D
 
-    # TODO : Refacotr into a util
-
     signal_to_noise_map = data / noise_map
     signal_to_noise_map[signal_to_noise_map < 0] = 0
 
@@ -517,12 +519,14 @@ 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(
-        shape_native=data.shape_native, pixel_scales=data.pixel_scales
+        shape_native=data.shape_native,
+        pixel_scales=data.pixel_scales,
+        origin=data.origin,
     )
 
     if len(noise_map.native) == 1:

autoarray/dataset/preprocess.py

@@ -86,7 +86,7 @@ def chi_squared(self) -> float:
         """
         Returns the chi-squared terms of the model data's fit to an dataset, by summing the chi-squared-map.
         """
-        return fit_util.chi_squared_from(chi_squared_map=self.chi_squared_map)
+        return fit_util.chi_squared_from(chi_squared_map=self.chi_squared_map.array)
 
     @property
     def noise_normalization(self) -> float:
@@ -95,7 +95,7 @@ def noise_normalization(self) -> float:
 
         [Noise_Term] = sum(log(2*pi*[Noise]**2.0))
         """
-        return fit_util.noise_normalization_from(noise_map=self.noise_map)
+        return fit_util.noise_normalization_from(noise_map=self.noise_map.array)
 
     @property
     def log_likelihood(self) -> float: