# 3D stable fluids algorithm based on FFT

This code is 3D extension of the code from my answer here. As well known the stably fluids algorithm is some kind of predictor corrector algorithm - see my answer here. This algorithm includes 3 steps - advection, diffusion and projection. In Fourier space the diffusion and projection can be combine in one step as follows

$$(u-u_n)/dt+(u.\nabla) u=0, (u_{n+1}-u)/dt+\nabla p-\mu \nabla^2 u=0$$

Apply $$\nabla .$$ to the last equation and use $$\nabla .u_n=0$$, then we have

$$\nabla^2 p-\nabla. u/dt-\mu \nabla^2( \nabla .u)=0$$

Using FFT we can transform last 2 linear equations to the system of algebraic equations and express $$u_{n+1}$$ Fourier image as

$$\vec{u}_{n+1}=(\vec{u}-(\vec{k}.\vec{u})\vec{k}/k^2)(1-\mu dtk^2)$$

where $$\vec{k}=(k_x,k_y,k_z),k^2=k_x^2+k_y^2+k_z^2$$. The code in a case of 3D flow around cylinder can be written as follows

n = 64; dt = 0.3; mu = 0.0001; nt = 300; mat =
Table[E^(
2 \[Pi] I (r - 1) (s - 1)/n), {r, 1, n}, {s, 1,
n}]; r = (Log[Flatten[mat]]/I) // DeleteDuplicates//N;
v = Table[{0., 0., 0.}, {n}, {n}, {n}];u0=.1;

Do[Do[If[i < 1 + n/16, v[[i, j, jz]] = {u0, 0., 0.}];
If[(i - n/4)^2 + (j - n/2)^2 < (n/16)^2,
v[[i, j, jz]] = {0., 0., 0}], {i, n}, {j, n}, {jz, n}];
{ui, vi, wi} =
Table[ListInterpolation[v[[All, All, All, i]]], {i, 3}];
v = Table[{i2, j2, jz2} = {i, j, jz} - n dt v[[i, j, jz]];
{ui[i2, j2, jz2], vi[i2, j2, jz2], wi[i2, j2, jz2]} // Quiet, {i,
n}, {j, n}, {jz, n}];
{uf, vf, wf} = Table[Fourier[v[[All, All, All, i]]], {i, 3}];
v = Table[{uf[[i, j, jz]], vf[[i, j, jz]], wf[[i, j, jz]]}, {i,
n}, {j, n}, {jz, n}];
v = Table[x = r[[i]];
y = r[[j]]; z = r[[jz]];
k = x^2 + y^2 + z^2;
If[k >
0, (v[[i, j,
jz]] - (v[[i, j, jz]] . {x, y, z}) {x, y, z}/k) (1 -
mu dt k) , v[[i, j, jz]]], {i, n}, {j, n}, {jz, n}];
{ur, vr, wr} =
Table[InverseFourier[v[[All, All, All, i]]] // Re, {i, 3}];
v = Table[{ur[[i, j, jz]], vr[[i, j, jz]], wr[[i, j, jz]]}, {i,
n}, {j, n}, {jz, n}]; vs[t] = v;, {t, 1, nt}]; // AbsoluteTiming


Flow visualization in 2D at z=n/2

lst2D = Table[
ImageRotate[
Show[ListDensityPlot[vs[t][[All, All, n/2, 1]],
ColorFunction -> Hue, PlotRange -> All, Frame -> False,
ImageSize -> Tiny],
Graphics[{Blue, Disk[{n/2, n/4}, n/16]}]], -Pi/2], {t, 15, 300,
5}];ListAnimate[lst2D]


3D flow visualization

