FFTW.jl for 2D array: Diffusion only happening in 1D

Question

From what I have read, using FFTW.jl / AbstractFFTs.jl's fft(A) when A is a 2D array should perform fft in 2D, not column-wise. Any idea why I am seeing only column-wise diffusion when (I think) I'm adding scaled second spatial derivative to u(t,x), as if using explicit solver for time?

Thank you! I am quite new to this.

code and heatmap screenshot

using Random
using FFTW
using Plots

gr()

N = (100,100)

# initialize with gaussian noise
u = randn(Float16, (N[1], N[2])).*0.4.+0.4
# include square of high concentration to observe diffusion clearly
u[40:50,40:50] .= 3

N = size(x)
L = 100
k1 = fftfreq(51)
k2 = fftfreq(51)
lap_mat = -(k1.^2 + k2.^2)

function lap_fft(x)
    lapF = rfft(x)
    lap = irfft(lap_mat.*lapF, 100)
    return lap
end

# ode stepper or Implicit-Explicit solver
for i in 1:100000
    u+=lap_fft(u)*0.0001
end
# plot state
heatmap(u)

Answer 1

Just because you are performing a real FFT, doesn't mean that you can real inverse fft the result. rfft goes from R -> C. What you can however do is the following:

function lap_fft(x)
    lapF = complex(zeros(100,100));        # only upper half filled
    lapF[1:51,1:100] = rfft(x) .* lap_mat; # R -> C
    return abs.(ifft(lapF));               # C -> R
end

Real FFT to complex frequency domain (only upper half filled because of data redundancy), multiply your filter in frequency domain, inverse FFT into complex image domain and obtain the magnitude abs.(), real part real.() etc.
But honestly, why the hassle with the real fft?

using Random
using FFTW
using Plots

gr()

N = (100,100)

# initialize with gaussian noise
u = randn(Float16, (N[1], N[2])).*0.4.+0.4;
# include square of high concentration to observe diffusion clearly
u[40:50,40:50] .= 3;

N = size(u);
L = 100;
k1 = fftfreq(100);
k2 = fftfreq(100);
tmp = -(k1.^2 + k2.^2);
lap_mat = sqrt.(tmp.*reshape(tmp,1,100));


function lap_fft(x)
    return abs.(ifftshift(ifft(fftshift(ifftshift(fft(fftshift(x))).*lap_mat))));
end

# ode stepper or Implicit-Explicit solver
for i in 1:100000
    u+=lap_fft(u)*0.001;
end
# plot state
heatmap(u)