Find gradient of function in 2-dim and 3-dim via FFT and via `np.gradient`. Then compare their results to analytic solution.

2d-gradient.py

# Run with:
# >>> python 2d-gradient.py <num_sample in one direction> <padding zero>
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import sys

# Plot Settings
fig = plt.figure(figsize=plt.figaspect(0.5))

# Test Settings
num_sample = int(sys.argv[1])   # sample number in x and y direction.
pad = int(sys.argv[2])          # how much times of row and column to pad 0 in FFT.
r_min, r_max = 0.0, 11.0
dx = (r_max - r_min) / (num_sample - 1)
x = np.linspace(r_min, r_max, num_sample, endpoint=True)
y = np.linspace(r_min, r_max, num_sample, endpoint=True)
xx, yy = np.meshgrid(x, y, indexing='ij')

# Test Function and its analytic solution of grad(f)
#f = np.sin(3.0 * xx) + np.sin(5.0 * yy)
#grad_fx_analytic = 3.0 * np.cos(3.0 * xx)
#grad_fy_analytic = 5.0 * np.cos(5.0 * yy)

f = np.sin(3.0 * xx * yy) + np.cos(5.0 * xx + 4.0 * yy)
grad_fx_analytic = 3.0 * yy * np.cos(3.0 * xx * yy) - 5.0 * np.sin(5.0 * xx + 4.0 * yy)
grad_fy_analytic = 3.0 * xx * np.cos(3.0 * xx * yy) - 4.0 * np.sin(5.0 * xx + 4.0 * yy)

# Do 2D - FFT 
#    Step1: padding 0 outside.
pad_array = np.zeros((f.shape[0], 2 * pad * f.shape[1]))
f_pad = np.hstack((f, pad_array))
pad_array = np.zeros((2 * pad * f.shape[0], f_pad.shape[1]))
f_pad = np.vstack((f_pad, pad_array))
f_pad = np.roll(f_pad, (pad * f.shape[0], pad * f.shape[1]), axis=(0, 1))

#    Step2: fft in 0, 1 axes
f_pad_FFT_x = np.fft.rfftn(f_pad, axes=[0])
f_pad_FFT_y = np.fft.rfftn(f_pad, axes=[1])

#    Step3: construct k array to find j.k.FFT(f), which is FFT of grad(f).
k = 2.0 * np.pi * np.fft.rfftfreq(num_sample + 2 * pad * f.shape[0], d=dx) # since f is a cube here, so sample in x and y direction are the same.
kx = np.tile(k, (f_pad_FFT_x.shape[1], 1)).T
ky = np.tile(k, (f_pad_FFT_y.shape[0], 1))
grad_f_pad_FFT_x = 1j * kx * f_pad_FFT_x
grad_f_pad_FFT_y = 1j * ky * f_pad_FFT_y

#    Step4: do inverse FFT, and dig out the interest section.
grad_fx_pad_FFT_IFFT = np.fft.irfftn(grad_f_pad_FFT_x, axes=[0])
grad_fy_pad_FFT_IFFT = np.fft.irfftn(grad_f_pad_FFT_y, axes=[1])

grad_fx_FFT_IFFT = np.roll(grad_fx_pad_FFT_IFFT, (-pad * f.shape[0], -pad * f.shape[1]), axis=(0,1))[0:f.shape[0], 0:f.shape[1]]
grad_fy_FFT_IFFT = np.roll(grad_fy_pad_FFT_IFFT, (-pad * f.shape[0], -pad * f.shape[1]), axis=(0,1))[0:f.shape[0], 0:f.shape[1]]

# Get grad(f) through np.gradient
grad_f = np.gradient(f)

# Compare
error_x = np.sum(np.absolute(grad_f[0] - grad_fx_FFT_IFFT)) / num_sample**2
error_y = np.sum(np.absolute(grad_f[1] - grad_fy_FFT_IFFT)) / num_sample**2
print("NPY_Gradient - NPY_FFT      / SamplePoints (x, y) = %lf, %lf" % (error_x, error_y))

