# Example python program to smooth a 3D function using the convolution theorem

import numpy as np
import matplotlib.pyplot as plt

# Parameters
lgrid = 100. # Length of box
ngrid = 32   # Grid size (should be power of 2)
sig = 5.     # Standard deviation of smoothing

# Co-ordinate values on the array
dgrid = lgrid/ngrid
xgrid = np.linspace(0.,lgrid-dgrid,ngrid)

# Set up original function to be smoothed (top-hat example)
f1 = np.zeros((ngrid,ngrid,ngrid))
for i in range(ngrid):
  for j in range(ngrid):
    for k in range(ngrid):
      if ((xgrid[i] > 0.4*lgrid) & (xgrid[i] < 0.6*lgrid) &
          (xgrid[j] > 0.4*lgrid) & (xgrid[j] < 0.6*lgrid) &
          (xgrid[k] > 0.4*lgrid) & (xgrid[k] < 0.6*lgrid)):
        f1[i,j,k] = 1.

# Set up Gaussian smoothing function in 3D, with periodic boundary
xgrid2 = xgrid
xgrid2[ngrid//2+1:] -= lgrid
f2 = np.zeros((ngrid,ngrid,ngrid))
for i in range(ngrid):
  for j in range(ngrid):
    for k in range(ngrid):
      f2[i,j,k] = np.exp(-(xgrid2[i]**2 + xgrid2[j]**2 + xgrid2[k]**2)/(2.*(sig**2)))

# Normalize Gaussian so it sums to 1.0
f2 /= np.sum(f2)

# Apply smoothing using the convolution theorem
f1fft = np.fft.fftn(f1)
f2fft = np.fft.fftn(f2)
f3 = np.fft.ifftn(f1fft*f2fft)
f3 = np.real(f3)

# Plot a cross-section through the smoothed image
fig = plt.figure()
plt.imshow(f3[ngrid//2,:,:],origin='lower',interpolation='nearest',clim=(0.,1.),cmap='Greys',extent=[0.,lgrid,0.,lgrid])
plt.show()