Angular spectrum method using python

Question

I am trying to numerically propagate a given (electric) field using the angular spectrum method. For this I am following "Principles and Applications of Fourier Optics" (Robert K. Tyson) Chapter 3, Page 2

I tried to recreate the maths using the following code

import numpy as np
import imageio

U = imageio.imread("ap.png")[:,:, 0] # load circular aperture
_lambda = 800e-9

def propagate2(self,z):
    A = np.fft.fft2(U, norm="ortho") # 2D FFT 
    alpha = np.fft.fftfreq(U.shape[0])*_lambda # direction cosine in x direction
    beta = np.fft.fftfreq(U.shape[1])*_lambda # direction cosine in x direction
    gamma = np.zeros([alpha.shape[0], beta.shape[0]])
    k = 2*np.pi/_lambda # wavevector

    for i,j in itertools.product(range(alpha.shape[0]), range(beta.shape[0])): # determine phase value for each (i,j)
        if alpha[i]**2+beta[j]**2 < 1:
            gamma[i,j] = np.sqrt(1-alpha[i]**2-beta[j]**2)
        else:
            gamma[i,j] = 1j*np.sqrt(np.abs(1-alpha[i]**2-beta[j]**2))
    phi = np.exp(1j*k*z*gamma)
    field = np.fft.ifft2(A*phi, norm="ortho") # 2D IFFT
    return field

This code should produce the usual double slit diffraction pattern, however (as to be seen below) won't produce and diffraction at all.

I am fairly certain that there is some problem with my alpha and beta values, however I can't seem to find it. Any help is highly appreciated.

ap.png:

Answer 1

I made a complete implementation of the Angular Spectrum Method using Python: https://github.com/rafael-fuente/Diffraction-Simulations--Angular-Spectrum-Method

To use it for your example, you need to clone the repository and put your double slit image in the path ./apertures/double_slit.png

Then, from the repository folder run the following script:

from diffractsim import MonochromaticField, mm, nm, cm, ApertureFromImage


F = MonochromaticField(
    wavelength=632.8 * nm, extent_x=5.6 * mm, extent_y=5.6 * mm, Nx=1024, Ny=1024
)

F.add(ApertureFromImage(
   "./apertures/double_slit.jpg", image_size=(4.0* mm, 4.0 * mm), simulation = F))
rgb = F.compute_colors_at(20*cm)
F.plot_colors(rgb, xlim=[-4.5, 4.5], ylim=[-4.5, 4.5])

This is the result of the script of the near field (20 cm) for different wavelengths:

As we can see in the plots, the higher the light's wavelength, the wider the length of the interference fringes.

You can also compute the diffraction pattern with a broad spectrum, like white light, with the PolychromaticField class:

from diffractsim import PolychromaticField, cf, mm, cm, ApertureFromImage


F = PolychromaticField(
    spectrum=2 * cf.illuminant_d65, extent_x=5.6 * mm, extent_y=5.6 * mm, Nx=1024, Ny=1024
)

F.add(ApertureFromImage(
    "./apertures/double_slit.png", image_size=(4.0* mm, 4.0 * mm), simulation = F))

rgb = F.compute_colors_at(20*cm)
F.plot_colors(rgb, xlim=[-4.5, 4.5], ylim=[-4.5, 4.5])

which results in:

The problem with your code is that you are not using the Fast Fourier Transform correctly.

The core of my implementation (propagation of the Angular Spectrum) is in the MonochromaticField propagate method, in monochromatic_simulator.py:

def propagate(self, z):
    self.z += z

    # compute angular spectrum
    fft_c = fft2(self.E)
    c = fftshift(fft_c)

    kx = 2*bd.pi*bd.fft.fftshift(bd.fft.fftfreq(simulation.Nx, d = simulation.dx))
    ky = 2*bd.pi*bd.fft.fftshift(bd.fft.fftfreq(simulation.Ny, d = simulation.dy))
    kx, ky = np.meshgrid(kx, ky)
    kz = np.sqrt((2 * np.pi / self.?) ** 2 - kx ** 2 - ky ** 2)

    # propagate the angular spectrum a distance z
    E = ifft2(ifftshift(c * np.exp(1j * kz * z)))

    # compute Field Intensity
    self.I = np.real(E * np.conjugate(E)) 

You can see that I used the method fftshift after performing the FFT to match the definition of the Fourier Transform.

Also, if you want to discard the evanescent fields as you seem you want to do, it's a much better approach to use numpy.where instead of looping every pixel, since Python loops are much slower:

    # propagate the angular spectrum a distance z
    mask = (2*np.pi/self.?)**2 - kx**2 - ky**2 > 0
    A = np.where(mask, c*np.exp(1j*kz * z),  0)
    E = ifft2(ifftshift(A))

This code should replace the last lines in the compute_colors_at method. Hope that this helps!

Answer 2

It can be tricky to get this right. Here it is:

u = ... # this is your 2D complex field that you wish to propagate
z = ... # this is the distance by which you wish to propagate

dx, dy = 1, 1 # or whatever

wavelen = 1 # or whatever
wavenum = 2 * np.pi / wavelen
wavenum_sq = wavenum * wavenum

kx = np.fft.fftfreq(u.shape[0], dx / (2 * np.pi))
ky = np.fft.fftfreq(u.shape[1], dy / (2 * np.pi))

# this is just for broadcasting, the "indexing" argument prevents NumPy
# from doing a transpose (default meshgrid behaviour)
kx, ky = np.meshgrid(kx, ky, indexing = 'ij', sparse = True)

kz_sq = kx * kx + ky * ky

# we zero out anything for which this is not true, see many textbooks on
# optics for an explanation
mask = wavenum * wavenum > kz_sq

g = np.zeros((len(kx), len(ky)), dtype = np.complex_)
g[mask] = np.exp(1j * np.sqrt(wavenum_sq - kz_sq[mask]) * z)

res = np.fft.ifft2(g * np.fft.fft2(u)) # this is the result

You probably want to pad u to prevent wrap around. In that case, simply compute g with the shape doubled then slice the result.