walberla: More refactoring

src/python/espressomd/detail/walberla.py

@@ -66,6 +66,15 @@ def get_node_indices_inside_shape(self, shape):
             if shape.is_inside(position=pos):
                 yield idx
 
+    def get_shape_bitmask(self, shape):
+        """Create a bitmask for the given shape."""
+        if not isinstance(shape, espressomd.shapes.Shape):
+            raise ValueError(
+                "Parameter 'shape' must be derived from espressomd.shapes.Shape")
+        mask_flat = shape.call_method("rasterize", grid_size=self.shape,
+                                      grid_spacing=self.agrid, grid_offset=0.5)
+        return np.reshape(mask_flat, self.shape).astype(bool)
+
 
 class LatticeModel:
 
@@ -77,6 +86,31 @@ def save_checkpoint(self, path, binary):
     def load_checkpoint(self, path, binary):
         return self.call_method("load_checkpoint", path=path, mode=int(binary))
 
+    def get_nodes_inside_shape(self, shape=None):
+        """
+        Provide a generator for iterating over all nodes inside the given shape.
+
+        Parameters
+        ----------
+        shape : :class:`espressomd.shapes.Shape`
+            Shape to use as filter.
+
+        """
+        for idx in self.lattice.get_node_indices_inside_shape(shape):
+            yield self[idx]
+
+    def get_shape_bitmask(self, shape=None):
+        """
+        Create a bitmask for the given shape.
+
+        Parameters
+        ----------
+        shape : :class:`espressomd.shapes.Shape`
+            Shape to rasterize.
+
+        """
+        return self.lattice.get_shape_bitmask(shape=shape)
+
 
 def get_slice_bounding_box(slices, grid_size):
     shape = []
@@ -111,32 +145,6 @@ def get_slice_bounding_box(slices, grid_size):
             "shape": shape}
 
 
-class VTKRegistry:
-
-    def __init__(self):
-        if not espressomd.code_features.has_features("WALBERLA"):
-            raise NotImplementedError("Feature WALBERLA not compiled in")
-        self.map = {}
-
-    def __getitem__(self, vtk_uid):
-        return self.map[vtk_uid]
-
-    def __contains__(self, vtk_uid):
-        return vtk_uid in self.map
-
-    def _register_vtk_object(self, vtk_obj):
-        vtk_uid = vtk_obj.vtk_uid
-        self.map[vtk_uid] = vtk_obj
-
-    def __getstate__(self):
-        return self.map
-
-    def __setstate__(self, active_vtk_objects):
-        self.map = {}
-        for vtk_uid, vtk_obj in active_vtk_objects.items():
-            self.map[vtk_uid] = vtk_obj
-
-
 class VTKOutputBase(ScriptInterfaceHelper):
 
     def __init__(self, *args, **kwargs):
@@ -150,11 +158,14 @@ def __init__(self, *args, **kwargs):
             super().__init__(*args, **params)
         else:
             super().__init__(**kwargs)
-        self._add_to_registry()
 
     def valid_observables(self):
         return set(self.call_method("get_valid_observable_names"))
 
+    def valid_keys(self):
+        return {"delta_N", "execution_count", "observables", "identifier",
+                "base_folder", "prefix", "enabled"}
+
     def default_params(self):
         return {"delta_N": 0, "enabled": True, "execution_count": 0,
                 "base_folder": "vtk_out", "prefix": "simulation_step"}

src/python/espressomd/electrokinetics.py

@@ -21,8 +21,8 @@
 import numpy as np
 
 from . import utils
-from .detail.walberla import VTKRegistry, VTKOutputBase, LatticeWalberla  # pylint: disable=unused-import
-from .script_interface import ScriptInterfaceHelper, script_interface_register, ScriptObjectList
+from .detail.walberla import VTKOutputBase, LatticeWalberla  # pylint: disable=unused-import
+from .script_interface import ScriptInterfaceHelper, script_interface_register, ScriptObjectList, array_variant
 import espressomd.detail.walberla
 import espressomd.shapes
 import espressomd.code_features
