From 4f2a320646ec0d5ac0e6e6a21a1f7871973e77e6 Mon Sep 17 00:00:00 2001 From: Sergei Date: Tue, 16 Jan 2024 13:45:45 +0100 Subject: [PATCH] Provided documentation, made tests of several PSFs --- src/zernpy/calculations/calc_psfs.py | 385 ++++++++++++++++++++++++--- 1 file changed, 342 insertions(+), 43 deletions(-) diff --git a/src/zernpy/calculations/calc_psfs.py b/src/zernpy/calculations/calc_psfs.py index 84a9842..34714ba 100644 --- a/src/zernpy/calculations/calc_psfs.py +++ b/src/zernpy/calculations/calc_psfs.py @@ -40,67 +40,209 @@ pixel2um_coeff_plot = k*NA*(pixel_size/10.0) # coefficient used for better plotting with the reduced pixel size for preventing pixelated -# %% Integral Functions - integration on the radial external, angular - internal (np.sum) -def phase_integral_part_ang(zernike_pol: ZernPol, alpha: float, phi: np.array, p: float, theta: float, r: float) -> np.array: +# %% Integral Functions - integration by trapezoidal rule first on radius (normalized on the pupil), after - on angle +def diffraction_integral_r(zernike_pol: ZernPol, alpha: float, phi: float, p: np.array, theta: float, r: float) -> np.array: + """ + Diffraction integral function (see the references for the source of the equation). + + Parameters + ---------- + zernike_pol : ZernPol + Zernike polynomial definition. + alpha : float + Amplitude of the polynomial. + phi : float + Angle on the pupil coordinates. + p : np.array + Integration interval on the pupil coordinates. + theta : float + Angle on the image coordinates. + r : float + Radius on the image coordinates. + + References + ---------- + [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf + [2] https://opg.optica.org/ao/abstract.cfm?uri=ao-52-10-2062 (DOI: 10.1364/AO.52.002062) + + Returns + ------- + numpy.ndarray + Values of the diffraction integral. + + """ + # Multiplication by phi (in r*p*np.cos(...) only guides to scaling of the resulting PSF phase_arg = (alpha*zernike_pol.polynomial_value(p, phi) - r*p*np.cos(phi - theta))*1j return np.exp(phase_arg)*p -def angular_integral(zernike_pol: ZernPol, r: float, theta: float, p: float, alpha: float) -> complex: +def radial_integral(zernike_pol: ZernPol, r: float, theta: float, phi: float, alpha: float) -> complex: + """ + Make integration of the diffraction integral on the radius. + + Parameters + ---------- + zernike_pol : ZernPol + Zernike polynomial definition. + r : float + Radius on the image coordinates. + theta : float + Angle on the image coordinates. + phi : float + Angle on the pupil coordinates. + alpha : float + Amplitude of the polynomial. + + Returns + ------- + complex + Complex amplitude of the field as result of integration on the pupil radius coordinate. + + """ # Integration on the pupil angle. Vectorized form of the trapezoidal rule - h_phi = pi/n_phi_points; phi = np.arange(start=h_phi, stop=2.0*pi - h_phi, step=h_phi) - fa = phase_integral_part_ang(zernike_pol, alpha, 0.0, p, theta, r); fb = phase_integral_part_ang(zernike_pol, alpha, 2.0*pi, p, theta, r) - ang_int = np.sum(phase_integral_part_ang(zernike_pol, alpha, phi, p, theta, r)) + 0.5*(fa + fb) - return h_phi*ang_int + h_p = 1.0/n_p_points; p = np.arange(start=h_p, stop=1.0 - h_p, step=h_p) + fa = diffraction_integral_r(zernike_pol, alpha, phi, 0.0, theta, r); fb = diffraction_integral_r(zernike_pol, alpha, phi, 1.0, theta, r) + ang_int = np.sum(diffraction_integral_r(zernike_pol, alpha, phi, p, theta, r)) + 0.5*(fa + fb) + return h_p*ang_int + + +def get_psf_point_r(zernike_pol, r: float, theta: float, alpha: float = 1.0) -> float: + """ + Get the point for calculation of PSF. + + Parameters + ---------- + zernike_pol : tuple (m, n) or ZernPol + Zernike polynomial definition. + r : float + Radius on the image coordinates. + theta : float + Angle on the image coordinates. + alpha : float, optional + Amplitude of the polynomial. The default is 1.0. + Returns + ------- + float + |U|**2 - the module and square of the amplitude, intensity as the PSF value. -def get_psf_point(zernike_pol, r: float, theta: float, alpha: float = 1.0) -> float: + """ # Check and initialize Zernike polynomial if provided only orders (m, n) = define_orders(zernike_pol) # get polynomial orders if not isinstance(zernike_pol, ZernPol): zernike_pol = ZernPol(m=m, n=n) # Integration on the pupil radius using Simpson equation - n_integral_points = n_p_points; h_p = 1.0/n_integral_points + n_integral_points = n_phi_points; h_phi = 2.0*pi/n_integral_points even_sum = 0.0j; odd_sum = 0.0j for i in range(2, n_integral_points-2, 2): - p = i*h_p; even_sum += angular_integral(zernike_pol, r, theta, p, alpha) + phi = i*h_phi; even_sum += radial_integral(zernike_pol, r, theta, phi, alpha) for i in range(1, n_integral_points-1, 2): - p = i*h_p; odd_sum += angular_integral(zernike_pol, r, theta, p, alpha) - yA = angular_integral(zernike_pol, r, theta, 0.0, alpha); yB = angular_integral(zernike_pol, r, theta, 1.0, alpha) - integral_sum = (h_p/3.0)*(yA + yB + 2.0*even_sum + 4.0*odd_sum) - return np.power(np.abs(integral_sum), 2)/(4.0*pi*pi) + phi = i*h_phi; odd_sum += radial_integral(zernike_pol, r, theta, phi, alpha) + yA = radial_integral(zernike_pol, r, theta, 0.0, alpha); yB = radial_integral(zernike_pol, r, theta, 2.0*pi, alpha) + integral_sum = (h_phi/3.0)*(yA + yB + 2.0*even_sum + 4.0*odd_sum); integral_normalization = 1.0/(pi*pi) + return np.power(np.abs(integral_sum), 2)*integral_normalization + +# %% Integral Functions - integration on the radial external, angular - internal by the trapezoidal rule +def diffraction_integral_ang(zernike_pol: ZernPol, alpha: float, phi: np.array, p: float, theta: float, r: float) -> np.array: + """ + Diffraction integral function (see the references for the source of the equation). -# %% Integral Functions - vice versa (order of integration on the radius r and angle theta) -def phase_integral_part_rad(zernike_pol: ZernPol, alpha: float, phi: float, p: np.array, theta: float, r: float) -> np.array: - # !!! multiplication by phi (in r*p*np.cos(...) only guides to scaling of the resulting PSF + Parameters + ---------- + zernike_pol : ZernPol + Zernike polynomial definition. + alpha : float + Amplitude of the polynomial. + phi : float + Integration interval of angles on the pupil coordinates. + p : float + Radius on the pupil coordinates. + theta : float + Angle on the image coordinates. + r : float + Radius on the image coordinates. + + References + ---------- + [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf + [2] https://opg.optica.org/ao/abstract.cfm?uri=ao-52-10-2062 (DOI: 10.1364/AO.52.002062) + + Returns + ------- + numpy.ndarray + Values of the diffraction integral. + + """ phase_arg = (alpha*zernike_pol.polynomial_value(p, phi) - r*p*np.cos(phi - theta))*1j return np.exp(phase_arg)*p -def radial_integral(zernike_pol: ZernPol, r: float, theta: float, phi: float, alpha: float) -> complex: +def angular_integral(zernike_pol: ZernPol, r: float, theta: float, p: float, alpha: float) -> complex: + """ + Make integration of the diffraction integral on the angle. + + Parameters + ---------- + zernike_pol : ZernPol + Zernike polynomial definition. + r : float + Radius on the image coordinates. + theta : float + Angle on the image coordinates. + p : float + Radius on the pupil coordinates. + alpha : float + Amplitude of the polynomial. + + Returns + ------- + complex + Complex amplitude of the field as result of integration on the pupil radius coordinate. + + """ # Integration on the pupil angle. Vectorized form of the trapezoidal rule - h_p = 1.0/n_p_points; p = np.arange(start=h_p, stop=1.0 - h_p, step=h_p) - fa = phase_integral_part_rad(zernike_pol, alpha, phi, 0.0, theta, r); fb = phase_integral_part_rad(zernike_pol, alpha, phi, 1.0, theta, r) - ang_int = np.sum(phase_integral_part_rad(zernike_pol, alpha, phi, p, theta, r)) + 0.5*(fa + fb) - return h_p*ang_int + h_phi = pi/n_phi_points; phi = np.arange(start=h_phi, stop=2.0*pi - h_phi, step=h_phi) + fa = diffraction_integral_ang(zernike_pol, alpha, 0.0, p, theta, r); fb = diffraction_integral_ang(zernike_pol, alpha, 2.0*pi, p, theta, r) + ang_int = np.sum(diffraction_integral_ang(zernike_pol, alpha, phi, p, theta, r)) + 0.5*(fa + fb) + return h_phi*ang_int -def get_psf_point_r(zernike_pol, r: float, theta: float, alpha: float = 1.0) -> float: +def get_psf_point_ang(zernike_pol, r: float, theta: float, alpha: float = 1.0) -> float: + """ + Get the point for calculation of PSF. + + Parameters + ---------- + zernike_pol : tuple (m, n) or ZernPol + Zernike polynomial definition. + r : float + Radius on the image coordinates. + theta : float + Angle on the image coordinates. + alpha : float, optional + Amplitude of the polynomial. The default is 1.0. + + Returns + ------- + float + |U|**2 - the module and square of the amplitude, intensity as the PSF value. + + """ # Check and initialize Zernike polynomial if provided only orders (m, n) = define_orders(zernike_pol) # get polynomial orders if not isinstance(zernike_pol, ZernPol): zernike_pol = ZernPol(m=m, n=n) # Integration on the pupil radius using Simpson equation - n_integral_points = n_phi_points; h_phi = 2.0*pi/n_integral_points + n_integral_points = n_p_points; h_p = 1.0/n_integral_points even_sum = 0.0j; odd_sum = 0.0j for i in range(2, n_integral_points-2, 2): - phi = i*h_phi; even_sum += radial_integral(zernike_pol, r, theta, phi, alpha) + p = i*h_p; even_sum += angular_integral(zernike_pol, r, theta, p, alpha) for i in range(1, n_integral_points-1, 2): - phi = i*h_phi; odd_sum += radial_integral(zernike_pol, r, theta, phi, alpha) - yA = radial_integral(zernike_pol, r, theta, 0.0, alpha); yB = radial_integral(zernike_pol, r, theta, 2.0*pi, alpha) - integral_sum = (h_phi/3.0)*(yA + yB + 2.0*even_sum + 4.0*odd_sum); integral_normalization = 1.0/(pi*pi) - return np.power(np.abs(integral_sum), 2)*integral_normalization + p = i*h_p; odd_sum += angular_integral(zernike_pol, r, theta, p, alpha) + yA = angular_integral(zernike_pol, r, theta, 0.0, alpha); yB = angular_integral(zernike_pol, r, theta, 1.0, alpha) + integral_sum = (h_p/3.0)*(yA + yB + 2.0*even_sum + 4.0*odd_sum) + return np.power(np.abs(integral_sum), 2)/(4.0*pi*pi) # %% Nijboer's Thesis Functions @@ -135,6 +277,27 @@ def radial_func(n: int, r: float) -> float: def radial_ampl_func(n: int, r: np.ndarray) -> np.ndarray: + """ + Calculate the radial amplitude functions used in the Nijboer's thesis as Jv(r)/r. + + Parameters + ---------- + n : int + Radial order of the Zernike polynomial. + r : np.ndarray + Radii on the pupil coordinates. + Note that the array can only start with the special zero point (r[0] = 0.0). + + Reference + --------- + [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf + + Returns + ------- + radial_ampl_f : np.ndarray + Jv(r)/r function values. + + """ # Expanse the calculation for using the numpy array, assuming it starting with 0.0 radial_ampl_f = None print(type(r)) @@ -148,6 +311,30 @@ def radial_ampl_func(n: int, r: np.ndarray) -> np.ndarray: def get_aberrated_psf(zernike_pol, r: float, theta: float, alpha: float = 1.0) -> float: + """ + Get the precalculated expansion of diffraction integral based on the Nijboer's thesis. + + Parameters + ---------- + zernike_pol : tuple (m, n) or ZernPol + Zernike polynomial definition. + r : float + Radius on the image coordinates. + theta : float + Angle on the image coordinates. + alpha : float, optional + Amplitude of the polynomial. The default is 1.0. + + Reference + --------- + [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf + + Returns + ------- + float + PSF point as |U|**2. + + """ (m, n) = define_orders(zernike_pol) # get polynomial orders # The analytical equation could be found in the Nijboer's thesis if m == 0 and n == 0: # piston value @@ -176,12 +363,36 @@ def get_aberrated_psf(zernike_pol, r: float, theta: float, alpha: float = 1.0) - else: return b*b - a*a else: - return 0.0 # !!! Should be exchanged to the integral or precalculated equations + return 0.0 # %% Another attempt to use Nijboer's Thesis Functions -# Tested, this approximation somehow doens't work +# Tested, this approximation somehow doens't work (maybe, the amplitudes of polynomial aren't small enough) def radial_integral_nijboer(zernike_pol: ZernPol, r: float, order: int, power: int) -> float: + """ + Calculate the radial integrals used in the Nijboer's thesis. + + Parameters + ---------- + zernike_pol : tuple (m, n) or ZernPol + Zernike polynomial definition. + r : float + Radius on the image coordinates. + order : int + Of the Bessel function Jv(r). + power : int + Power of the polynomial radial component. + + Reference + --------- + [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf + + Returns + ------- + float + Result of integration. + + """ integral_sum = 0.0; (m, n) = define_orders(zernike_pol) # get polynomial orders # Integration on the pupil radius. Vectorized form of simple integration equation h_p = 1.0/1000; p = np.arange(start=h_p, stop=1.0-h_p, step=h_p) @@ -190,7 +401,32 @@ def radial_integral_nijboer(zernike_pol: ZernPol, r: float, order: int, power: i return integral_sum -def psf_point_approx_sum(zernike_pol: ZernPol, r: float, theta: float, alpha: float) -> complex: +def psf_point_approx_sum(zernike_pol: ZernPol, r: float, theta: float, alpha: float): + """ + Get the point for calculation of PSF as the expansion of the diffraction integral acquired in the Nijboer's thesis. + + Parameters + ---------- + zernike_pol : tuple (m, n) or ZernPol + Zernike polynomial definition. + r : float + Radius on the image coordinates. + theta : float + Angle on the image coordinates. + alpha : float, optional + Amplitude of the polynomial. The default is 1.0. + + Reference + --------- + [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf + + + Returns + ------- + float + |U|**2 - the module and square of the amplitude, intensity as the PSF value. + + """ (m, n) = define_orders(zernike_pol) # get polynomial orders if not isinstance(zernike_pol, ZernPol): zernike_pol = ZernPol(m=m, n=n) @@ -262,6 +498,7 @@ def get_psf_kernel(zernike_pol, calibration_coefficient: float, alpha: float) -> (m, n) = define_orders(zernike_pol) # get polynomial orders # Define the kernel size just imperically max_size = int(round(20.0*(1.0/calibration_coefficient), 0)) + 1; size = max_size + # Auto definition of the required PSF size is failed for the different PSFs forms (e.g., vertical coma with different amplitudes) # Make kernel with odd sizes for precisely centering the kernel if size % 2 == 0: size += 1 @@ -274,8 +511,8 @@ def get_psf_kernel(zernike_pol, calibration_coefficient: float, alpha: float) -> # The PSF also has the angular dependency, not only the radial one theta = np.arctan2((i - i_center), (j - j_center)) theta += np.pi # shift angles to the range [0, 2pi] - # kernel[i, j] = get_aberrated_psf(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha) - kernel[i, j] = get_psf_point_r(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha*calibration_coefficient) # ??? Scaling alpha + # The scaling below is not needed because the Zernike polynomial is scaled as the RMS values + kernel[i, j] = get_psf_point_r(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha) return kernel @@ -309,10 +546,7 @@ def show_ideal_psf(zernike_pol, size: int, calibration_coefficient: float, alpha # The PSF also has the angular dependency, not only the radial one theta = np.arctan2((i - i_center), (j - j_center)) theta += np.pi # shift angles to the range [0, 2pi] - # img[i, j] = get_aberrated_psf(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha) - # img[i, j] = get_psf_point(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha) img[i, j] = get_psf_point_r(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha) - # img[i, j] = psf_point_approx_sum(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha) if img[0, 0] > np.max(img)/100: __warn_message = f"The provided size for plotting PSF ({size}) isn't sufficient for proper representation" warnings.warn(__warn_message) @@ -325,7 +559,34 @@ def show_ideal_psf(zernike_pol, size: int, calibration_coefficient: float, alpha def plot_correlation(zernike_pol, size: int, calibration_coefficient: float, alpha: float, title: str = None, - show_original: bool = True, show_psf: bool = True): + show_original: bool = True, show_psf: bool = True, R: int = 1): + """ + Plot result of correlation with the object shown as the pixelated circle. + + Parameters + ---------- + zernike_pol : tuple (m, n) or ZernPol + Zernike polynomial definition. + size : int + Size of the picture with the object. + calibration_coefficient : float + Coefficient between physical values and pixels. + alpha : float + Amplitude of the polynomial. + title : str, optional + Title of plots. The default is None. + show_original : bool, optional + Flag for plotting the original object. The default is True. + show_psf : bool, optional + Flag for plotting the used PSF. The default is True. + R : int, optional + Radius of the circle. The default is 1. + + Returns + ------- + None. + + """ if size % 2 == 0: size += 1 # make the image with odd sizes img = np.zeros((size, size), dtype=float) @@ -354,6 +615,31 @@ def plot_correlation(zernike_pol, size: int, calibration_coefficient: float, alp def plot_correlation_photo(zernike_pol, calibration_coefficient: float, alpha: float, title: str = None, show_original: bool = True, show_psf: bool = False): + """ + Plot result of convolution of PSF with the sample image. + + Parameters + ---------- + zernike_pol : tuple (m, n) or ZernPol + Zernike polynomial definition. + size : int + Size of the picture with the object. + calibration_coefficient : float + Coefficient between physical values and pixels. + alpha : float + Amplitude of the polynomial. + title : str, optional + Title of plots. The default is None. + show_original : bool, optional + Flag for plotting the original object. The default is True. + show_psf : bool, optional + Flag for plotting the used PSF. The default is True. + + Returns + ------- + None. + + """ try: sample = img_as_ubyte(io.imread(os.path.join(os.getcwd(), "nesvizh_grey.jpg"), as_gray=True)) if show_original: @@ -371,9 +657,10 @@ def plot_correlation_photo(zernike_pol, calibration_coefficient: float, alpha: f # %% Tests if __name__ == '__main__': orders1 = (0, 2); orders2 = (0, 0); orders3 = (-1, 1); orders4 = (-3, 3); plot_pure_psfs = False; plot_photo_convolution = False + plot_photo_convolution_row = False; figsizes = (6.5, 6.5) # Plotting - plt.ion(); plt.close('all'); conv_pic_size = 14; detailed_plots_sizes = 22; calibration_coeff = pixel2um_coeff/1.5; alpha = 1.5 + plt.ion(); plt.close('all'); conv_pic_size = 14; detailed_plots_sizes = 24; calibration_coeff = pixel2um_coeff/1.5; alpha = 2.0 if plot_pure_psfs: t1 = time.time() p_img = show_ideal_psf(orders2, detailed_plots_sizes, calibration_coeff, alpha, "Piston") @@ -392,10 +679,16 @@ def plot_correlation_photo(zernike_pol, calibration_coefficient: float, alpha: f # p_img2 = show_ideal_psf(orders1, detailed_plots_sizes, calibration_coeff, -alpha, "Defocus -") # show_ideal_psf(orders4, detailed_plots_sizes, calibration_coeff, alpha, "+ Vertical Trefoil") # show_ideal_psf(orders4, detailed_plots_sizes, calibration_coeff, -alpha, "- Vertical Trefoil") - show_ideal_psf((-2, 2), detailed_plots_sizes, calibration_coeff, alpha, "+ Obliq. Astigmatism") - show_ideal_psf((-2, 2), detailed_plots_sizes, calibration_coeff, -alpha, "- Obliq. Astigmatism") - show_ideal_psf((-1, 3), detailed_plots_sizes, calibration_coeff, alpha, "+ Vert. Coma") - show_ideal_psf((-1, 3), detailed_plots_sizes, calibration_coeff, -alpha, "- Vert. Coma") + # show_ideal_psf((-2, 2), detailed_plots_sizes, calibration_coeff, alpha, "+ Obliq. Astigmatism") + # show_ideal_psf((-2, 2), detailed_plots_sizes, calibration_coeff, -alpha, "- Obliq. Astigmatism") + show_ideal_psf((0, 4), detailed_plots_sizes, calibration_coeff, alpha, f"{alpha} Primary Spherical") + show_ideal_psf((0, 4), detailed_plots_sizes, calibration_coeff, -alpha, f"{-alpha} Primary Spherical") + + # Testing the kernel calculation + # v_coma_m = get_psf_kernel((-1, 3), calibration_coeff, -1.0) + # plt.figure(figsize=figsizes); plt.imshow(v_coma_m, cmap=plt.cm.viridis); plt.tight_layout() + # v_coma_p = get_psf_kernel((-1, 3), calibration_coeff, 1.0) + # plt.figure(figsize=figsizes); plt.imshow(v_coma_p, cmap=plt.cm.viridis); plt.tight_layout() # Plot convolution of the sample photo with the various psfs if plot_photo_convolution: @@ -404,6 +697,12 @@ def plot_correlation_photo(zernike_pol, calibration_coefficient: float, alpha: f plot_correlation_photo(orders1, pixel2um_coeff/1.5, 1.0, "Defocus", show_psf=True, show_original=True) plot_correlation_photo(orders4, pixel2um_coeff/1.5, 1.0, "Vertical Trefoil", show_psf=True, show_original=True) + # Testing of the convolution results with different amplitudes and signs + if plot_photo_convolution_row: + amplitudes = [-2.0, 0.0, 2.0]; vert_coma = (-1, 3) + for amplitude in amplitudes: + plot_correlation_photo(vert_coma, pixel2um_coeff/1.5, amplitude, f"Y Coma {amplitude}", show_psf=False, show_original=False) + # Plot convolution of point objet with the various psfs # plot_correlation(orders2, conv_pic_size, pixel2um_coeff, 0.85, "Piston", show_psf=True) # plot_correlation(orders3, conv_pic_size, pixel2um_coeff/1.75, 0.5, "Y Tilt", True, show_psf=True)