Spaces:
Sleeping
Sleeping
| import matplotlib.pyplot as plt | |
| import cupy as cp | |
| import numpy as np | |
| import math | |
| from numpy.fft import fft2, fftshift, ifftshift | |
| from cupyx.scipy import fft as fft_gpu | |
| from zern_generator import Zernike | |
| def denoise(screen: np.ndarray): | |
| return np.where(np.abs(screen) < 1e-4, 0, screen) | |
| class FraunhoferGPU: | |
| def __init__(self, zernike: Zernike): | |
| self.focal_distance = 100 * 10 ** -3 | |
| self.distance = 100 * 10 ** -3 | |
| self.radius = 5 * 10 ** -3 | |
| self.span = (-(10 ** -2), 10 ** -2) | |
| self.wavelength = 500 * 10 ** -9 | |
| self.k = 2*math.pi/self.wavelength | |
| self.zernike = zernike | |
| def get_focal_point(self,input_screen): | |
| input_screen_gpu = cp.asarray(input_screen) | |
| xv = cp.asarray(self.zernike.xv) | |
| yv = cp.asarray(self.zernike.yv) | |
| phase_mul = ((cp.exp(1j*self.k*self.distance) * cp.exp(1j*self.k*(xv**2 + yv**2)/(2*self.distance)))/(1j*self.wavelength*self.distance)) | |
| U_fraun = fft_gpu.fft2(fft_gpu.ifftshift(input_screen_gpu), overwrite_x=True) | |
| U_fraun = fftshift(U_fraun * phase_mul) | |
| I_fraun = cp.abs(U_fraun)**2 | |
| return I_fraun / cp.max(I_fraun) | |
| def generate_diff(self, defocus_amount, *zernikes): | |
| zernikes = list(zernikes) | |
| if len(zernikes) < 5: | |
| zernikes += [0] * (5-len(zernikes)) | |
| wf = self.zernike.generate_zern_screen([1], self.radius) | |
| f_plus = self.zernike.generate_zern_screen(self.zernike.add_defocus(defocus_amount,zernikes), norm_radius=self.radius) | |
| f_minus = self.zernike.generate_zern_screen(self.zernike.add_defocus(-defocus_amount,zernikes), norm_radius=self.radius) | |
| wf_c = cp.asarray(wf) | |
| f_plus_c = cp.asarray(f_plus) | |
| f_minus_c = cp.asarray(f_minus) | |
| focal_image_plus = denoise(self.get_focal_point(wf_c*cp.exp(f_plus_c * 1j))) | |
| focal_image_minus = denoise(self.get_focal_point(wf_c*cp.exp(f_minus_c * 1j))) | |
| return ( | |
| cp.asnumpy(focal_image_plus), | |
| cp.asnumpy(focal_image_minus), | |
| cp.asnumpy(focal_image_plus - focal_image_minus) | |
| ) | |
| class Fraunhofer: | |
| def __init__(self, zernike: Zernike): | |
| self.focal_distance = 100 * 10 ** -3 | |
| self.distance = 100 * 10 ** -3 | |
| self.radius = 5 * 10 ** -3 | |
| self.span = (-(10 ** -2), 10 ** -2) | |
| self.wavelength = 500 * 10 ** -9 | |
| self.k = 2*math.pi/self.wavelength | |
| self.zernike = zernike | |
| # def _denoise(screen: np.ndarray): | |
| # return np.where(np.abs(screen) < 1e-4, 0, screen) | |
| def get_focal_point(self,input_screen): | |
| input_screen_gpu = np.asarray(input_screen) | |
| phase_mul = ((np.exp(1j*self.k*self.distance) * np.exp(1j*self.k*(self.zernike.xv**2 + self.zernike.yv**2)/(2*self.distance)))/(1j*self.wavelength*self.distance)) | |
| U_fraun = fft2(ifftshift(input_screen_gpu)) | |
| U_fraun = fftshift(U_fraun * phase_mul) | |
| #U_fraun = fftshift(U_fraun) | |
| I_fraun = np.abs(U_fraun)**2 | |
| return I_fraun / np.max(I_fraun) | |
| def generate_diff(self, defocus_amount, *zernikes): | |
| zernikes = list(zernikes) | |
| if len(zernikes) < 5: | |
| zernikes += [0] * (5-len(zernikes)) | |
| wf = self.zernike.generate_zern_screen([1], self.radius) | |
| f_plus = self.zernike.generate_zern_screen( | |
| self.zernike.add_defocus(defocus_amount,zernikes), | |
| norm_radius=self.radius | |
| ) | |
| f_minus = self.zernike.generate_zern_screen( | |
| self.zernike.add_defocus(-defocus_amount,zernikes), | |
| norm_radius=self.radius | |
| ) | |
| wf_c = np.asarray(wf) | |
| f_plus_c = np.asarray(f_plus) | |
| f_minus_c = np.asarray(f_minus) | |
| focal_image_plus = denoise( | |
| self.get_focal_point(wf_c*np.exp(f_plus_c * 1j)) | |
| ) | |
| focal_image_minus = denoise( | |
| self.get_focal_point(wf_c*np.exp(f_minus_c * 1j)) | |
| ) | |
| diff = focal_image_plus - focal_image_minus | |
| normalized_diff = diff / cp.max(diff) | |
| return ( | |
| cp.asnumpy(focal_image_plus), | |
| cp.asnumpy(focal_image_minus), | |
| cp.asnumpy(normalized_diff) | |
| ) | |
| if __name__ == "__main__": | |
| zernike = Zernike() | |
| zernike.set_image_params(npix=2 ** 8) | |
| zoom_factor = 20 | |
| f = Fraunhofer(zernike=zernike) | |
| ratio = (f.wavelength*f.focal_distance)/f.radius | |
| zernikes = [[1 if i == j else 0 for i in range(0,24)] for j in range(20,24)] | |
| print(zernikes) | |
| fig, ax = plt.subplots(4,1) | |
| for i, z in enumerate(zernikes): | |
| print(zernike._osa_to_noll_list(z)) | |
| wf_a = zernike.generate_zern_screen([1], 0.01/zoom_factor) | |
| wf_c = zernike.generate_zern_screen(z, 0.01/zoom_factor) | |
| wf_c_nz = zernike.generate_zern_screen(z, 0.01) | |
| wf = wf_a * np.exp(wf_c*1j) | |
| focal = f.get_focal_point(wf) | |
| zoom = ratio*zoom_factor | |
| #ax[i][0].imshow(np.abs(wf_c_nz)) | |
| ax[i].imshow(np.abs(focal)) | |
| ax[i].imshow(np.abs(focal)) | |
| ax[0].set_title("Zernike 5 - DEFOCUS") | |
| ax[1].set_title("Zernike 6 - V-PR-ASTIGMATISM") | |
| ax[2].set_title("Zernike 7 - V-TREFOIL") | |
| ax[3].set_title("Zernike 8 - V-PR_COMA") | |
| ax[0].axis('off') | |
| ax[1].axis('off') | |
| ax[2].axis('off') | |
| ax[3].axis('off') | |
| plt.show() | |