lst3D =
Table[Show[
ListDensityPlot3D[vs[t][[All, All, All, 1]], ColorFunction -> Hue,
PlotRange -> All, Boxed -> False, ImageSize -> Tiny,
Axes -> False, ViewPoint -> {-2., 1., 1.},
OpacityFunction -> 0.05],
Graphics3D[{Blue,
Cylinder[{{0, n/2, n/4}, {n, n/2, n/4}}, n/16]}]], {t, 30, 300,
3}];ListAnimate[lst3D]


The code is working fine, but very slow. I try to compile code but without success. How can we improve computation time?

Update 1. As it proposed in the comment by yarchik we can use trilinear interpolation instead of ListInterpolation. The corresponding module advect was made for 3D flow simulation here. The code with this module is follows

Clear["Global*"]

n = 64; dt = 0.3; mu = 0.0001; nt = 100; mat =
Table[E^(
2 \[Pi] I (r - 1) (s - 1)/n), {r, 1, n}, {s, 1,
n}]; r = (Log[Flatten[mat]]/I) // DeleteDuplicates;
wr = vr = ur = Table[0, {n}, {n}, {n}]; u0 = .1;

advect[n_, d0_, u1_, v1_, w1_, dt_] :=
Module[{x, y, z, d1, dt0, i, j, k, i0, i1, j0, j1, k0, k1, s0, s1,
t0, t1, p1, p0, d00, d10, d01, d11, cd0, cd1, xd, yd, zd, nx, ny,
nz}, nx = n; ny = n; nz = n; d1 = Table[0, {nx}, {ny}, {nz}];
dt0 = dt n;
Do[Do[
Do[x = i - dt0 u1[[i, j, k]]; y = j - dt0 v1[[i, j, k]];
z = k - dt0 w1[[i, j, k]];
i0 =
Which[x <= 1, 1, 1 < x < nx - 1, Floor[x], True, nx - 1];
i1 = i0 + 1;
j0 =
Which[y <= 1, 1, 1 < y < ny - 1, Floor[y], True, ny - 1];
j1 = j0 + 1;
k0 = Which[z <= 1, 1, 1 < z < nz - 1, Floor[z], True,
nz - 1];
k1 = k0 + 1;(*Trilinear interpolation*)xd = x - i0;
yd = y - j0; zd = z - k0;
d00 = d0[[i0, j0, k0]] (1 - xd) + d0[[i1, j0, k0]] xd;
d01 = d0[[i0, j0, k1]] (1 - xd) + d0[[i1, j0, k1]] xd;
d10 = d0[[i0, j1, k0]] (1 - xd) + d0[[i1, j1, k0]] xd;
d11 = d0[[i0, j1, k1]] (1 - xd) + d0[[i1, j1, k1]] xd;
cd0 = d00 (1 - yd) + d10 yd; cd1 = d01 (1 - yd) + d11 yd;
d1[[i, j, k]] = cd0 (1 - zd) + cd1 zd;
, {k, 2, nz - 1}];, {j, 2, ny - 1}];, {i, 1, nx}]; d1];
Do[Do[If[
i < 1 + n/16, {ur[[i, j, jz]], vr[[i, j, jz]],
wr[[i, j, jz]]} = {u0, 0, 0}];
If[(i - n/4)^2 + (j - n/2)^2 < (n/16)^2, {ur[[i, j, jz]],
vr[[i, j, jz]], wr[[i, j, jz]]} = {0, 0, 0}], {i, n}, {j,
n}, {jz, n}];
ui = advect[n, ur, ur, vr, wr, dt];
vi = advect[n, vr, ur, vr, wr, dt];
wi = advect[n, wr, ur, vr, wr, dt];
uf = Fourier[ui]; vf = Fourier[vi]; wf = Fourier[wi];
v = Table[{uf[[i, j, jz]], vf[[i, j, jz]], wf[[i, j, jz]]}, {i,
n}, {j, n}, {jz, n}];
v = Table[x = r[[i]];
y = r[[j]]; z = r[[jz]];
k = x^2 + y^2 + z^2;
If[k >
0, (v[[i, j,
jz]] - (v[[i, j, jz]] . {x, y, z}) {x, y, z}/k) (1 -
mu dt k) , v[[i, j, jz]]], {i, n}, {j, n}, {jz, n}];
{ur, vr, wr} =
Table[InverseFourier[v[[All, All, All, i]]] // Re, {i, 3}];
us[t] = ur; vs[t] = vr; ws[t] = wr;, {t, 1, nt}] // AbsoluteTiming



Visualization

Show[ListDensityPlot3D[
Table[Norm[{ur[[i, j, jz]], vr[[i, j, jz]], wr[[i, j, jz]]}], {i,
n}, {j, n}, {jz, n}], AxesLabel -> {"z", "y", "x"},
ColorFunction -> Hue, OpacityFunction -> 0.05,
ViewPoint -> {-2., 1., 1.}],
Graphics3D[{Blue, Cylinder[{{0, n/2, n/4}, {n, n/2, n/4}}, n/16]}]]


This code also is very slow and can't be compiled due to mixture complex and real variables.

Update 2 Nevertheless we can compile part of code using idea from xzczd answer. In this code we add separate module bcu for boundary condition, and onestep to make advection step with boundary conditions

Clear["Global*"]

mu = 1./10000; U0 = 0.; V0 = 0.; W0 = 0.; n = 64; nx = n; ny =
nz = n; {nx0, ny0, R0} = {n/4, n/2, n/16}; dt = .3; uinfl = .1;
n1 = n + 1; nt = 100;
u0 = Table[U0, {nx}, {ny}, {nz}];
v0 = Table[V0, {nx}, {ny}, {nz}]; w0 =
Table[W0, {nx}, {ny}, {nz}]; mat =
Table[E^(
2 \[Pi] I (r - 1) (s - 1)/n), {r, 1, n}, {s, 1,
n}]; r = (Log[Flatten[mat]]/I) // DeleteDuplicates // N;
Do[u0[[i, j, jz]] = If[i < 1 + n/16, uinfl, 0];, {i, n}, {j, n}, {jz,
n}];

bcu[nx_, ny_, nz_, in_, up_, ud_, ul_, ur_, ub_] :=
Module[{bd = ub},
Do[bd[[nx, i, j]] = bd[[nx - 1, i, j]];
bd[[1, i, j]] = bd[[2, i, j]];, {i, 2, ny - 1}, {j, 2, nz - 1}];
Do[bd[[i, 1, j]] = ud;
bd[[i, ny, j]] = up; bd[[i, j, 1]] = ul;
bd[[i, j, nz]] = ur;, {i, 1, nx}, {j, 1, ny}];
bd];
advect[n_, nx_, ny_, nz_, d0_, u_, v_, w_, dt_] :=
Module[{x, y, z, d1, dt0, i0, i1, j0, j1, k0, k1, s0, s1, t0, t1,
p1, p0, d00, d10, d01, d11, cd0, cd1, xd, yd, zd},
d1 = ConstantArray[0, {nx, ny, nz}]; dt0 = dt n;
Do[Do[
Do[x = i - dt0 u[[i, j, k]]; y = j - dt0 v[[i, j, k]];
z = k - dt0 w[[i, j, k]];
i0 =
Which[x <= 1, 1, 1 < x < nx - 1, Floor[x], True, nx - 1];
i1 = i0 + 1;
j0 =
Which[y <= 1, 1, 1 < y < ny - 1, Floor[y], True, ny - 1];
j1 = j0 + 1;
k0 = Which[z <= 1, 1, 1 < z < nz - 1, Floor[z], True,
nz - 1];
k1 = k0 + 1;(*Trilinear interpolation*)xd = x - i0;
yd = y - j0; zd = z - k0;
d00 = d0[[i0, j0, k0]] (1 - xd) + d0[[i1, j0, k0]] xd;
d01 = d0[[i0, j0, k1]] (1 - xd) + d0[[i1, j0, k1]] xd;
d10 = d0[[i0, j1, k0]] (1 - xd) + d0[[i1, j1, k0]] xd;
d11 = d0[[i0, j1, k1]] (1 - xd) + d0[[i1, j1, k1]] xd;
cd0 = d00 (1 - yd) + d10 yd; cd1 = d01 (1 - yd) + d11 yd;
d1[[i, j, k]] = cd0 (1 - zd) + cd1 zd;
, {k, 2, nz - 1}];, {j, 2, ny - 1}];, {i, 1, nx}]; d1];

onestep[n_, nx_, ny_, nz_, nx0_, ny0_, R0_, uin_, vin_, win_, dt_,
uinfl_] := Module[{u0, v0, w0},
u0 = Table[0., {nx}, {ny}, {nz}];
v0 = Table[0., {nx}, {ny}, {nz}];
w0 = Table[0., {nx}, {ny}, {nz}];
u0 = uin; v0 = vin; w0 = win;
Do[u0[[i, j, jz]] = If[i < 1 + n/16, uinfl, u0[[i, j, jz]]];, {i,
n}, {j, n}, {jz, n}];
Do[Do[u0[[i, j, k]] = 0;, {i,
nx0 - Round[Sqrt[R0^2 - (j - ny0)^2]],
nx0 + Round[Sqrt[R0^2 - (j - ny0)^2]], 1}], {j, ny0 - R0,
ny0 + R0, 1}, {k, 1, nz, 1}];
u0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, u0];
v0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, v0];
w0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, w0];
u0 = advect[n, nx, ny, nz, u0, u0, v0, w0, dt];
v0 = advect[n, nx, ny, nz, v0, u0, v0, w0, dt];
w0 = advect[n, nx, ny, nz, w0, u0, v0, w0, dt];
Do[Do[u0[[i, j, k]] = 0; v0[[i, j, k]] = 0;
w0[[i, j, k]] = 0;, {i, nx0 - Round[Sqrt[R0^2 - (j - ny0)^2]],
nx0 + Round[Sqrt[R0^2 - (j - ny0)^2]], 1}], {j, ny0 - R0,
ny0 + R0, 1}, {k, 1, nz, 1}];
u0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, u0];
v0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, v0];
w0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, w0];

