# Stable fluids code for electromagnetic mixture application

This code has been translated from the original Jos Stam code and improved with some Mathematica functions. It solves problem of viscous incompressible flow with electromagnetic force in a rectangle with periodic boundary condition on one side and with Dirichlet condition on the other side. In the initial condition fluid velocity is zero, but density has unit step like distribution (for instance, liquids gold and silver). First we define initial data and boundary condition module:

ClearAll["Global*"]

dif = 1/500; pec = 20; U0 = 0; V0 = 0; F0 = 3; dn0 = 1; kap = 5; n = \
31; n1 = n + 1; dt = 40./n^2; sm = 200; r = 20; a = dt dif n n; c =
ConstantArray[0, {n1, n1}]; d = ConstantArray[0, {n1, n1}]; den =
ConstantArray[0, {n1, n1}]; c0 = ConstantArray[0, {n1, n1}]; u0 =
ConstantArray[0., {n1, n1}]; v0 = ConstantArray[0., {n1, n1}]; u =
ConstantArray[0, {n1, n1}]; v = ConstantArray[0, {n1, n1}]; Do[
den[[All, j]] = dn0 (1 + .3 Tanh[kap (n1/2 - j)]);, {j, n1}]; dnup =
den[[1, n1]]; dnd = den[[1, 1]];

periodic[n_, up_, ud_, ub_] :=
Module[{bd = ub},
Do[bd[[1, i]] = .5 (bd[[n, i]] + bd[[2, i]]);
bd[[n + 1, i]] = bd[[1, i]]; bd[[i, 1]] = ud;
bd[[i, n + 1]] = up;, {i, 2, n}];
bd[[1, 1]] = .5 (bd[[2, 1]] + bd[[1, 2]]);
bd[[n + 1, n + 1]] = .5 (bd[[n, n + 1]] + bd[[n + 1, n]]);
bd[[n + 1, 1]] = .5 (bd[[n, 1]] + bd[[n + 1, 2]]);
bd[[1, n + 1]] = .5 (bd[[1, n]] + bd[[2, n + 1]]); bd];


Second, diffusion step module with Gauss-Seidel relaxation algorithm:

diffuse[n_, r_, a_, c_, c0_] :=
Module[{c1 = c}, c1 = ConstantArray[0, {n + 1, n + 1}];
Do[Do[Do[c1[[i,
j]] = (c0[[i, j]] +
a (c1[[i - 1, j]] + c1[[i + 1, j]] + c1[[i, j - 1]] +
c1[[i, j + 1]]))/(1 + 4 a);, {j, 2, n}];, {i, 2, n}];
Do[c1[[1, i]] = c1[[n, i]]; c1[[n + 1, i]] = c1[[2, i]];
c1[[i, 1]] = c0[[i, 1]];
c1[[i, n + 1]] = c0[[i, n + 1]];, {i, 2, n}];
c1[[1, 1]] = .5 (c1[[2, 1]] + c1[[1, 2]]);
c1[[n + 1, n + 1]] = .5 (c1[[n, n + 1]] + c1[[n + 1, n]]);
c1[[n + 1, 1]] = .5 (c1[[n, 1]] + c1[[n + 1, 2]]);
c1[[1, n + 1]] = .5 (c1[[1, n]] + c1[[2, n + 1]]);, {k, 0, r}];
c1];


advect[n_, d_, d0_, u_, v_, dt_] :=
Module[{x, y, d1, dt0, i0, i1, j0, j1, s0, s1, t0, t1},
d1 = ConstantArray[0, {n + 1, n + 1}]; dt0 = dt n;
Do[Do[x = i - dt0 u[[i, j]]; y = j - dt0 v[[i, j]];
i0 = Which[x <= 1, 1, 1 < x < n, Floor[x], x >= n, n];
i1 = i0 + 1;
j0 = Which[y <= 1, 1, 1 < y < n, Floor[y], y >= n, n];
j1 = j0 + 1; s1 = x - i0; s0 = 1 - s1; t1 = y - j0; t0 = 1 - t1;
d1[[i, j]] =
s0 (t0 d0[[i0, j0]] + t1 d0[[i0, j1]]) +
s1 (t0 d0[[i1, j0]] + t1 d0[[i1, j1]]);, {j, 1, n + 1}];, {i,
1, n + 1}]; d1];


