add velocity penalization

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

@@ -244,3 +244,186 @@ class PythonPenalizeVorticity(HostOperator):
         # Y direction
         self.tmp = self.stencil[1](a=self.tmp_v[0], out=self.tmp, axis=0)
         self.W[0][...] += -self.tmp
+
+
+class PythonPenalizeVelocity(HostOperator):
+    """
+
+    """
+    __default_method = {
+        SpaceDiscretization: SpaceDiscretization.FDC4
+    }
+
+    __available_methods = {
+        SpaceDiscretization: (InstanceOf(SpaceDiscretization), InstanceOf(int))
+    }
+
+    @classmethod
+    def default_method(cls):
+        dm = super(PythonPenalizeVelocity, cls).default_method()
+        dm.update(cls.__default_method)
+        return dm
+
+    @classmethod
+    def available_methods(cls):
+        am = super(PythonPenalizeVelocity, cls).available_methods()
+        am.update(cls.__available_methods)
+        return am
+
+    @debug
+    def __init__(self, obstacles, variables,
+                 velocity,
+                 dt, coeff=None, ubar=None, formulation=None, **kwds):
+        check_instance(velocity, Field)
+        check_instance(variables, dict, keys=Field,
+                       values=CartesianTopologyDescriptors)
+        check_instance(dt, ScalarParameter)
+        check_instance(coeff, (ScalarParameter, float, type(lambda x: x)), allow_none=True)
+        check_instance(ubar, TensorParameter, allow_none=True)
+        check_instance(formulation, PenalizationFormulation, allow_none=True)
+        check_instance(obstacles, (tuple, dict), values=Field,
+                       keys=(ScalarParameter, float, type(lambda x: x)), check_kwds=False)
+
+        input_fields = {velocity: variables[velocity], }
+        output_fields = {velocity: variables[velocity]}
+        input_params = {dt.name: dt}
+
+        if isinstance(coeff, ScalarParameter):
+            self.coeff = lambda o: coeff()*o
+        elif isinstance(coeff, type(lambda x: x)):
+            self.coeff = coeff
+        elif not coeff is None:
+            c = ScalarParameter("penal_coeff", initial_value=coeff)
+            self.coeff = lambda o: c()*o
+
+        if isinstance(obstacles, dict):
+            obs = {}
+            for c, o in obstacles.iteritems():
+                if isinstance(c, ScalarParameter):
+                    obs[lambda x: c()*x] = o
+                elif isinstance(c, type(lambda x: x)):
+                    obs[c] = o
+                elif not c is None:
+                    _ = ScalarParameter("penal_coeff", initial_value=c)
+                    obs[lambda x: _()*x] = o
+        else:
+            obs = {}
+            for o in obstacles:
+                obs[self.coeff] = o
+        for o in obs.values():
+            assert o.nb_components == 1
+            input_fields[o] = variables[o]
+        self._ubar = ubar
+        if ubar is None:
+            self._ubar = TensorParameter(
+                name="ubar", shape=(velocity.dim,),
+                quiet=True, dtype=velocity.dtype,
+                initial_value=(0.,)*velocity.dim)
+        if formulation is None or formulation is PenalizationFormulation.IMPLICIT:
+            self._compute_penalization = self._compute_penalization_implicit
+        else:
+            self._compute_penalization = self._compute_penalization_exact
+
+        self.velocity = velocity
+        self.obstacles = obs
+        self.dt = dt
+        super(PythonPenalizeVelocity, self).__init__(
+            input_fields=input_fields, output_fields=output_fields,
+            input_params=input_params, **kwds)
+
+    @debug
+    def handle_method(self, method):
+        super(PythonPenalizeVelocity, self).handle_method(method)
+        self.space_discretization = method.pop(SpaceDiscretization)
+        csg = CenteredStencilGenerator()
+        csg.configure(dtype=MPQ, dim=1)
+        try:
+            self.order = int(str(self.space_discretization)[3:])
+        except:
+            self.order = self.space_discretization
+        self.stencil = tuple(csg.generate_exact_stencil(
+            derivative=1,
+            order=self.order) for _ in range(3))
+
+    @debug
+    def get_field_requirements(self):
+        requirements = super(PythonPenalizeVelocity,
+                             self).get_field_requirements()
+        stencil = self.stencil[0]
+        G = max(max(a, b) for a, b in zip(stencil.L, stencil.R))
+        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 = (tuple(range(field.dim)), )
+        return requirements
+
+    @debug
+    def discretize(self):
+        if self.discretized:
+            return
+        super(PythonPenalizeVelocity, self).discretize()
+        dim = self.velocity.dim
+        self.dvelocity = self.input_discrete_tensor_fields[self.velocity]
+
+        dv = self.dvelocity
+        stencil = self.stencil[0]
+        G = max(max(a, b) for a, b in zip(stencil.L, stencil.R))
+        view = dv.local_slices(ghosts=(G, )*dim)
+        V = tuple(Vi[view] for Vi in dv.buffers)
+        self.V = 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, )*dim)]
+
+        for s, dx in zip(self.stencil, dv.space_step):
+            s.replace_symbols({s.dx: dx})
+
+        self.v_ghost_exchanger = dv.build_ghost_exchanger()
+
+    @debug
+    def get_work_properties(self):
+        requests = super(PythonPenalizeVelocity, self).get_work_properties()
+        buffers = MemoryRequest.empty_like(
+            a=self.V[0], nb_components=4, backend=self.dvelocity.backend)
+        requests.push_mem_request('Wtmp', buffers)
+        return requests
+
+    def setup(self, work):
+        super(PythonPenalizeVelocity, self).setup(work=work)
+        Wtmp = work.get_buffer(self, 'Wtmp', handle=True)
+        self.tmp_v = Wtmp[0:3]
+        self.tmp = Wtmp[-1]
+
+    @classmethod
+    def supported_dimensions(cls):
+        return (3, 2)
+
+    @classmethod
+    def supports_mpi(cls):
+        return True
+
+    def _compute_penalization_implicit(self):
+        dt = self.dt()
+        self.tmp[...] = 0.
+        for c, o in self.dobstacles.iteritems():
+            self.tmp[...] += (-dt) * c(o) / (1.0 + dt * c(o))
+
+    def _compute_penalization_exact(self):
+        dt = self.dt()
+        self.tmp[...] = 1.
+        for c, o in self.dobstacles.iteritems():
+            self.tmp[...] *= np.exp(-dt*c(o))
+        self.tmp[...] -= 1.
+
+    @op_apply
+    def apply(self, **kwds):
+        super(PythonPenalizeVelocity, self).apply(**kwds)
+        self.v_ghost_exchanger()
+
+        # Penalize velocity
+        self._compute_penalization()
+        for ubar, v, tmp in zip(self._ubar(), self.V, self.tmp_v):
+            tmp[...] = (v - ubar) * self.tmp
+        self.v_ghost_exchanger()

