dataset generation

.gitignore

@@ -1,3 +1,4 @@
 venv
 .idea
-__pycache__
+__pycache__
+screens

fraunhofer_test.py

@@ -1,7 +1,8 @@
 import matplotlib.pyplot as plt
+import cupy as cp
 import numpy as np
 import math
-from scipy.fft import fft2, fftshift, ifftshift
+from cupyx.scipy.fft import fft2, fftshift, ifftshift
 from zern_generator import Zernike
 
 
@@ -17,40 +18,45 @@ class Fraunhofer:
         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 _denoise(screen: np.ndarray):
+    #     return np.where(np.abs(screen) < 1e-4, 0, screen)
+
     def get_focal_point(self,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), overwrite_x=True)
-        U_fraun = fftshift(U_fraun * phase_mul)
+        input_screen_gpu = cp.asarray(input_screen)
+        #phase_mul = ((cp.exp(1j*self.k*self.distance) * cp.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), overwrite_x=True)
+        #U_fraun = fftshift(U_fraun * phase_mul)
+        U_fraun = fftshift(U_fraun)
+
+
+        I_fraun = cp.abs(U_fraun)**2
+        return I_fraun / cp.max(I_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)
 
-        focal_image_plus = denoise(self.get_focal_point(wf*np.exp(f_plus * 1j)))
-        focal_image_minus = denoise(self.get_focal_point(wf*np.exp(f_minus * 1j)))
+        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)
+        )
 
-        return focal_image_plus, focal_image_minus, focal_image_plus-focal_image_minus
-        
 
 
-    

generate_dataset.py

@@ -0,0 +1,113 @@
+from zern_generator import Zernike
+from PIL import Image
+from fraunhofer_test import Fraunhofer
+import cupy as cp
+import matplotlib.pyplot as plt
+import itertools
+from time import time
+from pprint import pprint
+import os
+from multiprocessing import Pool, cpu_count
+import numpy as np
+
+zernike = Zernike()
+zernike.set_image_params(npix=2 ** 8)
+fraunhofer = Fraunhofer(zernike)
+
+
+def generate_image(z):
+
+
+    start_time = time()
+    arr = fraunhofer.generate_diff(1, 0, *z)
+    a = cp.min(arr[2])
+    b = cp.max(arr[2])
+    img_arr = ((arr[2] - a) / (b - a) * 255).astype(cp.uint8)
+    img = Image.fromarray(np.stack([img_arr] * 3, axis=-1), mode='RGB')
+    generation_end_time = time()
+    if True:
+        coeffs = '_'.join([f'{x:.2f}' for x in z[2:]])
+        if not os.path.exists(f'./screens_gpu/{z[0]}/{z[1]}'):
+            try:
+                os.makedirs(f'./screens_gpu/{z[0]}/{z[1]}')
+            except:
+                pass
+        filename = f'./screens_gpu/{z[0]}/{z[1]}/{coeffs}_.png'
+        img.save(filename)
+        save_end_time = time()
+    else:
+        plt.imshow(img)
+        plt.colorbar()
+        plt.show()
+    return generation_end_time - start_time, save_end_time - generation_end_time
+
+
+def generate_dataset(zs):
+    c = 0
+    window_size = 50
+    times = []
+    start_time = time()
+    total_images = len(zs)
+    num_workers = 24
+    if num_workers is None:
+        num_workers = cpu_count()
+
+    with Pool(num_workers) as pool:
+        results = []
+
+        for result in pool.imap_unordered(generate_image, zs):
+            c += 1
+
+            results.append(result)
+
+            current_time = time()
+            iteration_time = current_time - start_time
+            times.append(iteration_time)
+            if len(times) > window_size:
+                times.pop(0)
+            avg_iteration_time = sum(times) / len(times)
+
+            # if c % 10 == 0:
+            #     os.system('cls')
+            #     print(f"Generated {c}/{total_images}")
+            #     elapsed_time = current_time - start_time
+            #     print(f"Elapsed time: {elapsed_time:.2f} s")
+            #     remaining_images = total_images - c
+            #     estimated_remaining_time = avg_iteration_time * remaining_images
+            #     print(f"Estimated remaining time: {estimated_remaining_time:.2f} s")
+        for result in results:
+            result.get()
+
+
+def generate_dataset_gpu(zs):
+    c = 0
+    window_size = 50
+    times = []
+    start_time = time()
+    total_images = len(zs)
+    num_workers = 24
+    if num_workers is None:
+        num_workers = cpu_count()
+
+    for z in zs:
+
+        generate_image(z)
+
+
+if __name__ == '__main__':
+    rng = (-5, 5)
+    step = 0.5
+    points = tuple(np.linspace(*rng, int((rng[1] - rng[0]) / step + 1)))
+    points_as_tuples = tuple(float(x) for x in points)
+
+    zernikes = np.array(list(set(itertools.product(points_as_tuples, repeat=3))))
+
+    # print("Total images: ", len(zernikes))
+    # total_gen_time, total_save_time = 0, 0
+    # for z in zernikes:
+    #     gen_time, save_time = generate_image(z)
+    #     total_save_time += save_time
+    #     total_gen_time += gen_time
+    # print("Average generation time: ", total_gen_time/len(zernikes))
+    # print("Average save time: ", total_save_time/len(zernikes))
+    generate_dataset(zernikes)

