Source code for aiapy.psf.deconvolve

"""
Deconvolve an AIA image with the channel point spread function.
"""
import copy
import warnings

import numpy as np
from sunpy import log

try:
    import cupy

    HAS_CUPY = True
except ImportError:
    HAS_CUPY = False

from aiapy.util import AiapyUserWarning

from .psf import psf as calculate_psf

__all__ = ["deconvolve"]


[docs] def deconvolve(smap, *, psf=None, iterations=25, clip_negative=True, use_gpu=True): """ Deconvolve an AIA image with the point spread function. Perform image deconvolution on an AIA image with the instrument point spread function using the Richardson-Lucy deconvolution algorithm [1]_. .. note:: If the `~cupy` package is installed and your machine has an NVIDIA GPU, the deconvolution will automatically be accelerated with CUDA. This can lead to more than an order of magnitude in performance increase compared to pure `numpy` on a CPU. For more information on PSF deconvolution on a GPU, see [2]_. Parameters ---------- smap : `~sunpy.map.Map` An AIA image psf : `~numpy.ndarray`, optional The point spread function. If None, it will be calculated iterations : `int`, optional Number of iterations in the Richardson-Lucy algorithm clip_negative : `bool`, optional If the image has negative intensity values, set them to zero. use_gpu : `bool`, optional If True and `~cupy` is installed, do PSF deconvolution on the GPU with `~cupy`. Returns ------- `~sunpy.map.Map` Deconvolved AIA image See Also -------- psf References ---------- .. [1] https://en.wikipedia.org/wiki/Richardson%E2%80%93Lucy_deconvolution .. [2] Cheung, M., 2015, *GPU Technology Conference Silicon Valley*, `GPU-Accelerated Image Processing for NASA's Solar Dynamics Observatory <https://on-demand.gputechconf.com/gtc/2015/presentation/S5209-Mark-Cheung.pdf>`__ """ # TODO: do we need a check to make sure this is a full-frame image? img = smap.data if clip_negative: img = np.where(img < 0, 0, img) if np.any(img < 0): warnings.warn( "Image contains negative intensity values. Consider setting clip_negative to True", AiapyUserWarning, stacklevel=3, ) if psf is None: psf = calculate_psf(smap.wavelength) if use_gpu and not HAS_CUPY: log.info("cupy not installed or working, falling back to CPU") if HAS_CUPY and use_gpu: img = cupy.array(img) psf = cupy.array(psf) # Center PSF at pixel (0,0) psf = np.roll(np.roll(psf, psf.shape[0] // 2, axis=0), psf.shape[1] // 2, axis=1) # Convolution requires FFT of the PSF psf = np.fft.rfft2(psf) psf_conj = psf.conj() img_decon = np.copy(img) for _ in range(iterations): ratio = img / np.fft.irfft2(np.fft.rfft2(img_decon) * psf) img_decon = img_decon * np.fft.irfft2(np.fft.rfft2(ratio) * psf_conj) return smap._new_instance( cupy.asnumpy(img_decon) if (HAS_CUPY and use_gpu) else img_decon, copy.deepcopy(smap.meta), plot_settings=copy.deepcopy(smap.plot_settings), mask=smap.mask, )