structured-optics 0.2__tar.gz → 0.3__tar.gz

{structured_optics-0.2 → structured_optics-0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: structured_optics
-Version: 0.2
+Version: 0.3
 Summary: Package for Structured Light simulation
 Author: Altilano C. Barbosa
 Requires-Python: >=3.9

{structured_optics-0.2 → structured_optics-0.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "structured_optics"
-version = "0.2"
+version = "0.3"
 description = "Package for Structured Light simulation"
 requires-python = ">=3.9"
 authors = [

{structured_optics-0.2 → structured_optics-0.3}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='structured_optics',
-    version='0.2',
+    version='0.3',
     packages=find_packages()
 
 )

structured_optics-0.3/structured_optics/__init__.py

@@ -0,0 +1,4 @@
+from .struct_opt import *
+from .modes import *
+
+

{structured_optics-0.2 → structured_optics-0.3}/structured_optics/algebra_utils.py

@@ -3,7 +3,7 @@ import numpy as np
 
 
 
-def hg_proj(Beam,N, completeness = False):
+def hg_proj(Beam,N):
     overlaps = np.zeros((N+1, N+1), dtype='complex')
     auxiliary_beam  = Beam.copy_clean()
     for n in range(N+1):
@@ -11,14 +11,10 @@ def hg_proj(Beam,N, completeness = False):
             auxiliary_beam.hg(n,m)
             overlaps[n,m] = overlap(Beam, auxiliary_beam)
     n, m =  np.meshgrid(np.arange(0, N+1, 1), np.arange(0, N+1, 1))
-    comp = np.sum(np.abs(overlaps)**2)
-    if completeness == True:
-        return n, m, overlaps, comp
-    else:
-        return n, m, overlaps
+    return n, m, overlaps
     
 
-def lg_proj(Beam, N, completeness = False):
+def lg_proj(Beam, N):
     overlaps = np.zeros((int(2*N)+1, int(N/2)+1), dtype='complex')
     auxiliary_beam  = Beam.copy_clean()
     for l in np.arange(-N, N+1, 1):
@@ -26,27 +22,23 @@ def lg_proj(Beam, N, completeness = False):
             auxiliary_beam.lg(l,p)
             overlaps[l+N,p] = overlap(Beam, auxiliary_beam)
     p, l = np.meshgrid(np.arange(0, int(N/2)+1, 1), np.arange(-N, N+1, 1))
-    comp = np.sum(np.abs(overlaps)**2)
-    if completeness == True:
-        return l, p, overlaps, comp
-    else:
-        return l, p, overlaps
+    return l, p, overlaps
     
 def hg_basis(Beam, N, waist = None, norm1 = False):
-    aux = Beam.copy_clean()
+    aux = Beam.copy_clean(1)
     if waist != None:
         aux.waist = waist
     basis = np.empty((N+1, N+1, Beam.Dy, Beam.Dx), dtype = 'complex')
     for n in range(N+1):
         for m in range(N+1):
             aux.hg(n, m)
-            basis[n, m] = aux.field
+            basis[n, m] = aux.Ex
     if norm1 == True:
         basis = basis/np.max(basis)
     return basis
 
 def lg_basis(Beam, N, waist=None, norm1=False):
-    aux = Beam.copy_clean()
+    aux = Beam.copy_clean(1)
     if waist != None:
         aux.waist = waist
     basis = np.empty((N+1, N+1, Beam.Dy, Beam.Dx), dtype = 'complex')
@@ -55,7 +47,7 @@ def lg_basis(Beam, N, waist=None, norm1=False):
             l = m-n
             p = min(n,m)
             aux.lg(l, p)
-            basis[n, m] = aux.field
+            basis[n, m] = aux.Ex
     if norm1 == True:
         basis = basis/np.max(basis)
     return basis
@@ -73,18 +65,12 @@ def bessel_basis(Beam, Nmax, waist=None, norm1=False):
         basis = basis / np.max(np.abs(basis))
     return basis
 
-def build_from_coefs_and_basis(Beam, coefs, basis):
+def build_from_coefs_and_basis(Beam, coefs, basis, pol_index = 0):
     N = len(coefs)
     aux = Beam.copy_clean()
     for i in range(N):
         for j in range(N):
-            aux.field = aux.field + coefs[i,j] * basis[i,j]
+            aux.field[pol_index] = aux.field[pol_index] + coefs[i,j] * basis[i,j]
     Beam.field = aux.field
     return Beam
 
-def astigmatic_mode_converter(Beam, N, theta): 
-    n, m, overlaps = Beam.hg_projector(N)
-    converter = np.exp(1j*(m-n)*theta)
-    overlaps = overlaps*converter
-    Beam.build_from_coefs_and_basis(overlaps, Beam.hg_basis(N))
-    return Beam

{structured_optics-0.2 → structured_optics-0.3}/structured_optics/hologram.py

@@ -6,17 +6,14 @@ import time
 
 def slm_hologram(Beam, x_grating, y_grating, method, input_beam, eps, max_range):
     #generate a hologram for a slm
-    if Beam.pol_dim != 1:
-        print('Beam must be scalar (pol_dim=1) to generate holograms')
-        return None
     
     dx = (Beam.x[0,1] - Beam.x[0,0])
     dy = (Beam.y[1,0] - Beam.y[0,0])
     if input_beam is None:
         input_field = np.exp(-((Beam.x**2 + Beam.y**2)/(2*np.max(Beam.x))**2))
     else:
-        input_field = input_beam.field
-    desired = Beam.field/np.max(np.abs(Beam.field))
+        input_field = input_beam.Ex
+    desired = Beam.Ex/np.max(np.abs(Beam.Ex))
     lamb_x = dx*x_grating
     lamb_y = dy*y_grating
     Ain = np.abs(input_field)
@@ -49,10 +46,7 @@ def slm_hologram(Beam, x_grating, y_grating, method, input_beam, eps, max_range)
 
 def dmd_hologram(Beam, cx, cy, sign):
     #generates a hologram for a dmd
-    if Beam.pol_dim != 1:
-        print('Beam must be scalar (pol_dim=1) to generate holograms')
-        return None
-    U = np.abs(Beam.field)
+    U = np.abs(Beam.Ex)
     U = U/np.amax(U)
     phi = Beam.phase()
     A = np.arcsin(U)

structured_optics-0.3/structured_optics/modes.py

@@ -0,0 +1,299 @@
+from structured_optics.utils import *
+import numpy as np
+from scipy.special import factorial, jn_zeros, jv
+from scipy.optimize import brentq
+
+
+
+#modes implementations
+
+@rotated_mode
+def hg(Beam, N, M, z=0):                                                        
+    #get a HG mode at distance z of order N+M
+    zr = Beam.zr()
+    q0 = 1j*zr
+    q = -z + 1j*zr
+    k = 2*np.pi/Beam.lamb
+    w = Beam.waist*np.sqrt(1 + (z/zr)**2)
+    Cn = np.sqrt(np.sqrt(2/np.pi)*q0/(2**N * factorial(N) * q * Beam.waist))
+    un = Cn*(-np.conjugate(q)/q)**(N/2)*herm(np.sqrt(2)*(Beam.x-Beam.x0)/w, N)*np.exp(-1j*k*(Beam.x-Beam.x0)**2/(2*q))
+    Cm = np.sqrt(np.sqrt(2/np.pi)*q0/(2**M * factorial(M) * q * Beam.waist))
+    um = Cm*(-np.conjugate(q)/q)**(M/2)*herm(np.sqrt(2)*(Beam.y-Beam.y0)/w, M)*np.exp(-1j*k*(Beam.y-Beam.y0)**2/(2*q))
+    F = un*um*np.exp(1j*k*z)
+    return F
+
+@rotated_mode
+def hg_astigmatic(Beam, N, M, wx, wy, z=0):                                                        
+    #get a astigmatic HG mode at distance z of order N+M
+    zrx = np.pi*wx**2/Beam.lamb
+    q0x = 1j*zrx
+    qx = -z + q0x
+    k = 2*np.pi/Beam.lamb
+    wxz = wx*np.sqrt(1 + (z/zrx)**2)
+    Cn = np.sqrt(np.sqrt(2/np.pi)*q0x/(2**N * factorial(N) * qx * wx))
+    un = Cn*(-np.conjugate(qx)/qx)**(N/2)*herm(np.sqrt(2)*(Beam.x-Beam.x0)/wxz, N)*np.exp(-1j*k*(Beam.x-Beam.x0)**2/(2*qx))
+    
+    zry = np.pi*wy**2/Beam.lamb
+    q0y = 1j*zry
+    qy = -z + q0y
+    wyz = wy*np.sqrt(1 + (z/zry)**2)
+    Cm = np.sqrt(np.sqrt(2/np.pi)*q0y/(2**M * factorial(M) * qy * wy))
+    um = Cm*(-np.conjugate(qy)/qy)**(M/2)*herm(np.sqrt(2)*(Beam.y-Beam.y0)/wyz, M)*np.exp(-1j*k*(Beam.y-Beam.y0)**2/(2*qy))
+    F = un*um*np.exp(1j*k*z)
+    return F
+
+@rotated_mode
+def lg(Beam,l,p, z=0):                                    
+    #get a LG mode at distance z of order abs(N) + 2M
+    zr = Beam.zr()
+    k = 2*np.pi/Beam.lamb
+    w = Beam.waist*np.sqrt(1 + (z/zr)**2)
+    R = (z**2 + zr**2)
+    gouy = (np.abs(l) + 2*p + 1)*np.arctan(z/zr)
+    r = np.sqrt((Beam.x-Beam.x0)**2 + (Beam.y-Beam.y0)**2)
+    Cte = np.sqrt(2*factorial(p) / (np.pi*factorial(p+np.abs(l))))
+    F = Cte*(1/w)*(np.sqrt(2)*r/w)**(np.abs(l)) * np.exp(-(r/w)**2) * laguerre(2*(r/w)**2, np.abs(l), p) *\
+        np.exp(1j*(k*z + k*z*(r**2)/(2*R) + l*np.arctan2(Beam.y-Beam.y0,Beam.x-Beam.x0) - gouy))
+    return F
+
+@rotated_mode
+def nbessel(Beam, N, z=0):                       
+    #get a Bessel mode of order N at distance z
+    k = 2*np.pi / Beam.lamb        # wavenumber
+    j01 = jn_zeros(abs(N), 1)[0]  # first zero of J_N
+    # enforce the first zero at r = Beam.waist
+    kr = j01 / Beam.waist
+    kz = np.sqrt(k**2 - kr**2)    
+    r = np.sqrt((Beam.x-Beam.x0)**2 + (Beam.y-Beam.y0)**2)
+    phi = np.arctan2(Beam.y-Beam.y0, Beam.x-Beam.x0)
+    F = (np.exp(1j*kz*z) * jv(abs(N), kr*r) * np.exp(-1j*N*phi))
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def gbessel(Beam, N, r0):                                    
+    #get a gaussian bessel beam of order N at z=0
+    rad = 2*np.pi*Beam.waist**2/r0
+    r = np.sqrt((Beam.x-Beam.x0)**2 + (Beam.y-Beam.y0)**2)
+    F = jv(np.abs(N),rad*r/Beam.waist**2)*\
+        np.exp(-1j*N*np.arctan2(Beam.y-Beam.y0, Beam.x-Beam.x0))*np.exp(-(r/Beam.waist)**2)
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def lg_prod(Beam, N, ls, centers):
+    if centers is None:
+        centers = np.zeros((N, 2))
+    centers = np.asarray(centers)
+    if ls is None:
+        ls = np.ones(N, dtype=int)
+    ls = np.asarray(ls)
+
+    F = np.ones((Beam.Dy, Beam.Dx), dtype='complex128')
+    for i in range(N):
+        x0 = centers[i, 0]
+        y0 = centers[i, 1]
+        l = ls[i]
+        sign = -1 if l < 0 else 1
+        exp_pow = int(np.abs(l))
+        arg = (Beam.x - x0) + 1j * sign * (Beam.y - y0)
+        r2 = (Beam.x - x0)**2 + (Beam.y - y0)**2
+        lg = arg**exp_pow * np.exp(-r2/N / ((Beam.waist)**2))
+        F *= lg
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def frac_oam(Beam, Ma, n_modes, beta, theta_0, z=0):        
+    #get a fractional OAM beam, with OAM Ma (!= integer), by the method of LG supperpositions.
+    mu = Ma%1
+    m = Ma-mu
+    if mu >=0.5:
+        n_min = np.ceil(Ma - n_modes / 2)
+    else: 
+        n_min = np.floor(Ma - n_modes / 2)
+    n_max = n_min + n_modes - 1
+    F = np.zeros((Beam.Dy, Beam.Dx), dtype='complex128')
+    for l in np.arange(int(n_min), int(n_max)+1):
+        coef = np.exp(-1j*mu*beta)*1j*np.exp(1j*(Ma-l)*theta_0)/(2*np.pi*(Ma-l))*np.exp(1j*(m-l)*beta)*(1-np.exp(1j*mu*2*np.pi))
+        F += coef*(lg(Beam,l, 0, z = z)) 
+
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def frac_oam_qs(Beam, Ma, n_modes, beta, theta_0, z=0):
+    #get a quasi-stable fractional oam mode.
+    n_min = np.round(Ma-n_modes/2) 
+    n_max = n_min + n_modes - 1
+    mu = Ma%1
+    m = Ma-mu
+    F = np.zeros((Beam.Dy, Beam.Dx), dtype='complex128')
+    for l in np.arange(int(n_min), int(n_max)+1):
+        coef = np.exp(-1j*mu*beta)*1j*np.exp(1j*(Ma-l)*theta_0)/(2*np.pi*(Ma-l))*np.exp(1j*(m-l)*beta)*(1-np.exp(1j*mu*2*np.pi))
+        p = np.floor((np.abs(Ma) + n_modes/2 -np.abs(l))/2)
+        F += coef*(lg(Beam, l, p, z = z)) 
+
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+
+@rotated_mode
+def IG_even(Beam, p, m, q, z=0):
+    #Even Ince-Gaussian mode IG_p,m^e(x,y)
+    
+    xi, eta = cartesian_to_elliptic(Beam.x-Beam.x0, Beam.y-Beam.y0, q, Beam.waist, z, Beam.lamb)
+
+    Ce_xi  = C_ince(1j*xi,  p, m, q)
+    Ce_eta = C_ince(eta, p, m, q)
+
+    zr = np.pi * Beam.waist**2 / Beam.lamb
+    k = 2*np.pi/Beam.lamb
+    wz = Beam.waist * np.sqrt(1 + (z * Beam.lamb / (np.pi * Beam.waist**2))**2)
+    Rz = (z**2 + zr**2)
+    r2 = (Beam.x-Beam.x0)**2 + (Beam.y-Beam.y0)**2
+    gouy = (p+1)*np.arctan(z / zr)
+    F = Ce_xi * Ce_eta * np.exp(-r2 / (wz**2))*np.exp(1j*(k*z + k*z*(r2)/(2*Rz) - gouy))
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def IG_odd(Beam, p, m, q, z=0):
+    #Odd Ince-Gaussian mode IG_p,m^o(x,y)
+    
+    xi, eta = cartesian_to_elliptic(Beam.x-Beam.x0, Beam.y-Beam.y0, q, Beam.waist, z, Beam.lamb)
+    So_xi  = S_ince(1j*xi,  p, m, q)
+    So_eta = S_ince(eta, p, m, q)
+
+    zr = np.pi * Beam.waist**2 / Beam.lamb
+    k = 2*np.pi/Beam.lamb
+    wz = Beam.waist * np.sqrt(1 + (z * Beam.lamb / (np.pi * Beam.waist**2))**2)
+    Rz = (z**2 + zr**2)
+    r2 = (Beam.x-Beam.x0)**2 + (Beam.y-Beam.y0)**2
+    gouy = (p+1)*np.arctan(z / zr)
+    F = So_xi * So_eta * np.exp(-r2 / (wz**2))*np.exp(1j*(k*z + k*z*(r2)/(2*Rz) - gouy))
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def HInceG(Beam, p, m, q, helicity, z=0):
+    #Hermite-Ince-Gaussian mode HIG_p,m^e(x,y), combining even and odd Ince-Gaussian modes with given helicity.
+    #For some reason IG_odd has a 3*pi/2 phase shift wrt IG_even, so the helicity sign is inverted here.
+    xi, eta = cartesian_to_elliptic(Beam.x-Beam.x0, Beam.y-Beam.y0, q, Beam.waist, z, Beam.lamb)
+
+    Ce_xi  = C_ince(1j*xi,  p, m, q)
+    Ce_eta = C_ince(eta, p, m, q)
+
+    So_xi  = S_ince(1j*xi,  p, m, q)
+    So_eta = S_ince(eta, p, m, q)
+
+    zr = np.pi * Beam.waist**2 / Beam.lamb
+    k = 2*np.pi/Beam.lamb
+    wz = Beam.waist * np.sqrt(1 + (z * Beam.lamb / (np.pi * Beam.waist**2))**2)
+    Rz = (z**2 + zr**2)
+    r2 = (Beam.x-Beam.x0)**2 + (Beam.y-Beam.y0)**2
+    gouy = (p+1)*np.arctan(z / zr)
+    F = (Ce_xi * Ce_eta - helicity*So_xi * So_eta) * np.exp(-r2 / (wz**2))*np.exp(1j*(k*z + k*z*(r2)/(2*Rz) - gouy))
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def circle(Beam, center, radius):
+    if radius == None:
+        radius = Beam.waist
+    #Create a circular mode
+    center_x, center_y = center
+    # Calculate distance from center for all points
+    distances = np.sqrt((Beam.x - center_x)**2 + (Beam.y - center_y)**2)
+    # Set points within radius to 1
+    mask = distances <= radius
+    F = np.zeros((Beam.Dy, Beam.Dx), dtype='complex')
+    F[mask] = 1
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def square(Beam, center, side_length):
+    if side_length == None:
+        side_length = Beam.waist
+    center_x, center_y = center
+    
+    # Calculate square boundaries
+    half_side = side_length / 2.0
+    
+    # Create mask for points inside square
+    mask = (np.abs(Beam.x - center_x) <= half_side) & (np.abs(Beam.y - center_y) <= half_side)
+    
+    # Set points inside square to 1
+    F = np.zeros((Beam.Dy, Beam.Dx), dtype='complex')
+    F[mask] = 1
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def triangle(Beam, center, side_length):
+    if side_length == None:
+        side_length = Beam.waist
+    center_x, center_y = center
+    
+    # Height of equilateral triangle
+    height = side_length * np.sqrt(3) / 2
+    
+    # Vertices for upward-pointing equilateral triangle
+    # Top vertex
+    v0_x = center_x
+    v0_y = center_y + height * 2/3
+    
+    # Bottom left vertex
+    v1_x = center_x - side_length / 2
+    v1_y = center_y - height * 1/3
+    
+    # Bottom right vertex
+    v2_x = center_x + side_length / 2
+    v2_y = center_y - height * 1/3
+    
+    # Use barycentric coordinates to determine if point is inside triangle
+    def sign(px, py, v1_x, v1_y, v2_x, v2_y):
+        return (px - v2_x) * (v1_y - v2_y) - (v1_x - v2_x) * (py - v2_y)
+    
+    # Compute signs for all points
+    d1 = sign(Beam.x, Beam.y, v0_x, v0_y, v1_x, v1_y)
+    d2 = sign(Beam.x, Beam.y, v1_x, v1_y, v2_x, v2_y)
+    d3 = sign(Beam.x, Beam.y, v2_x, v2_y, v0_x, v0_y)
+    
+    # Point is inside if all signs are same
+    has_neg = (d1 < 0) | (d2 < 0) | (d3 < 0)
+    has_pos = (d1 > 0) | (d2 > 0) | (d3 > 0)
+    
+    mask = ~(has_neg & has_pos)
+    F = np.zeros((Beam.Dy, Beam.Dx), dtype='complex')
+    F[mask] = 1
+    F = F/np.sqrt(np.sum(np.abs(F**2))*(4*Beam.nix/Beam.Dx)*(Beam.niy/Beam.Dy))
+    return F
+
+@rotated_mode
+def lp(Beam, l, m, n_core, n_clad, parity="cos"):
+    """Evaluate the scalar LP_l field Psi(r,phi)"""
+    params = get_LP_params(l,m, n_core, n_clad, Beam.waist, Beam.lamb)
+
+    r = np.sqrt(Beam.x**2 + Beam.y**2)
+    PHI = np.arctan2(Beam.y, Beam.x)
+    a = Beam.waist
+    inside = r <= a
+    F = np.zeros_like(r)
+    F[inside] = jv(l, params['u'] * r[inside] / a)
+    scale = jv(l, params['u']) / kv(l, params['v'])
+    F[~inside] = scale * kv(l, params['v'] * r[~inside] / a)
+    ang = np.cos(l * PHI) if parity == "cos" else np.sin(l * PHI)
+    return F * ang
+
+@rotated_mode
+def lp_hel(Beam, l, m, n_core, n_clad):
+    co = lp(Beam, np.abs(l), m, n_core, n_clad, parity = 'cos')
+    si = lp(Beam, np.abs(l), m, n_core, n_clad, parity = 'sin')
+    if l>=0:
+        F = (co + 1j*si)/np.sqrt(2)
+    else:
+        F = (co - 1j*si)/np.sqrt(2)
+    return F
+

structured_optics-0.3/structured_optics/polarization.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+
+HPROJ = np.array([[1,0],[0,0]], dtype='complex')
+VPROJ = np.array([[0,0],[0,1]], dtype='complex')
+DPROJ = np.array([[1,1],[1,1]], dtype='complex')/2
+APROJ = np.array([[1,-1],[-1,1]], dtype='complex')/2
+RPROJ = np.array([[1,1j],[-1j,1]], dtype='complex')/2
+LPROJ = np.array([[1,-1j],[1j,1]], dtype='complex')/2
+
+
+def J_rot(matrix, ang):
+    rt = np.array([[np.cos(ang), -np.sin(ang)],[np.sin(ang), np.cos(ang)]], dtype='complex')
+    return rt@matrix@np.conjugate(rt.T)
+
+def J_phase_retarder(delta):
+    return np.array([[np.exp(1j*delta/2), 0],[0, np.exp(-1j*delta/2)]])
+
+def J_hwp(ang):
+    return J_rot(J_phase_retarder(np.pi/2), ang)
+
+def J_qwp(ang):
+    return J_rot(J_phase_retarder(np.pi/4), ang)
+
+

{structured_optics-0.2 → structured_optics-0.3}/structured_optics/prop_methods.py

@@ -6,8 +6,8 @@ def propagate_fresnel(Beam, z):
     #Propagate the field by a distance z using fresnel integral. convolution method
     k = 2*np.pi/Beam.lamb
     prop = np.exp(-1j * z * (Beam.kx ** 2 + Beam.ky ** 2) / (2 * k) + 1j*k*z)
-    Beam.fourier_field = fft.fft2(Beam.field)
-    Beam.field = fft.ifft2(prop*Beam.fourier_field)
+    Beam.fourier_field = fft.fft2(Beam.field, axes=(-2,-1))
+    Beam.field = fft.ifft2(prop*Beam.fourier_field, axes=(-2,-1))
     return Beam
 
 def propagate_incoherent(Beam, z):
@@ -41,6 +41,6 @@ def propagate_fraunhofer(Beam, z):
         Beam.nix = np.max(Beam.x)
         Beam.niy = np.max(Beam.y)
         C = np.exp(1j*k*z)/(1j*Beam.lamb*z)*np.exp(1j*k*(Beam.x**2 + Beam.y**2)/(2*z))
-        Beam.fourier_field = fft.fft2(Beam.field)
+        Beam.fourier_field = fft.fft2(Beam.field, axes=(-2,-1))
         Beam.field = C*fft.fftshift(Beam.fourier_field)*dx*dy
     return Beam