test__convolve_image

Author: Jammy2211

autoarray/structures/arrays/kernel_2d.py

@@ -8,7 +8,6 @@
 
 from autoconf.fitsable import header_obj_from
 
-from autoarray.mask.mask_2d import Mask2D
 from autoarray.structures.arrays.uniform_2d import AbstractArray2D
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.grids.uniform_2d import Grid2D
@@ -17,6 +16,7 @@
 from autoarray import exc
 from autoarray import type as ty
 from autoarray.structures.arrays import array_2d_util
+from autoarray.mask.mask_2d import mask_2d_util
 
 
 class Kernel2D(AbstractArray2D):
@@ -483,28 +483,74 @@ def convolved_array_from(self, array: Array2D) -> Array2D:
 
         return Array2D(values=convolved_array_1d, mask=array_2d.mask)
 
-    def jax_convolve(self, image, blurring_image, method="auto"):
+    def convolve_image(self, image, blurring_image, jax_method="fft"):
+        """
+        For a given 1D array and blurring array, convolve the two using this convolver.
+
+        Parameters
+        ----------
+        image
+            1D array of the values which are to be blurred with the convolver's PSF.
+        blurring_image
+            1D array of the blurring values which blur into the array after PSF convolution.
+        jax_method
+            If JAX is enabled this keyword will indicate what method is used for the PSF
+            convolution. Can be either `direct` to calculate it in real space or `fft`
+            to calculated it via a fast Fourier transform. `fft` is typically faster for
+            kernels that are more than about 5x5. Default is `fft`.
+        """
+
+        if self.mask.shape[0] % 2 == 0 or self.mask.shape[1] % 2 == 0:
+            raise exc.KernelException("Kernel2D Kernel2D must be odd")
+
+        print(type(image.native + blurring_image.native))
+        print(type(self.native))
+
+        convolved_array_2d = scipy.signal.convolve2d((image.native + blurring_image.native)._array, self.native._array, mode="same")
 
+        convolved_array_1d = array_2d_util.array_2d_slim_from(
+            mask_2d=np.array(image.mask),
+            array_2d_native=convolved_array_2d,
+        )
+
+        return Array2D(values=convolved_array_1d, mask=image.mask)
+
+    def convolve_image_jax_from(self, array, blurring_array, method="auto"):
+        """
+        For a given 1D array and blurring array, convolve the two using this convolver.
+
+        Parameters
+        ----------
+        array
+            1D array of the values which are to be blurred with the convolver's PSF.
+        blurring_array
+            1D array of the blurring values which blur into the array after PSF convolution.
+        jax_method
+            If JAX is enabled this keyword will indicate what method is used for the PSF
+            convolution. Can be either `direct` to calculate it in real space or `fft`
+            to calculated it via a fast Fourier transform. `fft` is typically faster for
+            kernels that are more than about 5x5. Default is `fft`.
+        """
         slim_to_native = jnp.nonzero(
-            jnp.logical_not(self.mask.array), size=image.shape[0]
+            jnp.logical_not(self.mask.array), size=array.shape[0]
         )
         slim_to_native_blurring = jnp.nonzero(
-            jnp.logical_not(self.blurring_mask), size=blurring_image.shape[0]
+            jnp.logical_not(self.blurring_mask), size=blurring_array.shape[0]
         )
 
-        expanded_image_native = jnp.zeros(self.mask.shape)
+        expanded_array_native = jnp.zeros(self.mask.shape)
 
-        expanded_image_native = expanded_image_native.at[slim_to_native].set(
-            image.array
+        expanded_array_native = expanded_array_native.at[slim_to_native].set(
+            array.array
         )
-        expanded_image_native = expanded_image_native.at[slim_to_native_blurring].set(
-            blurring_image.array
+        expanded_array_native = expanded_array_native.at[slim_to_native_blurring].set(
+            blurring_array.array
         )
 
         kernel = np.array(self.kernel.native.array)
 
         convolve_native = jax.scipy.signal.convolve(
-            expanded_image_native, kernel, mode="same", method=method
+            expanded_array_native, kernel, mode="same", method=method
         )
 
         convolve_slim = convolve_native[slim_to_native]

test_autoarray/structures/arrays/test_kernel_2d.py

@@ -1,11 +1,10 @@
-from astropy.io import fits
 from astropy import units
 from astropy.modeling import functional_models
 from astropy.coordinates import Angle
+import jax.numpy as jnp
 import numpy as np
 import pytest
 from os import path
-import os
 
 import autoarray as aa
 from autoarray import exc
@@ -359,6 +358,36 @@ def test__convolved_array_from():
     ).all()
 
 
+def test__convolved_array_from__input_jax_array():
+
+    array_2d = jnp.array(
+        [
+            [0.0, 0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 1.0],
+            [0.0, 0.0, 0.0, 0.0],
+        ])
+
+    kernel_2d = aa.Kernel2D.no_mask(
+        values=[[1.0, 1.0, 1.0], [2.0, 2.0, 1.0], [1.0, 3.0, 3.0]], pixel_scales=1.0
+    )
+
+    blurred_array_2d = kernel_2d.convolved_array_from(array_2d)
+
+    assert (
+        blurred_array_2d.native
+        == np.array(
+            [
+                [1.0, 1.0, 0.0, 0.0],
+                [2.0, 1.0, 1.0, 1.0],
+                [3.0, 3.0, 2.0, 2.0],
+                [0.0, 0.0, 1.0, 3.0],
+            ]
+        )
+    ).all()
+
+
+
 def test__convolve_mapping_matrix():
     mask = np.array(
         [
@@ -482,19 +511,19 @@ def test__convolve_mapping_matrix():
     )
 
 
-def test__compare_to_full_2d_convolution():
-    # Setup a blurred data, using the PSF to perform the convolution in 2D, then masks it to make a 1d array.
-
-    import scipy.signal
+def test__convolve_image():
 
     mask = aa.Mask2D.circular(
         shape_native=(30, 30), pixel_scales=(1.0, 1.0), radius=4.0
     )
-    kernel = aa.Kernel2D.no_mask(values=np.arange(49).reshape(7, 7), pixel_scales=1.0)
-    image = aa.Array2D.no_mask(values=np.arange(900).reshape(30, 30), pixel_scales=1.0)
+
+    import scipy.signal
+
+    kernel = np.arange(49).reshape(7, 7)
+    image = np.arange(900).reshape(30, 30)
 
     blurred_image_via_scipy = scipy.signal.convolve2d(
-        image.native, kernel.native, mode="same"
+        image, kernel, mode="same"
     )
     blurred_image_via_scipy = aa.Array2D.no_mask(
         values=blurred_image_via_scipy, pixel_scales=1.0
@@ -503,7 +532,10 @@ def test__compare_to_full_2d_convolution():
         values=blurred_image_via_scipy.native, mask=mask
     )
 
-    # Now reproduce this data using the frame convolver_image
+    # Now reproduce this data using the convolve_image function
+
+    image = aa.Array2D.no_mask(values=np.arange(900).reshape(30, 30), pixel_scales=1.0)
+    kernel = aa.Kernel2D.no_mask(values=np.arange(49).reshape(7, 7), pixel_scales=1.0)
 
     masked_image = aa.Array2D(values=image.native, mask=mask)