add diff

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from zern_generator import generate_zern_wavefront_fig, generate_batch
+from fraunhofer_test import generate_diff
 
 css = """
 .container {
@@ -31,7 +32,9 @@ with gr.Blocks(theme=theme, css=css) as demo:
 
             resolution_num = gr.Number(value=256, label="Разрешение")
             num_zernikes = gr.Number(label="Количество коэффициентов Цернике", value=2, interactive=True)
+            defocus_amt = gr.Number(label="Величина дефокуса", value=2, interactive=True)
             sliders_container = gr.Column(elem_classes=["container"])
+           
             # file_inp.upload(get_zern_from_file, inputs=[file_inp], outputs=[*sliders_container.children])
 
 
@@ -56,19 +59,25 @@ with gr.Blocks(theme=theme, css=css) as demo:
             with gr.Tab(label="3D"):
                 plot_3d = gr.Plot()
             with gr.Tab(label="Batch"):
-
                 plot_batch = gr.Plot(label="Тестовые экраны (NOLL scheme)")
+            with gr.Tab(label="Diff"):
+                plot_diff = gr.Plot(label="Разница в фокусе")
 
             with gr.Row():
                 btn = gr.Button("Рассчитать")
                 btn_batch = gr.Button("Рассчитать тестовые экраны")
+                btn_diff = gr.Button("Рассчитать разницу в фокусе")
 
 
             def on_button_click(resolution):
                 return generate_zern_wavefront_fig(resolution, *[zerns[key] for key in sorted(zerns)])
+            
+            def on_diff_button_click(defocus_amount):
+                return generate_diff(defocus_amount, *[zerns[key] for key in sorted(zerns)])
 
 
             btn.click(on_button_click, inputs=[resolution_num], outputs=[plot_2d, plot_3d])
             btn_batch.click(generate_batch, inputs=[resolution_num], outputs=[plot_batch])
+            btn_diff.click(on_diff_button_click, inputs=[defocus_amt], outputs=[plot_diff])
 
 demo.launch()

fraunhofer_test.py

@@ -13,58 +13,73 @@ from multiprocessing import Pool
 from zern_generator import generate_zern_wavefront
 from scipy.ndimage import zoom
 
-fig, axs = plt.subplots(2, 2)
-
 
 f = 100 * 10 ** -3  # Форкусное расстояние линзы
 
 
 radius = 5 * 10 ** -3  # Радиус линзы
-
 span = (-(10 ** -2), 10 ** -2)
 lens = po.ThinLens(radius, f)
 z = 100 * 10 ** -3  # Расстояние
 wavelength = 500 * 10 ** -9  # Длина волны
-
 pixel_scale = 250 * 10 ** -4  # Размер пикселя
 npix = 2**9  # Размер сетки
+k = 2*math.pi/wavelength
 
+def gen_focal_image(xv, yv, input_wf):
+    phase_mul = ((np.exp(1j*k*z) * np.exp(1j*k*(xv**2 + yv**2)/(2*z)))/(1j*wavelength*z))
+    U_fraun = fft2(ifftshift(input_wf), overwrite_x=True)
+    U_fraun = fftshift(U_fraun * phase_mul)
 
-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, 3, 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
+    I_fraun = np.abs(U_fraun)**2
+    return I_fraun / np.max(I_fraun)
+
+
+def generate_zernike_screen(xv, yv, coeffs, norm_radius):
+    dist_from_center = np.sqrt(xv**2 + yv**2)
+    mask = dist_from_center <= norm_radius
+    f2 = generate_zern_wavefront(xv, yv, coeffs, norm_radius=radius)
+    f2[mask == 0] = 0
+    f2[np.isnan(f2)] = 0
+    return f2
+
+
+def circle(xv, yv, radius, resolution):
+    dist_from_center = np.sqrt(xv**2 + yv**2)
+    mask = dist_from_center <= radius
+    wf = np.zeros((resolution, resolution), dtype=np.complex64)
+    wf[mask] = 1 + 0j
+    return wf
 
-wf = np.zeros((npix, npix), dtype=np.complex64)
-wf[mask] = 1 + 0j
-wf += f2 * 1j
 
 
-D = 2e-6
-fscale = wavelength*z
 
-target_min, target_max = -50e-6, 50e-6
-v = np.fft.fftfreq(npix, d=D/npix)
-v = fscale * np.fft.fftshift(v)
-extent = [1e3 * v[0], 1e3 * v[-1], 1e3 * v[0], 1e3 * v[-1]]
+def generate_diff(defocus_amount, *zernikes):
+    zernikes = list(zernikes)
+    if len(zernikes) < 5:
+        zernikes += [0] * (5-len(zernikes))
+    fig, axs = plt.subplots(1, 3, figsize = (30, 15))
 
-phase_mul = ((np.exp(1j*k*z) * np.exp(1j*k*(xv**2 + yv**2)/(2*z)))/(1j*wavelength*z))
-U_fraun = fft2(ifftshift(wf), overwrite_x=True)
-U_fraun = fftshift(U_fraun * phase_mul)
+    ddx = np.linspace(*span, npix)
+    ddy = np.linspace(*span, npix)
+    xv, yv = np.meshgrid(ddx, ddy) 
+    wf = circle(xv, yv, radius, npix)
+    zernikes[4] += defocus_amount
+    f_plus = generate_zernike_screen(xv, yv, zernikes, norm_radius=radius)
+    zernikes[4] -= 2*defocus_amount
+    f_minus = generate_zernike_screen(xv, yv, zernikes, norm_radius=radius)
 
+    focal_image_plus = gen_focal_image(xv, yv, wf*np.exp(f_plus * 1j))
+    
+    focal_image_minus = gen_focal_image(xv, yv, wf*np.exp(f_minus * 1j))
 
-I_fraun = np.abs(U_fraun)**2
-I_fraun = I_fraun / np.max(I_fraun)
+    axs[0].imshow(focal_image_plus, cmap='grey', extent=(*span, *span))
+    axs[1].imshow(focal_image_minus, cmap='grey', extent=(*span, *span))
+    axs[2].imshow((focal_image_plus-focal_image_minus), cmap='grey', extent=(*span, *span))
+    axs[0].axis('off')
+    axs[1].axis('off')
+    axs[2].axis('off')
+    fig.show()
+    return fig
 
-axs[0][1].imshow(I_fraun, cmap='grey', extent=(*span, *span))
-axs[1][1].plot(np.linspace(extent[0], extent[1], I_fraun.shape[0]),
-               I_fraun[:, I_fraun.shape[0]//2], antialiased='none')
-axs[0][0].imshow(np.abs(wf)**2, cmap='grey', extent=(*span, *span))
-axs[1][0].imshow(f2, cmap='grey', extent=(*span, *span), interpolation='none')
-plt.show()

zern_generator.py

@@ -9,6 +9,10 @@ span = (-1, 1)
 
 
 def generate_zern_wavefront(xv, yv, zerns, norm_radius = 1):
+    indexed_zerns = sorted(map(lambda x: (_mn_to_noll(*_osa_to_mn(x[0])), x[1]), enumerate(zerns)), key=lambda x: x[0])
+    print(indexed_zerns)
+    zerns = [x[1] for x in indexed_zerns]
+    print(zerns)
     num_zerns = len(zerns)
 
     degree = math.ceil((-3 + math.sqrt(8 * num_zerns + 1)) / 2)
@@ -58,3 +62,25 @@ def generate_batch(resolution):
             zern_index += 1
     plt.show()
     return fig
+
+
+def _osa_to_mn(idx):
+    n = math.ceil((-3 + math.sqrt((9 + 8 * idx))) / 2)
+    m = 2 * idx - n * (n + 2)
+    return m, n
+
+def _mn_to_noll(m, n):
+    first = (n * (n + 1)) // 2
+    second = abs(m)
+    third = 0
+    # I wanted to simplify the logic but it seems to work like a colander
+    if m > 0 and n % 4 in [0, 1]:
+        third = 0
+    elif m < 0 and n % 4 in [2, 3]:
+        third = 0
+    elif m >= 0 and n % 4 in [2, 3]:
+        third = 1
+    elif m <= 0 and n % 4 in [0, 1]:
+        third = 1
+    j = first + second + third
+    return j