optic.py

@@ -0,0 +1,129 @@
+import astropy.units as u
+import pyoptica as po
+import matplotlib.pyplot as plt
+import numpy as np
+import math
+import os
+from scipy.fft import fft2, ifft2, fftshift, ifftshift
+import aotools.opticalpropagation as op
+import scipy.special as sc
+import time
+import multiprocessing
+from multiprocessing import Pool
+from zern_generator import generate_zern_wavefront, generate_zern_wavefront_fig
+from scipy.ndimage import zoom
+
+fig, (ax1, ax2) = plt.subplots(2,1)
+
+
+#f2.show()
+
+f = 100 * 10 ** -3 # Форкусное расстояние линзы
+
+
+
+radius = 5 * 10 ** -3 # Радиус линзы
+span = (-(2*radius), 2*radius)
+lens = po.ThinLens(radius, f)
+
+z = 100 * 10 ** -3 # Расстояние 
+
+wavelength = 500 * 10 ** -9 # Длина волны
+pixel_scale = 250 * 10 ** -4 # Размер пикселя
+npix = 2**9 # Размер сетки
+
+
+ddx = np.linspace(*span, npix)
+ddy = np.linspace(*span, npix)
+xv, yv = np.meshgrid(ddx, ddy)
+f2 = generate_zern_wavefront(xv, yv, [0,0,0,1,0,0], norm_radius=radius)
+dist_from_center = np.sqrt(xv**2 + yv**2)
+mask = dist_from_center <= radius
+f2[mask==0] = 0
+f2[np.isnan(f2)] = 0
+k = 2*math.pi/wavelength
+
+
+wf = np.zeros((npix, npix), dtype=np.complex64)
+
+wf[mask] = 1 + 0j
+
+wf += f2 * 1j
+
+wf = fftshift(fft2(wf))
+
+phase_mul = ((np.exp(1j*k*z) * np.exp(1j*k*(xv**2 + yv**2)/(2*z)))/(1j*wavelength*z))
+
+#wf = wf * phase_mul
+wf= np.abs(wf)**2
+#wf = np.log10(wf)
+
+crop = 20
+center = npix//2
+vx = np.linspace(*span, npix)
+ax1.imshow(wf, extent = (*span, *span), cmap="gray")
+#ax2.plot(vx, wf[npix//2])
+ax2.imshow(f2, extent = (*span, *span),cmap="gray")
+print(np.ptp(wf))
+plt.show()
+
+
+
+
+
+# amp = np.array(wf.amplitude)
+
+# a = fftshift(fft2(wf.amplitude))
+# a = np.abs(a)**2
+# print(a)
+# ax.imshow(a, extent=(-npix*pixel_scale/2, npix*pixel_scale/2,-npix*pixel_scale/2, npix*pixel_scale/2))
+# plt.show()
+
+
+# # Двойной интеграл из 4.13
+# def calc_point_fourier(x,y, uxy):
+#     res = 0 + 0j
+#     for x0 in range(npix):
+#         x_r = to_real(x0)
+#         for y0 in range(npix):       
+#             y_r = to_real(y0)
+#             res += uxy * np.exp(-1j*k*(x*x_r + y*y_r)/z)
+#     return res
+
+
+# # Считает точку по 4.13
+# def calc_point(x,y):
+#     x_r = to_real(x)
+#     y_r = to_real(y)
+#     if y%100 == 1:
+#         os.system('cls')
+#         print(f"{npix*x + y}/{npix**2}")
+#     return ((np.exp(1j*k*z) * np.exp(1j*k*(x_r**2 + y_r**2)/(2*z)))/(1j*wavelength*z)) * calc_point_fourier(x_r, y_r, amp[x][y])
+
+
+    
+
+# if __name__ == "__main__":
+#     pairs = []
+#     for x in range(npix):
+#         for y in range(npix):
+#             pairs.append((x,y))
+#     multiprocessing.freeze_support()
+#     with Pool(20) as p:
+#         res = p.starmap_async(calc_point, pairs)
+#         a = res.get()
+#         a = np.abs(a)**2
+#         a = a.reshape((npix, npix))
+#         a /= np.max(a)
+#         ax.imshow(a, extent=(-npix*pixel_scale/2, npix*pixel_scale/2,-npix*pixel_scale/2, npix*pixel_scale/2))
+#         plt.show()
+
+
+
+
+
+
+try:
+    pass
+except KeyError:
+    

