Adapted EK kernel generation to new interface and added 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 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,46 @@ 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 subsitution 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 +183,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 +227,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: