src/lbmpy/__init__.py

+from .creationfunctions import (
+    create_lb_ast,
+    create_lb_collision_rule,
+    create_lb_function,
+    create_lb_method,
+    create_lb_update_rule,
+    LBMConfig,
+    LBMOptimisation,
+)
+from .enums import Stencil, Method, ForceModel, CollisionSpace, SubgridScaleModel
+from .lbstep import LatticeBoltzmannStep
+from .macroscopic_value_kernels import (
+    pdf_initialization_assignments,
+    macroscopic_values_getter,
+    strain_rate_tensor_getter,
+    compile_macroscopic_values_getter,
+    compile_macroscopic_values_setter,
+    create_advanced_velocity_setter_collision_rule,
+)
+from .maxwellian_equilibrium import get_weights
+from .relaxationrates import (
+    relaxation_rate_from_lattice_viscosity,
+    lattice_viscosity_from_relaxation_rate,
+    relaxation_rate_from_magic_number,
+)
+from .scenarios import create_lid_driven_cavity, create_fully_periodic_flow
+from .stencils import LBStencil
+
+
+__all__ = [
+    "create_lb_ast",
+    "create_lb_collision_rule",
+    "create_lb_function",
+    "create_lb_method",
+    "create_lb_update_rule",
+    "LBMConfig",
+    "LBMOptimisation",
+    "Stencil",
+    "Method",
+    "ForceModel",
+    "CollisionSpace",
+    "SubgridScaleModel",
+    "LatticeBoltzmannStep",
+    "pdf_initialization_assignments",
+    "macroscopic_values_getter",
+    "strain_rate_tensor_getter",
+    "compile_macroscopic_values_getter",
+    "compile_macroscopic_values_setter",
+    "create_advanced_velocity_setter_collision_rule",
+    "get_weights",
+    "relaxation_rate_from_lattice_viscosity",
+    "lattice_viscosity_from_relaxation_rate",
+    "relaxation_rate_from_magic_number",
+    "create_lid_driven_cavity",
+    "create_fully_periodic_flow",
+    "LBStencil",
+]
+
+
+from . import _version
+__version__ = _version.get_versions()['version']

lbmpy/advanced_streaming/__init__.py → src/lbmpy/advanced_streaming/__init__.py

-from .indexing import BetweenTimestepsIndexing, NeighbourOffsetArrays
+from .indexing import BetweenTimestepsIndexing
 from .communication import get_communication_slices, LBMPeriodicityHandling
 from .utility import Timestep, get_accessor, is_inplace, get_timesteps, \
     numeric_index, numeric_offsets, inverse_dir_index, AccessPdfValues
 