@@ -104,15 +104,15 @@ class EKSpecies(ScriptInterfaceHelper,
         Species diffusion coefficient.
     valency : :obj:`float`
         Species valency.
+    advection : :obj:`bool`
+        Whether to enable advection.
+    friction_coupling : :obj:`bool`
+        Whether to enable friction coupling.
     ext_efield : (3,) array_like of :obj:`float`, optional
         External electrical field.
     kT : :obj:`float`, optional
         Thermal energy of the simulated heat bath (for thermalized species).
         Set it to 0 for an unthermalized species.
-    advection : :obj:`bool`, optional
-        Whether to enable advection.
-    friction_coupling : :obj:`bool`, optional
-        Whether to enable friction coupling.
 
     Methods
     -------
@@ -163,6 +163,10 @@ def __init__(self, *args, **kwargs):
 
         super().__init__(*args, **kwargs)
 
+    def default_params(self):
+        return {"single_precision": False, "kT": 0.,
+                "ext_efield": [0.0, 0.0, 0.0]}
+
     def __getitem__(self, key):
         if isinstance(key, (tuple, list, np.ndarray)) and len(key) == 3:
             if any(isinstance(item, slice) for item in key):
@@ -211,39 +215,15 @@ def add_boundary_from_shape(self, shape, value, boundary_type):
                 raise ValueError(
                     f"Cannot process density value grid of shape {np.shape(value)}")
 
-        value_flat = value.reshape((-1,))
-        mask_flat = shape.call_method('rasterize', grid_size=self.shape,
-                                      grid_spacing=self.lattice.agrid,
-                                      grid_offset=0.5)
-
-        value_view = np.ascontiguousarray(value_flat, dtype=np.double)
-        raster_view = np.ascontiguousarray(mask_flat, dtype=np.int32)
+        mask = self.get_shape_bitmask(shape=shape).astype(int)
         if issubclass(boundary_type, FluxBoundary):
-            self.call_method(
-                "update_flux_boundary_from_shape",
-                raster_view=raster_view,
-                value_view=value_view)
-        if issubclass(boundary_type, DensityBoundary):
-            self.call_method(
-                "update_density_boundary_from_shape",
-                raster_view=raster_view,
-                value_view=value_view)
-
-    def get_nodes_inside_shape(self, shape):
-        """
-        Provide a generator for iterating over all nodes inside the given shape.
-        """
-        for idx in self.lattice.get_node_indices_inside_shape(shape):
-            yield self[idx]
-
-    def get_shape_bitmask(self, shape):
-        """Create a bitmask for the given shape."""
-        utils.check_type_or_throw_except(
-            shape, 1, espressomd.shapes.Shape, "expected a espressomd.shapes.Shape")
-        mask_flat = shape.call_method("rasterize", grid_size=self.shape,
-                                      grid_spacing=self.lattice.agrid,
-                                      grid_offset=0.5)
-        return np.reshape(mask_flat, self.shape).astype(type(True))
+            boundaries_update_method = "update_flux_boundary_from_shape"
+        else:
+            boundaries_update_method = "update_density_boundary_from_shape"
+        self.call_method(
+            boundaries_update_method,
+            raster=array_variant(mask.flatten()),
+            values=array_variant(value.flatten()))
 
 
 class FluxBoundary:
@@ -272,10 +252,6 @@ def __init__(self, density):
         self.density = density
 
 
-if espressomd.code_features.has_features("WALBERLA"):
-    _walberla_vtk_registry = VTKRegistry()
-
-
 @script_interface_register
 class VTKOutput(VTKOutputBase):
     """
@@ -313,13 +289,6 @@ class VTKOutput(VTKOutputBase):
     _so_creation_policy = "GLOBAL"
     _so_bind_methods = ("enable", "disable", "write")
 
-    def _add_to_registry(self):
-        _walberla_vtk_registry._register_vtk_object(self)
-
-    def valid_keys(self):
-        return {'species', 'delta_N', 'execution_count', 'observables',
-                'identifier', 'base_folder', 'prefix', 'enabled'}
-
     def required_keys(self):
         return self.valid_keys() - self.default_params().keys()
 

src/python/espressomd/lb.py

@@ -21,7 +21,7 @@
 import numpy as np
 
 from . import utils
-from .detail.walberla import VTKRegistry, VTKOutputBase, LatticeWalberla
+from .detail.walberla import VTKOutputBase, LatticeWalberla
 from .script_interface import ScriptInterfaceHelper, script_interface_register, array_variant
 import espressomd.detail.walberla
 import espressomd.shapes
@@ -60,11 +60,11 @@ def _deactivate(self):
         self._deactivate_method()
 
     def _activate_method(self):
-        self.call_method('activate')
+        self.call_method("activate")
         utils.handle_errors("HydrodynamicInteraction activation failed")
 
     def _deactivate_method(self):
-        self.call_method('deactivate')
+        self.call_method("deactivate")
         utils.handle_errors("HydrodynamicInteraction deactivation failed")
 
     def validate_params(self, params):
@@ -111,7 +111,7 @@ def _check_mach_limit(cls, velocities):
         if np.any(np.linalg.norm(velocities, axis=1) > vel_max):
             speed_of_sound = 1. / np.sqrt(3.)
             mach_number = vel_max / speed_of_sound
-            raise ValueError(f'Slip velocity exceeds Mach {mach_number:.2f}')
+            raise ValueError(f"Slip velocity exceeds Mach {mach_number:.2f}")
 
     @property
     def pressure_tensor(self):
@@ -123,10 +123,6 @@ def pressure_tensor(self, value):
         raise RuntimeError(f"Property 'pressure_tensor' is read-only")
 
 
-if espressomd.code_features.has_features("WALBERLA"):
-    _walberla_vtk_registry = VTKRegistry()
-
-
 @script_interface_register
 class VTKOutput(VTKOutputBase):
     """
@@ -164,13 +160,6 @@ class VTKOutput(VTKOutputBase):
     _so_creation_policy = "GLOBAL"
     _so_bind_methods = ("enable", "disable", "write")
 
-    def _add_to_registry(self):
-        _walberla_vtk_registry._register_vtk_object(self)
-
-    def valid_keys(self):
-        return {'lb_fluid', 'delta_N', 'execution_count', 'observables',
-                'identifier', 'base_folder', 'prefix', 'enabled'}
-
     def required_keys(self):
         return self.valid_keys() - self.default_params().keys()
 
@@ -297,7 +286,7 @@ def validate_params(self, params):
                 raise ValueError("missing argument 'lattice' or 'agrid'")
             params["lattice"] = LatticeWalberla(
                 agrid=params.pop("agrid"), n_ghost_layers=1)
-        elif 'agrid' in params:
+        elif "agrid" in params:
             raise ValueError("cannot provide both 'lattice' and 'agrid'")
 
         utils.check_required_keys(self.required_keys(), params.keys())
@@ -350,80 +339,16 @@ def add_boundary_from_shape(self, shape,
                 f'Cannot process velocity value grid of shape {np.shape(velocity)}')
 
         # range checks
-        lattice_speed = self.call_method('get_lattice_speed')
+        lattice_speed = self.call_method("get_lattice_speed")
         velocity = np.array(velocity, dtype=float).reshape((-1, 3))
         velocity *= 1. / lattice_speed
         self._check_mach_limit(velocity)
 
-        velocity_flat = velocity.reshape((-1,))
-        mask_flat = shape.call_method("rasterize", grid_size=self.shape,
-                                      grid_spacing=self.agrid, grid_offset=0.5)
-
+        mask = self.get_shape_bitmask(shape=shape).astype(int)
         self.call_method(
             "add_boundary_from_shape",
-            raster=array_variant(mask_flat),
-            velocity=array_variant(velocity_flat))
-
-    def add_boundary_from_list(self, nodes,
-                               velocity=np.zeros(3, dtype=float),
-                               boundary_type=VelocityBounceBack):
-        """
-        Set boundary conditions from a list of node indices.
-
-        Parameters
-        ----------
-        nodes : (N, 3) array_like of :obj:`int`
-            List of node indices to update. If they were originally not
-            boundary nodes, they will become boundary nodes.
-        velocity : (3,) or (N, 3) array_like of :obj:`float`, optional
-            Slip velocity. By default no-slip boundary conditions are used.
-            If a vector of 3 values, a uniform slip velocity is used,
-            otherwise ``N`` must be identical to the ``N`` of ``nodes``.
-        boundary_type : Union[:class:`~espressomd.lb.VelocityBounceBack`] (optional)
-            Type of the boundary condition.
-
-        """
-        if not issubclass(boundary_type, VelocityBounceBack):
-            raise ValueError(
-                "boundary_type must be a subclass of VelocityBounceBack")
-
-        nodes = np.array(nodes, dtype=int)
-        velocity = np.array(velocity, dtype=float)
-        if len(nodes.shape) != 2 or nodes.shape[1] != 3:
-            raise ValueError(
-                f'Cannot process node list of shape {nodes.shape}')
-        if velocity.shape == (3,):
-            velocity = np.tile(velocity, (nodes.shape[0], 1))
-        elif nodes.shape != velocity.shape:
-            raise ValueError(
-                f'Node indices and velocities must have the same shape, got {nodes.shape} and {velocity.shape}')
-
-        # range checks
-        lattice_speed = self.call_method('get_lattice_speed')
-        velocity *= 1. / lattice_speed
-        self._check_mach_limit(velocity)
-        if np.any(np.logical_or(np.max(nodes, axis=0) >= self.shape,
-                                np.min(nodes, axis=0) < 0)):
-            raise ValueError(
-                f'Node indices must be in the range (0, 0, 0) to {self.shape}')
-
-        self.call_method('add_boundary_from_list', nodes=array_variant(
-            nodes.reshape((-1,))), velocity=array_variant(velocity.reshape((-1,))))
-
-    def get_nodes_inside_shape(self, shape):
-        """
-        Provide a generator for iterating over all nodes inside the given shape.
-        """
-        for idx in self.lattice.get_node_indices_inside_shape(shape):
-            yield self[idx]
-
-    def get_shape_bitmask(self, shape):
-        """Create a bitmask for the given shape."""
-        utils.check_type_or_throw_except(
-            shape, 1, espressomd.shapes.Shape, "expected a espressomd.shapes.Shape")
-        mask_flat = shape.call_method('rasterize', grid_size=self.shape,
-                                      grid_spacing=self.agrid, grid_offset=0.5)
-        return np.reshape(mask_flat, self.shape).astype(type(True))
+            raster=array_variant(mask.flatten()),
+            values=array_variant(velocity.flatten()))
 
 
 class LBFluidWalberlaGPU(HydrodynamicInteraction):
@@ -764,7 +689,7 @@ def edge_detection(boundary_mask, periodicity):
     edge[1, 1, 1] = -np.sum(edge)
 
     # periodic convolution
-    wrapped_mask = np.pad(fluid_mask.astype(int), 3 * [(2, 2)], mode='wrap')
+    wrapped_mask = np.pad(fluid_mask.astype(int), 3 * [(2, 2)], mode="wrap")
     if not periodicity[0]:
         wrapped_mask[:2, :, :] = 0
         wrapped_mask[-2:, :, :] = 0
@@ -775,7 +700,7 @@ def edge_detection(boundary_mask, periodicity):
         wrapped_mask[:, :, :2] = 0
         wrapped_mask[:, :, -2:] = 0
     convolution = scipy.signal.convolve(
-        wrapped_mask, edge, mode='same', method='direct')[2:-2, 2:-2, 2:-2]
+        wrapped_mask, edge, mode="same", method="direct")[2:-2, 2:-2, 2:-2]
     convolution = np.multiply(convolution, boundary_mask)
 
     return np.array(np.nonzero(convolution < 0)).T