Merge pull request #225 from punch-mission/viz-kernels

add visualization for kernels
punch-mission · Nov 2, 2024 · 543a5be · 543a5be
2 parents ca4ffb4 + 2d63339
commit 543a5be
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 57 deletions.
diff --git a/docs/source/example.ipynb b/docs/source/example.ipynb
diff --git a/regularizepsf/psf.py b/regularizepsf/psf.py
@@ -337,21 +337,23 @@ def visualize_psfs(self,
                        colorbar_label="Normalized brightness",
                        imshow_args=imshow_args)
 
-    def visualize_kernels(self,
+    def visualize_ffts(self,
                   fig: mpl.figure.Figure | None = None,
                   fig_scale: int = 1,
                   all_patches: bool = False, imshow_args: dict | None = None) -> None:  # noqa: ANN002, ANN003
-        """Visualize the transfer kernels."""
+        """Visualize the fft of the PSF."""
         imshow_args = KERNEL_IMSHOW_ARGS_DEFAULT if imshow_args is None else imshow_args
 
-        extent = np.max(np.abs(self._fft_cube.values))
+        arr = np.abs(np.fft.fftshift(np.fft.ifft2(self._fft_cube.values)))
+        extent = np.max(np.abs(arr))
         if 'vmin' not in imshow_args:
-            imshow_args['vmin'] = extent
+            imshow_args['vmin'] = -extent
         if 'vmax' not in imshow_args:
             imshow_args['vmax'] = extent
 
         return visualize_grid(
-            self._fft_cube, all_patches=all_patches, fig=fig,
+            IndexedCube(self._fft_cube.coordinates, arr),
+            all_patches=all_patches, fig=fig,
             fig_scale=fig_scale, colorbar_label="Transfer kernel amplitude",
             imshow_args=imshow_args)
 

diff --git a/regularizepsf/transform.py b/regularizepsf/transform.py
@@ -5,12 +5,14 @@
 from typing import TYPE_CHECKING
 
 import h5py
+import matplotlib as mpl
 import numpy as np
 import scipy
 from astropy.io import fits
 
 from regularizepsf.exceptions import InvalidCoordinateError
 from regularizepsf.util import IndexedCube
+from regularizepsf.visualize import KERNEL_IMSHOW_ARGS_DEFAULT, visualize_grid
 
 if TYPE_CHECKING:
     import pathlib
@@ -138,14 +140,25 @@ def slice_padded_image(coordinate: tuple[int, int]) -> tuple[slice, slice]:
             2 * self.psf_shape[1] : image.shape[1] + 2 * self.psf_shape[1],
         ]
 
-    def visualize(self) -> None:
-        """Visualize the PSFTransform.
-
-        Returns
-        -------
-        None
-
-        """
+    def visualize(self,
+                          fig: mpl.figure.Figure | None = None,
+                          fig_scale: int = 1,
+                          all_patches: bool = False, imshow_args: dict | None = None) -> None:  # noqa: ANN002, ANN003
+        """Visualize the transfer kernels."""
+        imshow_args = KERNEL_IMSHOW_ARGS_DEFAULT if imshow_args is None else imshow_args
+
+        arr = np.abs(np.fft.fftshift(np.fft.ifft2(self._transfer_kernel.values)))
+        extent = np.max(np.abs(arr))
+        if 'vmin' not in imshow_args:
+            imshow_args['vmin'] = -extent
+        if 'vmax' not in imshow_args:
+            imshow_args['vmax'] = extent
+
+        return visualize_grid(
+            IndexedCube(self._transfer_kernel.coordinates, arr),
+            all_patches=all_patches, fig=fig,
+            fig_scale=fig_scale, colorbar_label="Transfer kernel amplitude",
+            imshow_args=imshow_args)
 
     def save(self, path: pathlib.Path) -> None:
         """Save a PSFTransform to a file. Supports h5 and FITS.
@@ -173,6 +186,7 @@ def save(self, path: pathlib.Path) -> None:
                                             name="transfer_imag", quantize_level=32)]).writeto(path)
         else:
             raise NotImplementedError(f"Unsupported file type {path.suffix}. Change to .h5 or .fits.")
+
     @classmethod
     def load(cls, path: pathlib.Path) -> ArrayPSFTransform:
         """Load a PSFTransform object. Supports h5 and FITS.