-__all__ = ['BetweenTimestepsIndexing', 'NeighbourOffsetArrays',
+__all__ = ['BetweenTimestepsIndexing',
            'get_communication_slices', 'LBMPeriodicityHandling',
            'Timestep', 'get_accessor', 'is_inplace', 'get_timesteps',
            'numeric_index', 'numeric_offsets', 'inverse_dir_index', 'AccessPdfValues']

lbmpy/advanced_streaming/communication.py → src/lbmpy/advanced_streaming/communication.py

 import itertools
 from pystencils import CreateKernelConfig, Field, Assignment, AssignmentCollection
-from pystencils.slicing import shift_slice, get_slice_before_ghost_layer, normalize_slice
-from lbmpy.advanced_streaming.utility import is_inplace, get_accessor, numeric_index,\
-    Timestep, get_timesteps, numeric_offsets
+from pystencils.slicing import (
+    shift_slice,
+    get_slice_before_ghost_layer,
+    normalize_slice,
+)
+from lbmpy.advanced_streaming.utility import (
+    is_inplace,
+    get_accessor,
+    numeric_index,
+    Timestep,
+    get_timesteps,
+    numeric_offsets,
+)
 from pystencils.datahandling import SerialDataHandling
 from pystencils.enums import Target
 from itertools import chain
@@ -10,22 +20,28 @@ from itertools import chain
 
 class LBMPeriodicityHandling:
 
-    def __init__(self, stencil, data_handling, pdf_field_name,
-                 streaming_pattern='pull', ghost_layers=1,
-                 pycuda_direct_copy=True):
+    def __init__(
+        self,
+        stencil,
+        data_handling,
+        pdf_field_name,
+        streaming_pattern="pull",
+        ghost_layers=1,
+        cupy_direct_copy=True,
+    ):
         """
-            Periodicity Handling for Lattice Boltzmann Streaming.
-
-            **On the usage with cuda:**
-            - pycuda allows the copying of sliced arrays within device memory using the numpy syntax,
-            e.g. `dst[:,0] = src[:,-1]`. In this implementation, this is the default for periodicity
-            handling. Alternatively, if you set `pycuda_direct_copy=False`, GPU kernels are generated and
-            compiled. The compiled kernels are almost twice as fast in execution as pycuda array copying,
-            but especially for large stencils like D3Q27, their compilation can take up to 20 seconds.
-            Choose your weapon depending on your use case.
+        Periodicity Handling for Lattice Boltzmann Streaming.
+
+        **On the usage with cuda:**
+        - cupy allows the copying of sliced arrays within device memory using the numpy syntax,
+        e.g. `dst[:,0] = src[:,-1]`. In this implementation, this is the default for periodicity
+        handling. Alternatively, if you set `cupy_direct_copy=False`, GPU kernels are generated and
+        compiled. The compiled kernels are almost twice as fast in execution as cupy array copying,
+        but especially for large stencils like D3Q27, their compilation can take up to 20 seconds.
+        Choose your weapon depending on your use case.
         """
         if not isinstance(data_handling, SerialDataHandling):
-            raise ValueError('Only serial data handling is supported!')
+            raise ValueError("Only serial data handling is supported!")
 
         self.stencil = stencil
         self.dim = stencil.D
@@ -40,7 +56,7 @@ class LBMPeriodicityHandling:
         self.inplace_pattern = is_inplace(streaming_pattern)
 
         self.cpu = self.target == Target.CPU
-        self.pycuda_direct_copy = self.target == Target.GPU and pycuda_direct_copy
+        self.cupy_direct_copy = self.target == Target.GPU and cupy_direct_copy
 
         def is_copy_direction(direction):
             s = 0
@@ -56,14 +72,18 @@ class LBMPeriodicityHandling:
         self.comm_slices = []
         timesteps = get_timesteps(streaming_pattern)
         for timestep in timesteps:
-            slices_per_comm_dir = get_communication_slices(stencil=stencil,
-                                                           comm_stencil=copy_directions,
-                                                           streaming_pattern=streaming_pattern,
-                                                           prev_timestep=timestep,
-                                                           ghost_layers=ghost_layers)
-            self.comm_slices.append(list(chain.from_iterable(v for k, v in slices_per_comm_dir.items())))
-
-        if self.target == Target.GPU and not pycuda_direct_copy:
+            slices_per_comm_dir = get_communication_slices(
+                stencil=stencil,
+                comm_stencil=copy_directions,
+                streaming_pattern=streaming_pattern,
+                prev_timestep=timestep,
+                ghost_layers=ghost_layers,
+            )
+            self.comm_slices.append(
+                list(chain.from_iterable(v for k, v in slices_per_comm_dir.items()))
+            )
+
+        if self.target == Target.GPU and not cupy_direct_copy:
             self.device_copy_kernels = list()
             for timestep in timesteps:
                 self.device_copy_kernels.append(self._compile_copy_kernels(timestep))
@@ -81,16 +101,16 @@ class LBMPeriodicityHandling:
             arr[dst] = arr[src]
 
     def _compile_copy_kernels(self, timestep):
+        assert self.target == Target.GPU
         pdf_field = self.dh.fields[self.pdf_field_name]
         kernels = []
         for src, dst in self.comm_slices[timestep.idx]:
-            kernels.append(
-                periodic_pdf_copy_kernel(pdf_field, src, dst, target=self.target))
+            kernels.append(periodic_pdf_gpu_copy_kernel(pdf_field, src, dst))
         return kernels
 
     def _periodicity_handling_gpu(self, prev_timestep):
         arr = self.dh.gpu_arrays[self.pdf_field_name]
-        if self.pycuda_direct_copy:
+        if self.cupy_direct_copy:
             for src, dst in self.comm_slices[prev_timestep.idx]:
                 arr[dst] = arr[src]
         else:
@@ -100,7 +120,12 @@ class LBMPeriodicityHandling:
 
 
 def get_communication_slices(
-        stencil, comm_stencil=None, streaming_pattern='pull', prev_timestep=Timestep.BOTH, ghost_layers=1):
+    stencil,
+    comm_stencil=None,
+    streaming_pattern="pull",
+    prev_timestep=Timestep.BOTH,
+    ghost_layers=1,
+):
     """
     Return the source and destination slices for periodicity handling or communication between blocks.
 
@@ -116,7 +141,9 @@ def get_communication_slices(
     if comm_stencil is None:
         comm_stencil = itertools.product(*([-1, 0, 1] for _ in range(stencil.D)))
 
-    pdfs = Field.create_generic('pdfs', spatial_dimensions=len(stencil[0]), index_shape=(stencil.Q,))
+    pdfs = Field.create_generic(
+        "pdfs", spatial_dimensions=len(stencil[0]), index_shape=(stencil.Q,)
+    )
     write_accesses = get_accessor(streaming_pattern, prev_timestep).write(pdfs, stencil)
     slices_per_comm_direction = dict()
 
@@ -130,7 +157,9 @@ def get_communication_slices(
             d = stencil.index(streaming_dir)
             write_index = numeric_index(write_accesses[d])[0]
 
-            origin_slice = get_slice_before_ghost_layer(comm_dir, ghost_layers=ghost_layers, thickness=1)
+            origin_slice = get_slice_before_ghost_layer(
+                comm_dir, ghost_layers=ghost_layers, thickness=1
+            )
             src_slice = _fix_length_one_slices(origin_slice)
 
             write_offsets = numeric_offsets(write_accesses[d])
@@ -138,13 +167,15 @@ def get_communication_slices(
 
             # TODO: this is just a hotfix. _trim_slice_in_direction breaks FreeSlip BC with adjacent periodic side
             if streaming_pattern != "pull":
-                src_slice = shift_slice(_trim_slice_in_direction(src_slice, tangential_dir), write_offsets)
+                src_slice = shift_slice(
+                    _trim_slice_in_direction(src_slice, tangential_dir), write_offsets
+                )
 
             neighbour_transform = _get_neighbour_transform(comm_dir, ghost_layers)
             dst_slice = shift_slice(src_slice, neighbour_transform)
 
-            src_slice = src_slice + (write_index, )
-            dst_slice = dst_slice + (write_index, )
+            src_slice = src_slice + (write_index,)
+            dst_slice = dst_slice + (write_index,)
 
             slices_for_dir.append((src_slice, dst_slice))
 
@@ -152,10 +183,10 @@ def get_communication_slices(
     return slices_per_comm_direction
 
 
-def periodic_pdf_copy_kernel(pdf_field, src_slice, dst_slice,
-                             domain_size=None, target=Target.GPU):
-    """Copies a rectangular array slice onto another non-overlapping array slice"""
-    from pystencils.gpucuda.kernelcreation import create_cuda_kernel
+def periodic_pdf_gpu_copy_kernel(pdf_field, src_slice, dst_slice, domain_size=None):
+    """Generate a GPU kernel which copies all values from one slice of a field
+    to another non-overlapping slice."""
+    from pystencils import create_kernel
 
     pdf_idx = src_slice[-1]
     assert isinstance(pdf_idx, int), "PDF index needs to be an integer constant"
@@ -163,6 +194,7 @@ def periodic_pdf_copy_kernel(pdf_field, src_slice, dst_slice,
     src_slice = src_slice[:-1]
     dst_slice = dst_slice[:-1]
 
+    # TODO this is the domain_size with GL
     if domain_size is None:
         domain_size = pdf_field.spatial_shape
 
@@ -175,18 +207,28 @@ def periodic_pdf_copy_kernel(pdf_field, src_slice, dst_slice,
     def _stop(s):
         return s.stop if isinstance(s, slice) else s
 
-    offset = [_start(s1) - _start(s2) for s1, s2 in zip(normalized_from_slice, normalized_to_slice)]
-    assert offset == [_stop(s1) - _stop(s2) for s1, s2 in zip(normalized_from_slice, normalized_to_slice)], \
-        "Slices have to have same size"
-
-    copy_eq = AssignmentCollection(main_assignments=[Assignment(pdf_field(pdf_idx), pdf_field[tuple(offset)](pdf_idx))])
-    config = CreateKernelConfig(iteration_slice=dst_slice, skip_independence_check=True)
-    ast = create_cuda_kernel(copy_eq, config=config)
-    if target == Target.GPU:
-        from pystencils.gpucuda import make_python_function
-        return make_python_function(ast)
-    else:
-        raise ValueError('Invalid target:', target)
+    offset = [
+        _start(s1) - _start(s2)
+        for s1, s2 in zip(normalized_from_slice, normalized_to_slice)
+    ]
+    assert offset == [
+        _stop(s1) - _stop(s2)
+        for s1, s2 in zip(normalized_from_slice, normalized_to_slice)
+    ], "Slices have to have same size"
+
+    copy_eq = AssignmentCollection(
+        main_assignments=[
+            Assignment(pdf_field(pdf_idx), pdf_field[tuple(offset)](pdf_idx))
+        ]
+    )
+    config = CreateKernelConfig(
+        iteration_slice=dst_slice,
+        skip_independence_check=True,
+        target=Target.GPU,
+    )
+
+    ast = create_kernel(copy_eq, config=config)
+    return ast.compile()
 
 
 def _extend_dir(direction):
@@ -195,10 +237,10 @@ def _extend_dir(direction):
     elif direction[0] == 0:
         for d in [-1, 0, 1]:
             for rest in _extend_dir(direction[1:]):
-                yield (d, ) + rest
+                yield (d,) + rest
     else:
         for rest in _extend_dir(direction[1:]):
-            yield (direction[0], ) + rest
+            yield (direction[0],) + rest
 
 
 def _get_neighbour_transform(direction, ghost_layers):

lbmpy/advanced_streaming/indexing.py → src/lbmpy/advanced_streaming/indexing.py

@@ -3,17 +3,12 @@ import sympy as sp
 import pystencils as ps
 
 from pystencils.typing import TypedSymbol, create_type
-from pystencils.backends.cbackend import CustomCodeNode
-
 from lbmpy.advanced_streaming.utility import get_accessor, inverse_dir_index, is_inplace, Timestep
+from lbmpy.custom_code_nodes import TranslationArraysNode
 
 from itertools import product
 
 
-def _array_pattern(dtype, name, content):
-    return f"const {str(dtype)} {name} [] = {{ {','.join(str(c) for c in content)} }}; \n"
-
-
 class BetweenTimestepsIndexing:
 
     #   ==============================================
@@ -30,7 +25,7 @@ class BetweenTimestepsIndexing:
 
     @property
     def inverse_dir_symbol(self):
-        """Symbol denoting the inversion of a PDF field index. 
+        """Symbol denoting the inversion of a PDF field index.
         Use only at top-level of index to f_out or f_in, otherwise it can't be correctly replaced."""
         return sp.IndexedBase('invdir')
 