Projection step module:

project[n_, r_, u0_, v0_, u_, v_] :=
Module[{ux = u, vy = v, div, p},
p = ConstantArray[0, {n + 1, n + 1}];
div = ConstantArray[0, {n + 1, n + 1}];
ux = ConstantArray[0, {n + 1, n + 1}];
vy = ConstantArray[0, {n + 1, n + 1}];
Do[div[[i,
j]] = -.5 /
n (u0[[i + 1, j]] - u0[[i - 1, j]] + v0[[i, 1 + j]] -
v0[[i, j - 1]]);, {i, 2, n}, {j, 2, n}];
Do[Do[Do[
p[[i, j]] = (div[[i,
j]] + (p[[i - 1, j]] + p[[i + 1, j]] + p[[i, j - 1]] +
p[[i, j + 1]]))/4;, {j, 2, n}], {i, 2, n}];, {k, 0, r}];
Do[ux[[i, j]] = u0[[i, j]] - .5 n (p[[i + 1, j]] - p[[i - 1, j]]);
vy[[i, j]] =
v0[[i, j]] - .5 n (p[[i, j + 1]] - p[[i, j - 1]]);, {i, 2,
n}, {j, 2, n}]; {ux, vy}];


Electromagnetic force:

Fx[t_, x_, y_] :=
F0 ((y - .5) Sin[2 Pi t]^2 - (x - 0.5) Sin[2 Pi t] Cos[2 Pi t]);
Fy[t_, x_, y_] :=
F0 (-(x - .5) Cos[2 Pi t]^2 + (y - .5) Sin[2 Pi t] Cos[2 Pi t]);


One step module:

onestep[n_, step_, r_, a_, uin_, vin_, dt_] :=
Module[{u1, v1, f1, f2, c, u, v, u0, v0},
f1 = ConstantArray[0, {n + 1, n + 1}];
f2 = ConstantArray[0., {n + 1, n + 1}];
u0 = ConstantArray[0., {n + 1, n + 1}];
v0 = ConstantArray[0., {n + 1, n + 1}];
u = ConstantArray[0., {n + 1, n + 1}];
v = ConstantArray[0., {n + 1, n + 1}];
u1 = ConstantArray[0., {n + 1, n + 1}];
v1 = ConstantArray[0., {n + 1, n + 1}]; u0 = uin; v0 = vin;
u0 = advect[n, c, u0, u0, v0, dt]; v0 = advect[n, c, v0, u0, v0, dt];
u0 = periodic[n, 0, 0, u0]; v0 = periodic[n, 0, 0, v0];
u0 = diffuse[n, r, a, c, u0]; v0 = diffuse[n, r, a, c, v0];
u0 = periodic[n, 0, 0, u0]; v0 = periodic[n, 0, 0, v0];
Do[f1[[i, j]] =
Fx[dt (step + .5), (i - 1)/n, (j - 1)/n]/den[[i, j]];
f2[[i, j]] =
Fy[dt (step + .5), (i - 1)/n, (j - 1)/n]/den[[i, j]];, {i, 2,
n + 1}, {j, 2, n + 1}]; u0 += f1 dt;
v0 += f2 dt; {u1, v1} = project[n, r, u0, v0, u, v];
u0 = periodic[n, 0, 0, u1]; v0 = periodic[n, 0, 0, v1]; {u0, v0}]


Numerical solution for the flow and density:

Do[{u0, v0} = onestep[n, step, r, a, u0, v0, dt]; uu[step] = u0;
vv[step] = v0; den = diffuse[n, r, a/pec, c, den];
den = periodic[n, dnup, dnd, den];
den = advect[n, c, den, u0, v0, dt];
den = periodic[n, dnup, dnd, den];
dd[step] = den;, {step, 0, sm}] // AbsoluteTiming


Visualization:

Do[lstu[k] =
Flatten[Table[{{(i - 1)/n, (j - 1)/n}, uu[k][[i, j]]}, {i, n1}, {j,
n1}], 1];
lstv[k] =
Flatten[Table[{{(i - 1)/n, (j - 1)/n}, vv[k][[i, j]]}, {i, n1}, {j,
n1}], 1];, {k, 0, sm}];
Do[Uvel[i] = Interpolation[lstu[i], InterpolationOrder -> 3];, {i, 1,
sm}];
Do[Vvel[i] = Interpolation[lstv[i], InterpolationOrder -> 3];, {i, 1,
sm}]; Do[lst4[k] =
Flatten[Table[{{(i - 1)/n, (j - 1)/n}, dd[k][[i, j]]}, {i, n1}, {j,
n1}], 1];, {k, 0, sm}];
Do[rh[k] = Interpolation[lst4[k], InterpolationOrder -> 3];, {k, 0,
sm}];

