walberla: add ek thermalization

maintainer/walberla_kernels/custom_additional_extensions.py

@@ -24,6 +24,7 @@
 import pystencils as ps
 import pystencils_walberla
 import sympy as sp
+import pystencils_walberla.utility
 
 
 class Dirichlet_Custom(ps.boundaries.Dirichlet):
@@ -221,7 +222,7 @@ def generate_boundary(
 ):
     struct_name = "IndexInfo"
 
-    config = pystencils_walberla.codegen.config_from_context(
+    config = pystencils_walberla.utility.config_from_context(
         generation_context,
         target=target,
         data_type=data_type,
@@ -231,6 +232,8 @@ def generate_boundary(
     create_kernel_params = config.__dict__
     del create_kernel_params["target"]
     del create_kernel_params["index_fields"]
+    del create_kernel_params["default_number_int"]
+    del create_kernel_params["skip_independence_check"]
 
     coordinate_names = ("x", "y", "z")[:dim]
 
@@ -255,7 +258,13 @@ def generate_boundary(
         index_fields=[index_field], target=target, **create_kernel_params
     )
 
-    kernel = ps.kernelcreation.create_kernel(assignment, config=kernel_config)
+    elements = [ps.boundaries.boundaryhandling.BoundaryOffsetInfo(stencil)]
+    # dummy read, such that it recognizes the field....
+    elements += [ps.astnodes.SympyAssignment(
+        ps.TypedSymbol("dummy", np.int32), index_field[0]("x"))]
+    elements += assignment
+
+    kernel = ps.kernelcreation.create_kernel(elements, config=kernel_config)
 
     if isinstance(kernel, ps.astnodes.KernelFunction):
         kernel.function_name = f"boundary_{class_name}"

maintainer/walberla_kernels/ekin.py

@@ -139,16 +139,18 @@ def flux(self, include_vof: bool = False,
             stencil_offsets = list(
                 map(lambda d: ps.stencil.direction_string_to_offset(d), stencil))
 
+            A0 = sum([sp.Matrix(ps.stencil.direction_string_to_offset(d)).norm()
+                      for d in self.flux_field.staggered_stencil]) / self.dim
+
             for i, (val, d, rng_symb) in enumerate(
                     zip(stencil, stencil_offsets, rng_symbol_gen)):
                 assert _flux_collection.main_assignments[i].lhs == self.flux_field.staggered_access(
                     val)
                 _flux_collection.main_assignments[i] = ps.Assignment(
                     self.flux_field.staggered_access(val),
                     _flux_collection.main_assignments[i].rhs + sp.sqrt(
-                        2 * self.diffusion * discretize(self.density_field.center, d)) / sp.Matrix(
-                        d).norm() * rng_symb * sp.sqrt(
-                        3) / 4)
+                        2 * self.diffusion * discretize(self.density_field.center, d) / sp.Matrix(
+                            d).norm() / A0) * (rng_symb - 0.5) * sp.sqrt(12))
 
         if include_vof:
             assert self.velocity_field is not None, "velocity field is not provided!"

maintainer/walberla_kernels/generate_ek_kernels.py

@@ -23,6 +23,7 @@
 import lbmpy
 import argparse
 import packaging.specifiers
+import numpy as np
 
 import pystencils_espresso
 import code_generation_context
@@ -37,7 +38,7 @@
 args = parser.parse_args()
 
 # Make sure we have the correct versions of the required dependencies
-for module, requirement in [(ps, "==1.2"), (lbmpy, "==1.2")]:
+for module, requirement in [(ps, "==1.3.3"), (lbmpy, "==1.3.3")]:
     assert packaging.specifiers.SpecifierSet(requirement).contains(module.__version__), \
         f"{module.__name__} version {module.__version__} " \
         f"doesn't match requirement {requirement}"
@@ -47,16 +48,20 @@
 data_type_cpp = "double" if double_precision else "float"
 data_type_np = pystencils_espresso.data_type_np[data_type_cpp]
 precision_suffix = pystencils_espresso.precision_suffix[double_precision]
-precision_rng = pystencils_espresso.precision_rng[double_precision]
+precision_rng = pystencils_espresso.precision_rng_modulo[double_precision]
 
 
-def replace_getData_with_uncheckedFastGetData(filename: str) -> None:
+def replace_macros(filename: str) -> None:
     with open(filename, "r+") as f:
         content = f.read()
         f.seek(0)
         f.truncate(0)
+        # replace getData with uncheckedFastGetData
         content = content.replace("block->getData<IndexVectors>(indexVectorID);",
                                   "block->uncheckedFastGetData<IndexVectors>(indexVectorID);")
+        # remove dummy assignment
+        content = content.replace(
+            r"const int32_t dummy = *((int32_t *  )(& _data_indexVector[12*ctr_0]));", "")
         f.write(content)
 
 
@@ -117,34 +122,47 @@ def replace_getData_with_uncheckedFastGetData(filename: str) -> None:
     rate_coef=rate_coef,
 )
 
+block_offsets = tuple(
+    ps.TypedSymbol(f"block_offset_{i}", np.uint32)
+    for i in range(3))
+
 params = {
     "target": target,
-    "cpu_vectorize_info": {"assume_inner_stride_one": False}}
+    "cpu_vectorize_info": {"assume_inner_stride_one": False}, }
 
 with code_generation_context.CodeGeneration() as ctx:
     ctx.double_accuracy = double_precision
 
     # codegen configuration
     config = pystencils_espresso.generate_config(ctx, params)
 