@@ -168,90 +163,21 @@ class BetweenTimestepsIndexing:
         return trivial_index_translations, trivial_offset_translations
 
     def create_code_node(self):
-        return BetweenTimestepsIndexing.TranslationArraysNode(self)
-
-    class TranslationArraysNode(CustomCodeNode):
-
-        def __init__(self, indexing):
-            code = ''
-            symbols_defined = set()
-
-            for f_dir, inv in indexing._required_index_arrays:
-                indices, offsets = indexing._get_translated_indices_and_offsets(f_dir, inv)
-                index_array_symbol = indexing._index_array_symbol(f_dir, inv)
-                symbols_defined.add(index_array_symbol)
-                code += _array_pattern(indexing._index_dtype, index_array_symbol.name, indices)
-
-            for f_dir, inv in indexing._required_offset_arrays:
-                indices, offsets = indexing._get_translated_indices_and_offsets(f_dir, inv)
-                offset_array_symbols = indexing._offset_array_symbols(f_dir, inv)
-                symbols_defined |= set(offset_array_symbols)
-                for d, arrsymb in enumerate(offset_array_symbols):
-                    code += _array_pattern(indexing._offsets_dtype, arrsymb.name, offsets[d])
-
-            super(BetweenTimestepsIndexing.TranslationArraysNode, self).__init__(
-                code, symbols_read=set(), symbols_defined=symbols_defined)
+        array_content = list()
+        symbols_defined = set()
+        for f_dir, inv in self._required_index_arrays:
+            indices, offsets = self._get_translated_indices_and_offsets(f_dir, inv)
+            index_array_symbol = self._index_array_symbol(f_dir, inv)
+            symbols_defined.add(index_array_symbol)
+            array_content.append((self._index_dtype, index_array_symbol.name, indices))
 