frames = Table[
ContourPlot[rh[i][x, y], {x, 0, 1}, {y, 0, 1},
ColorFunction -> "BlueGreenYellow", Frame -> False,
PlotRange -> All, ImageSize -> Small, MaxRecursion -> 2,
Contours -> 5, ContourStyle -> Yellow], {i, 0, sm}];

Export["D:\\Animation\\Periodic.gif", frames,
AnimationRepetitions -> Infinity]


The question is about code improvement. How we can reduce computation time?

• You've localized c in definition of onestep, is it a typo? Also, you write Module[{c1 = c}, c1 = ConstantArray[0, {n + 1, n + 1}];  in definition of diffuse, is it a quick repair? May 16, 2021 at 4:54
• @xzczd Thank you very much for your remarks. I don't remember my logic when I created this code. But code has been tested and it works not as a stable fluid only. Also improvements you made are very nice. May 16, 2021 at 10:20
• I’m unfamiliar with this physics. Where does the force expression come from? May 16, 2021 at 19:48
• @Gilbert This force generated in the conducting fluid by the rotating magnetic field. May 16, 2021 at 20:39

Let's Compile.

dif = 1/500; pec = 20; U0 = 0; V0 = 0; F0 = 3; dn0 = 1; kap = 5; n = 63(*31*);
n1 = n + 1;(*dt=40./n^2;*)sm = 4 200; r = 20;(*a=dt dif n n;*)
(*c=ConstantArray[0.,{n1,n1}];*)
(*d=ConstantArray[0,{n1,n1}];*)
den = ConstantArray[dn0 (1 + .3 Tanh[-kap Range[-n1/2, n1/2]]), n1];
(*c0=ConstantArray[0,{n1,n1}];*)
u0 = ConstantArray[0., {n1, n1}];
v0 = ConstantArray[0., {n1, n1}];
(*u=ConstantArray[0,{n1,n1}];
v=ConstantArray[0,{n1,n1}];*)
(*Do[den[[All,j]]=dn0 (1+.3 Tanh[kap (n1/2-j)]);,{j,n1}];*)
(*dnup=den[[1,n1]];dnd=den[[1,1]];*)

periodic[n_, up_, ud_, ub_] :=
Module[{bd = ub}, Do[bd[[1, i]] = .5 (bd[[n, i]] + bd[[2, i]]);
bd[[n + 1, i]] = bd[[1, i]]; bd[[i, 1]] = ud;
bd[[i, n + 1]] = up;, {i, 2, n}];
bd[[1, 1]] = .5 (bd[[2, 1]] + bd[[1, 2]]);
bd[[n + 1, n + 1]] = .5 (bd[[n, n + 1]] + bd[[n + 1, n]]);
bd[[n + 1, 1]] = .5 (bd[[n, 1]] + bd[[n + 1, 2]]);
bd[[1, n + 1]] = .5 (bd[[1, n]] + bd[[2, n + 1]]); bd];

diffuse[n_, r_, a_, c_, c0_] :=
Module[{c1 = c},(*c1=ConstantArray[0,{n+1,n+1}];*)
Do[Do[Do[
c1[[i, j]] = (c0[[i, j]] +
a (c1[[i - 1, j]] + c1[[i + 1, j]] + c1[[i, j - 1]] + c1[[i, j + 1]]))/(1 +
4 a);, {j, 2, n}];, {i, 2, n}];
Do[c1[[1, i]] = c1[[n, i]]; c1[[n + 1, i]] = c1[[2, i]];
c1[[i, 1]] = c0[[i, 1]];
c1[[i, n + 1]] = c0[[i, n + 1]];, {i, 2, n}];
c1[[1, 1]] = .5 (c1[[2, 1]] + c1[[1, 2]]);
c1[[n + 1, n + 1]] = .5 (c1[[n, n + 1]] + c1[[n + 1, n]]);
c1[[n + 1, 1]] = .5 (c1[[n, 1]] + c1[[n + 1, 2]]);
c1[[1, n + 1]] = .5 (c1[[1, n]] + c1[[2, n + 1]]);, {k, 0, r}];
c1];