zern_generator.py

@@ -2,6 +2,7 @@ import math
 
 import matplotlib.pyplot as plt
 import numpy as np
+import cupy as cp
 from zernike import RZern
 import utils
 
@@ -9,7 +10,6 @@ SAMPLE_SIZE = 256
 span = (-1, 1)
 
 
-
 class Zernike:
     def __init__(self):
         self.set_image_params()
@@ -29,33 +29,37 @@ class Zernike:
             "P-SPHERICAL",
             "V-SC-ASTIGMATISM",
             "V-QUADRAFOIL",
-            
+
             # "V-SC-COMA",
             # "H-SC-COMA",
             # "V-SC-TREFOIL",
             # "O-SC-TREFOIL",
             # "O-PENTAFOIL",
             # "V-PENTAFOIL"
-            ]
-        
-    def set_image_params(self, span = (-(10 ** -2), 10 ** -2), npix = 2**8):
+        ]
+
+    def set_image_params(self, span=(-(10 ** -2), 10 ** -2), npix=2 ** 8):
         self.span = span
         self.npix = npix
         ddx = np.linspace(*self.span, self.npix)
         ddy = np.linspace(*self.span, self.npix)
-        self.xv, self.yv = np.meshgrid(ddx, ddy) 
-    
-    def generate_zern_screen(self, zerns, norm_radius = 1):
+        self.xv, self.yv = np.meshgrid(ddx, ddy)
+
+    def generate_zern_screen(self, zerns, norm_radius=1):
+
+        if isinstance(zerns, cp.ndarray):
+            zerns = zerns.get()
         zerns = self._osa_to_noll_list(zerns)
         num_zerns = len(zerns)
         mask = utils.circle_mask(self.xv, self.yv, norm_radius)
-        
-        
+
         degree = math.ceil((-3 + math.sqrt(8 * num_zerns + 1)) / 2)
         zerns = list(zerns) + [0] * int((((degree + 1) * (degree + 2) / 2) - len(zerns)))
         cart = RZern(degree)
-        cart.make_cart_grid(self.xv/norm_radius, self.yv/norm_radius)
-        screen =  np.array(cart.eval_grid(np.array(zerns), matrix=True))
+
+        cart.make_cart_grid(self.xv / norm_radius, self.yv / norm_radius)
+
+        screen = np.array(cart.eval_grid(np.array(zerns), matrix=True))
         screen[mask == 0] = 0
         screen[np.isnan(screen)] = 0
         return screen
@@ -70,17 +74,17 @@ class Zernike:
         fig_2d.colorbar(a)
         plt.show()
         return fig_2d, fig_3d
-    
+
     def add_defocus(self, defocus_value, zernikes):
+
         zernikes = list(zernikes)
         if len(zernikes) < 5:
-            zernikes += [0] * (5-len(zernikes))
+            zernikes += [0] * (5 - len(zernikes))
         zernikes[4] += defocus_value
         return zernikes
 
-
     def generate_batch(self, resolution):