hysop/operator/penalization.py

@@ -15,7 +15,7 @@ from hysop.topology.cartesian_descriptor import CartesianTopologyDescriptors
 from hysop.parameters.scalar_parameter import ScalarParameter
 from hysop.parameters.tensor_parameter import TensorParameter
 from hysop.core.graph.computational_node_frontend import ComputationalGraphNodeFrontend
-from hysop.backend.host.python.operator.penalization import PythonPenalizeVorticity
+from hysop.backend.host.python.operator.penalization import PythonPenalizeVorticity, PythonPenalizeVelocity
 
 
 class PenalizeVorticity(ComputationalGraphNodeFrontend):
@@ -113,3 +113,91 @@ class PenalizeVorticity(ComputationalGraphNodeFrontend):
             variables=variables,
             implementation=implementation,
             **kwds)
+
+
+class PenalizeVelocity(ComputationalGraphNodeFrontend):
+    """
+    Solve
+    \f{eqnarray*}
+    \frac{\partial w}{\partial t} &=& \lambda\chi_s\nabla\times(v_D - v)
+    \f}
+    using penalization.
+    """
+    __implementations = {
+        Implementation.PYTHON: PythonPenalizeVelocity
+    }
+
+    @classmethod
+    def implementations(cls):
+        return cls.__implementations
+
+    @classmethod
+    def default_implementation(cls):
+        return Implementation.PYTHON
+
+    @debug
+    def __init__(self, obstacles, variables,
+                 velocity,
+                 dt, coeff=None, ubar=None, formulation=None,
+                 implementation=None, **kwds):
+        """
+        Parameters
+        ----------
+        obstacles : dict or list of :class:`~hysop.Field`
+            sets of geometries on which penalization must be applied
+        velocity: field
+            input velocity
+        coeff : ScalarParameter, optional
+            penalization factor (\f\lambda\f) applied to all geometries.
+        ubar : TensorParameter, optional
+            Solid velocity (default to 0)
+        formulation: PenalizationFormulation
+            Solving penalization either with IMPLICIT scheme or EXACT solution
+        variables: dict
+            dictionary of fields as keys and topologies as values.
+        dt: ScalarParameter
+            Timestep parameter that will be used for time integration.
+        implementation: Implementation, optional, defaults to None
+            target implementation, should be contained in available_implementations().
+            If None, implementation will be set to default_implementation().
+        kwds:
+            Keywords arguments that will be passed towards implementation
+            poisson operator __init__.
+
+        Set::
+
+        obstacles = {coeff1: obs1, coeff2: obs2, ...}
+        coeff = None
+
+        to apply a different coefficient on each subset.
+        Set::
+
+        obstacles = [obs1, obs2, ...]
+        coeff = some_value
+
+        Warning : coeff as a function is not yet implemented!!
+        """
+        if not isinstance(obstacles, dict):
+            obstacles = to_list(obstacles)
+            assert len(set(obstacles)) == len(obstacles)
+            obstacles = tuple(obstacles)
+
+        check_instance(velocity, Field)
+        check_instance(variables, dict, keys=Field,
+                       values=CartesianTopologyDescriptors)
+        check_instance(dt, ScalarParameter)
+        check_instance(coeff, (ScalarParameter, float, type(lambda x: x)), allow_none=True)
+        check_instance(formulation, PenalizationFormulation, allow_none=True)
+        check_instance(ubar, TensorParameter, allow_none=True)
+        check_instance(obstacles, (tuple, dict), values=Field,
+                       keys=(ScalarParameter, float, type(lambda x: x)), check_kwds=False)
+        super(PenalizeVelocity, self).__init__(
+            velocity=velocity,
+            coeff=coeff,
+            ubar=ubar,
+            obstacles=obstacles,
+            dt=dt,
+            formulation=formulation,
+            variables=variables,
+            implementation=implementation,
+            **kwds)

hysop/operators.py

@@ -9,7 +9,7 @@ from hysop.operator.poisson import Poisson
 from hysop.operator.poisson_curl import PoissonCurl
 from hysop.operator.diffusion import Diffusion  # FFTW diffusion
 from hysop.operator.advection import Advection  # Scales fortran advection
-from hysop.operator.penalization import PenalizeVorticity
+from hysop.operator.penalization import PenalizeVorticity, PenalizeVelocity
 from hysop.operator.flowrate_correction import FlowRateCorrection
 from hysop.operator.vorticity_absorption import VorticityAbsorption
 from hysop.operator.transpose import Transpose