-        def __str__(self):
-            return "Variable PDF Access Translation Arrays"
+        for f_dir, inv in self._required_offset_arrays:
+            indices, offsets = self._get_translated_indices_and_offsets(f_dir, inv)
+            offset_array_symbols = self._offset_array_symbols(f_dir, inv)
+            symbols_defined |= set(offset_array_symbols)
+            for d, arrsymb in enumerate(offset_array_symbols):
+                array_content.append((self._offsets_dtype, arrsymb.name, offsets[d]))
 
-        def __repr__(self):
-            return "Variable PDF Access Translation Arrays"
+        return TranslationArraysNode(array_content, symbols_defined)
 
 #   end class AdvancedStreamingIndexing
-
-
-class NeighbourOffsetArrays(CustomCodeNode):
-
-    @staticmethod
-    def neighbour_offset(dir_idx, stencil):
-        if isinstance(sp.sympify(dir_idx), sp.Integer):
-            return stencil[dir_idx]
-        else:
-            return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx]
-                         for symbol in NeighbourOffsetArrays._offset_symbols(len(stencil[0]))])
-
-    @staticmethod
-    def _offset_symbols(dim):
-        return [TypedSymbol(f"neighbour_offset_{d}", create_type(np.int32)) for d in ['x', 'y', 'z'][:dim]]
-
-    def __init__(self, stencil, offsets_dtype=np.int32):
-        offsets_dtype = create_type(offsets_dtype)
-        dim = len(stencil[0])
-
-        array_symbols = NeighbourOffsetArrays._offset_symbols(dim)
-        code = "\n"
-        for i, arrsymb in enumerate(array_symbols):
-            code += _array_pattern(offsets_dtype, arrsymb.name, (d[i] for d in stencil))
-
-        offset_symbols = NeighbourOffsetArrays._offset_symbols(dim)
-        super(NeighbourOffsetArrays, self).__init__(code, symbols_read=set(),
-                                                    symbols_defined=set(offset_symbols))
-
-
-class MirroredStencilDirections(CustomCodeNode):
-
-    @staticmethod
-    def mirror_stencil(direction, mirror_axis):
-        assert mirror_axis <= len(direction), f"only {len(direction)} axis available for mirage"
-        direction = list(direction)
-        direction[mirror_axis] = -direction[mirror_axis]
-
-        return tuple(direction)
-
-    @staticmethod
-    def _mirrored_symbol(mirror_axis):
-        axis = ['x', 'y', 'z']
-        return TypedSymbol(f"{axis[mirror_axis]}_axis_mirrored_stencil_dir", create_type(np.int32))
-
-    def __init__(self, stencil, mirror_axis, dtype=np.int32):
-        offsets_dtype = create_type(dtype)
-
-        mirrored_stencil_symbol = MirroredStencilDirections._mirrored_symbol(mirror_axis)
-        mirrored_directions = [stencil.index(MirroredStencilDirections.mirror_stencil(direction, mirror_axis))
-                               for direction in stencil]
-        code = "\n"
-        code += _array_pattern(offsets_dtype, mirrored_stencil_symbol.name, mirrored_directions)
-
-        super(MirroredStencilDirections, self).__init__(code, symbols_read=set(),
-                                                        symbols_defined={mirrored_stencil_symbol})

