Feature/convolve fft

README.rst

@@ -75,7 +75,7 @@ API Overview
 ------------
 
 Galaxy morphology calculations are performed in **PyAutoGalaaxy** by building a ``Plane`` object from ``LightProfile``
-and ``Galaxy`` objects. Below, we create a simple galaxy system where a redshift 0.5
+and ``Galaxy`` objects. We create a simple galaxy system where a redshift 0.5
 ``Galaxy`` with an ``Sersic`` ``LightProfile`` representing a bulge and an ``Exponential`` ``LightProfile``
 representing a disk.
 

autogalaxy/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: true
+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.
 grid:
   max_evaluation_grid_size: 1000   # An evaluation grid whose shape is adaptive chosen is used to compute quantities like critical curves, this integer is the max size of the grid ensuring faster run times.
 adapt:

autogalaxy/cosmology/__init__.py

@@ -1,3 +1,3 @@
 from .lensing import LensingCosmology
 from .wrap import Planck15
-from .model import FlatwCDMWrap, FlatLambdaCDMWrap
+from .model import FlatwCDMWrap, FlatLambdaCDMWrap

autogalaxy/operate/image.py

@@ -37,7 +37,7 @@ def _blurred_image_2d_from(
         psf: aa.Kernel2D,
     ) -> aa.Array2D:
 
-        values = psf.convolve_image(
+        values = psf.convolved_image_from(
             image=image_2d,
             blurring_image=blurring_image_2d,
         )
@@ -141,6 +141,9 @@ def unmasked_blurred_image_2d_from(self, grid, psf):
             grid=padded_grid, operated_only=False
         )
 