{u0, v0, w0}];

cf = With[{cg = CompileGetElement, hp = HoldPattern,
dv = DownValues},
Hold@Compile[{{u0argu, _Real, 3}, {v0argu, _Real,
3}, {w0argu, _Real,
3}, {n, _Integer}, {nx, _Integer}, {ny, _Integer}, {nz, \
_Integer}, {nx0, _Integer}, {ny0, _Integer}, {R0, _Integer}, {dt, \
_Real}, {uinfl, _Real}},
Module[{u0 = u0argu, v0 = v0argu, w0 = w0argu, uu, vv, ww},
{u0, v0, w0} =
onestep[n, nx, ny, nz, nx0, ny0, R0, u0, v0, w0, dt,
uinfl];
{u0, v0, w0}], CompilationTarget -> "C",
RuntimeOptions -> "Speed"] /. dv@onestep /.
Flatten[dv /@ {advect, bcu}] /.
hp@ConstantArray[c_, {i_, j_, kc_}] :>
Table[0., {i}, {j}, {kc}] /. hp@Part[a__] :> cg[a] /.
hp[cg[a__] = rhs_] :> (Part[a] = rhs) //
ReleaseHold];


With this compilation code about 7 times faster

 Do[{ui, vi, wi} =
cf[u0, v0, w0, n, nx, ny, nz, nx0, ny0, R0, dt, uinfl];
uf = Fourier[ui]; vf = Fourier[vi]; wf = Fourier[wi];
v = Table[{uf[[i, j, jz]], vf[[i, j, jz]], wf[[i, j, jz]]}, {i,
n}, {j, n}, {jz, n}];
v = Table[x = r[[i]];
y = r[[j]]; z = r[[jz]];
k = x^2 + y^2 + z^2;
If[k >
0, (v[[i, j,
jz]] - (v[[i, j, jz]] . {x, y, z}) {x, y, z}/k) (1 -
mu dt k) , v[[i, j, jz]]], {i, n}, {j, n}, {jz, n}];
{u0, v0, w0} =
Table[InverseFourier[v[[All, All, All, i]]] // Re, {i, 3}];
us[t] = u0; vs[t] = v0; ws[t] = w0;, {t, 1, nt}] // AbsoluteTiming


Update 3. We also can compile the complex part of the code as follows

Clear["Global*"]

mu = 1./10000; U0 = 0.; V0 = 0.; W0 = 0.; n = 64; nx = n; ny =
nz = n; {nx0, ny0, R0} = {n/4, n/2, n/16}; dt = .3; uinfl = .1;
n1 = n + 1; nt = 300;
u0 = Table[U0, {nx}, {ny}, {nz}];
v0 = Table[V0, {nx}, {ny}, {nz}]; w0 =
Table[W0, {nx}, {ny}, {nz}]; mat =
Table[E^(
2 \[Pi] I (r - 1) (s - 1)/n), {r, 1, n}, {s, 1,
n}]; r = (Log[Flatten[mat]]/I) // DeleteDuplicates // N;
Do[u0[[i, j, jz]] = If[i < 1 + n/16, uinfl, 0];, {i, n}, {j, n}, {jz,
n}];

bcu[nx_, ny_, nz_, in_, up_, ud_, ul_, ur_, ub_] :=
Module[{bd = ub},
Do[bd[[nx, i, j]] = bd[[nx - 1, i, j]];
bd[[1, i, j]] = bd[[2, i, j]];, {i, 2, ny - 1}, {j, 2, nz - 1}];
Do[bd[[i, 1, j]] = ud;
bd[[i, ny, j]] = up; bd[[i, j, 1]] = ul;
bd[[i, j, nz]] = ur;, {i, 1, nx}, {j, 1, ny}];
bd];
advect[n_, nx_, ny_, nz_, d0_, u_, v_, w_, dt_] :=
Module[{x, y, z, d1, dt0, i0, i1, j0, j1, k0, k1, s0, s1, t0, t1,
p1, p0, d00, d10, d01, d11, cd0, cd1, xd, yd, zd},
d1 = ConstantArray[0, {nx, ny, nz}]; dt0 = dt n;
Do[Do[
Do[x = i - dt0 u[[i, j, k]]; y = j - dt0 v[[i, j, k]];
z = k - dt0 w[[i, j, k]];
i0 =
Which[x <= 1, 1, 1 < x < nx - 1, Floor[x], True, nx - 1];
i1 = i0 + 1;
j0 =
Which[y <= 1, 1, 1 < y < ny - 1, Floor[y], True, ny - 1];
j1 = j0 + 1;
k0 = Which[z <= 1, 1, 1 < z < nz - 1, Floor[z], True,
nz - 1];
k1 = k0 + 1;(*Trilinear interpolation*)xd = x - i0;
yd = y - j0; zd = z - k0;
d00 = d0[[i0, j0, k0]] (1 - xd) + d0[[i1, j0, k0]] xd;
d01 = d0[[i0, j0, k1]] (1 - xd) + d0[[i1, j0, k1]] xd;
d10 = d0[[i0, j1, k0]] (1 - xd) + d0[[i1, j1, k0]] xd;
d11 = d0[[i0, j1, k1]] (1 - xd) + d0[[i1, j1, k1]] xd;
cd0 = d00 (1 - yd) + d10 yd; cd1 = d01 (1 - yd) + d11 yd;
d1[[i, j, k]] = cd0 (1 - zd) + cd1 zd;
, {k, 2, nz - 1}];, {j, 2, ny - 1}];, {i, 1, nx}]; d1];

onestep[n_, nx_, ny_, nz_, nx0_, ny0_, R0_, uin_, vin_, win_, dt_,
uinfl_] := Module[{u0, v0, w0},
u0 = Table[0., {nx}, {ny}, {nz}];
v0 = Table[0., {nx}, {ny}, {nz}];
w0 = Table[0., {nx}, {ny}, {nz}];
u0 = uin; v0 = vin; w0 = win;
Do[u0[[i, j, jz]] = If[i < 1 + n/16, uinfl, u0[[i, j, jz]]];, {i,
n}, {j, n}, {jz, n}];
Do[Do[u0[[i, j, k]] = 0;, {i,
nx0 - Round[Sqrt[R0^2 - (j - ny0)^2]],
nx0 + Round[Sqrt[R0^2 - (j - ny0)^2]], 1}], {j, ny0 - R0,
ny0 + R0, 1}, {k, 1, nz, 1}];
u0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, u0];
v0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, v0];
w0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, w0];
u0 = advect[n, nx, ny, nz, u0, u0, v0, w0, dt];
v0 = advect[n, nx, ny, nz, v0, u0, v0, w0, dt];
w0 = advect[n, nx, ny, nz, w0, u0, v0, w0, dt];
Do[Do[u0[[i, j, k]] = 0; v0[[i, j, k]] = 0;
w0[[i, j, k]] = 0;, {i, nx0 - Round[Sqrt[R0^2 - (j - ny0)^2]],
nx0 + Round[Sqrt[R0^2 - (j - ny0)^2]], 1}], {j, ny0 - R0,
ny0 + R0, 1}, {k, 1, nz, 1}];
u0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, u0];
v0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, v0];
w0 = bcu[nx, ny, nz, 0, 0, 0, 0, 0, w0];

