xP = yP = np.linspace(-b, +b, Npixels)
XP, YP = np.meshgrid(xP, yP)
kx = XP * k / f
ky = YP * k / f
# kx = XP / a * k_cut_off # other way to calculate kx and ky
# ky = YP / a * k_cut_off
# (kx,ky) in radial coordinates
k_rho = np.sqrt(kx**2 + ky**2)
k_theta = np.arctan2(ky, kx)
N = 3 # Zernike radial order
M = 1 # Zernike azimutal frequency
phase = np.pi * nm_polynomial(
N, M, k_rho / k_cut_off, k_theta, normalized=False)
weight = 1 # weight of the polynomials in units of lambda (weight 0.5 means wavefront abberated of lamba/2)
ATF = np.exp(1.j * weight * phase) # Amplitude Transfer Function
mask_idx = (k_rho > k_cut_off)
ATF[mask_idx] = 0 # Creates a circular mask
ASF = ifftshift(ifft2(fftshift(ATF))) # Amplitude Spread Function
PSF = np.abs(ASF)**2 # Point Spread Function
PSF = PSF / np.sum(PSF) # PSF normalized with its area
OTF = fftshift(fft2(ifftshift(PSF)))
def add_Zernike_aberration(self, N, M, weight):
self.phase += weight * nm_polynomial(
N, M, self.k_rho / self.k_cut_off, self.k_theta, normalized=True)