add 2D support for python penalization and dummy operators

hysop/backend/host/python/operator/penalization.py

@@ -124,7 +124,7 @@ class PythonPenalizeVorticity(HostOperator):
         for is_input, (field, td, req) in requirements.iter_requirements():
             min_ghosts = (max(g, G) for g in req.min_ghosts.copy())
             req.min_ghosts = min_ghosts
-            req.axes = ((0, 1, 2), )
+            req.axes = (tuple(range(field.dim)), )
         return requirements
 
     @debug
@@ -132,22 +132,26 @@ class PythonPenalizeVorticity(HostOperator):
         if self.discretized:
             return
         super(PythonPenalizeVorticity, self).discretize()
+        dim = self.velocity.dim
         self.dvelocity = self.input_discrete_tensor_fields[self.velocity]
-        self.dvorticity = self.input_discrete_tensor_fields[self.vorticity]
+        if dim == 2:
+            self.dvorticity = self.input_discrete_fields[self.vorticity]
+        if dim == 3:
+            self.dvorticity = self.input_discrete_tensor_fields[self.vorticity]
 
         dv, dw = self.dvelocity, self.dvorticity
         stencil = self.stencil[0]
         G = max(max(a, b) for a, b in zip(stencil.L, stencil.R))
-        view = dv.local_slices(ghosts=(G, G, G))
+        view = dv.local_slices(ghosts=(G, )*dim)
         V = tuple(Vi[view] for Vi in dv.buffers)
-        view = dw.local_slices(ghosts=(G, G, G))
+        view = dw.local_slices(ghosts=(G, )*dim)
         W = tuple(Wi[view] for Wi in dw.buffers)
         self.W, self.V = W, V
         self.dobstacles = {}
         for c, o in self.obstacles.iteritems():
             o_df = self.input_discrete_fields[o]
             self.dobstacles[c] = o_df.data[0][o_df.local_slices(
-                ghosts=(G, G, G))]
+                ghosts=(G, )*dim)]
 
         for s, dx in zip(self.stencil, dw.space_step):
             s.replace_symbols({s.dx: dx})
@@ -168,10 +172,14 @@ class PythonPenalizeVorticity(HostOperator):
         Wtmp = work.get_buffer(self, 'Wtmp', handle=True)
         self.tmp_v = Wtmp[0:3]
         self.tmp = Wtmp[-1]
+        if self.velocity.dim == 2:
+            self._compute_vorticity = self._compute_vorticity_2d
+        if self.velocity.dim == 3:
+            self._compute_vorticity = self._compute_vorticity_3d
 
     @classmethod
     def supported_dimensions(cls):
-        return (3, )
+        return (3, 2)
 
     @classmethod
     def supports_mpi(cls):
@@ -193,34 +201,46 @@ class PythonPenalizeVorticity(HostOperator):
     @op_apply
     def apply(self, **kwds):
         super(PythonPenalizeVorticity, self).apply(**kwds)
-        V, W = self.V, self.W
         self.v_ghost_exchanger()
 
-        # # Penalize velocity
+        # Penalize velocity
         self._compute_penalization()
-        for ubar, v, tmp in zip(self._ubar(), V, self.tmp_v):
+        for ubar, v, tmp in zip(self._ubar(), self.V, self.tmp_v):
             tmp[...] = (v - ubar) * self.tmp
 
+        self._compute_vorticity()
+        self.w_ghost_exchanger()
+
+    def _compute_vorticity_3d(self):
         # compute penalized vorticity:
         #     (dvz/dy - dvy/dz)
         # W = (dvx/dz - dvz/dx)
         #     (dvy/dx - dvx/dy)
         # X direction
         self.tmp = self.stencil[0](a=self.tmp_v[2], out=self.tmp, axis=2)
-        W[1][...] += -self.tmp
+        self.W[1][...] += -self.tmp
         self.tmp = self.stencil[0](a=self.tmp_v[1], out=self.tmp, axis=2)
-        W[2][...] += self.tmp
+        self.W[2][...] += self.tmp
 
         # Y direction
         self.tmp = self.stencil[1](a=self.tmp_v[0], out=self.tmp, axis=1)
-        W[2][...] += -self.tmp
+        self.W[2][...] += -self.tmp
         self.tmp = self.stencil[1](a=self.tmp_v[2], out=self.tmp, axis=1)
-        W[0][...] += self.tmp
+        self.W[0][...] += self.tmp
 
         # Z direction
         self.tmp = self.stencil[2](a=self.tmp_v[1], out=self.tmp, axis=0)
-        W[0][...] += -self.tmp
+        self.W[0][...] += -self.tmp
         self.tmp = self.stencil[2](a=self.tmp_v[0], out=self.tmp, axis=0)
-        W[1][...] += self.tmp
+        self.W[1][...] += self.tmp
 
-        self.w_ghost_exchanger()
+    def _compute_vorticity_2d(self):
+        # compute penalized vorticity:
+        # W = (dvy/dx - dvx/dy)
+        # X direction
+        self.tmp = self.stencil[0](a=self.tmp_v[1], out=self.tmp, axis=1)
+        self.W[0][...] += self.tmp
+
+        # Y direction
+        self.tmp = self.stencil[1](a=self.tmp_v[0], out=self.tmp, axis=0)
+        self.W[0][...] += -self.tmp

hysop/fields/field_requirements.py

@@ -134,6 +134,10 @@ class DiscreteFieldRequirements(object):
 
     def set_axes(self, axes):
         check_instance(axes, tuple, values=tuple, allow_none=True)
+        if axes:
+            if not all([len(_) == self._dim for _ in axes]):
+                msg = 'all given axis should be of length {}, given {}'.format(self._dim, axes)
+                assert False, msg
         self._axes = axes
 
     def get_tstates(self):

hysop/operator/dummy.py

@@ -31,7 +31,7 @@ class PythonDummy(HostOperator):
             if (reqs is None):
                 continue
             (field, td, req) = reqs
-            req.axes = ((0, 1, 2), )
+            req.axes = (tuple(range(field.dim)), )
         return requirements
 
     @op_apply
@@ -43,10 +43,11 @@ class PythonDummy(HostOperator):
     def supports_mpi(cls):
         return True
 
+
 class Dummy(ComputationalGraphNodeFrontend):
 
     __implementations = {
-            Implementation.PYTHON: PythonDummy
+        Implementation.PYTHON: PythonDummy
     }
 
     @classmethod