-        #term = ["PISTON","V-TILT","H-TILT","DEFOCUS","O-PR-ASTIGMATISM","V-PR-ASTIGMATISM","V-PR-COMA","H-PR-COMA","V-TREFOIL","O-TREFOIL","P-SPHERICAL","O-SC-ASTIGMATISM","V-SC-ASTIGMATISM","V-QUADRAFOIL","O-QUADRAFOIL","V-SC-COMA","H-SC-COMA","V-SC-TREFOIL","O-SC-TREFOIL","O-PENTAFOIL","V-PENTAFOIL"]
+        # term = ["PISTON","V-TILT","H-TILT","DEFOCUS","O-PR-ASTIGMATISM","V-PR-ASTIGMATISM","V-PR-COMA","H-PR-COMA","V-TREFOIL","O-TREFOIL","P-SPHERICAL","O-SC-ASTIGMATISM","V-SC-ASTIGMATISM","V-QUADRAFOIL","O-QUADRAFOIL","V-SC-COMA","H-SC-COMA","V-SC-TREFOIL","O-SC-TREFOIL","O-PENTAFOIL","V-PENTAFOIL"]
         fig = plt.figure()
         fig.set_figwidth(30)
         fig.set_figheight(20)
@@ -92,15 +96,14 @@ class Zernike:
 
                 zern = [0] * (zern_index - 1) + [1]
                 ax = plt.subplot2grid((5, 6), (i, j))
-                Phi = self.generate_zern_screen(zern, norm_radius=self.span[1]/2)
-                ax.set_title(f" \n Zernike {zern_index-1} — {self.term[zern_index-1]} ")
+                Phi = self.generate_zern_screen(zern, norm_radius=self.span[1] / 2)
+                ax.set_title(f" \n Zernike {zern_index - 1} — {self.term[zern_index - 1]} ")
                 ax.imshow(Phi, extent=(*span, *span), cmap='gray')
 
                 zern_index += 1
         plt.show()
         return fig
 
-
     def _osa_to_mn(self, idx):
         n = math.ceil((-3 + math.sqrt((9 + 8 * idx))) / 2)
         m = 2 * idx - n * (n + 2)
@@ -122,7 +125,6 @@ class Zernike:
         j = first + second + third
         return j
 
-
     def _osa_to_noll(self, idx):
         return self._mn_to_noll(*self._osa_to_mn(idx))
 
@@ -130,8 +132,11 @@ class Zernike:
         indexed_zerns = sorted(map(lambda x: (self._osa_to_noll(x[0]), x[1]), enumerate(zernikes)), key=lambda x: x[0])
         used_indeces = [x for x, _ in indexed_zerns]
         mn, mx = min(used_indeces), max(used_indeces)
-        missing_indeces = set(range(mn, mx+1)) - set(used_indeces)
+        missing_indeces = set(range(mn, mx + 1)) - set(used_indeces)
         return [x[1] for x in list(sorted(indexed_zerns + [(x, 0) for x in missing_indeces]))]
 
-term = ["PISTON","V-TILT","H-TILT","DEFOCUS","O-PR-ASTIGMATISM","V-PR-ASTIGMATISM","V-PR-COMA","H-PR-COMA","V-TREFOIL","O-TREFOIL","P-SPHERICAL","O-SC-ASTIGMATISM","V-SC-ASTIGMATISM","V-QUADRAFOIL","O-QUADRAFOIL","V-SC-COMA","H-SC-COMA","V-SC-TREFOIL","O-SC-TREFOIL","O-PENTAFOIL","V-PENTAFOIL"]
-print(Zernike()._osa_to_noll_list(list(range(0,15))))
+
+term = ["PISTON", "V-TILT", "H-TILT", "DEFOCUS", "O-PR-ASTIGMATISM", "V-PR-ASTIGMATISM", "V-PR-COMA", "H-PR-COMA",
+        "V-TREFOIL", "O-TREFOIL", "P-SPHERICAL", "O-SC-ASTIGMATISM", "V-SC-ASTIGMATISM", "V-QUADRAFOIL", "O-QUADRAFOIL",
+        "V-SC-COMA", "H-SC-COMA", "V-SC-TREFOIL", "O-SC-TREFOIL", "O-PENTAFOIL", "V-PENTAFOIL"]
+print(Zernike()._osa_to_noll_list(list(range(0, 15))))