lbmpy/advanced_streaming/utility.py → src/lbmpy/advanced_streaming/utility.py

@@ -58,24 +58,27 @@ odd_accessors = {
 }
 
 
+def is_inplace(streaming_pattern):
+    if streaming_pattern not in streaming_patterns:
+        raise ValueError('Invalid streaming pattern', streaming_pattern)
+
+    return streaming_pattern in ['aa', 'esotwist', 'esopull', 'esopush']
+
+
 def get_accessor(streaming_pattern: str, timestep: Timestep) -> PdfFieldAccessor:
     if streaming_pattern not in streaming_patterns:
         raise ValueError(
             "Invalid value of parameter 'streaming_pattern'.", streaming_pattern)
 
+    if is_inplace(streaming_pattern) and (timestep == Timestep.BOTH):
+        raise ValueError(f"Invalid timestep for streaming pattern {streaming_pattern}: {str(timestep)}")
+
     if timestep == Timestep.EVEN:
         return even_accessors[streaming_pattern]
     else:
         return odd_accessors[streaming_pattern]
 
 
-def is_inplace(streaming_pattern):
-    if streaming_pattern not in streaming_patterns:
-        raise ValueError('Invalid streaming pattern', streaming_pattern)
-
-    return streaming_pattern in ['aa', 'esotwist', 'esopull', 'esopush']
-
-
 def get_timesteps(streaming_pattern):
     return (Timestep.EVEN, Timestep.ODD) if is_inplace(streaming_pattern) else (Timestep.BOTH, )
 

src/lbmpy/analytical_solutions.py

+from typing import Union
+from numpy.typing import NDArray
+
+
+def poiseuille_flow(middle_distance: Union[float, NDArray], height,
+                    ext_force_density: float, dyn_visc: float) -> Union[float, NDArray]:
+    """
+    Analytical solution for plane Poiseuille flow.
+
+    Args:
+        middle_distance: Distance to the middle plane of the channel.
+        height: Distance between the boundaries.
+        ext_force_density: Force density on the fluid normal to the boundaries.
+        dyn_visc: dyn_visc
+
+    Returns:
+        A numpy array of the poiseuille profile if middle_distance is given as array otherwise of velocity of
+        the position given with middle_distance
+    """
+    return ext_force_density * 1. / (2 * dyn_visc) * (height**2.0 / 4.0 - middle_distance**2.0)

src/lbmpy/boundaries/__init__.py

+from lbmpy.boundaries.boundaryconditions import (
+    UBB, FixedDensity, DiffusionDirichlet, SimpleExtrapolationOutflow, WallFunctionBounce,
+    ExtrapolationOutflow, NeumannByCopy, NoSlip, NoSlipLinearBouzidi, QuadraticBounceBack, StreamInConstant, FreeSlip)
+from lbmpy.boundaries.boundaryhandling import LatticeBoltzmannBoundaryHandling
+from lbmpy.boundaries.wall_function_models import MoninObukhovSimilarityTheory, LogLaw, MuskerLaw, SpaldingsLaw
+
+__all__ = ['NoSlip', 'NoSlipLinearBouzidi', 'QuadraticBounceBack', 'FreeSlip', 'WallFunctionBounce',
+           'UBB', 'FixedDensity',
+           'SimpleExtrapolationOutflow', 'ExtrapolationOutflow',
+           'DiffusionDirichlet', 'NeumannByCopy', 'StreamInConstant',
+           'LatticeBoltzmannBoundaryHandling',
+           'MoninObukhovSimilarityTheory', 'LogLaw', 'MuskerLaw', 'SpaldingsLaw']

lbmpy/boundaries/boundaries_in_kernel.py → src/lbmpy/boundaries/boundaries_in_kernel.py

 import sympy as sp
 
-from lbmpy.boundaries.boundaryhandling import LbmWeightInfo
 from lbmpy.advanced_streaming.indexing import BetweenTimestepsIndexing
 from lbmpy.advanced_streaming.utility import Timestep, get_accessor
