Merge pull request #189 from Jammy2211/feature/convolve_fft

Feature/convolve fft

Author: Jammy2211

autoarray/config/general.yaml

@@ -2,6 +2,8 @@ jax:
   use_jax: true                       # If True, uses JAX internally, whereas False uses normal Numpy.
 fits:
   flip_for_ds9: false                 # If True, the image is flipped before output to a .fits file, which is useful for viewing in DS9.
+psf:
+  use_fft_default: true              # If True, PSFs are convolved using FFTs by default, which is faster and uses less memory in all cases except for very small PSFs, False uses direct convolution.
 inversion:
   check_reconstruction: true          # If True, the inversion's reconstruction is checked to ensure the solution of a meshs's mapper is not an invalid solution where the values are all the same.
   use_positive_only_solver: true      # If True, inversion's use a positive-only linear algebra solver by default, which is slower but prevents unphysical negative values in the reconstructed solutuion.

autoarray/dataset/imaging/dataset.py

@@ -4,6 +4,7 @@
 from typing import Optional, Union
 
 from autoconf import cached_property
+from autoconf import instance
 
 from autoarray.dataset.abstract.dataset import AbstractDataset
 from autoarray.dataset.grids import GridsDataset
@@ -29,7 +30,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_psf: bool = False,
+        disable_fft_pad: bool = True,
         use_normalized_psf: Optional[bool] = True,
         check_noise_map: bool = True,
     ):
@@ -76,63 +77,59 @@ 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_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.
+        disable_fft_pad
+            The FFT PSF convolution is optimal for a certain 2D FFT padding or trimming, which places the fewest zeros
+            around the image. If this is set to `True`, this optimal padding is not performed and the image is used
+            as-is.
         use_normalized_psf
             If `True`, the PSF kernel values are rescaled such that they sum to 1.0. This can be important for ensuring
             the PSF kernel does not change the overall normalization of the image when it is convolved with it.
         check_noise_map
             If True, the noise-map is checked to ensure all values are above zero.
         """
 
-        self.unmasked = None
+        self.disable_fft_pad = disable_fft_pad
 
-        self.pad_for_psf = pad_for_psf
+        if 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
-                )
-            except exc.MaskException:
-                over_sample_size_lp = (
-                    over_sample_util.over_sample_size_convert_to_array_2d_from(
-                        over_sample_size=over_sample_size_lp, mask=data.mask
-                    )
-                )
-                over_sample_size_lp = (
-                    over_sample_size_lp.padded_before_convolution_from(
-                        kernel_shape=psf.shape_native, mask_pad_value=1
-                    )
-                )
+            full_shape, fft_shape, mask_shape = psf.fft_shape_from(mask=data.mask)
 
-                over_sample_size_pixelization = (
-                    over_sample_util.over_sample_size_convert_to_array_2d_from(
-                        over_sample_size=over_sample_size_pixelization, mask=data.mask
-                    )
-                )
-                over_sample_size_pixelization = (
-                    over_sample_size_pixelization.padded_before_convolution_from(
-                        kernel_shape=psf.shape_native, mask_pad_value=1
-                    )
+        if psf is not None and not disable_fft_pad and data.mask.shape != fft_shape:
+
+            # If using real-space convolution instead of FFT, enforce odd-odd shapes
+            if not psf.use_fft:
+                fft_shape = tuple(s + 1 if s % 2 == 0 else s for s in fft_shape)
+
+            logger.info(
+                f"Imaging data has been trimmed or padded for FFT convolution.\n"
+                f"  - Original shape : {data.mask.shape}\n"
+                f"  - FFT shape    : {fft_shape}\n"
+                f"Padding ensures accurate PSF convolution in Fourier space. "
+                f"Set `disable_fft_pad=True` in Imaging object to turn off automatic padding."
+            )
+
+            over_sample_size_lp = (
+                over_sample_util.over_sample_size_convert_to_array_2d_from(
+                    over_sample_size=over_sample_size_lp, mask=data.mask
                 )
+            )
+            over_sample_size_lp = over_sample_size_lp.resized_from(
+                new_shape=fft_shape, mask_pad_value=1
+            )
 
-                data = data.padded_before_convolution_from(
-                    kernel_shape=psf.shape_native, mask_pad_value=1
+            over_sample_size_pixelization = (
+                over_sample_util.over_sample_size_convert_to_array_2d_from(
+                    over_sample_size=over_sample_size_pixelization, mask=data.mask
                 )
-                if noise_map is not None:
-                    noise_map = noise_map.padded_before_convolution_from(
-                        kernel_shape=psf.shape_native, mask_pad_value=1
-                    )
-                logger.info(
-                    f"The image and noise map of the `Imaging` objected have been padded to the dimensions"
-                    f"{data.shape}. This is because the blurring region around the mask (which defines where"
-                    f"PSF flux may be convolved into the masked region) extended beyond the edge of the image."
-                    f""
-                    f"This can be prevented by using a smaller mask, smaller PSF kernel size or manually padding"
-                    f"the image and noise-map yourself."
+            )
+            over_sample_size_pixelization = over_sample_size_pixelization.resized_from(
+                new_shape=fft_shape, mask_pad_value=1
+            )
+
+            data = data.resized_from(new_shape=fft_shape, mask_pad_value=1)
+            if noise_map is not None:
+                noise_map = noise_map.resized_from(
+                    new_shape=fft_shape, mask_pad_value=1
                 )
 
         super().__init__(
@@ -179,6 +176,9 @@ def __init__(
                 normalize=use_normalized_psf,
                 image_mask=image_mask,
                 blurring_mask=blurring_mask,
+                mask_shape=mask_shape,
+                full_shape=full_shape,
+                fft_shape=fft_shape,
             )
 
         self.psf = psf
@@ -337,31 +337,34 @@ def from_fits(
             over_sample_size_pixelization=over_sample_size_pixelization,
         )
 
-    def apply_mask(self, mask: Mask2D) -> "Imaging":
+    def apply_mask(self, mask: Mask2D, disable_fft_pad: bool = False) -> "Imaging":
         """
         Apply a mask to the imaging dataset, whereby the mask is applied to the image data, noise-map and other
         quantities one-by-one.
 
