add fraunhofer diffraction test

fraunhofer_test.py

@@ -0,0 +1,70 @@
+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
+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  # Размер сетки
+
+
+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
+
+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]]
+
+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)
+
+
+I_fraun = np.abs(U_fraun)**2
+I_fraun = I_fraun / np.max(I_fraun)
+
+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

@@ -8,14 +8,14 @@ SAMPLE_SIZE = 256
 span = (-1, 1)
 
 
-def generate_zern_wavefront(xv, yv, zerns):
+def generate_zern_wavefront(xv, yv, zerns, norm_radius = 1):
     num_zerns = len(zerns)
 
     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(xv, yv)
+    cart.make_cart_grid(xv/norm_radius, yv/norm_radius)
     return np.array(cart.eval_grid(np.array(zerns), matrix=True))