advect[n_, d0_, u_, v_, dt_] :=
Module[{x, y, d1, dt0, i0, i1, j0, j1, s0, s1, t0, t1},
d1 = ConstantArray[0, {n + 1, n + 1}]; dt0 = dt n;
Do[Do[x = i - dt0 u[[i, j]]; y = j - dt0 v[[i, j]];
i0 = Which[x <= 1, 1, 1 < x < n, Floor[x], True(*x≥n*), n];
i1 = i0 + 1;
j0 = Which[y <= 1, 1, 1 < y < n, Floor[y], True(*y≥n*), n];
j1 = j0 + 1; s1 = x - i0; s0 = 1 - s1; t1 = y - j0; t0 = 1 - t1;
d1[[i, j]] =
s0 (t0 d0[[i0, j0]] + t1 d0[[i0, j1]]) +
s1 (t0 d0[[i1, j0]] + t1 d0[[i1, j1]]);, {j, 1, n + 1}];, {i, 1, n + 1}]; d1];

project[n_, r_, u0_, v0_, u_, v_] :=
Module[{ux = u, vy = v, div, p}, p = ConstantArray[0, {n + 1, n + 1}];
div = ConstantArray[0, {n + 1, n + 1}];
ux = ConstantArray[0, {n + 1, n + 1}];
vy = ConstantArray[0, {n + 1, n + 1}];
Do[div[[i, j]] = -.5/
n (u0[[i + 1, j]] - u0[[i - 1, j]] + v0[[i, 1 + j]] - v0[[i, j - 1]]);, {i, 2,
n}, {j, 2, n}];
Do[Do[Do[
p[[i, j]] = (div[[i, j]] + (p[[i - 1, j]] + p[[i + 1, j]] + p[[i, j - 1]] +
p[[i, j + 1]]))/4;, {j, 2, n}], {i, 2, n}];, {k, 0, r}];
Do[ux[[i, j]] = u0[[i, j]] - .5 n (p[[i + 1, j]] - p[[i - 1, j]]);
vy[[i, j]] = v0[[i, j]] - .5 n (p[[i, j + 1]] - p[[i, j - 1]]);, {i, 2, n}, {j, 2,
n}]; {ux, vy}];

Fx[t_, x_, y_] := F0 ((y - .5) Sin[2 Pi t]^2 - (x - 0.5) Sin[2 Pi t] Cos[2 Pi t]);
Fy[t_, x_, y_] := F0 (-(x - .5) Cos[2 Pi t]^2 + (y - .5) Sin[2 Pi t] Cos[2 Pi t]);

onestep[n_, step_, r_, a_, uin_, vin_, dt_, c_] :=
Module[{u1, v1, f1, f2(*,c*), u, v, u0, v0}, f1 = ConstantArray[0, {n + 1, n + 1}];
f2 = ConstantArray[0., {n + 1, n + 1}];
u0 = ConstantArray[0., {n + 1, n + 1}];
v0 = ConstantArray[0., {n + 1, n + 1}];
u = ConstantArray[0., {n + 1, n + 1}];
v = ConstantArray[0., {n + 1, n + 1}];
u1 = ConstantArray[0., {n + 1, n + 1}];
v1 = ConstantArray[0., {n + 1, n + 1}]; u0 = uin; v0 = vin;
u0 = advect[n, u0, u0, v0, dt]; v0 = advect[n, v0, u0, v0, dt];
u0 = periodic[n, 0, 0, u0]; v0 = periodic[n, 0, 0, v0];
u0 = diffuse[n, r, a, c, u0]; v0 = diffuse[n, r, a, c, v0];
u0 = periodic[n, 0, 0, u0]; v0 = periodic[n, 0, 0, v0];
Do[f1[[i, j]] = Fx[dt (step + .5), (i - 1)/n, (j - 1)/n]/den[[i, j]];
f2[[i, j]] = Fy[dt (step + .5), (i - 1)/n, (j - 1)/n]/den[[i, j]];,
{i, 2, n + 1}, {j, 2, n + 1}]; u0 += f1 dt;
v0 += f2 dt; {u1, v1} = project[n, r, u0, v0, u, v];
u0 = periodic[n, 0, 0, u1]; v0 = periodic[n, 0, 0, v1]; {u0, v0}]