+        # Required to make sure right 2D indexes are computed with update mask
+        psf.slim_to_native_tuple = None
+
         padded_image_2d = padded_grid.mask.unmasked_blurred_array_from(
             padded_array=padded_image_2d_not_operated,
             psf=psf,
@@ -187,9 +190,7 @@ def visibilities_from(
 
         if jnp.any(image_2d.array):
             return transformer.visibilities_from(image=image_2d)
-        return aa.Visibilities.zeros(
-            shape_slim=(transformer.uv_wavelengths.shape[0],)
-        )
+        return aa.Visibilities.zeros(shape_slim=(transformer.uv_wavelengths.shape[0],))
 
         # return jax.lax.cond(
         #     jnp.any(image_2d.array),
@@ -411,7 +412,7 @@ def galaxy_blurred_image_2d_dict_from(
                 galaxy_key
             ]
 
-            blurred_image_2d = psf.convolve_image(
+            blurred_image_2d = psf.convolved_image_from(
                 image=image_2d_not_operated,
                 blurring_image=blurring_image_2d_not_operated,
             )

autogalaxy/profiles/light/linear/abstract.py

@@ -273,6 +273,55 @@ def mapping_matrix(self) -> np.ndarray:
 
         return jnp.stack(image_2d_list, axis=1)
 
+    @cached_property
+    def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]:
+        """
+        The inversion object takes the `mapping_matrix` of each linear object and combines it with the PSF
+        operator to perform a 2D convolution and compute the `operated_mapping_matrix`.
+
+        If this property is overwritten this operation is not performed, with the `operated_mapping_matrix` output this
+        property automatically used instead.
+
+        This is used for a linear light profile because the in-built mapping matrix convolution does not account for
+        how light profile images have flux outside the masked region which is blurred into the masked region. This
+        flux is outside the region that defines the `mapping_matrix` and thus this override is required to properly
+        incorporate it.
+
+        Returns
+        -------
+        A blurred mapping matrix of dimensions (total_mask_pixels, 1) which overrides the mapping matrix calculations
+        performed in the linear equation solvers.
+        """
+
+        if isinstance(self.light_profile_list[0], LightProfileOperated):
+            return self.mapping_matrix
+
+        # number of source pixels = number of light profiles
+        n_src = len(self.light_profile_list)
+
+        # allocate slim-form arrays for mapping matrices
+        mapping_matrix = jnp.zeros((self.grid.shape_slim, n_src))
+        blurring_mapping_matrix = jnp.zeros((self.blurring_grid.shape_slim, n_src))
+
+        # build each column
+        for pixel, light_profile in enumerate(self.light_profile_list):
+            # main grid mapping for this light profile
+            mapping_matrix = mapping_matrix.at[:, pixel].set(
+                light_profile.image_2d_from(grid=self.grid).array
+            )
+
+            # blurring grid mapping for this light profile
+            blurring_mapping_matrix = blurring_mapping_matrix.at[:, pixel].set(
+                light_profile.image_2d_from(grid=self.blurring_grid).array
+            )
+
+        return self.psf.convolved_mapping_matrix_from(
+            mapping_matrix=mapping_matrix,
+            mask=self.grid.mask,
+            blurring_mapping_matrix=blurring_mapping_matrix,
+            blurring_mask=self.blurring_grid.mask,
+        )
+
     @cached_property
     def operated_mapping_matrix_override(self) -> Optional[np.ndarray]:
         """
@@ -303,7 +352,7 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]:
 
             blurring_image_2d = light_profile.image_2d_from(grid=self.blurring_grid)
 
-            blurred_image_2d = self.psf.convolve_image(
+            blurred_image_2d = self.psf.convolved_image_from(
                 image=image_2d, blurring_image=blurring_image_2d
             )
 

autogalaxy/profiles/light/snr/abstract.py

@@ -84,9 +84,7 @@ def set_intensity_from(
 
         image_2d = self.image_2d_from(grid=grid)
         if psf is not None:
-            image_2d = psf.convolve_image_no_blurring(
-                image=image_2d, mask=image_2d.mask
-            )
+            image_2d = psf.convolved_image_from(image=image_2d, blurring_image=None)
 
         brightest_value = np.max(image_2d)
 

autogalaxy/profiles/mass/dark/nfw.py

@@ -412,4 +412,6 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
     @staticmethod
     def potential_func_sph(eta):
-        return ((np.log(eta.array / 2.0)) ** 2) - (np.arctanh(np.sqrt(1 - eta.array**2))) ** 2
+        return ((np.log(eta.array / 2.0)) ** 2) - (
+            np.arctanh(np.sqrt(1 - eta.array**2))
+        ) ** 2

test_autogalaxy/config/general.yaml

@@ -2,6 +2,8 @@ analysis:
   n_cores: 1
 fits:
   flip_for_ds9: true
+psf:
+  use_fft_default: false              # 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. Real space used for unit tests.
 grid:
   remove_projected_centre: false
 inversion:

test_autogalaxy/conftest.py

@@ -7,7 +7,6 @@ def pytest_configure():
 
 import os
 from os import path
-
 import pytest
 from matplotlib import pyplot
 

test_autogalaxy/imaging/test_simulator.py

@@ -150,6 +150,6 @@ def test__simulator__simulate_imaging_from_galaxy__source_galaxy__compare_to_ima
     )
 
     assert dataset.shape_native == (11, 11)
-    assert dataset.data == pytest.approx(imaging_via_image.data, 1.0e-4)
+    assert dataset.data.array == pytest.approx(imaging_via_image.data.array, 1.0e-4)
     assert (dataset.psf == imaging_via_image.psf).all()
     assert dataset.noise_map == pytest.approx(imaging_via_image.noise_map, 1.0e-4)

test_autogalaxy/operate/test_image.py

@@ -18,7 +18,7 @@ def test__blurred_image_2d_from(
     image_2d = lp.image_2d_from(grid=grid_2d_7x7)
     blurring_image_2d = lp.image_2d_from(grid=blurring_grid_2d_7x7)
 
-    blurred_image_2d_manual = psf_3x3.convolve_image(
+    blurred_image_2d_manual = psf_3x3.convolved_image_from(
         image=image_2d, blurring_image=blurring_image_2d
     )
 
@@ -55,7 +55,7 @@ def test__blurred_image_2d_from(
         grid=grid_2d_7x7, psf=psf_3x3, blurring_grid=blurring_grid_2d_7x7
     )
 
-    blurred_image_2d_manual_not_operated = psf_3x3.convolve_image(
+    blurred_image_2d_manual_not_operated = psf_3x3.convolved_image_from(
         image=image_2d_not_operated,
         blurring_image=blurring_image_2d_not_operated,
     )
@@ -152,7 +152,9 @@ def test__unmasked_blurred_image_2d_from():
         grid=grid, psf=psf
     )
 
-    assert unmasked_blurred_image_2d == pytest.approx(image_2d_manual, 1.0e-4)
+    assert unmasked_blurred_image_2d.array == pytest.approx(
+        image_2d_manual.array, 1.0e-4
+    )
 
 
 def test__visibilities_from_grid_and_transformer(grid_2d_7x7, transformer_7x7_7):
@@ -255,23 +257,27 @@ def test__unmasked_blurred_image_2d_list_from():
     padded_grid = grid.padded_grid_from(kernel_shape_native=psf.shape_native)
 
     manual_blurred_image_0 = lp_0.image_2d_from(grid=padded_grid)
-    manual_blurred_image_0 = psf.convolved_array_from(array=manual_blurred_image_0)
+    manual_blurred_image_0 = psf.convolved_image_from(
+        image=manual_blurred_image_0, blurring_image=None
+    )
 
     manual_blurred_image_1 = lp_1.image_2d_from(grid=padded_grid)
-    manual_blurred_image_1 = psf.convolved_array_from(array=manual_blurred_image_1)
+    manual_blurred_image_1 = psf.convolved_image_from(
+        image=manual_blurred_image_1, blurring_image=None
+    )
 
     gal = ag.Galaxy(redshift=0.5, lp_0=lp_0, lp_1=lp_1)
 
     unmasked_blurred_image_2d_list = gal.unmasked_blurred_image_2d_list_from(
         grid=grid, psf=psf
     )
 
-    assert unmasked_blurred_image_2d_list[0].native == pytest.approx(
-        manual_blurred_image_0.native[1:4, 1:4], 1.0e-4
+    assert unmasked_blurred_image_2d_list[0].native.array == pytest.approx(
+        manual_blurred_image_0.native.array[1:4, 1:4], 1.0e-4
     )
 
-    assert unmasked_blurred_image_2d_list[1].native == pytest.approx(
-        manual_blurred_image_1.native[1:4, 1:4], 1.0e-4
+    assert unmasked_blurred_image_2d_list[1].native.array == pytest.approx(
+        manual_blurred_image_1.native.array[1:4, 1:4], 1.0e-4
     )
 
 
@@ -332,7 +338,7 @@ def test__galaxy_blurred_image_2d_dict_from(grid_2d_7x7, blurring_grid_2d_7x7, p
     image_2d = lp_0.image_2d_from(grid=grid_2d_7x7)
     blurring_image_2d = lp_0.image_2d_from(grid=blurring_grid_2d_7x7)
 
-    image_2d_convolved = psf_3x3.convolve_image(
+    image_2d_convolved = psf_3x3.convolved_image_from(
         image=image_2d, blurring_image=blurring_image_2d
     )