from .._ndarray_backend import xp
from .. import pad
from .base import Refocus
[docs]
class RefocusNumpy(Refocus):
"""Refocusing with numpy-based Fourier transform
.. versionadded:: 0.3.0
"""
backend_expected = "numpy"
# cupy doesn't work due to padding
backend_incompatible = "cupy"
def _init_fft(self, field, padding):
"""Perform initial Fourier transform of the input field
Parameters
----------
field: 2d complex-valued ndarray
Input field to be refocused
padding: bool
Whether to perform boundary-padding with linear ramp
Returns
-------
fft_field0: 2d complex-valued ndarray
Fourier transform the initial field
"""
if padding:
field = pad.pad_add(field)
return xp.fft.fft2(field)
[docs]
def propagate(self, distance):
fft_kernel = self.get_kernel(distance=distance)
refoc = xp.fft.ifft2(self.fft_origin * fft_kernel)
if self.padding:
refoc = pad.pad_rem(refoc)
return refoc