+from lbmpy.custom_code_nodes import LbmWeightInfo
 from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo
 from pystencils.assignment import Assignment
 from pystencils.astnodes import Block, Conditional, LoopOverCoordinate, SympyAssignment
@@ -67,7 +67,7 @@ def boundary_conditional(boundary, direction, streaming_pattern, prev_timestep,
 
     assignments = []
     for direction_idx in dir_indices:
-        rule = boundary(f_out, f_in, direction_idx, inv_dir, lb_method, index_field=None)
+        rule = boundary(f_out, f_in, direction_idx, inv_dir, lb_method, index_field=None, force_vector=None)
 
         #   rhs: replace f_out by post collision symbols.
         rhs_substitutions = {f_out(i): sym for i, sym in enumerate(lb_method.post_collision_pdf_symbols)}

lbmpy/boundaries/boundaryhandling.py → src/lbmpy/boundaries/boundaryhandling.py

+from dataclasses import replace
 import numpy as np
-import sympy as sp
-from lbmpy.advanced_streaming.indexing import BetweenTimestepsIndexing
-from lbmpy.advanced_streaming.utility import is_inplace, Timestep, AccessPdfValues
-from pystencils import Field, Assignment, TypedSymbol, create_kernel
-from pystencils.stencil import inverse_direction
-from pystencils import CreateKernelConfig, Target
+
+from pystencils import Assignment, CreateKernelConfig, create_kernel, Field, Target
 from pystencils.boundaries import BoundaryHandling
 from pystencils.boundaries.createindexlist import numpy_data_type_for_boundary_object
-from pystencils.backends.cbackend import CustomCodeNode
+from pystencils.field import FieldType
+from pystencils.simp import add_subexpressions_for_field_reads
+from pystencils.stencil import inverse_direction
+
+from lbmpy.advanced_streaming.indexing import BetweenTimestepsIndexing
+from lbmpy.advanced_streaming.utility import is_inplace, Timestep, AccessPdfValues
 
 
 class LatticeBoltzmannBoundaryHandling(BoundaryHandling):
     """
-    Enables boundary handling for LBM simulations with advanced streaming patterns. 
-    For the in-place patterns AA and EsoTwist, two kernels are generated for a boundary 
+    Enables boundary handling for LBM simulations with advanced streaming patterns.
+    For the in-place patterns AA and EsoTwist, two kernels are generated for a boundary
     object and the right one selected depending on the time step.
     """
 
@@ -154,53 +156,39 @@ class LatticeBoltzmannBoundaryHandling(BoundaryHandling):
 
 # end class LatticeBoltzmannBoundaryHandling
 
-
-class LbmWeightInfo(CustomCodeNode):
-    def __init__(self, lb_method, data_type='double'):
-        self.weights_symbol = TypedSymbol("weights", data_type)
-        data_type_string = "double" if self.weights_symbol.dtype.numpy_dtype == np.float64 else "float"
-
-        weights = [str(w.evalf(17)) for w in lb_method.weights]
-        if data_type_string == "float":
-            weights = "f, ".join(weights)
-            weights += "f"  # suffix for the last element
-        else:
-            weights = ", ".join(weights)
-        w_sym = self.weights_symbol
-        code = f"const {data_type_string} {w_sym.name} [] = {{{ weights }}};\n"
-        super(LbmWeightInfo, self).__init__(code, symbols_read=set(), symbols_defined={w_sym})
-
-    def weight_of_direction(self, dir_idx, lb_method=None):
-        if isinstance(sp.sympify(dir_idx), sp.Integer):
-            return lb_method.weights[dir_idx].evalf(17)
-        else:
-            return sp.IndexedBase(self.weights_symbol, shape=(1,))[dir_idx]
-
-
-# end class LbmWeightInfo
-
-
 def create_lattice_boltzmann_boundary_kernel(pdf_field, index_field, lb_method, boundary_functor,
                                              prev_timestep=Timestep.BOTH, streaming_pattern='pull',
-                                             target=Target.CPU, **kernel_creation_args):
+                                             target=Target.CPU, force_vector=None, **kernel_creation_args):
 
     indexing = BetweenTimestepsIndexing(
         pdf_field, lb_method.stencil, prev_timestep, streaming_pattern, np.int32, np.int32)
 
+    dim = lb_method.stencil.D
     f_out, f_in = indexing.proxy_fields
     dir_symbol = indexing.dir_symbol
     inv_dir = indexing.inverse_dir_symbol
 
-    boundary_assignments = boundary_functor(f_out, f_in, dir_symbol, inv_dir, lb_method, index_field)
+    config = CreateKernelConfig(target=target, default_number_int="int32",
+                                skip_independence_check=True, **kernel_creation_args)
+
+    default_data_type = config.data_type.default_factory()
+
+    if force_vector is None:
+        force_vector_type = np.dtype([(f"F_{i}", default_data_type.c_name) for i in range(dim)], align=True)
+        force_vector = Field.create_generic('force_vector', spatial_dimensions=1,
+                                            dtype=force_vector_type, field_type=FieldType.INDEXED)
+
+    config = replace(config, index_fields=[index_field, force_vector])
+
+    boundary_assignments = boundary_functor(f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector)
     boundary_assignments = indexing.substitute_proxies(boundary_assignments)
 