{u0, v0, w0}];

cf = With[{cg = CompileGetElement, hp = HoldPattern,
dv = DownValues},
Hold@Compile[{{u0argu, _Real, 3}, {v0argu, _Real,
3}, {w0argu, _Real,
3}, {n, _Integer}, {nx, _Integer}, {ny, _Integer}, {nz, \
_Integer}, {nx0, _Integer}, {ny0, _Integer}, {R0, _Integer}, {dt, \
_Real}, {uinfl, _Real}},
Module[{u0 = u0argu, v0 = v0argu, w0 = w0argu, uu, vv, ww},
{u0, v0, w0} =
onestep[n, nx, ny, nz, nx0, ny0, R0, u0, v0, w0, dt,
uinfl];
{u0, v0, w0}], CompilationTarget -> "C",
RuntimeOptions -> "Speed"] /. dv@onestep /.
Flatten[dv /@ {advect, bcu}] /.
hp@ConstantArray[c_, {i_, j_, kc_}] :>
Table[0., {i}, {j}, {kc}] /. hp@Part[a__] :> cg[a] /.
hp[cg[a__] = rhs_] :> (Part[a] = rhs) //
ReleaseHold];
cf1 = With[{cg = CompileGetElement, hp = HoldPattern},
Hold@Compile[{{vargu, _Complex, 4}, {u0argu, _Complex,
3}, {v0argu, _Complex, 3}, {w0argu, _Complex,
3}, {n, _Integer}, {nx, _Integer}, {ny, _Integer}, {nz, \
_Integer}, {dt, _Real}, {mu, _Real}, {r, _Real, 1}},
Module[{uf = u0argu, vf = v0argu, wf = w0argu, v = vargu, x,
y, z, k2},
v = Table[{uf[[i, j, jz]], vf[[i, j, jz]],
wf[[i, j, jz]]}, {i, n}, {j, n}, {jz, n}];
v = Table[x = r[[i]];
y = r[[j]]; z = r[[jz]];
k2 = x^2 + y^2 + z^2;

If[k2 > 0, (v[[i, j,
jz]] - (v[[i, j, jz]] . {x, y, z}) {x, y, z}/k2) (1 -
mu dt k2) , v[[i, j, jz]]], {i, n}, {j, n}, {jz, n}];
v], CompilationTarget -> "C", RuntimeOptions -> "Speed"] /.
hp@Part[a__] :> cg[a] /. hp[cg[a__] = rhs_] :> (Part[a] = rhs) //
ReleaseHold];