cf = With[{cg = CompileGetElement, hp = HoldPattern, dv = DownValues},
Hold@Compile[{{u0argu, _Real, 2}, {v0argu, _Real, 2}, {denargu, _Real,
2}, {sm, _Integer}, {n, _Integer}, {r, _Integer}, dif, pec, F0},
Module[{u0 = u0argu, v0 = v0argu, uu, vv, dd, den = denargu,
c = Table[0., {n + 1}, {n + 1}], dt = 40./n^2, a, dnup = den[[1, n + 1]],
dnd = den[[1, 1]]}, a = dt dif n n;
uu = vv = dd = Table[0., {sm + 1}, {n + 1}, {n + 1}];
Do[{u0, v0} = onestep[n, step, r, a, u0, v0, dt, c];
uu[[step + 1]] = u0;
vv[[step + 1]] = v0;
den = diffuse[n, r, a/pec, c, den];
den = periodic[n, dnup, dnd, den];
den = advect[n, den, u0, v0, dt];
den = periodic[n, dnup, dnd, den];
dd[[step + 1]] = den;, {step, 0, sm}]; {uu, vv, dd}],
CompilationTarget -> C, RuntimeOptions -> "Speed"] /. dv@onestep /.
Flatten[dv /@ {Fx, Fy, advect, diffuse, periodic, project}] /.
hp@ConstantArray[c_, {i_, j_}] :> Table[0., {i}, {j}] /.
hp@Part[a__] :> cg[a] /. hp[cg[a__] = rhs_] :> (Part[a] = rhs) //
ReleaseHold // Last]; // AbsoluteTiming
(* {1.375986, Null} *)

rst = cf[u0, v0, den, sm, n, r, dif, pec, F0]; // AbsoluteTiming
(* {3.036883, Null} *)

(*
lst =
ListContourPlot[Transpose[#], ColorFunction -> "BlueGreenYellow", Contours -> 5,
PlotRange -> All, Frame -> False] & /@
rst[[-1]]; // AbsoluteTiming
*)
(* {28.801388, Null} *)

lst =
ArrayPlot[Transpose[#], ColorFunction -> "BlueGreenYellow", DataReversed -> True,
Frame -> None, PlotRangePadding -> None] & /@ rst[[-1]]; // AbsoluteTiming
(* {6.044640, Null} *)

lst // ListAnimate


Notice I've increased sm and n. For sm = 200; n = 31; it only takes 0.2 second to calculate rst on my laptop, which is a 350x speed-up.

To understand why the code is modified in this manner, you may want to read

How to make the code inside Compile conciser without hurting performance?

How to define a function inside a Compile?

Why is CompilationTarget -> C slower than directly writing with C?

If doesn't compile if its output is the output of Compile?

Memory usage of compiled function: C vs WVM

• Last step is very advanced (+1). Now with this improvement we can try to use this code with real time force like in CUDAFluidDynamics[]. May 16, 2021 at 10:06
• @AlexTrounev I'm not sure what it is, but something seems to be wrong with the algorithm. Try ArrayPlot[Transpose[Join @@ ConstantArray[# // Most, 2]], ColorFunction -> "BlueGreenYellow", DataReversed -> True, Frame -> None, PlotRangePadding -> None] &@rst[[-1, Floor[sm/2]]]. (When time gets large e.g. rst[[-1,-1]], the discontinuity becomes unobvious though. ) May 17, 2021 at 4:55
• Do you mean in your code with n = 63; sm=800 on the border? Even it is the stable fluid algorithm there are also instabilities happen due to higher Peclet number (in this case 20000). I will run several tests with this code to fix this problem. May 17, 2021 at 11:12
• Ah, I forgot that force has discontinuity on x for periodic solution. This discontinuity is not eliminated by viscosity on the first steps since we use explicit first order time integration step u0 += f1 dt; v0 += f2 dt. Therefore, if we need periodic solution without discontinuity, we should update force so that it has no discontinuity as well. May 17, 2021 at 13:13
• Jos Stam recommended to use u0 += f1 dt; v0 += f2 dt as a first step, and then advection, diffusion, projection. But this algorithm is less stable (not in this case). May 17, 2021 at 14:36