error_x = np.sum(np.absolute(grad_fx_analytic - grad_fx_FFT_IFFT)) / num_sample**2
error_y = np.sum(np.absolute(grad_fy_analytic - grad_fy_FFT_IFFT)) / num_sample**2
print("Analytic     - NPY_FFT      / SamplePoints  (x, y) = %lf , %lf" % (error_x, error_y))

error_x = np.sum(np.absolute(grad_fx_analytic - grad_f[0])) / num_sample**2
error_y = np.sum(np.absolute(grad_fy_analytic - grad_f[1])) / num_sample**2
print("Analytic     - NPY_Gradient / SamplePoints (x, y) = %lf, %lf" % (error_x, error_y))

# Plot result.
ax = fig.add_subplot(1,3,3)
q = ax.quiver(x, y, grad_fx_analytic, grad_fy_analytic )
ax.quiverkey(q, X=0.3, Y=0.3, U=1, label=r"$\nabla f$", labelpos='E', coordinates="figure")
ax.set_title("grad(f) analytic")
ax.set_xlabel("x")
ax.set_ylabel("y")

ax = fig.add_subplot(1,3,2)
q = ax.quiver(x, y, grad_fx_FFT_IFFT,grad_fy_FFT_IFFT )
ax.quiverkey(q, X=0.3, Y=0.3, U=1, label=r"$\nabla f$", labelpos='E', coordinates="figure")
ax.set_title("grad(f) result from FFT")
ax.set_xlabel("x")
ax.set_ylabel("y")

ax = fig.add_subplot(1,3,1)
q = ax.quiver(x, y, grad_f[0], grad_f[1])
ax.quiverkey(q, X=0.3, Y=0.3, U=1, label=r"$\nabla f$", labelpos='E', coordinates="figure")
ax.set_title("grad(f) result from np.gradient")
ax.set_xlabel("x")
ax.set_ylabel("y")

plt.show()

3d-dft-gradient.py

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import sys

# Plot Settings
fig = plt.figure(figsize=plt.figaspect(0.5))

# Test Settings
num_sample = int(sys.argv[1])   # sample number in x and y direction.
pad = int(sys.argv[2])          # how much times of row and column to pad 0 in FFT.
r_min, r_max = 0.0, 11.0
dx = (r_max - r_min) / (num_sample - 1)
x = np.linspace(r_min, r_max, num_sample, endpoint=True)
y = np.linspace(r_min, r_max, num_sample, endpoint=True)
z = np.linspace(r_min, r_max, num_sample, endpoint=True)
xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')

# Test Function and its analytic solution of grad(f)
f = np.sin(3.0 * xx) + np.sin(4.0 * yy) + np.sin(5.0 * zz)
grad_fx_analytic = 3.0 * np.cos(3.0 * xx)
grad_fy_analytic = 4.0 * np.cos(4.0 * yy)
grad_fz_analytic = 5.0 * np.cos(5.0 * zz)

# Do 3D - FFT 
#    Step1: padding 0 outside.
pad_array = np.zeros((2 * pad * f.shape[0], f.shape[1], f.shape[2] ))
f_pad = np.concatenate((f, pad_array), axis=0)
pad_array = np.zeros((f_pad.shape[0], 2 * pad * f.shape[1], f.shape[2]))
f_pad = np.concatenate((f_pad, pad_array), axis=1)
pad_array = np.zeros((f_pad.shape[0], f_pad.shape[1], 2 * pad * f.shape[2]))
f_pad = np.concatenate((f_pad, pad_array), axis=2)
f_pad = np.roll(f_pad, (pad * f.shape[0], pad * f.shape[1], pad * f.shape[2]), axis=(0, 1, 2))

#    Step2: fft in 0, 1, 2 axes
f_pad_FFT_x = np.fft.rfftn(f_pad, axes=[0])
f_pad_FFT_y = np.fft.rfftn(f_pad, axes=[1])
f_pad_FFT_z = np.fft.rfftn(f_pad, axes=[2])