-        The original unmasked imaging data is stored as the `self.unmasked` attribute. This is used to ensure that if
-        the `apply_mask` function is called multiple times, every mask is always applied to the original unmasked
-        imaging dataset.
+        The `apply_mask` function cannot be called multiple times, if it is a mask may remove data, therefore
+        an exception is raised. If you wish to apply a new mask, reload the dataset from .fits files.
 
         Parameters
         ----------
         mask
             The 2D mask that is applied to the image.
         """
-        if self.data.mask.is_all_false:
-            unmasked_dataset = self
-        else:
-            unmasked_dataset = self.unmasked
+        invalid = np.logical_and(self.data.mask, np.logical_not(mask))
+
+        if np.any(invalid):
+            raise exc.DatasetException(
+                "The new mask overlaps with pixels that are already unmasked in the dataset. "
+                "You cannot apply a new mask on top of an existing one. "
+                "If you wish to apply a different mask, please reload the dataset from .fits files."
+            )
 
-        data = Array2D(values=unmasked_dataset.data.native, mask=mask)
+        data = Array2D(values=self.data.native, mask=mask)
 
-        noise_map = Array2D(values=unmasked_dataset.noise_map.native, mask=mask)
+        noise_map = Array2D(values=self.noise_map.native, mask=mask)
 
-        if unmasked_dataset.noise_covariance_matrix is not None:
-            noise_covariance_matrix = unmasked_dataset.noise_covariance_matrix
+        if self.noise_covariance_matrix is not None:
+            noise_covariance_matrix = self.noise_covariance_matrix
 
             noise_covariance_matrix = np.delete(
                 noise_covariance_matrix, mask.derive_indexes.masked_slim, 0
@@ -385,11 +388,9 @@ 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_psf=True,
+            disable_fft_pad=disable_fft_pad,
         )
 
-        dataset.unmasked = unmasked_dataset
-
         logger.info(
             f"IMAGING - Data masked, contains a total of {mask.pixels_in_mask} image-pixels"
         )
@@ -400,6 +401,7 @@ def apply_noise_scaling(
         self,
         mask: Mask2D,
         noise_value: float = 1e8,
+        disable_fft_pad: bool = False,
         signal_to_noise_value: Optional[float] = None,
         should_zero_data: bool = True,
     ) -> "Imaging":
@@ -455,18 +457,6 @@ def apply_noise_scaling(
         else:
             data = self.data.native.array
 
-        data_unmasked = Array2D.no_mask(
-            values=data,
-            shape_native=self.data.shape_native,
-            pixel_scales=self.data.pixel_scales,
-        )
-
-        noise_map_unmasked = Array2D.no_mask(
-            values=noise_map,
-            shape_native=self.noise_map.shape_native,
-            pixel_scales=self.noise_map.pixel_scales,
-        )
-
         data = Array2D(values=data, mask=self.data.mask)
 
         noise_map = Array2D(values=noise_map, mask=self.data.mask)
@@ -478,15 +468,10 @@ 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_psf=False,
+            disable_fft_pad=disable_fft_pad,
             check_noise_map=False,
         )
 
-        if self.unmasked is not None:
-            dataset.unmasked = self.unmasked
-            dataset.unmasked.data = data_unmasked
-            dataset.unmasked.noise_map = noise_map_unmasked
-
         logger.info(
             f"IMAGING - Data noise scaling applied, a total of {mask.pixels_in_mask} pixels were scaled to large noise values."
         )
@@ -497,6 +482,7 @@ def apply_over_sampling(
         self,
         over_sample_size_lp: Union[int, Array2D] = None,
         over_sample_size_pixelization: Union[int, Array2D] = None,
+        disable_fft_pad: bool = False,
     ) -> "AbstractDataset":
         """
         Apply new over sampling objects to the grid and grid pixelization of the dataset.