With 2 compiled functions cf,cf1 the code is in 86 times faster than without compilation. Also, due to bcu we can simulate 3D flow in a channel.

v = Table[{0., 0., 0.}, {nx}, {ny}, {nz}];
Do[{ui, vi, wi} =
cf[u0, v0, w0, n, nx, ny, nz, nx0, ny0, R0, dt, uinfl];
uf = Fourier[ui]; vf = Fourier[vi]; wf = Fourier[wi];
v1 = cf1[v, uf, vf, wf, n, nx, ny, nz, dt, mu, r];
{u0, v0, w0} =
Table[InverseFourier[v1[[All, All, All, i]]] // Re, {i, 3}];
us[t] = u0; vs[t] = v0;
ws[t] = w0;, {t, 1, nt}] // AbsoluteTiming


Visualization

lst3D = Table[
Show[ListDensityPlot3D[us[t], ColorFunction -> Hue,
PlotRange -> All, Boxed -> False, ImageSize -> Tiny,
Axes -> False, ViewPoint -> {-2., 1., 1.},
OpacityFunction -> 0.05],
Graphics3D[{Blue,
Cylinder[{{0, n/2, n/4}, {n, n/2, n/4}}, n/16]}]], {t, 20, nt,
4}]; ListAnimate[lst3D]


• Maybe linear interpolation is sufficient ? Commented Oct 30, 2022 at 9:40
• @yarchik Do you mean trilinear interpolation as in Update 2 on mathematica.stackexchange.com/questions/261185/… ? Commented Oct 31, 2022 at 3:38
• Yes, that was my first idea. I thought it might be too expensive and not compilable to construct interpolation function each time-step. Commented Oct 31, 2022 at 7:49
• @yarchik Yes, you are right. With trilinear interpolation the compiled code in 100 times faster than with ListInterpolation without compilation. Commented Nov 1, 2022 at 21:49
• Impressive speed up! Commented Nov 1, 2022 at 21:53

Here's a 3D version of my answer:

SetAttributes[compile, HoldAll];
compile[argu__] :=
With[{cg = CompileGetElement},
Hold@Compile[argu, RuntimeOptions -> "Speed", CompilationTarget -> "C"] /.
Part -> cg /. HoldPattern[cg[a__] = rhs_] :> (Part[a] = rhs)] // ReleaseHold //
Last

interrule = inter[valueL_, valueR_, scale_] :> (1 - scale) valueL + scale valueR;

Hold@compile[{{arg, _Real, 4}, {dt, _Real}},
Module[{u, v, w, unew, vnew, wnew, nx, ny, nz, inew, jnew, knew, iL, jL, kL,
iR, jR, kR},
{u, v, w} = arg; {nx, ny, nz} = Dimensions@v;
unew = vnew = wnew = Table[0., {nx}, {ny}, {nz}];
Do[If[i < 1 + nx/16.,
u[[i, j, k]] = 0.1; v[[i, j, k]] = 0.; w[[i, j, k]] = 0.];
If[(i - nx/4.)^2 + (j - ny/2.)^2 < (nx/16.)^2,
u[[i, j, k]] = 0.; v[[i, j, k]] = 0.; w[[i, j, k]] = 0.], {i, nx}, {j,
ny}, {k, nz}];
Do[inew = Mod[i - nx dt u[[i, j, k]], nx - 1, 1];
jnew = Mod[j - nx dt v[[i, j, k]], ny - 1, 1];
knew = Mod[k - nx dt w[[i, j, k]], nz - 1, 1];
iL = Floor@inew; jL = Floor@jnew; kL = Floor@knew;
iR = iL + 1; jR = jL + 1; kR = kL + 1;
unew[[i, j, k]] = interfunc[u];
vnew[[i, j, k]] = interfunc[v];
wnew[[i, j, k]] = interfunc[w], {i, nx}, {j, ny}, {k, nz}];
{unew, vnew, wnew}
]] /. interfunc[v_] :> inter[
inter[
inter[v[[iL, jL, kL]], v[[iL, jR, kL]], jnew - jL],
inter[v[[iR, jL, kL]], v[[iR, jR, kL]], jnew - jL], inew - iL],
inter[
inter[v[[iL, jL, kR]], v[[iL, jR, kR]], jnew - jL],
inter[v[[iR, jL, kR]], v[[iR, jR, kR]], jnew - jL], inew - iL],
knew - kL] //. interrule // ReleaseHold;

viscosity…and…conservation =
compile[{{arg, _Complex, 4}, dt, mu},
Module[{u, v, w, x, y, z, nx, ny, nz, k, k2},
{u, v, w} = arg; {nx, ny, nz} = Dimensions@v;
Do[
x = Mod[i - 1, nx, -nx/2.];
y = Mod[j - 1, ny, -ny/2.];
z = Mod[k - 1, nz, -nz/2.];
k2 = x^2 + y^2 + z^2;
With[{coef = (1 - k2 mu dt (Pi/(nx/2.))^2),
mid = x u[[i, j, k]] + y v[[i, j, k]] + z w[[i, j, k]]},
If[k2 > 0,
u[[i, j, k]] = coef (u[[i, j, k]] - x mid/k2);
v[[i, j, k]] = coef (v[[i, j, k]] - y mid/k2);
w[[i, j, k]] = coef (w[[i, j, k]] - z mid/k2), 0.]], {i, nx}, {j, ny}, {k,
nz}];
{u, v, w}]];

nx = 64; ny = nx; nz = nx; dt = 0.3; mu = 0.001; nt = 300;
v = Table[0., {3}, {nx}, {ny}, {nz}];

vlst = Table[v = force…and…advection[v, dt];
v = Fourier /@ v;
v = viscosity…and…conservation[v, dt, mu];
v = InverseFourier /@ v // Re, {t, nt}]; // AbsoluteTiming
(* {40.3973, Null} *)

arrayplot = ArrayPlot[#, DataReversed -> True, ColorFunction -> "Rainbow"] &;

vlst[[-1]] // First // #[[All, All, nz/2]]\[Transpose] & // arrayplot


vlst[[-1]] // First //
ListDensityPlot3D[#, ColorFunction -> "Rainbow",
PlotRange -> All] & // AbsoluteTiming


For 2GHz dual core laptop, the timing for 300×64×64×64 grid is about 40 seconds. Limited by the RAM of my laptop, it's a bit troublesome for me to make better visualization (the GIF in question isn't obtained with 300 time steps, I guess?), so I'd like to stop here.

• +1 for the answer, but a question: What is the black magic with the three dots in the "force…and…advection" variable? How's this possible??? Commented Oct 30, 2022 at 10:11
• This answer is very useful (+1). Commented Oct 30, 2022 at 10:19
• @Hans It's actually \[Ellipsis]` :D . In Chinese this symbol is used for ellipsis, so it's relatively easy for me to input this symbol using Chinese IME. Commented Oct 30, 2022 at 10:31
• Oh I see! I always envied other programming languages, e.g. Python etc, where you can use underscores etc in the variables, but now I'm doing this :D Commented Oct 30, 2022 at 10:36
• @Hans Some other alternatives of underscore are mentioned here: mathematica.stackexchange.com/q/111337/1871 :) Commented Oct 30, 2022 at 10:51