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 convpsf(psf1, psf2):
'''Convolves two PSFs.
Args:
psf1 (prysm.PSF): first PSF.
psf2 (prysm.PSF): second PSF.
Returns:
PSF: A new `PSF` that is the convolution of psf1 and psf2.
Notes:
The PSF with the lower nyquist frequency defines the sampling of the
output. The PSF with a higher nyquist will be truncated in the
frequency domain (without aliasing) and projected onto the
sampling grid of the PSF with a lower nyquist.
'''
if psf2.samples_x == psf1.samples_x and \
psf2.samples_y == psf1.samples_y and \
psf2.sample_spacing == psf1.sample_spacing:
# no need to interpolate, use FFTs to convolve
psf3 = PSF(data=abs(ifftshift(ifft2(fft2(psf1.data) * fft2(psf2.data)))),
sample_spacing=psf1.sample_spacing)
return psf3._renorm()
else:
# need to interpolate, suppress all frequency content above nyquist for the less sampled psf
if psf1.sample_spacing > psf2.sample_spacing:
# psf1 has the lower nyquist, resample psf2 in the fourier domain to match psf1
return _unequal_spacing_conv_core(psf1, psf2)
else:
# psf2 has lower nyquist, resample psf1 in the fourier domain to match psf2
return _unequal_spacing_conv_core(psf2, psf1)
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 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 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 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 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