@@ -526,7 +512,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_psf=False,
+            disable_fft_pad=disable_fft_pad,
             check_noise_map=False,
         )
 

autoarray/dataset/imaging/simulator.py

@@ -126,7 +126,7 @@ def via_image_from(
             pixel_scales=image.pixel_scales,
         )
 
-        image = self.psf.convolved_array_from(array=image)
+        image = self.psf.convolved_image_from(image=image, blurring_image=None)
 
         image = image + background_sky_map
 
@@ -169,12 +169,16 @@ def via_image_from(
         image = Array2D(values=image, mask=mask)
 
         dataset = Imaging(
-            data=image, psf=self.psf, noise_map=noise_map, check_noise_map=False
+            data=image,
+            psf=self.psf,
+            noise_map=noise_map,
+            check_noise_map=False,
+            disable_fft_pad=True,
         )
 
         if over_sample_size is not None:
             dataset = dataset.apply_over_sampling(
-                over_sample_size_lp=over_sample_size.native
+                over_sample_size_lp=over_sample_size.native, disable_fft_pad=True
             )
 
         return dataset

autoarray/geometry/geometry_1d.py

@@ -1,11 +1,6 @@
 from __future__ import annotations
 import logging
-import numpy as np
-from typing import TYPE_CHECKING, List, Tuple, Union
-
-if TYPE_CHECKING:
-    from autoarray.structures.grids.uniform_1d import Grid1D
-    from autoarray.mask.mask_2d import Mask2D
+from typing import Tuple
 
 from autoarray import type as ty
 

autoarray/inversion/inversion/imaging/abstract.py

@@ -92,7 +92,7 @@ def operated_mapping_matrix_list(self) -> List[np.ndarray]:
 
         return [
             (
-                self.psf.convolve_mapping_matrix(
+                self.psf.convolved_mapping_matrix_from(
                     mapping_matrix=linear_obj.mapping_matrix, mask=self.mask
                 )
                 if linear_obj.operated_mapping_matrix_override is None
@@ -134,7 +134,7 @@ def linear_func_operated_mapping_matrix_dict(self) -> Dict:
             if linear_func.operated_mapping_matrix_override is not None:
                 operated_mapping_matrix = linear_func.operated_mapping_matrix_override
             else:
-                operated_mapping_matrix = self.psf.convolve_mapping_matrix(
+                operated_mapping_matrix = self.psf.convolved_mapping_matrix_from(
                     mapping_matrix=linear_func.mapping_matrix,
                     mask=self.mask,
                 )
@@ -215,7 +215,7 @@ def mapper_operated_mapping_matrix_dict(self) -> Dict:
         mapper_operated_mapping_matrix_dict = {}
 
         for mapper in self.cls_list_from(cls=AbstractMapper):
-            operated_mapping_matrix = self.psf.convolve_mapping_matrix(
+            operated_mapping_matrix = self.psf.convolved_mapping_matrix_from(
                 mapping_matrix=mapper.mapping_matrix,
                 mask=self.mask,
             )

autoarray/inversion/inversion/imaging/mapping.py

@@ -73,7 +73,7 @@ def _data_vector_mapper(self) -> np.ndarray:
             mapper = mapper_list[i]
             param_range = mapper_param_range_list[i]
 
-            operated_mapping_matrix = self.psf.convolve_mapping_matrix(
+            operated_mapping_matrix = self.psf.convolved_mapping_matrix_from(
                 mapping_matrix=mapper.mapping_matrix, mask=self.mask
             )
 
@@ -132,7 +132,7 @@ def _curvature_matrix_mapper_diag(self) -> Optional[np.ndarray]:
             mapper_i = mapper_list[i]
             mapper_param_range_i = mapper_param_range_list[i]
 
-            operated_mapping_matrix = self.psf.convolve_mapping_matrix(
+            operated_mapping_matrix = self.psf.convolved_mapping_matrix_from(
                 mapping_matrix=mapper_i.mapping_matrix, mask=self.mask
             )
 

autoarray/inversion/inversion/imaging/w_tilde.py

@@ -518,8 +518,13 @@ def mapped_reconstructed_data_dict(self) -> Dict[LinearObj, Array2D]:
                     reconstruction=np.array(reconstruction),
                 )
 
-                mapped_reconstructed_image = self.psf.convolve_image_no_blurring(
-                    image=mapped_reconstructed_image, mask=self.mask
+                mapped_reconstructed_image = Array2D(
+                    values=mapped_reconstructed_image, mask=self.mask
+                )
+
+                mapped_reconstructed_image = self.psf.convolved_image_from(
+                    image=mapped_reconstructed_image,
+                    blurring_image=None,
                 ).array
 
                 mapped_reconstructed_image = Array2D(

autoarray/mask/derive/mask_2d.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 import logging
-import copy
 import numpy as np
 from typing import TYPE_CHECKING, Tuple
 
@@ -10,7 +9,6 @@
 from autoarray import exc
 from autoarray.mask.derive.indexes_2d import DeriveIndexes2D
 
-from autoarray.structures.arrays import array_2d_util
 from autoarray.mask import mask_2d_util
 
 logging.basicConfig()

autoarray/mask/mask_2d.py

@@ -653,7 +653,9 @@ def unmasked_blurred_array_from(self, padded_array, psf, image_shape) -> Array2D
             The 1D unmasked image which is blurred.
         """
 
-        blurred_image = psf.convolved_array_from(array=padded_array)
+        blurred_image = psf.convolved_image_from(
+            image=padded_array, blurring_image=None
+        )
 
         return self.trimmed_array_from(
             padded_array=blurred_image, image_shape=image_shape

autoarray/operators/mock/mock_psf.py

@@ -2,5 +2,5 @@ class MockPSF:
     def __init__(self, operated_mapping_matrix=None):
         self.operated_mapping_matrix = operated_mapping_matrix
 
-    def convolve_mapping_matrix(self, mapping_matrix, mask):
+    def convolved_mapping_matrix_from(self, mapping_matrix, mask):
         return self.operated_mapping_matrix