def single_analytical_ft_convolution(without_analytic, with_analytic):
"""Convolves two convolvable objects utilizing their analytic fourier transforms.
Parameters
----------
without_analytic : `Convolvable`
A Convolvable object which lacks an analytic fourier transform
with_analytic : `Convolvable`
A Convolvable object which has an analytic fourier transform
Returns
-------
`ConvolutionResult`
A convolution result
"""
fourier_data = m.fftshift(m.fft2(m.fftshift(without_analytic.data)))
fourier_unit_x = forward_ft_unit(without_analytic.sample_spacing,
without_analytic.samples_x)
fourier_unit_y = forward_ft_unit(without_analytic.sample_spacing,
without_analytic.samples_y)
a_ft = with_analytic.analytic_ft(fourier_unit_x, fourier_unit_y)
result = _conv_result_core(fourier_data, a_ft)
return Convolvable(result, without_analytic.unit_x,
without_analytic.unit_y, False)
def _unequal_spacing_conv_core(psf1, psf2):
'''Interpolates psf2 before using fft-based convolution
Args:
psf1 (prysm.PSF): PSF. This one defines the sampling of the output.
psf2 (prysm.PSF): PSF. This one will have its frequency response
truncated.
Returns:
PSF: a new `PSF` that is the convolution of psf1 and psf2.
'''
# map psf1 into the fourier domain
ft1 = fft2(fftshift(psf1.data))
unit1x = forward_ft_unit(psf1.sample_spacing, psf1.samples_x)
unit1y = forward_ft_unit(psf1.sample_spacing, psf1.samples_y)
# map psf2 into the fourier domain
ft2 = fft2(fftshift(psf2.data))
unit2x = forward_ft_unit(psf2.sample_spacing, psf2.samples_x)
unit2y = forward_ft_unit(psf2.sample_spacing, psf2.samples_y)
ft3 = ifftshift(resample_2d_complex(fftshift(ft2), (unit2y, unit2x), (unit1y, unit1x)))
psf3 = PSF(data=abs(ifftshift(ifft2(ft1 * ft3))),
sample_spacing=psf1.sample_spacing)
return psf3._renorm()
def prop_pupil_plane_to_psf_plane(wavefunction, Q, incoherent=True, norm=None):
"""Propagate a pupil plane to a PSF plane and compute the grid along which the PSF exists.
Parameters
----------
wavefunction : `numpy.ndarray`
the pupil wavefunction
Q : `float`
oversampling / padding factor
incoherent : `bool`, optional
whether to return the incoherent (real valued) PSF, or the
coherent (complex-valued) PSF. Incoherent = |coherent|^2
norm : `str`, {None, 'ortho'}
normalization parameter passed directly to numpy/cupy fft
Returns
-------
psf : `numpy.ndarray`
incoherent point spread function
"""
padded_wavefront = pad2d(wavefunction, Q)
impulse_response = m.ifftshift(
m.fft2(m.fftshift(padded_wavefront), norm=norm))
if incoherent:
return abs(impulse_response)**2
else:
return impulse_response
def conv(self, psf2):
'''Convolves this PSF with another
Args:
psf2 (`PSF`): PSf to convolve with this one.
Returns:
PSF: A new `PSF` that is the convolution of these two PSFs.
Notes:
output PSF has equal sampling to whichever PSF has a lower nyquist
frequency.
'''
try:
psf_ft = fftshift(fft2(self.data))
psf_unit_x = forward_ft_unit(self.sample_spacing, self.samples_x)
psf_unit_y = forward_ft_unit(self.sample_spacing, self.samples_y)
psf2_ft = psf2.analytic_ft(psf_unit_x, psf_unit_y)
psf3 = PSF(data=abs(ifft2(psf_ft * psf2_ft)),
sample_spacing=self.sample_spacing)
return psf3._renorm()
except AttributeError: # no analytic FT on the PSF/subclass
print('No analytic FT, falling back to numerical approach.')
return convpsf(self, psf2)
def forward_ft_unit(sample_spacing, samples, shift=True):
"""Compute the units resulting from a fourier transform.
Parameters
----------
sample_spacing : `float`
center-to-center spacing of samples in an array
samples : `int`
number of samples in the data
shift : `bool`, optional
whether to shift the output. If True, first element is a negative freq
if False, first element is 0 freq.
Returns
-------
`numpy.ndarray`
array of sample frequencies in the output of an fft
"""
unit = m.fftfreq(samples, sample_spacing)
if shift:
return m.fftshift(unit)
else:
return unit
def from_pupil(pupil, efl, padding=1):
''' Uses scalar diffraction propogation to generate a PSF from a pupil.
Args:
pupil (prysm.Pupil): Pupil, with OPD data and wavefunction.
efl (float): effective focal length of the optical system.
padding (number): number of pupil widths to pad each side of the
pupil with during computation.
Returns:
PSF. A new PSF instance.
'''
# padded pupil contains 1 pupil width on each side for a width of 3
psf_samples = (pupil.samples * padding) * 2 + pupil.samples
sample_spacing = pupil_sample_to_psf_sample(pupil_sample=pupil.sample_spacing * 1000,
num_samples=psf_samples,
wavelength=pupil.wavelength,
efl=efl)
padded_wavefront = pad2d(pupil.fcn, padding)
impulse_response = ifftshift(fft2(fftshift(padded_wavefront)))
psf = abs(impulse_response)**2
return PSF(psf / np.max(psf), sample_spacing)
def _numerical_ft_convolution_core_equalspacing_unequalsamplecount(
more_samples, less_samples):
# compute the ordinate axes of the input and output
in_x = forward_ft_unit(less_samples.sample_spacing,
less_samples.data.shape[0])
in_y = forward_ft_unit(less_samples.sample_spacing,
less_samples.data.shape[1])
output_x = forward_ft_unit(more_samples.sample_spacing,
more_samples.data.shape[0])
output_y = forward_ft_unit(more_samples.sample_spacing,
more_samples.data.shape[1])
# FFT the less sampled one and map it onto the denser grid
less_fourier = m.fftshift(m.fft2(m.fftshift(less_samples.data)))
interpf = interp2d(in_x, in_y, less_fourier, kind='linear')
resampled_less = interpf(output_x, output_y)
# FFT convolve the two convolvables
more_fourier = m.fftshift(m.fft2(m.fftshift(more_samples.data)))
data = _conv_result_core(resampled_less, more_fourier)
return Convolvable(data, more_samples.unit_x, more_samples.unit_y, False)
def bandreject_filter(array, sample_spacing, wllow, wlhigh):
sy, sx = array.shape
# compute the bandpass in sample coordinates
ux, uy = forward_ft_unit(sample_spacing,
sx), forward_ft_unit(sample_spacing, sy)
fhigh, flow = 1 / wllow, 1 / wlhigh
# make an ordinate array in frequency space and use it to make a mask
uxx, uyy = m.meshgrid(ux, uy)
highpass = ((uxx < -fhigh) | (uxx > fhigh)) | ((uyy < -fhigh) |
(uyy > fhigh))
lowpass = ((uxx > -flow) & (uxx < flow)) & ((uyy > -flow) & (uyy < flow))
mask = highpass | lowpass
# adjust NaNs and FFT
work = array.copy()
work[~m.isfinite(work)] = 0
fourier = m.fftshift(m.fft2(m.ifftshift(work)))
fourier[mask] = 0
out = m.fftshift(m.ifft2(m.ifftshift(fourier)))
return out.real
def from_psf(psf):
''' Generates an MTF from a PSF.
Args:
psf (:class:`PSF`): PSF to compute an MTF from.
Returns:
:class:`MTF`: A new MTF instance.
'''
dat = abs(fftshift(fft2(psf.data)))
unit_x = forward_ft_unit(psf.sample_spacing, psf.samples_x)
unit_y = forward_ft_unit(psf.sample_spacing, psf.samples_y)
return MTF(dat / dat[psf.center_x, psf.center_y], unit_x, unit_y)
def _numerical_ft_convolution_core_unequalspacing(finer_sampled,
coarser_sampled):
# compute the ordinate axes of the input of each
in_x_more = forward_ft_unit(finer_sampled.sample_spacing,
finer_sampled.data.shape[0])
in_y_more = forward_ft_unit(finer_sampled.sample_spacing,
finer_sampled.data.shape[1])
in_x_less = forward_ft_unit(coarser_sampled.sample_spacing,
coarser_sampled.data.shape[0])
in_y_less = forward_ft_unit(coarser_sampled.sample_spacing,
coarser_sampled.data.shape[1])
# fourier-space interpolate the larger bandwidth signal onto the grid defined by the lower
# bandwidth signal. This assumes the lower bandwidth signal is Nyquist sampled, which is
# not necessarily the case. The accuracy of this method depends on the quality of the input.
more_fourier = m.fftshift(m.fft2(m.fftshift(finer_sampled.data)))
interpf = interp2d(in_x_more, in_y_more, more_fourier, kind='linear')
resampled_more = interpf(in_x_less, in_y_less)
# FFT the less well sampled input and perform the Fourier based convolution.
less_fourier = m.fftshift(m.fft2(m.fftshift(coarser_sampled.data)))
data = _conv_result_core(resampled_more, less_fourier)
return Convolvable(data, in_x_less, in_y_less, False)
def forward_ft_unit(sample_spacing, samples):
"""Compute the units resulting from a fourier transform.
Parameters
----------
sample_spacing : `float`
center-to-center spacing of samples in an array
samples : `int`
number of samples in the data
Returns
-------
`numpy.ndarray`
array of sample frequencies in the output of an fft
"""
return m.fftshift(m.fftfreq(samples, sample_spacing))
def prop_pupil_plane_to_psf_plane(wavefunction, Q, norm=None):
"""Propagate a pupil plane to a PSF plane and compute the grid along which the PSF exists.
Parameters
----------
wavefunction : `numpy.ndarray`
the pupil wavefunction
Q : `float`
oversampling / padding factor
norm : `str`, {None, 'ortho'}
normalization parameter passed directly to numpy/cupy fft
Returns
-------
psf : `numpy.ndarray`
incoherent point spread function
"""
padded_wavefront = pad2d(wavefunction, Q)
impulse_response = m.ifftshift(m.fft2(m.fftshift(padded_wavefront), norm=norm))
return abs(impulse_response) ** 2
def from_psf(psf):
"""Generate an MTF from a PSF.
Parameters
----------
psf : `PSF`
PSF to compute an MTF from
Returns
-------
`MTF`
A new MTF instance
"""
if getattr(psf, '_mtf', None) is not None:
return psf._mtf
else:
dat = abs(m.fftshift(m.fft2(psf.data)))
unit_x = forward_ft_unit(psf.sample_spacing / 1e3, psf.samples_x) # 1e3 for microns => mm
unit_y = forward_ft_unit(psf.sample_spacing / 1e3, psf.samples_y)
return MTF(dat / dat[psf.center_y, psf.center_x], unit_x, unit_y)
def psd(height, sample_spacing, window=None):
"""Compute the power spectral density of a signal.
Parameters
----------
height : `numpy.ndarray`
height or phase data
sample_spacing : `float`
spacing of samples in the input data
window : {'welch', 'hann'} or ndarray, optional
Returns
-------
unit_x : `numpy.ndarray`
ordinate x frequency axis
unit_y : `numpy.ndarray`
ordinate y frequency axis
psd : `numpy.ndarray`
power spectral density
Notes
-----
See GH_FFT for a rigorous treatment of FFT scalings
https://holometer.fnal.gov/GH_FFT.pdf
"""
window = make_window(height, sample_spacing, window)
fft = m.ifftshift(m.fft2(m.fftshift(height * window)))
psd = abs(fft)**2 # mag squared first as per GH_FFT
fs = 1 / sample_spacing
S2 = (window**2).sum()
coef = S2 * fs * fs
psd /= coef
ux = forward_ft_unit(sample_spacing, height.shape[1])
uy = forward_ft_unit(sample_spacing, height.shape[0])
return ux, uy, psd
def show_fourier(self, interp_method='lanczos', fig=None, ax=None):
''' Displays the fourier transform of the image.
Args:
interp_method (`string`): method used to interpolate the data
for display.
fig (`matplotlib.figure`): figure to plot in.
ax (`matplotlib.axis`): axis to plot in.
Returns:
`tuple` containing:
`matplotlib.figure`: figure containing the plot.
`matplotlib.axis`: axis containing the plot.
'''
dat = abs(fftshift(fft2(pad2d(self.data))))
dat /= dat.max()
unit_x = forward_ft_unit(self.sample_spacing, self.samples_x)
unit_y = forward_ft_unit(self.sample_spacing, self.samples_y)
xmin, xmax = unit_x[0], unit_x[-1]
ymin, ymax = unit_y[0], unit_y[-1]
fig, ax = share_fig_ax(fig, ax)
im = ax.imshow(dat**0.1,
extent=[xmin, xmax, ymin, ymax],
cmap='Greys_r',
interpolation=interp_method,
origin='lower')
fig.colorbar(im)
ax.set(xlabel='Spatial Frequency X [cy/mm]',
ylabel='Spatial Frequency Y [cy/mm]',
title='Normalized FT of image, to 0.1 power')
return fig, ax
def synthesize_surface_from_psd(psd, nu_x, nu_y):
"""Synthesize a surface height map from PSD data.
Parameters
----------
psd : `numpy.ndarray`
PSD data, units nm²/(cy/mm)²
nu_x : `numpy.ndarray`
x spatial frequency, cy/mm
nu_y : `numpy.ndarray`
y spatial frequency, cy_mm
"""
# generate a random phase to be matched to the PSD
randnums = m.rand(*psd.shape)
randfft = m.fft2(randnums)
phase = m.angle(randfft)
# calculate the output window
# the 0th element of nu_y has the greatest frequency in magnitude because of
# the convention to put the nyquist sample at -fs instead of +fs for even-size arrays
fs = -2 * nu_y[0]
dx = dy = 1 / fs
ny, nx = psd.shape
x, y = m.arange(nx) * dx, m.arange(ny) * dy
# calculate the area of the output window, "S2" in GH_FFT notation
A = x[-1] * y[-1]
# use ifft to compute the PSD
signal = m.exp(1j * phase) * m.sqrt(A * psd)
coef = 1 / dx / dy
out = m.ifftshift(m.ifft2(m.fftshift(signal))) * coef
out = out.real
return x, y, out
def show_fourier(self,
freq_x=None,
freq_y=None,
interp_method='lanczos',
fig=None,
ax=None):
'''Display the fourier transform of the image.
Parameters
----------
interp_method : `string`
method used to interpolate the data for display.
freq_x : iterable
x frequencies to use for convolvable with analytical FT and no data
freq_y : iterable
y frequencies to use for convolvable with analytic FT and no data
fig : `matplotlib.figure.Figure`
Figure containing the plot
ax : `matplotlib.axes.Axis`
Axis containing the plot
Returns
-------
fig : `matplotlib.figure.Figure`
Figure containing the plot
ax : `matplotlib.axes.Axis`
Axis containing the plot
Notes
-----
freq_x and freq_y are unused when the convolvable has a .data field.
'''
if self.has_analytic_ft:
if self.data is None:
if freq_x is None or freq_y is None:
raise ValueError(
'Convolvable has analytic FT and no data, must provide x and y coordinates'
)
else:
lx, ly, ss = len(self.unit_x), len(
self.unit_y), self.sample_spacing
freq_x, freq_y = forward_ft_unit(ss,
lx), forward_ft_unit(ss, ly)
data = self.analytic_ft(freq_x, freq_y)
else:
data = abs(m.fftshift(m.fft2(pad2d(self.data))))
data /= data.max()
freq_x = forward_ft_unit(self.sample_spacing, self.samples_x)
freq_y = forward_ft_unit(self.sample_spacing, self.samples_y)
xmin, xmax = freq_x[0], freq_x[-1]
ymin, ymax = freq_y[0], freq_y[-1]
fig, ax = share_fig_ax(fig, ax)
im = ax.imshow(data,
extent=[xmin, xmax, ymin, ymax],
cmap='Greys_r',
interpolation=interp_method,
origin='lower')
fig.colorbar(im, ax=ax, label='Normalized Spectral Intensity [a.u.]')
ax.set(xlim=(xmin, xmax),
xlabel=r' $\nu_x$ [cy/mm]',
ylim=(ymin, ymax),
ylabel=r' $\nu_y$ [cy/mm]')
return fig, ax
def _conv_result_core(data1, data2):
dat = abs(m.fftshift(m.ifft2(data1 * data2)))
return dat / dat.max()
def _numerical_ft_convolution_core_equalspacing(convolvable1, convolvable2):
# two are identically sampled; just FFT convolve them without modification
ft1 = m.fftshift(m.fft2(m.fftshift(convolvable1.data)))
ft2 = m.fftshift(m.fft2(m.fftshift(convolvable2.data)))
data = _conv_result_core(ft1, ft2)
return Convolvable(data, convolvable1.unit_x, convolvable1.unit_y, False)
