3D periodic jet with levelset

examples/bubble/periodic_jet_levelset.py

@@ -21,11 +21,11 @@ def init_rho(data, coords):
 def init_phi(data, coords, L):
     phi = data[0]
     phi[...] = np.inf
-    Di = np.empty_like(phi)
-    X,Y    = coords 
-    Ly, Lx = L
+    Di = np.ones_like(phi)
+    X,Y = coords[:2]
+    Ly, Lx = L[:2]
     R = Lx*(0.15 - 0.05*np.sin(2*np.pi*Y/Ly) + 0.001*np.cos(2*np.pi*10*Y/Ly))
-    Di[...] = (X-0.5*Lx)**2 - R**2
+    Di[...] *= (X-0.5*Lx)**2 - R**2
     mask = (np.abs(Di)<np.abs(phi))
     phi[mask] = Di[mask]
     assert np.isfinite(phi).all()
@@ -148,10 +148,32 @@ def compute(args):
     diffuse = DirectionalDiffusion(implementation=impl,
              name='diffuse_{}'.format(vorti.name),
              pretty_name=u'D{}'.format(vorti.pretty_name.decode('utf-8')),
-             coeffs = (args.mu, args.mu/10.0),
-             fields = (vorti, phi),
+             coeffs = (args.mu/10.0,),
+             fields = (phi,),
              variables = {vorti: npts , phi: npts},
              dt=dt, **extra_op_kwds)
+    #> Directional stretching + diffusion
+    if (dim==3):
+        stretch_diffuse = DirectionalStretchingDiffusion(implementation=impl,
+                 name='stretch_diffuse',
+                 pretty_name=u'SD{}'.format(vorti.pretty_name.decode('utf-8')),
+                 formulation = args.stretching_formulation,
+                 viscosity = args.mu,
+                 velocity  = velo,       
+                 vorticity = vorti,
+                 variables = {velo: npts, vorti: npts},
+                 dt=dt, **extra_op_kwds)
+    elif (dim==2):
+        stretch_diffuse = DirectionalDiffusion(implementation=impl,
+                 name='diffuse_{}'.format(vorti.name),
+                 pretty_name=u'D{}'.format(vorti.pretty_name.decode('utf-8')),
+                 coeffs = args.mu,
+                 fields  = vorti,
+                 variables = {vorti: npts},
+                 dt=dt, **extra_op_kwds)
+    else:
+        msg='Unsupported dimension {}.'.format(dim)
+        raise RuntimeError(msg)
     
     #> External force rot(-rho*g)
     Fext = np.zeros(shape=(dim,), dtype=object).view(SymbolicTensor)
@@ -169,7 +191,7 @@ def compute(args):
 
     #> Directional splitting operator subgraph
     splitting = StrangSplitting(splitting_dim=dim, order=args.strang_order)
-    splitting.push_operators(advec, diffuse, eval_fields, external_force)
+    splitting.push_operators(advec, diffuse, stretch_diffuse, eval_fields, external_force)
     
     ### Build standard operators
     #> Poisson operator to recover the velocity from the vorticity
@@ -310,8 +332,8 @@ if __name__=='__main__':
             
         def _setup_parameters(self, args):
             dim = args.ndim
-            if (dim not in (2,)):
-                msg='Domain should be 2D.'
+            if (dim not in (2,3)):
+                msg='Domain should be 2D or 3D.'
                 self.error(msg)