#    Step3: construct k array to find j.k.FFT(f), which is FFT of grad(f).
k = 2.0 * np.pi * np.fft.rfftfreq(num_sample + 2 * pad * f.shape[0], d=dx) # since f is a cube here, so sample in x,y,z direction are the same.
kx = np.tile( k.reshape((k.shape[0], 1, 1)), (1, f_pad_FFT_x.shape[1], f_pad_FFT_x.shape[2]))
ky = np.tile( k.reshape((1, k.shape[0], 1)), (f_pad_FFT_y.shape[0], 1, f_pad_FFT_y.shape[2]))
kz = np.tile( k.reshape((1, 1, k.shape[0])), (f_pad_FFT_z.shape[0], f_pad_FFT_z.shape[1], 1))

grad_f_pad_FFT_x = 1j * kx * f_pad_FFT_x
grad_f_pad_FFT_y = 1j * ky * f_pad_FFT_y
grad_f_pad_FFT_z = 1j * kz * f_pad_FFT_z

#    Step4: do inverse FFT, and dig out the interest section.
grad_fx_pad_FFT_IFFT = np.fft.irfftn(grad_f_pad_FFT_x, axes=[0])
grad_fy_pad_FFT_IFFT = np.fft.irfftn(grad_f_pad_FFT_y, axes=[1])
grad_fz_pad_FFT_IFFT = np.fft.irfftn(grad_f_pad_FFT_z, axes=[2])

grad_fx_FFT_IFFT = np.roll(grad_fx_pad_FFT_IFFT, (-pad * f.shape[0], -pad * f.shape[1], -pad * f.shape[2]), axis=(0,1,2))[0:f.shape[0], 0:f.shape[1], 0:f.shape[2]]
grad_fy_FFT_IFFT = np.roll(grad_fy_pad_FFT_IFFT, (-pad * f.shape[0], -pad * f.shape[1], -pad * f.shape[2]), axis=(0,1,2))[0:f.shape[0], 0:f.shape[1], 0:f.shape[2]]
grad_fz_FFT_IFFT = np.roll(grad_fz_pad_FFT_IFFT, (-pad * f.shape[0], -pad * f.shape[1], -pad * f.shape[2]), axis=(0,1,2))[0:f.shape[0], 0:f.shape[1], 0:f.shape[2]]

# Get grad(f) through np.gradient
grad_f = np.gradient(f)

# Compare
error_x = np.sum(np.absolute(grad_f[0] - grad_fx_FFT_IFFT)) / num_sample**3
error_y = np.sum(np.absolute(grad_f[1] - grad_fy_FFT_IFFT)) / num_sample**3
error_z = np.sum(np.absolute(grad_f[2] - grad_fz_FFT_IFFT)) / num_sample**3
print("NPY_Gradient - NPY_FFT      / SamplePoints  (x, y, z) = %.3lf , %.3lf , %.3lf" % (error_x, error_y, error_z))

error_x = np.sum(np.absolute(grad_fx_analytic - grad_fx_FFT_IFFT)) / num_sample**3
error_y = np.sum(np.absolute(grad_fy_analytic - grad_fy_FFT_IFFT)) / num_sample**3
error_z = np.sum(np.absolute(grad_fz_analytic - grad_fz_FFT_IFFT)) / num_sample**3
print("Analytic     - NPY_FFT      / SamplePoints  (x, y, z) = %.3lf , %.3lf , %.3lf" % (error_x, error_y, error_z))

error_x = np.sum(np.absolute(grad_fx_analytic - grad_f[0])) / num_sample**3
error_y = np.sum(np.absolute(grad_fy_analytic - grad_f[1])) / num_sample**3
error_z = np.sum(np.absolute(grad_fz_analytic - grad_f[2])) / num_sample**3
print("Analytic     - NPY_Gradient / SamplePoints  (x, y, z) = %.3lf , %.3lf, %.3lf" % (error_x, error_y, error_z))