-    pystencils_walberla.generate_sweep(
-        ctx,
-        f"DiffusiveFluxKernel_{precision_suffix}",
-        ek.flux(include_vof=False, include_fluctuations=False,
-                rng_node=precision_rng),
-        staggered=True,
-        **params)
-    pystencils_walberla.generate_sweep(
-        ctx,
-        f"DiffusiveFluxKernelWithElectrostatic_{precision_suffix}",
-        ek_electrostatic.flux(include_vof=False, include_fluctuations=False,
-                              rng_node=precision_rng),
-        staggered=True,
-        **params)
+    for midfix, fluctuation in (("", False), ("Thermalized", True)):
+        pystencils_walberla.generate_sweep(
+            ctx,
+            f"DiffusiveFluxKernel{midfix}_{precision_suffix}",
+            ek.flux(include_vof=False, include_fluctuations=fluctuation,
+                    rng_node=precision_rng),
+            staggered=True,
+            block_offset=block_offsets if fluctuation else None,
+            **params)
+        pystencils_walberla.generate_sweep(
+            ctx,
+            f"DiffusiveFluxKernelWithElectrostatic{midfix}_{precision_suffix}",
+            ek_electrostatic.flux(include_vof=False, include_fluctuations=fluctuation,
+                                  rng_node=precision_rng),
+            staggered=True,
+            block_offset=block_offsets if fluctuation else None,
+            **params)
+
+    # the substitution for field reads is necessary, because otherwise there are
+    # "ResolvedFieldAccess" nodes that fail in the code generation
+    flux_advection = ps.AssignmentCollection(ek.flux_advection())
+    flux_advection = ps.simp.add_subexpressions_for_field_reads(flux_advection)
+
     pystencils_walberla.generate_sweep(
         ctx,
         f"AdvectiveFluxKernel_{precision_suffix}",
-        ek.flux_advection(),
+        flux_advection,
         staggered=True,
         **params)
     pystencils_walberla.generate_sweep(
@@ -166,7 +184,7 @@ def replace_getData_with_uncheckedFastGetData(filename: str) -> None:
     dynamic_flux_additional_data = custom_additional_extensions.FluxAdditionalDataHandler(
         stencil=stencil, boundary_object=dynamic_flux)
 
-    pystencils_walberla.generate_staggered_flux_boundary(
+    pystencils_espresso.generate_staggered_flux_boundary(
         generation_context=ctx,
         class_name=f"FixedFlux_{precision_suffix}",
         boundary_object=dynamic_flux,
@@ -210,8 +228,7 @@ def replace_getData_with_uncheckedFastGetData(filename: str) -> None:
             dim=dim,
             target=target,
             assignment=assignments)
-        replace_getData_with_uncheckedFastGetData(
-            filename=f"{filename_stem}.cpp")
+        replace_macros(filename=f"{filename_stem}.cpp")
 
     # ek reactions helper functions
     custom_additional_extensions.generate_kernel_selector(

maintainer/walberla_kernels/pystencils_espresso.py

@@ -17,6 +17,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 
+import numpy as np
 import sympy as sp
 import lbmpy.fieldaccess
 import lbmpy.macroscopic_value_kernels
@@ -25,6 +26,9 @@
 import pystencils_walberla
 import pystencils_walberla.utility
 
+from pystencils.astnodes import LoopOverCoordinate
+from pystencils.backends.cbackend import CustomCodeNode
+
 
 def skip_philox_unthermalized(code, result_symbols, rng_name):
     for r in result_symbols:
@@ -45,12 +49,106 @@ def get_code(self, *args, **kwargs):
         return skip_philox_unthermalized(code, self.result_symbols, self._name)
 
 
+class PhiloxTwoDoublesModulo(ps.rng.PhiloxTwoDoubles):
+    def __init__(self, dim, time_step=ps.typing.TypedSymbol(
+            "time_step", np.uint32), *args, **kwargs):
+        super().__init__(dim, time_step=time_step, *args, **kwargs)
+
+        if kwargs.get("keys") is None:
+            keys = (0,) * self._num_keys
+        else:
+            keys = kwargs.get("keys")
+
+        if kwargs.get("offsets") is None:
+            offsets = (0,) * dim
+        else:
+            offsets = kwargs.get("offsets")
+
+        coordinates = [
+            LoopOverCoordinate.get_loop_counter_symbol(i) +
+            offsets[i] for i in range(dim)]
+        if dim < 3:
+            coordinates.append(0)
+
+        # add folding to coordinates with symbols
+        field_sizes = [
+            ps.typing.TypedSymbol(
+                f"field_size_{i}",
+                np.uint32) for i in range(dim)]
+
+        new_coordinates = [
+            (coord) %
+            field_size for coord,
+            field_size in zip(
+                coordinates,
+                field_sizes)]
+        self._args = sp.sympify([time_step, *new_coordinates, *keys])
+        symbols_read = set.union(*[s.atoms(sp.Symbol) for s in self.args])
+
+        headers = self.headers
+        # set headers again, since the constructor of CustomCodeNode resets the
+        # header-list
+        CustomCodeNode.__init__(
+            self,
+            "",
+            symbols_read=symbols_read,
+            symbols_defined=self.result_symbols)
+        self.headers = headers
+
+
 class PhiloxFourFloats(ps.rng.PhiloxFourFloats):
     def get_code(self, *args, **kwargs):
         code = super().get_code(*args, **kwargs)
         return skip_philox_unthermalized(code, self.result_symbols, self._name)
 
 
+class PhiloxFourFloatsModulo(ps.rng.PhiloxFourFloats):
+    def __init__(self, dim, time_step=ps.typing.TypedSymbol(
+            "time_step", np.uint32), *args, **kwargs):
+        super().__init__(dim, time_step=time_step, *args, **kwargs)
+
+        if kwargs.get("keys") is None:
+            keys = (0,) * self._num_keys
+        else:
+            keys = kwargs.get("keys")
+
+        if kwargs.get("offsets") is None:
+            offsets = (0,) * dim
+        else:
+            offsets = kwargs.get("offsets")
+
+        coordinates = [
+            LoopOverCoordinate.get_loop_counter_symbol(i) +
+            offsets[i] for i in range(dim)]
+        if dim < 3:
+            coordinates.append(0)
+
+        # add folding to coordinates with symbols
+        field_sizes = [
+            ps.typing.TypedSymbol(
+                f"field_size_{i}",
+                np.uint32) for i in range(dim)]
+
+        new_coordinates = [
+            (coord) %
+            field_size for coord,
+            field_size in zip(
+                coordinates,
+                field_sizes)]
+        self._args = sp.sympify([time_step, *new_coordinates, *keys])
+        symbols_read = set.union(*[s.atoms(sp.Symbol) for s in self.args])
+
+        headers = self.headers
+        # set headers again, since the constructor of CustomCodeNode resets the
+        # header-list
+        CustomCodeNode.__init__(
+            self,
+            "",
+            symbols_read=symbols_read,
+            symbols_defined=self.result_symbols)
+        self.headers = headers
+
+
 precision_prefix = {
     True: 'DoublePrecision',
     False: 'SinglePrecision'}
@@ -60,6 +158,9 @@ def get_code(self, *args, **kwargs):
 precision_rng = {
     True: PhiloxTwoDoubles,
     False: PhiloxFourFloats}
+precision_rng_modulo = {
+    True: PhiloxTwoDoublesModulo,
+    False: PhiloxFourFloatsModulo}
 data_type_np = {'double': 'float64', 'float': 'float32'}
 
 
@@ -161,3 +262,20 @@ def generate_setters(ctx, lb_method, params):
         fields['velocity'].center_vector,
         fields['pdfs'].center_vector)
     return pdfs_setter
+
+
+# this can be removed when https://i10git.cs.fau.de/walberla/walberla/-/issues/247 is fixed
+def generate_staggered_flux_boundary(generation_context, class_name, boundary_object,
+                                     dim, neighbor_stencil, index_shape, target=ps.Target.CPU, **kwargs):
+    assert dim == len(neighbor_stencil[0])
+    pystencils_walberla.boundary.generate_boundary(
+        generation_context=generation_context,
+        class_name=class_name,
+        boundary_object=boundary_object,
+        field_name='flux',
+        neighbor_stencil=neighbor_stencil,
+        index_shape=index_shape,
+        field_type=ps.FieldType.STAGGERED_FLUX,
+        kernel_creation_function=None,
+        target=target,
+        **kwargs)

maintainer/walberla_kernels/templates/Boundary.tmpl.h

@@ -59,6 +59,15 @@
 #pragma GCC diagnostic ignored "-Wunused-parameter"
 #endif
 
+#ifdef __GNUC__
+#define RESTRICT __restrict__
+#elif _MSC_VER
+#define RESTRICT __restrict
+#else
+#define RESTRICT
+
+#endif
+
 namespace walberla {
 namespace {{namespace}} {
 

src/core/unit_tests/ek_interface_test.cpp

@@ -148,7 +148,8 @@ BOOST_AUTO_TEST_CASE(ek_interface_walberla) {
     auto constexpr order = 2.;
     auto ek_species = walberla::new_ek_walberla(
         espresso::ek_lattice, params.diffusion, params.kT, params.valency,
-        params.ext_efield, params.density, false, false, single_precision);
+        params.ext_efield, params.density, false, false, single_precision,
+        false, 0u);
     auto ek_reactant = std::make_shared<EKReactant>(ek_species, stoich, order);
     auto ek_reaction = std::make_shared<walberla::EKReactionImpl>(
         espresso::ek_lattice,

src/python/espressomd/electrokinetics.py

@@ -94,6 +94,10 @@ class EKSpecies(ScriptInterfaceHelper,
         Set it to 0 for an unthermalized species.
     single_precision : :obj:`bool`, optional
         Use single-precision floating-point arithmetic.
+    thermalized : :obj:`bool`, optional
+        Whether the species is thermalized.
+    seed : :obj:`int`, optional
+        Seed for the random number generator.
 
     Methods
     -------
@@ -169,7 +173,8 @@ def __init__(self, *args, **kwargs):
 
     def default_params(self):
         return {"single_precision": False,
-                "kT": 0., "ext_efield": [0., 0., 0.]}
+                "kT": 0., "ext_efield": [0., 0., 0.],
+                "thermalized": False}
 
     def __getitem__(self, key):
         if isinstance(key, (tuple, list, np.ndarray)) and len(key) == 3: