performance tuning – 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 silber). 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-Zeidel 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); 

Advection step module

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)

Figure 1

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