get_ipython().run_line_magic('matplotlib', 'notebook')
import numpy as np
import matplotlib.pyplot as plt

# Create a 512x512 black grid
imag = np.zeros((512, 512))

# Add some slits with different intensities
slit = 3
while slit > 0:
    X1 = slit * 100 + np.random.randint(50)
    X2 = X1 + 10 + np.random.randint(50)
    Y1 = 100 + np.random.randint(300)
    Y2 = Y1 + 10 + np.random.randint(50)
    imag[X1:X2, Y1:Y2] = 100 + np.random.randint(156)
    slit -= 1

# Plot the image of the slits as a mesh
I = np.arange(512)
J = np.arange(512)

fig = plt.figure(figsize=(8, 16))
ax = fig.add_subplot(311, projection='3d')
X, Y = np.meshgrid(I, J)
Z = imag
ax.plot_surface(X, Y, Z, cmap='plasma')
ax.set_title('Image of the slits')

# Produce a Gaussian point spread function
spread = np.zeros((512, 512))
XS, YS = np.meshgrid(np.arange(61), np.arange(61))
XS = XS - 30.5
YS = YS - 30.5
gpeak = np.exp((-XS**2 - YS**2) / 128)
# Normalize
S = np.sum(gpeak)
gpeak = gpeak/S
spread[227:227 + gpeak.shape[0], 227:227 + gpeak.shape[1]] = gpeak

# Plot the point spread function
ax = fig.add_subplot(312, projection='3d')
Z = spread
ax.plot_surface(X, Y, Z, cmap='plasma')
#ax.plot_wireframe(X, Y, Z, color = "red")
ax.set_title('Image of the Gaussian point spread function')

# Find the Fourier transform of the image
imagfft = np.fft.fft2(imag)

# Find the Fourier transform of the point spread function
spreadfft = np.fft.fft2(spread)

# Find the Fourier transform of the convolution
convfft = np.multiply(imagfft, spreadfft)
convimag = np.fft.ifftshift(np.fft.ifft2(convfft))

# Plot the convolution result
ax = fig.add_subplot(313, projection='3d')
Z = abs(convimag)
ax.plot_surface(X, Y, Z, cmap='plasma')
ax.set_title('Convolution')
plt.show