-    #   Code Elements inside the loop
+    if pdf_field.dtype != default_data_type:
+        boundary_assignments = add_subexpressions_for_field_reads(boundary_assignments, data_type=default_data_type)
+
     elements = [Assignment(dir_symbol, index_field[0]('dir'))]
     elements += boundary_assignments.all_assignments
 
-    config = CreateKernelConfig(index_fields=[index_field], target=target, default_number_int="int32",
-                                skip_independence_check=True, **kernel_creation_args)
-
     kernel = create_kernel(elements, config=config)
 
     #   Code Elements ahead of the loop

src/lbmpy/boundaries/wall_function_models.py

+import sympy as sp
+from abc import ABC, abstractmethod
+
+from pystencils import Assignment
+
+
+class WallFunctionModel(ABC):
+    def __init__(self, name):
+        self._name = name
+
+    @abstractmethod
+    def shear_stress_assignments(self, density_symbol: sp.Symbol, shear_stress_symbol: sp.Symbol,
+                                 velocity_symbol: sp.Symbol, wall_distance, u_tau_target):
+        """
+        Computes a symbolic representation for the log law.
+
+        Args:
+            density_symbol: symbol density, should be provided by the LB method's conserved quantity computation
+            shear_stress_symbol: symbolic wall shear stress to which the calculated shear stress will be assigned
+            velocity_symbol: symbolic velocity that is taken as a reference in the wall functions
+            wall_distance: distance to the wall, equals to 0.5 in standard cell-centered LBM
+            u_tau_target: in implicit wall functions, a target friction velocity can be provided which will be used as
+                          initial guess in the Newton iteration. This target friction velocity can be obtained, e.g.,
+                          from the target friction Reynolds number
+        """
+        pass
+
+
+# end class WallFunctionModel
+
+
+class ExplicitWallFunctionModel(WallFunctionModel, ABC):
+    """
+    Abstract base class for explicit wall functions that can be solved directly for the wall shear stress.
+    """
+
+    def __init__(self, name):
+        super(ExplicitWallFunctionModel, self).__init__(name=name)
+
+
+class MoninObukhovSimilarityTheory(ExplicitWallFunctionModel):
+    def __init__(self, z0, kappa=0.41, phi=0, name="MOST"):
+        self.z0 = z0
+        self.kappa = kappa
+        self.phi = phi
+
+        super(MoninObukhovSimilarityTheory, self).__init__(name=name)
+
+    def shear_stress_assignments(self, density_symbol: sp.Symbol, shear_stress_symbol: sp.Symbol,
+                                 velocity_symbol: sp.Symbol, wall_distance, u_tau_target=None):
+        u_tau = velocity_symbol * self.kappa / sp.ln(wall_distance / self.z0 + self.phi)
+        return [Assignment(shear_stress_symbol, u_tau ** 2 * density_symbol)]
+
+
+class ImplicitWallFunctionModel(WallFunctionModel, ABC):
+    """
+    Abstract base class for implicit wall functions that require a Newton procedure to solve for the wall shear stress.
+    """
+
+    def __init__(self, name, newton_steps, viscosity):
+        self.newton_steps = newton_steps
+        self.u_tau = sp.symbols(f"wall_function_u_tau_:{self.newton_steps + 1}")
+        self.delta = sp.symbols(f"wall_function_delta_:{self.newton_steps}")
+
+        self.viscosity = viscosity
+
+        super(ImplicitWallFunctionModel, self).__init__(name=name)
+
+    def newton_iteration(self, wall_law):
+        m = -wall_law / wall_law.diff(self.u_tau[0])
+
+        assignments = []
+        for i in range(self.newton_steps):
+            assignments.append(Assignment(self.delta[i], m.subs({self.u_tau[0]: self.u_tau[i]})))
+            assignments.append(Assignment(self.u_tau[i + 1], self.u_tau[i] + self.delta[i]))
+
+        return assignments
+
+
+class LogLaw(ImplicitWallFunctionModel):
+    """
+    Analytical model for the velocity profile inside the boundary layer, obtained from the mean velocity gradient.
+    Only valid in the log-law region.
+    """
+
+    def __init__(self, viscosity, newton_steps=5, kappa=0.41, b=5.2, name="LogLaw"):
+        self.kappa = kappa
+        self.b = b
+
+        super(LogLaw, self).__init__(name=name, newton_steps=newton_steps, viscosity=viscosity)
+
+    def shear_stress_assignments(self, density_symbol: sp.Symbol, shear_stress_symbol: sp.Symbol,
+                                 velocity_symbol: sp.Symbol, wall_distance, u_tau_target=None):
+        def law(u_p, y_p):
+            return 1 / self.kappa * sp.ln(y_p) + self.b - u_p
+
+        u_plus = velocity_symbol / self.u_tau[0]
+        y_plus = (wall_distance * self.u_tau[0]) / self.viscosity
+
+        u_tau_init = u_tau_target if u_tau_target else velocity_symbol / sp.Float(100)
+
+        wall_law = law(u_plus, y_plus)
+        assignments = [Assignment(self.u_tau[0], u_tau_init),  # initial guess
+                       *self.newton_iteration(wall_law),  # newton iterations
+                       Assignment(shear_stress_symbol, self.u_tau[-1] ** 2 * density_symbol)]  # final result
+
+        return assignments
+
+
+class SpaldingsLaw(ImplicitWallFunctionModel):
+    """
+    Single formula for the velocity profile inside the boundary layer, proposed by Spalding :cite:`spalding1961`.
+    Valid in the inner and the outer layer.
+    """
+
+    def __init__(self, viscosity, newton_steps=5, kappa=0.41, b=5.5, name="Spalding"):
+        self.kappa = kappa
+        self.b = b
+
+        super(SpaldingsLaw, self).__init__(name=name, newton_steps=newton_steps, viscosity=viscosity)
+
+    def shear_stress_assignments(self, density_symbol: sp.Symbol, shear_stress_symbol: sp.Symbol,
+                                 velocity_symbol: sp.Symbol, wall_distance, u_tau_target=None):
+        def law(u_p, y_p):
+            k_times_u = self.kappa * u_p
+            frac_1 = (k_times_u ** 2) / sp.Float(2)
+            frac_2 = (k_times_u ** 3) / sp.Float(6)
+            return (u_p + sp.exp(-self.kappa * self.b) * (sp.exp(k_times_u) - sp.Float(1) - k_times_u - frac_1 - frac_2)
+                    - y_p)
+
+        u_plus = velocity_symbol / self.u_tau[0]
+        y_plus = (wall_distance * self.u_tau[0]) / self.viscosity
+
+        u_tau_init = u_tau_target if u_tau_target else velocity_symbol / sp.Float(100)
+
+        wall_law = law(u_plus, y_plus)
+        assignments = [Assignment(self.u_tau[0], u_tau_init),  # initial guess
+                       *self.newton_iteration(wall_law),  # newton iterations
+                       Assignment(shear_stress_symbol, self.u_tau[-1] ** 2 * density_symbol)]  # final result
+
+        return assignments
+
+
+class MuskerLaw(ImplicitWallFunctionModel):
+    """
+    Quasi-analytical model for the velocity profile inside the boundary layer, proposed by Musker. Valid in the inner
+    and the outer layer.
+    Formulation taken from :cite:`malaspinas2015`, Equation (59).
+    """
+
+    def __init__(self, viscosity, newton_steps=5, name="Musker"):
+
+        super(MuskerLaw, self).__init__(name=name, newton_steps=newton_steps, viscosity=viscosity)
+
+    def shear_stress_assignments(self, density_symbol: sp.Symbol, shear_stress_symbol: sp.Symbol,
+                                 velocity_symbol: sp.Symbol, wall_distance, u_tau_target=None):
+        def law(u_p, y_p):
+            arctan = sp.Float(5.424) * sp.atan(sp.Float(0.119760479041916168) * y_p - sp.Float(0.488023952095808383))
+            logarithm = (sp.Float(0.434) * sp.log((y_p + sp.Float(10.6)) ** sp.Float(9.6)
+                                                  / (y_p ** 2 - sp.Float(8.15) * y_p + sp.Float(86)) ** 2))
+            return (arctan + logarithm - sp.Float(3.50727901936264842)) - u_p
+
+        u_plus = velocity_symbol / self.u_tau[0]
+        y_plus = (wall_distance * self.u_tau[0]) / self.viscosity
+
+        u_tau_init = u_tau_target if u_tau_target else velocity_symbol / sp.Float(100)
+
+        wall_law = law(u_plus, y_plus)
+        assignments = [Assignment(self.u_tau[0], u_tau_init),  # initial guess
+                       *self.newton_iteration(wall_law),  # newton iterations
+                       Assignment(shear_stress_symbol, self.u_tau[-1] ** 2 * density_symbol)]  # final result
+
+        return assignments

lbmpy/chapman_enskog/chapman_enskog.py → src/lbmpy/chapman_enskog/chapman_enskog.py

@@ -370,7 +370,7 @@ def take_moments(eqn, pdf_to_moment_name=(('f', '\\Pi'), ('\\Omega f', '\\Upsilo
                 if new_f_index is None:
                     rest *= factor
                 else:
-                    assert not(new_f_index and f_index)
+                    assert not (new_f_index and f_index)
                     f_index = new_f_index
 
         moment_tuple = [0] * len(velocity_terms)