Compare revisions

Markus Holzer · Markus Holzer · Helen Schottenhamml · Helen Schottenhamml · Helen Schottenhamml · Markus Holzer
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -29,7 +29,7 @@ stages:
 tests-and-coverage:
  stage: pretest
  extends: .every-commit
-  image: i10git.cs.fau.de:5005/pycodegen/pycodegen/full
+  image: i10git.cs.fau.de:5005/pycodegen/pycodegen/full:cupy12.3
  script:
    # - pip install sympy --upgrade
    - export NUM_CORES=$(nproc --all)

--- a/AUTHORS.txt
+++ b/AUTHORS.txt
@@ -11,3 +11,4 @@ Contributors:
  - Rudolf Weeber <weeber@icp.uni-stuttgart.de>
  - Christian Godenschwager <christian.godenschwager@fau.de>
  - Jan Hönig <jan.hoenig@fau.de>
+  - Philipp Suffa <philipp.suffa@fau.de>
--- a/COPYING.txt
+++ b/COPYING.txt
------------------------   Important ---------------------------------
-lbmpy is under the following GNU AGPLv3 license. 
-This license holds for the sources of lbmpy itself as well 
-as for all kernels generated with lbmpy i.e. 
-the output of lbmpy is also protected by the GNU AGPLv3 license. 
----------------------------------------------------------------------
                    GNU AFFERO GENERAL PUBLIC LICENSE
                       Version 3, 19 November 2007

--- a/src/lbmpy/__init__.py
+++ b/src/lbmpy/__init__.py
@@ -12,6 +12,7 @@ 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,
@@ -42,6 +43,7 @@ __all__ = [
    "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",

--- a/src/lbmpy/advanced_streaming/communication.py
+++ b/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 pystencils.slicing import (
-from lbmpy.advanced_streaming.utility import is_inplace, get_accessor, numeric_index, \
+    shift_slice,
-    Timestep, get_timesteps, numeric_offsets
+    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,
+    def __init__(
-                 streaming_pattern='pull', ghost_layers=1,
+        self,
-                 cupy_direct_copy=True):
+        stencil,
+        data_handling,
+        pdf_field_name,
+        streaming_pattern="pull",
+        ghost_layers=1,
+        cupy_direct_copy=True,
+    ):
        """
-            Periodicity Handling for Lattice Boltzmann Streaming.
+        Periodicity Handling for Lattice Boltzmann Streaming.
-            **On the usage with cuda:**
+        **On the usage with cuda:**
-            - cupy allows the copying of sliced arrays within device memory using the numpy syntax,
+        - 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
+        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
+        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,
+        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.
+        but especially for large stencils like D3Q27, their compilation can take up to 20 seconds.
-            Choose your weapon depending on your use case.
+        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
@@ -56,12 +72,16 @@ class LBMPeriodicityHandling:
        self.comm_slices = []
        timesteps = get_timesteps(streaming_pattern)
        for timestep in timesteps:
-            slices_per_comm_dir = get_communication_slices(stencil=stencil,
+            slices_per_comm_dir = get_communication_slices(
-                                                           comm_stencil=copy_directions,
+                stencil=stencil,
-                                                           streaming_pattern=streaming_pattern,
+                comm_stencil=copy_directions,
-                                                           prev_timestep=timestep,
+                streaming_pattern=streaming_pattern,
-                                                           ghost_layers=ghost_layers)
+                prev_timestep=timestep,
-            self.comm_slices.append(list(chain.from_iterable(v for k, v in slices_per_comm_dir.items())))
+                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()
@@ -81,11 +101,11 @@ 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(
+            kernels.append(periodic_pdf_gpu_copy_kernel(pdf_field, src, dst))
-                periodic_pdf_copy_kernel(pdf_field, src, dst, target=self.target))
        return kernels
    def _periodicity_handling_gpu(self, prev_timestep):
@@ -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, )
+            src_slice = src_slice + (write_index,)
-            dst_slice = dst_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,
+def periodic_pdf_gpu_copy_kernel(pdf_field, src_slice, dst_slice, domain_size=None):
-                             domain_size=None, target=Target.GPU):
+    """Generate a GPU kernel which copies all values from one slice of a field
-    """Copies a rectangular array slice onto another non-overlapping array slice"""
+    to another non-overlapping slice."""
-    from pystencils.gpucuda.kernelcreation import create_cuda_kernel
+    from pystencils import create_kernel
    pdf_idx = src_slice[-1]
    assert isinstance(pdf_idx, int), "PDF index needs to be an integer constant"
@@ -176,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)]
+    offset = [
-    assert offset == [_stop(s1) - _stop(s2) for s1, s2 in zip(normalized_from_slice, normalized_to_slice)], \
+        _start(s1) - _start(s2)
-        "Slices have to have same size"
+        for s1, s2 in zip(normalized_from_slice, normalized_to_slice)
+    ]
-    copy_eq = AssignmentCollection(main_assignments=[Assignment(pdf_field(pdf_idx), pdf_field[tuple(offset)](pdf_idx))])
+    assert offset == [
-    config = CreateKernelConfig(iteration_slice=dst_slice, skip_independence_check=True)
+        _stop(s1) - _stop(s2)
-    ast = create_cuda_kernel(copy_eq, config=config)
+        for s1, s2 in zip(normalized_from_slice, normalized_to_slice)
-    if target == Target.GPU:
+    ], "Slices have to have same size"
-        from pystencils.gpucuda import make_python_function
-        return make_python_function(ast)
+    copy_eq = AssignmentCollection(
-    else:
+        main_assignments=[
-        raise ValueError('Invalid target:', target)
+            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):
@@ -196,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):

--- a/src/lbmpy/boundaries/boundaries_in_kernel.py
+++ b/src/lbmpy/boundaries/boundaries_in_kernel.py
@@ -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)}

--- a/src/lbmpy/boundaries/boundaryconditions.py
+++ b/src/lbmpy/boundaries/boundaryconditions.py
@@ -28,10 +28,11 @@ class LbBoundary(abc.ABC):
    inner_or_boundary = True
    single_link = False
-    def __init__(self, name=None):
+    def __init__(self, name=None, calculate_force_on_boundary=False):
        self._name = name
+        self.calculate_force_on_boundary = calculate_force_on_boundary
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        """
        This function defines the boundary behavior and must therefore be implemented by all boundaries.
        The boundary is defined through a list of sympy equations from which a boundary kernel is generated.
@@ -48,6 +49,8 @@ class LbBoundary(abc.ABC):
        lb_method:      an instance of the LB method used. Use this to adapt the boundary to the method
                        (e.g. compressibility)
        index_field:    the boundary index field that can be used to retrieve and update boundary data
+        force_vector:   vector to store the force on the boundary. Has the same size as the index field and
+                        D-entries per cell
        Returns:
            list of pystencils assignments, or pystencils.AssignmentCollection
@@ -109,14 +112,31 @@ class NoSlip(LbBoundary):
    Args:
        name: optional name of the boundary.
+        calculate_force_on_boundary: stores the force for each PDF at the boundary in a force vector
    """
-    def __init__(self, name=None):
+    def __init__(self, name=None, calculate_force_on_boundary=False):
        """Set an optional name here, to mark boundaries, for example for force evaluations"""
-        super(NoSlip, self).__init__(name)
+        super(NoSlip, self).__init__(name, calculate_force_on_boundary)
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def get_additional_code_nodes(self, lb_method):
-        return Assignment(f_in(inv_dir[dir_symbol]), f_out(dir_symbol))
+        if self.calculate_force_on_boundary:
+            return [NeighbourOffsetArrays(lb_method.stencil)]
+        else:
+            return []
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
+        if self.calculate_force_on_boundary:
+            force = sp.Symbol("f")
+            subexpressions = [Assignment(force, sp.Float(2.0) * f_out(dir_symbol))]
+            offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
+            for i in range(lb_method.stencil.D):
+                subexpressions.append(Assignment(force_vector[0](f'F_{i}'), force * offset[i]))
+        else:
+            subexpressions = []
+        boundary_assignments = [Assignment(f_in(inv_dir[dir_symbol]), f_out(dir_symbol))]
+        return AssignmentCollection(boundary_assignments, subexpressions=subexpressions)
 class NoSlipLinearBouzidi(LbBoundary):
@@ -135,11 +155,10 @@ class NoSlipLinearBouzidi(LbBoundary):
        data_type: data type of the wall distance q
    """
-    def __init__(self, name=None, init_wall_distance=None, data_type='double'):
+    def __init__(self, name=None, init_wall_distance=None, data_type='double', calculate_force_on_boundary=False):
        self.data_type = data_type
        self.init_wall_distance = init_wall_distance
+        super(NoSlipLinearBouzidi, self).__init__(name, calculate_force_on_boundary)
-        super(NoSlipLinearBouzidi, self).__init__(name)
    @property
    def additional_data(self):
@@ -147,6 +166,12 @@ class NoSlipLinearBouzidi(LbBoundary):
        direction is needed. This information is stored in the index vector."""
        return [('q', create_type(self.data_type))]
+    def get_additional_code_nodes(self, lb_method):
+        if self.calculate_force_on_boundary:
+            return [NeighbourOffsetArrays(lb_method.stencil)]
+        else:
+            return []
    @property
    def additional_data_init_callback(self):
        def default_callback(boundary_data, **_):
@@ -160,7 +185,7 @@ class NoSlipLinearBouzidi(LbBoundary):
                 "(init_wall_distance=None). The boundary condition will fall back to a simple NoSlip BC")
            return default_callback
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        f_xf = sp.Symbol("f_xf")
        f_xf_inv = sp.Symbol("f_xf_inv")
        d_x2f = sp.Symbol("d_x2f")
@@ -182,6 +207,13 @@ class NoSlipLinearBouzidi(LbBoundary):
                           (case_two, sp.And(sp.Gt(q, 0), sp.Lt(q, 0.5))),
                           (case_three, True))
+        if self.calculate_force_on_boundary:
+            force = sp.Symbol("f")
+            subexpressions.append(Assignment(force, f_xf + rhs))
+            offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
+            for i in range(lb_method.stencil.D):
+                subexpressions.append(Assignment(force_vector[0](f'F_{i}'), force * offset[i]))
        boundary_assignments = [Assignment(f_in(inv_dir[dir_symbol]), rhs)]
        return AssignmentCollection(boundary_assignments, subexpressions=subexpressions)
@@ -201,14 +233,15 @@ class QuadraticBounceBack(LbBoundary):
        data_type: data type of the wall distance q
    """
-    def __init__(self, relaxation_rate, name=None, init_wall_distance=None, data_type='double'):
+    def __init__(self, relaxation_rate, name=None, init_wall_distance=None, data_type='double',
+                 calculate_force_on_boundary=False):
        self.relaxation_rate = relaxation_rate
        self.data_type = data_type
        self.init_wall_distance = init_wall_distance
        self.equilibrium_values_name = "f_eq"
        self.inv_dir_symbol = TypedSymbol("inv_dir", create_type("int32"))
-        super(QuadraticBounceBack, self).__init__(name)
+        super(QuadraticBounceBack, self).__init__(name, calculate_force_on_boundary)
    @property
    def additional_data(self):
@@ -258,7 +291,7 @@ class QuadraticBounceBack(LbBoundary):
        return result
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        omega = self.relaxation_rate
        inv = sp.IndexedBase(self.inv_dir_symbol, shape=(1,))[dir_symbol]
        weight_info = LbmWeightInfo(lb_method, data_type=self.data_type)
@@ -311,6 +344,13 @@ class QuadraticBounceBack(LbBoundary):
        t2 = (q * (f_xf + f_xf_inv)) / (one + q)
        result = (one - q) / (one + q) * t1 * half + t2
+        if self.calculate_force_on_boundary:
+            force = sp.Symbol("f")
+            subexpressions.append(Assignment(force, f_xf + result))
+            offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
+            for i in range(lb_method.stencil.D):
+                subexpressions.append(Assignment(force_vector[0](f'F_{i}'), force * offset[i]))
        boundary_assignments = [Assignment(f_in(inv_dir[dir_symbol]), result)]
        return AssignmentCollection(boundary_assignments, subexpressions=subexpressions)
@@ -355,7 +395,7 @@ class FreeSlip(LbBoundary):
        if name is None and normal_direction:
            name = f"Free Slip : {offset_to_direction_string([-x for x in normal_direction])}"
-        super(FreeSlip, self).__init__(name)
+        super(FreeSlip, self).__init__(name, calculate_force_on_boundary=False)
    def init_callback(self, boundary_data, **_):
        if len(boundary_data.index_array) > 1e6:
@@ -425,7 +465,7 @@ class FreeSlip(LbBoundary):
        else:
            return [NeighbourOffsetArrays(lb_method.stencil)]
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        neighbor_offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
        if self.normal_direction:
            tangential_offset = tuple(offset + normal for offset, normal in zip(neighbor_offset, self.normal_direction))
@@ -559,13 +599,13 @@ class WallFunctionBounce(LbBoundary):
        if name is None:
            name = f"WFB : {offset_to_direction_string([-x for x in normal_direction])}"
-        super(WallFunctionBounce, self).__init__(name)
+        super(WallFunctionBounce, self).__init__(name, calculate_force_on_boundary=False)
    def get_additional_code_nodes(self, lb_method):
        return [MirroredStencilDirections(self.stencil, self.mirror_axis),
                NeighbourOffsetArrays(lb_method.stencil)]
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        # needed symbols for offsets and indices
        # neighbour offset symbols are basically the stencil directions defined in stencils.py:L130ff.
        neighbor_offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
@@ -715,7 +755,7 @@ class UBB(LbBoundary):
        self.dim = dim
        self.data_type = data_type
-        super(UBB, self).__init__(name)
+        super(UBB, self).__init__(name, calculate_force_on_boundary=False)
    @property
    def additional_data(self):
@@ -751,7 +791,7 @@ class UBB(LbBoundary):
        This is useful if the inflow velocity should have a certain profile for instance"""
        return callable(self._velocity)
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        vel_from_idx_field = callable(self._velocity)
        vel = [index_field(f'vel_{i}') for i in range(self.dim)] if vel_from_idx_field else self._velocity
@@ -827,7 +867,7 @@ class SimpleExtrapolationOutflow(LbBoundary):
        self.normal_direction = tuple([int(n) for n in normal_direction])
        assert all([n in [-1, 0, 1] for n in self.normal_direction]), \
            "Only -1, 0 and 1 allowed for defining the normal direction"
-        super(SimpleExtrapolationOutflow, self).__init__(name)
+        super(SimpleExtrapolationOutflow, self).__init__(name, calculate_force_on_boundary=False)
    def get_additional_code_nodes(self, lb_method):
        """Return a list of code nodes that will be added in the generated code before the index field loop.
@@ -841,7 +881,7 @@ class SimpleExtrapolationOutflow(LbBoundary):
        """
        return [NeighbourOffsetArrays(lb_method.stencil)]
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        neighbor_offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
        tangential_offset = tuple(offset - normal for offset, normal in zip(neighbor_offset, self.normal_direction))
@@ -867,7 +907,7 @@ class ExtrapolationOutflow(LbBoundary):
    Args:
        normal_direction: direction vector normal to the outflow
-        lb_method: the lattice boltzman method to be used in the simulation
+        lb_method: the lattice Boltzmann method to be used in the simulation
        dt: lattice time step size
        dx: lattice spacing distance
        name: optional name of the boundary.
@@ -920,7 +960,7 @@ class ExtrapolationOutflow(LbBoundary):
            self.equilibrium_calculation = calc_eq_pdfs
-        super(ExtrapolationOutflow, self).__init__(name)
+        super(ExtrapolationOutflow, self).__init__(name, calculate_force_on_boundary=False)
    def init_callback(self, boundary_data, **_):
        dim = boundary_data.dim
@@ -973,7 +1013,7 @@ class ExtrapolationOutflow(LbBoundary):
        """
        return [NeighbourOffsetArrays(lb_method.stencil)]
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        subexpressions = []
        boundary_assignments = []
        dtdx = sp.Rational(self.dt, self.dx)
@@ -1019,9 +1059,9 @@ class FixedDensity(LbBoundary):
            name = "Fixed Density " + str(density)
        self.density = density
-        super(FixedDensity, self).__init__(name)
+        super(FixedDensity, self).__init__(name, calculate_force_on_boundary=False)
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        def remove_asymmetric_part_of_main_assignments(assignment_collection, degrees_of_freedom):
            new_main_assignments = [Assignment(a.lhs, get_symmetric_part(a.rhs, degrees_of_freedom))
                                    for a in assignment_collection.main_assignments]
@@ -1083,7 +1123,7 @@ class DiffusionDirichlet(LbBoundary):
        self.concentration_is_callable = callable(self.concentration)
        self.velocity_field = velocity_field
-        super(DiffusionDirichlet, self).__init__(name)
+        super(DiffusionDirichlet, self).__init__(name, calculate_force_on_boundary=False)
    @property
    def additional_data(self):
@@ -1116,7 +1156,7 @@ class DiffusionDirichlet(LbBoundary):
        else:
            return [LbmWeightInfo(lb_method, self._data_type)]
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        assert lb_method.conserved_quantity_computation.zero_centered_pdfs is False, \
            "DiffusionDirichlet only works for methods with normal pdfs storage -> set zero_centered=False"
        weight_info = LbmWeightInfo(lb_method, self._data_type)
@@ -1159,7 +1199,7 @@ class NeumannByCopy(LbBoundary):
        """
        return [NeighbourOffsetArrays(lb_method.stencil)]
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        neighbour_offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
        return [Assignment(f_in(inv_dir[dir_symbol]), f_out(inv_dir[dir_symbol])),
                Assignment(f_out[neighbour_offset](dir_symbol), f_out(dir_symbol))]
@@ -1178,7 +1218,7 @@ class StreamInConstant(LbBoundary):
    """
    def __init__(self, constant, name=None):
-        super(StreamInConstant, self).__init__(name)
+        super(StreamInConstant, self).__init__(name, calculate_force_on_boundary=False)
        self.constant = constant
    def get_additional_code_nodes(self, lb_method):
@@ -1192,7 +1232,7 @@ class StreamInConstant(LbBoundary):
        """
        return [NeighbourOffsetArrays(lb_method.stencil)]
-    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field, force_vector):
        neighbour_offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
        return [Assignment(f_in(inv_dir[dir_symbol]), self.constant),
                Assignment(f_out[neighbour_offset](dir_symbol), self.constant)]

--- a/src/lbmpy/boundaries/boundaryhandling.py
+++ b/src/lbmpy/boundaries/boundaryhandling.py
+from dataclasses import replace
 import numpy as np
 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.field import FieldType
 from pystencils.simp import add_subexpressions_for_field_reads
 from pystencils.stencil import inverse_direction
@@ -156,22 +158,31 @@ class LatticeBoltzmannBoundaryHandling(BoundaryHandling):
 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",
-    boundary_assignments = indexing.substitute_proxies(boundary_assignments)
-    config = CreateKernelConfig(index_fields=[index_field], 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)
    if pdf_field.dtype != default_data_type:
        boundary_assignments = add_subexpressions_for_field_reads(boundary_assignments, data_type=default_data_type)

--- a/src/lbmpy/boundaries/wall_function_models.py
+++ b/src/lbmpy/boundaries/wall_function_models.py
@@ -157,7 +157,7 @@ class MuskerLaw(ImplicitWallFunctionModel):
        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, 10))
+                                                  / (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]

--- a/src/lbmpy/creationfunctions.py
+++ b/src/lbmpy/creationfunctions.py
@@ -67,7 +67,8 @@ from lbmpy.enums import Stencil, Method, ForceModel, CollisionSpace, SubgridScal
 import lbmpy.forcemodels as forcemodels
 from lbmpy.fieldaccess import CollideOnlyInplaceAccessor, PdfFieldAccessor, PeriodicTwoFieldsAccessor
 from lbmpy.fluctuatinglb import add_fluctuations_to_collision_rule
-from lbmpy.partially_saturated_cells import add_psm_to_collision_rule, PSMConfig
+from lbmpy.partially_saturated_cells import (replace_by_psm_collision_rule, PSMConfig,
+                                             add_psm_solid_collision_to_collision_rule)
 from lbmpy.non_newtonian_models import add_cassons_model, CassonsParameters
 from lbmpy.methods import (create_mrt_orthogonal, create_mrt_raw, create_central_moment,
                           create_srt, create_trt, create_trt_kbc)
@@ -376,8 +377,8 @@ class LBMConfig:
        if not self.compressible and self.method in (Method.MONOMIAL_CUMULANT, Method.CUMULANT):
            raise ValueError("Incompressible cumulant-based methods are not supported (yet).")
-        if self.zero_centered and (self.entropic or self.fluctuating):
+        if self.zero_centered and self.entropic:
-            raise ValueError("Entropic and fluctuating methods can only be created with `zero_centered=False`.")
+            raise ValueError("Entropic methods can only be created with `zero_centered=False`.")
        #   Check or infer delta-equilibrium
        if self.delta_equilibrium is not None:
@@ -468,7 +469,7 @@ class LBMConfig:
        }
        if self.psm_config is not None and self.psm_config.fraction_field is not None:
-            self.force = [(1.0 - self.psm_config.fraction_field.center) * f for f in self.force]
+            self.force = [(1.0 - self.psm_config.fraction_field_symbol) * f for f in self.force]
        if isinstance(self.force_model, str):
            new_force_model = ForceModel[self.force_model.upper()]
@@ -684,11 +685,6 @@ def create_lb_collision_rule(lb_method=None, lbm_config=None, lbm_optimisation=N
    else:
        collision_rule = lb_method.get_collision_rule(pre_simplification=pre_simplification)
-    if lbm_config.psm_config is not None:
-        if lbm_config.psm_config.fraction_field is None or lbm_config.psm_config.object_velocity_field is None:
-            raise ValueError("Specify a fraction and object velocity field in the PSM Config")
-        collision_rule = add_psm_to_collision_rule(collision_rule, lbm_config.psm_config)
    if lbm_config.galilean_correction:
        from lbmpy.methods.cumulantbased import add_galilean_correction
        collision_rule = add_galilean_correction(collision_rule)
@@ -706,6 +702,11 @@ def create_lb_collision_rule(lb_method=None, lbm_config=None, lbm_optimisation=N
                                                     bulk_relaxation_rate=lbm_config.relaxation_rates[1],
                                                     limiter=cumulant_limiter)
+    if lbm_config.psm_config is not None:
+        if lbm_config.psm_config.fraction_field is None or lbm_config.psm_config.object_velocity_field is None:
+            raise ValueError("Specify a fraction and object velocity field in the PSM Config")
+        collision_rule = replace_by_psm_collision_rule(collision_rule, lbm_config.psm_config)
    if lbm_config.entropic:
        if lbm_config.subgrid_scale_model or lbm_config.cassons:
            raise ValueError("Choose either entropic, subgrid-scale or cassons")
@@ -783,7 +784,7 @@ def create_lb_method(lbm_config=None, **params):
    if lbm_config.psm_config is None:
        fraction_field = None
    else:
-        fraction_field = lbm_config.psm_config.fraction_field
+        fraction_field = lbm_config.psm_config.fraction_field_symbol
    common_params = {
        'compressible': lbm_config.compressible,
@@ -869,49 +870,36 @@ def create_lb_method(lbm_config=None, **params):
 def create_psm_update_rule(lbm_config, lbm_optimisation):
-    node_collection = []
-    # Use regular lb update rule for no overlapping particles
+    if lbm_config.psm_config is None:
-    config_without_psm = copy.deepcopy(lbm_config)
+        raise ValueError("Specify a PSM Config in the LBM Config, when creating a psm update rule")
-    config_without_psm.psm_config = None
-    # TODO: the force is still multiplied by (1.0 - self.psm_config.fraction_field.center)
+    config_without_particles = copy.deepcopy(lbm_config)
-    #  (should not harm if memory bound since self.psm_config.fraction_field.center should always be 0.0)
+    config_without_particles.psm_config.max_particles_per_cell = 0
    lb_update_rule = create_lb_update_rule(
-        lbm_config=config_without_psm, lbm_optimisation=lbm_optimisation
+        lbm_config=config_without_particles, lbm_optimisation=lbm_optimisation)
-    )
-    node_collection.append(
+    node_collection = lb_update_rule.all_assignments
-        Conditional(
-            lbm_config.psm_config.fraction_field.center(0) <= 0.0,
-            Block(lb_update_rule.all_assignments),
-        )
-    )
-    # Only one particle, i.e., no individual_fraction_field is provided
    if lbm_config.psm_config.individual_fraction_field is None:
-        assert lbm_config.psm_config.MaxParticlesPerCell == 1
+        assert lbm_config.psm_config.max_particles_per_cell == 1
+        fraction_field = lbm_config.psm_config.fraction_field
+    else:
+        fraction_field = lbm_config.psm_config.individual_fraction_field
+    for p in range(lbm_config.psm_config.max_particles_per_cell):
+        psm_solid_collision = add_psm_solid_collision_to_collision_rule(lb_update_rule, lbm_config, p)
        psm_update_rule = create_lb_update_rule(
-            lbm_config=lbm_config, lbm_optimisation=lbm_optimisation
+            collision_rule=psm_solid_collision, lbm_config=lbm_config, lbm_optimisation=lbm_optimisation)
-        )
        node_collection.append(
            Conditional(
-                lbm_config.psm_config.fraction_field.center(0) > 0.0,
+                fraction_field.center(p) > 0.0,
                Block(psm_update_rule.all_assignments),
            )
        )
-    else:
-        for p in range(lbm_config.psm_config.MaxParticlesPerCell):
-            # Add psm update rule for p overlapping particles
-            config_with_p_particles = copy.deepcopy(lbm_config)
-            config_with_p_particles.psm_config.MaxParticlesPerCell = p + 1
-            psm_update_rule = create_lb_update_rule(
-                lbm_config=config_with_p_particles, lbm_optimisation=lbm_optimisation
-            )
-            node_collection.append(
-                Conditional(
-                    lbm_config.psm_config.individual_fraction_field.center(p) > 0.0,
-                    Block(psm_update_rule.all_assignments),
-                )
-            )
    return NodeCollection(node_collection)

--- a/src/lbmpy/custom_code_nodes.py
+++ b/src/lbmpy/custom_code_nodes.py
@@ -68,7 +68,7 @@ class LbmWeightInfo(CustomCodeNode):
        weights = [f"(({self.weights_symbol.dtype.c_name})({str(w.evalf(17))}))" for w in lb_method.weights]
        weights = ", ".join(weights)
        w_sym = self.weights_symbol
-        code = f"const {self.weights_symbol.dtype.c_name} {w_sym.name} [] = {{{ weights }}};\n"
+        code = f"const {self.weights_symbol.dtype.c_name} {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):

--- a/src/lbmpy/fluctuatinglb.py
+++ b/src/lbmpy/fluctuatinglb.py
@@ -19,9 +19,7 @@ def add_fluctuations_to_collision_rule(collision_rule, temperature=None, amplitu
    """"""
    if not (temperature and not amplitudes) or (temperature and amplitudes):
        raise ValueError("Fluctuating LBM: Pass either 'temperature' or 'amplitudes'.")
-    if collision_rule.method.conserved_quantity_computation.zero_centered_pdfs:
-        raise ValueError("The fluctuating LBM is not implemented for zero-centered PDF storage.")
    method = collision_rule.method
    if not amplitudes:
        amplitudes = fluctuation_amplitude_from_temperature(method, temperature, c_s_sq)
@@ -44,9 +42,7 @@ def fluctuation_amplitude_from_temperature(method, temperature, c_s_sq=sp.Symbol
    """Produces amplitude equations according to (2.60) and (3.54) in Schiller08"""
    normalization_factors = sp.matrix_multiply_elementwise(method.moment_matrix, method.moment_matrix) * \
        sp.Matrix(method.weights)
-    density = method.zeroth_order_equilibrium_moment_symbol
+    density = method._cqc.density_symbol
-    if method.conserved_quantity_computation.zero_centered_pdfs:
-        density += 1
    mu = temperature * density / c_s_sq
    return [sp.sqrt(mu * norm * (1 - (1 - rr) ** 2))
            for norm, rr in zip(normalization_factors, method.relaxation_rates)]

--- a/src/lbmpy/macroscopic_value_kernels.py
+++ b/src/lbmpy/macroscopic_value_kernels.py
 import functools
+import sympy as sp
 from copy import deepcopy
 from lbmpy.simplificationfactory import create_simplification_strategy
-from pystencils import create_kernel, CreateKernelConfig
+from pystencils import create_kernel, CreateKernelConfig, Assignment
 from pystencils.field import Field, get_layout_of_array
 from pystencils.enums import Target
 from lbmpy.advanced_streaming.utility import get_accessor, Timestep
+from lbmpy.relaxationrates import get_shear_relaxation_rate
+from lbmpy.utils import second_order_moment_tensor
-def pdf_initialization_assignments(lb_method, density, velocity, pdfs,
+def get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, pre_collision_pdfs):
-                                   streaming_pattern='pull', previous_timestep=Timestep.BOTH,
-                                   set_pre_collision_pdfs=False):
-    """Assignments to initialize the pdf field with equilibrium"""
    if isinstance(pdfs, Field):
        accessor = get_accessor(streaming_pattern, previous_timestep)
-        if set_pre_collision_pdfs:
+        if pre_collision_pdfs:
            field_accesses = accessor.read(pdfs, lb_method.stencil)
        else:
            field_accesses = accessor.write(pdfs, lb_method.stencil)
-    elif streaming_pattern == 'pull' and not set_pre_collision_pdfs:
+    elif streaming_pattern == 'pull' and not pre_collision_pdfs:
        field_accesses = pdfs
    else:
        raise ValueError("Invalid value of pdfs: A PDF field reference is required to derive "
                         + f"initialization assignments for streaming pattern {streaming_pattern}.")
+    return field_accesses
+def get_individual_or_common_fraction_field(psm_config):
+    if psm_config.individual_fraction_field is not None:
+        return psm_config.individual_fraction_field
+    else:
+        return psm_config.fraction_field
+def pdf_initialization_assignments(lb_method, density, velocity, pdfs, psm_config=None,
+                                   streaming_pattern='pull', previous_timestep=Timestep.BOTH,
+                                   set_pre_collision_pdfs=False):
+    """Assignments to initialize the pdf field with equilibrium"""
+    field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, set_pre_collision_pdfs)
    if isinstance(density, Field):
        density = density.center
    if isinstance(velocity, Field):
        velocity = velocity.center_vector
@@ -35,23 +49,40 @@ def pdf_initialization_assignments(lb_method, density, velocity, pdfs,
    setter_eqs = lb_method.get_equilibrium(conserved_quantity_equations=inp_eqs)
    setter_eqs = setter_eqs.new_with_substitutions({sym: field_accesses[i]
                                                    for i, sym in enumerate(lb_method.post_collision_pdf_symbols)})
+    if lb_method.fraction_field is not None:
+        if psm_config is None:
+            raise ValueError("If setting up LBM with PSM, please specify a PSM config in the macroscopic setter")
+        else:
+            for equ in setter_eqs:
+                if equ.lhs in lb_method.first_order_equilibrium_moment_symbols:
+                    pos = lb_method.first_order_equilibrium_moment_symbols.index(equ.lhs)
+                    new_rhs = 0
+                    if isinstance(equ.rhs, sp.core.Add):
+                        for summand in equ.rhs.args:
+                            if summand in velocity:
+                                new_rhs += (1.0 - psm_config.fraction_field.center) * summand
+                            else:
+                                new_rhs += summand.subs(lb_method.fraction_field, psm_config.fraction_field.center)
+                    else:
+                        new_rhs += (1.0 - psm_config.fraction_field.center) * equ.rhs
+                    fraction_field = get_individual_or_common_fraction_field(psm_config)
+                    for p in range(psm_config.max_particles_per_cell):
+                        new_rhs += psm_config.object_velocity_field(p * lb_method.dim + pos) * fraction_field.center(p)
+                    setter_eqs.subexpressions.remove(equ)
+                    setter_eqs.subexpressions.append(Assignment(equ.lhs, new_rhs))
    return setter_eqs
-def macroscopic_values_getter(lb_method, density, velocity, pdfs,
+def macroscopic_values_getter(lb_method, density, velocity, pdfs, psm_config=None,
                              streaming_pattern='pull', previous_timestep=Timestep.BOTH,
                              use_pre_collision_pdfs=False):
-    if isinstance(pdfs, Field):
-        accessor = get_accessor(streaming_pattern, previous_timestep)
+    field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, use_pre_collision_pdfs)
-        if use_pre_collision_pdfs:
-            field_accesses = accessor.read(pdfs, lb_method.stencil)
-        else:
-            field_accesses = accessor.write(pdfs, lb_method.stencil)
-    elif streaming_pattern == 'pull' and not use_pre_collision_pdfs:
-        field_accesses = pdfs
-    else:
-        raise ValueError("Invalid value of pdfs: A PDF field reference is required to derive "
-                         + f"getter assignments for streaming pattern {streaming_pattern}.")
    cqc = lb_method.conserved_quantity_computation
    assert not (velocity is None and density is None)
    output_spec = {}
@@ -59,12 +90,54 @@ def macroscopic_values_getter(lb_method, density, velocity, pdfs,
        output_spec['velocity'] = velocity
    if density is not None:
        output_spec['density'] = density
-    return cqc.output_equations_from_pdfs(field_accesses, output_spec)
+    getter_equ = cqc.output_equations_from_pdfs(field_accesses, output_spec)
+    if lb_method.fraction_field is not None:
+        if psm_config.fraction_field is None:
+            raise ValueError("If setting up LBM with PSM, please specify a PSM config in the macroscopic getter")
+        else:
+            if lb_method.force_model is not None:
+                for equ in getter_equ:
+                    if equ.lhs in lb_method.force_model.symbolic_force_vector:
+                        new_rhs = equ.rhs.subs(lb_method.fraction_field, psm_config.fraction_field.center)
+                        getter_equ.subexpressions.remove(equ)
+                        getter_equ.subexpressions.append(Assignment(equ.lhs, new_rhs))
+            for i, equ in enumerate(getter_equ.main_assignments[-lb_method.dim:]):
+                new_rhs = (1.0 - psm_config.fraction_field.center) * equ.rhs
+                fraction_field = get_individual_or_common_fraction_field(psm_config)
+                for p in range(psm_config.max_particles_per_cell):
+                    new_rhs += psm_config.object_velocity_field(p * lb_method.dim + i) * fraction_field.center(p)
+                getter_equ.main_assignments.remove(equ)
+                getter_equ.main_assignments.append(Assignment(equ.lhs, new_rhs))
+        getter_equ.topological_sort()
+    return getter_equ
 macroscopic_values_setter = pdf_initialization_assignments
+def strain_rate_tensor_getter(lb_method, strain_rate_tensor, pdfs, streaming_pattern='pull',
+                              previous_timestep=Timestep.BOTH, use_pre_collision_pdfs=False):
+    field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, use_pre_collision_pdfs)
+    if isinstance(strain_rate_tensor, Field):
+        strain_rate_tensor = strain_rate_tensor.center_vector
+    omega_s = get_shear_relaxation_rate(lb_method)
+    equilibrium = lb_method.equilibrium_distribution
+    rho = equilibrium.density if equilibrium.compressible else equilibrium.background_density
+    f_neq = sp.Matrix([field_accesses[i] for i in range(lb_method.stencil.Q)]) - lb_method.get_equilibrium_terms()
+    pi = second_order_moment_tensor(f_neq, lb_method.stencil)
+    strain_rate_tensor_equ = - 1.5 * (omega_s / rho) * pi
+    assignments = [Assignment(strain_rate_tensor[i * lb_method.stencil.D + j], strain_rate_tensor_equ[i, j])
+                   for i in range(lb_method.stencil.D) for j in range(lb_method.stencil.D)]
+    return assignments
 def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None,
                                      ghost_layers=1, iteration_slice=None,
                                      field_layout='numpy', target=Target.CPU,

--- a/src/lbmpy/methods/abstractlbmethod.py
+++ b/src/lbmpy/methods/abstractlbmethod.py
@@ -105,22 +105,27 @@ class AbstractLbMethod(abc.ABC):
        the subexpressions, that assign the number to the newly introduced symbol
        """
        rr = relaxation_rates if relaxation_rates is not None else self.relaxation_rates
-        unique_relaxation_rates = set()
        subexpressions = {}
+        symbolic_relaxation_rates = list()
        for relaxation_rate in rr:
            relaxation_rate = sp.sympify(relaxation_rate)
-            if relaxation_rate not in unique_relaxation_rates:
+            if isinstance(relaxation_rate, sp.Symbol):
-                # special treatment for zero, sp.Zero would be an integer ..
+                symbolic_relaxation_rate = relaxation_rate
+            else:
                if isinstance(relaxation_rate, Zero):
                    relaxation_rate = 0.0
-                if not isinstance(relaxation_rate, sp.Symbol):
+                if relaxation_rate in subexpressions:
-                    rt_symbol = sp.Symbol(f"rr_{len(subexpressions)}")
+                    symbolic_relaxation_rate = subexpressions[relaxation_rate]
-                    subexpressions[relaxation_rate] = rt_symbol
+                else:
-            unique_relaxation_rates.add(relaxation_rate)
+                    symbolic_relaxation_rate = sp.Symbol(f"rr_{len(subexpressions)}")
+                    subexpressions[relaxation_rate] = symbolic_relaxation_rate
+            symbolic_relaxation_rates.append(symbolic_relaxation_rate)
-        new_rr = [subexpressions[sp.sympify(e)] if sp.sympify(e) in subexpressions else e for e in rr]
        substitutions = [Assignment(e[1], e[0]) for e in subexpressions.items()]
        if relaxation_rates_modifier is not None:
-            new_rr = [r * relaxation_rates_modifier for r in new_rr]
+            symbolic_relaxation_rates = [r * relaxation_rates_modifier for r in symbolic_relaxation_rates]
+        else:
+            for srr in symbolic_relaxation_rates:
+                assert isinstance(srr, sp.Symbol)
-        return substitutions, sp.diag(*new_rr)
+        return substitutions, sp.diag(*symbolic_relaxation_rates)
--- a/src/lbmpy/methods/creationfunctions.py
+++ b/src/lbmpy/methods/creationfunctions.py
@@ -201,19 +201,23 @@ def create_from_equilibrium(stencil, equilibrium, conserved_quantity_computation
    if cspace.collision_space == CollisionSpace.POPULATIONS:
        return MomentBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
-                                   force_model=force_model, zero_centered=zero_centered, fraction_field=fraction_field,
+                                   force_model=force_model, zero_centered=zero_centered,
+                                   fraction_field=fraction_field,
                                   moment_transform_class=None)
    elif cspace.collision_space == CollisionSpace.RAW_MOMENTS:
        return MomentBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
-                                   force_model=force_model, zero_centered=zero_centered, fraction_field=fraction_field,
+                                   force_model=force_model, zero_centered=zero_centered,
+                                   fraction_field=fraction_field,
                                   moment_transform_class=cspace.raw_moment_transform_class)
    elif cspace.collision_space == CollisionSpace.CENTRAL_MOMENTS:
        return CentralMomentBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
                                          force_model=force_model, zero_centered=zero_centered,
+                                          fraction_field=fraction_field,
                                          central_moment_transform_class=cspace.central_moment_transform_class)
    elif cspace.collision_space == CollisionSpace.CUMULANTS:
        return CumulantBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
                                     force_model=force_model, zero_centered=zero_centered,
+                                     fraction_field=fraction_field,
                                     central_moment_transform_class=cspace.central_moment_transform_class,
                                     cumulant_transform_class=cspace.cumulant_transform_class)
@@ -334,7 +338,7 @@ def create_mrt_raw(stencil, relaxation_rates, continuous_equilibrium=True, conse
 def create_central_moment(stencil, relaxation_rates, nested_moments=None,
-                          continuous_equilibrium=True, conserved_moments=True, fraction_field=None, **kwargs):
+                          continuous_equilibrium=True, conserved_moments=True, **kwargs):
    r"""
    Creates moment based LB method where the collision takes place in the central moment space.
@@ -348,7 +352,6 @@ def create_central_moment(stencil, relaxation_rates, nested_moments=None,
        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
                        used to compute the equilibrium moments.
        conserved_moments: If lower order moments are conserved or not.
-        fraction_field: fraction field for the PSM method
    Returns:
        :class:`lbmpy.methods.momentbased.CentralMomentBasedLbMethod` instance
    """
@@ -371,8 +374,8 @@ def create_central_moment(stencil, relaxation_rates, nested_moments=None,
        nested_moments = cascaded_moment_sets_literature(stencil)
    rr_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments, stencil.D, conserved_moments)
-    if fraction_field is not None:
+    if 'fraction_field' in kwargs and kwargs['fraction_field'] is not None:
-        relaxation_rates_modifier = (1.0 - fraction_field.center)
+        relaxation_rates_modifier = (1.0 - kwargs['fraction_field'])
        rr_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments, stencil.D,
                                            relaxation_rates_modifier=relaxation_rates_modifier)
@@ -527,7 +530,7 @@ def create_mrt_orthogonal(stencil, relaxation_rates, continuous_equilibrium=True
 # ----------------------------------------- Cumulant method creators ---------------------------------------------------
-def create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments=True, fraction_field=None, **kwargs):
+def create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments=True, **kwargs):
    r"""Creates a cumulant-based lattice Boltzmann method.
    Args:
@@ -547,8 +550,8 @@ def create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moment
    """
    cumulant_to_rr_dict = _get_relaxation_info_dict(relaxation_rates, cumulant_groups, stencil.D, conserved_moments)
-    if fraction_field is not None:
+    if 'fraction_field' in kwargs and kwargs['fraction_field'] is not None:
-        relaxation_rates_modifier = (1.0 - fraction_field.center)
+        relaxation_rates_modifier = (1.0 - kwargs['fraction_field'])
        cumulant_to_rr_dict = _get_relaxation_info_dict(relaxation_rates, cumulant_groups, stencil.D,
                                                        relaxation_rates_modifier=relaxation_rates_modifier)

--- a/src/lbmpy/methods/cumulantbased/cumulantbasedmethod.py
+++ b/src/lbmpy/methods/cumulantbased/cumulantbasedmethod.py
@@ -47,7 +47,7 @@ class CumulantBasedLbMethod(AbstractLbMethod):
    def __init__(self, stencil, equilibrium, relaxation_dict,
                 conserved_quantity_computation=None,
-                 force_model=None, zero_centered=False,
+                 force_model=None, zero_centered=False, fraction_field=None,
                 central_moment_transform_class=BinomialChimeraTransform,
                 cumulant_transform_class=CentralMomentsToCumulantsByGeneratingFunc):
        assert isinstance(conserved_quantity_computation,
@@ -63,6 +63,7 @@ class CumulantBasedLbMethod(AbstractLbMethod):
        self._cqc = conserved_quantity_computation
        self._force_model = force_model
        self._zero_centered = zero_centered
+        self._fraction_field = fraction_field
        self._weights = None
        self._cumulant_transform_class = cumulant_transform_class
        self._central_moment_transform_class = central_moment_transform_class
@@ -72,6 +73,10 @@ class CumulantBasedLbMethod(AbstractLbMethod):
        """Force model employed by this method."""
        return self._force_model
+    @property
+    def fraction_field(self):
+        return self._fraction_field
    @property
    def relaxation_info_dict(self):
        """Dictionary mapping this method's cumulants to their relaxation rates and equilibrium values.

--- a/src/lbmpy/methods/momentbased/centralmomentbasedmethod.py
+++ b/src/lbmpy/methods/momentbased/centralmomentbasedmethod.py
@@ -55,7 +55,7 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
    def __init__(self, stencil, equilibrium, relaxation_dict,
                 conserved_quantity_computation=None,
-                 force_model=None, zero_centered=False,
+                 force_model=None, zero_centered=False, fraction_field=None,
                 central_moment_transform_class=BinomialChimeraTransform):
        assert isinstance(conserved_quantity_computation, AbstractConservedQuantityComputation)
        super(CentralMomentBasedLbMethod, self).__init__(stencil)
@@ -65,6 +65,7 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
        self._cqc = conserved_quantity_computation
        self._force_model = force_model
        self._zero_centered = zero_centered
+        self._fraction_field = fraction_field
        self._weights = None
        self._central_moment_transform_class = central_moment_transform_class
@@ -73,6 +74,10 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
        """Force model employed by this method."""
        return self._force_model
+    @property
+    def fraction_field(self):
+        return self._fraction_field
    @property
    def relaxation_info_dict(self):
        """Dictionary mapping this method's moments to their relaxation rates and equilibrium values.

--- a/src/lbmpy/methods/momentbased/momentbasedmethod.py
+++ b/src/lbmpy/methods/momentbased/momentbasedmethod.py
@@ -48,7 +48,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
        self._cqc = conserved_quantity_computation
        self._force_model = force_model
        self._zero_centered = zero_centered
-        self.fraction_field = fraction_field
+        self._fraction_field = fraction_field
        self._weights = None
        self._moment_transform_class = moment_transform_class
@@ -57,6 +57,10 @@ class MomentBasedLbMethod(AbstractLbMethod):
        """Force model employed by this method."""
        return self._force_model
+    @property
+    def fraction_field(self):
+        return self._fraction_field
    @property
    def relaxation_info_dict(self):
        """Dictionary mapping this method's moments to their relaxation rates and equilibrium values.
@@ -176,8 +180,8 @@ class MomentBasedLbMethod(AbstractLbMethod):
    def get_collision_rule(self, conserved_quantity_equations: AssignmentCollection = None,
                           pre_simplification: bool = True) -> LbmCollisionRule:
-        if self.fraction_field is not None:
+        if self._fraction_field is not None:
-            relaxation_rates_modifier = (1.0 - self.fraction_field.center)
+            relaxation_rates_modifier = (1.0 - self._fraction_field)
            rr_sub_expressions, d = self._generate_symbolic_relaxation_matrix(
                relaxation_rates_modifier=relaxation_rates_modifier)
        else:

--- a/src/lbmpy/moment_transforms/rawmomenttransforms.py
+++ b/src/lbmpy/moment_transforms/rawmomenttransforms.py
@@ -172,7 +172,7 @@ class PdfsToMomentsByMatrixTransform(AbstractRawMomentTransform):
    #   ----------------------------- Private Members -----------------------------
-    @ property
+    @property
    def _default_simplification(self):
        forward_simp = SimplificationStrategy()
        # forward_simp.add(substitute_moments_in_conserved_quantity_equations)
@@ -218,7 +218,7 @@ class PdfsToMomentsByChimeraTransform(AbstractRawMomentTransform):
                                                                                self.moment_polynomials)
        self.poly_to_mono_matrix = self.mono_to_poly_matrix.inv()
-    @ property
+    @property
    def absorbs_conserved_quantity_equations(self):
        return True
@@ -414,7 +414,7 @@ class PdfsToMomentsByChimeraTransform(AbstractRawMomentTransform):
    #   ----------------------------- Private Members -----------------------------
-    @ property
+    @property
    def _default_simplification(self):
        from lbmpy.methods.momentbased.momentbasedsimplifications import (
            substitute_moments_in_conserved_quantity_equations,

--- a/src/lbmpy/partially_saturated_cells.py
+++ b/src/lbmpy/partially_saturated_cells.py
 import sympy as sp
 from dataclasses import dataclass
+from lbmpy.enums import Method
 from lbmpy.methods.abstractlbmethod import LbmCollisionRule
 from pystencils import Assignment, AssignmentCollection
 from pystencils.field import Field
@@ -13,103 +14,156 @@ class PSMConfig:
    Fraction field for PSM 
    """
+    fraction_field_symbol = sp.Symbol('B')
+    """
+    Fraction field symbol used for simplification 
+    """
    object_velocity_field: Field = None
    """
    Object velocity field for PSM 
    """
-    SC: int = 1
+    solid_collision: int = 1
    """
    Solid collision option for PSM
    """
-    MaxParticlesPerCell: int = 1
+    max_particles_per_cell: int = 1
    """
    Maximum number of particles overlapping with a cell 
    """
    individual_fraction_field: Field = None
    """
-    Fraction field for each overlapping particle in PSM 
+    Fraction field for each overlapping object / particle in PSM 
    """
-    particle_force_field: Field = None
+    object_force_field: Field = None
    """
-    Force field for each overlapping particle in PSM 
+    Force field for each overlapping object / particle in PSM 
    """
-def add_psm_to_collision_rule(collision_rule, psm_config):
+def get_psm_solid_collision_term(collision_rule, psm_config, particle_per_cell_counter):
    if psm_config.individual_fraction_field is None:
-        psm_config.individual_fraction_field = psm_config.fraction_field
+        fraction_field = psm_config.fraction_field
+    else:
+        fraction_field = psm_config.individual_fraction_field
    method = collision_rule.method
    pre_collision_pdf_symbols = method.pre_collision_pdf_symbols
    stencil = method.stencil
-    # Get equilibrium from object velocity for solid collision
-    forces_rhs = [0] * psm_config.MaxParticlesPerCell * stencil.D
    solid_collisions = [0] * stencil.Q
-    for p in range(psm_config.MaxParticlesPerCell):
+    equilibrium_fluid = method.get_equilibrium_terms()
-        equilibrium_fluid = method.get_equilibrium_terms()
+    equilibrium_solid = []
-        equilibrium_solid = []
-        for eq in equilibrium_fluid:
+    # get equilibrium form object velocity
-            eq_sol = eq
+    for eq in equilibrium_fluid:
-            for i in range(stencil.D):
+        eq_sol = eq
-                eq_sol = eq_sol.subs(sp.Symbol("u_" + str(i)),
+        for i in range(stencil.D):
-                                     psm_config.object_velocity_field.center(p * stencil.D + i), )
+            eq_sol = eq_sol.subs(sp.Symbol("u_" + str(i)),
-            equilibrium_solid.append(eq_sol)
+                                 psm_config.object_velocity_field.center(particle_per_cell_counter * stencil.D + i), )
+        equilibrium_solid.append(eq_sol)
-        # Build solid collision
-        for i, (eqFluid, eqSolid, f, offset) in enumerate(
+    # Build solid collision
-                zip(equilibrium_fluid, equilibrium_solid, pre_collision_pdf_symbols, stencil)):
+    for i, (eqFluid, eqSolid, f, offset) in enumerate(
-            inverse_direction_index = stencil.stencil_entries.index(stencil.inverse_stencil_entries[i])
+            zip(equilibrium_fluid, equilibrium_solid, pre_collision_pdf_symbols, stencil)):
-            if psm_config.SC == 1:
+        inverse_direction_index = stencil.stencil_entries.index(stencil.inverse_stencil_entries[i])
-                solid_collision = psm_config.individual_fraction_field.center(p) * (
+        if psm_config.solid_collision == 1:
-                    (
+            solid_collision = (fraction_field.center(particle_per_cell_counter)
-                        pre_collision_pdf_symbols[inverse_direction_index]
+                               * ((pre_collision_pdf_symbols[inverse_direction_index]
-                        - equilibrium_fluid[inverse_direction_index]
+                                   - equilibrium_fluid[inverse_direction_index]) - (f - eqSolid)))
-                    )
+        elif psm_config.solid_collision == 2:
-                    - (f - eqSolid)
+            # TODO get relaxation rate vector from method and use the right relaxation rate [i]
-                )
+            solid_collision = (fraction_field.center(particle_per_cell_counter)
-            elif psm_config.SC == 2:
+                               * ((eqSolid - f) + (1.0 - method.relaxation_rates[0]) * (f - eqFluid)))
-                # TODO get relaxation rate vector from method and use the right relaxation rate [i]
+        elif psm_config.solid_collision == 3:
-                solid_collision = psm_config.individual_fraction_field.center(p) * (
+            solid_collision = (fraction_field.center(particle_per_cell_counter)
-                    (eqSolid - f) + (1.0 - method.relaxation_rates[0]) * (f - eqFluid)
+                               * ((pre_collision_pdf_symbols[inverse_direction_index]
-                )
+                                   - equilibrium_solid[inverse_direction_index]) - (f - eqSolid)))
-            elif psm_config.SC == 3:
+        else:
-                solid_collision = psm_config.individual_fraction_field.center(p) * (
+            raise ValueError("Only sc=1, sc=2 and sc=3 are supported.")
-                    (
-                        pre_collision_pdf_symbols[inverse_direction_index]
+        solid_collisions[i] += solid_collision
-                        - equilibrium_solid[inverse_direction_index]
-                    )
+    return solid_collisions
-                    - (f - eqSolid)
-                )
-            else:
+def get_psm_force_from_solid_collision(solid_collisions, stencil, object_force_field, particle_per_cell_counter):
-                raise ValueError("Only SC=1, SC=2 and SC=3 are supported.")
+    force_assignments = []
-            solid_collisions[i] += solid_collision
+    for d in range(stencil.D):
-            for j in range(stencil.D):
+        forces_rhs = 0
-                forces_rhs[p * stencil.D + j] -= solid_collision * int(offset[j])
+        for sc, offset in zip(solid_collisions, stencil):
+            forces_rhs -= sc * int(offset[d])
-    # Add solid collision to main assignments of collision rule
+        force_assignments.append(Assignment(
+            object_force_field.center(particle_per_cell_counter * stencil.D + d), forces_rhs
+        ))
+    return AssignmentCollection(force_assignments)
+def replace_fraction_symbol_with_field(assignments, fraction_field_symbol, fraction_field_access):
+    new_assignments = []
+    for ass in assignments:
+        rhs = ass.rhs.subs(fraction_field_symbol, fraction_field_access.center(0))
+        new_assignments.append(Assignment(ass.lhs, rhs))
+    return new_assignments
+def add_psm_solid_collision_to_collision_rule(collision_rule, lbm_config, particle_per_cell_counter):
+    method = collision_rule.method
+    solid_collisions = get_psm_solid_collision_term(collision_rule, lbm_config.psm_config, particle_per_cell_counter)
+    post_collision_pdf_symbols = method.post_collision_pdf_symbols
+    assignments = []
+    for sc, post in zip(solid_collisions, post_collision_pdf_symbols):
+        assignments.append(Assignment(post, post + sc))
+    if lbm_config.psm_config.object_force_field is not None:
+        assignments += get_psm_force_from_solid_collision(solid_collisions, method.stencil,
+                                                          lbm_config.psm_config.object_force_field,
+                                                          particle_per_cell_counter)
+    # exchanging rho with zeroth order moment symbol
+    if lbm_config.method in (Method.CENTRAL_MOMENT, Method.MONOMIAL_CUMULANT, Method.CUMULANT):
+        new_assignments = []
+        zeroth_moment_symbol = 'm_00' if lbm_config.stencil.D == 2 else 'm_000'
+        for ass in assignments:
+            new_assignments.append(ass.subs(sp.Symbol('rho'), sp.Symbol(zeroth_moment_symbol)))
+        assignments = new_assignments
+    collision_assignments = AssignmentCollection(assignments)
+    ac = LbmCollisionRule(method, collision_assignments, [],
+                          collision_rule.simplification_hints)
+    return ac
+def replace_by_psm_collision_rule(collision_rule, psm_config):
+    method = collision_rule.method
    collision_assignments = []
-    for main, sc in zip(collision_rule.main_assignments, solid_collisions):
+    solid_collisions = [0] * psm_config.max_particles_per_cell
-        collision_assignments.append(Assignment(main.lhs, main.rhs + sc))
+    for p in range(psm_config.max_particles_per_cell):
+        solid_collisions[p] = get_psm_solid_collision_term(collision_rule, psm_config, p)
-    # Add hydrodynamic force calculations to collision assignments if two-way coupling is used
-    # (i.e., force field is not None)
+        if psm_config.object_force_field is not None:
-    if psm_config.particle_force_field is not None:
+            collision_assignments += get_psm_force_from_solid_collision(solid_collisions[p], method.stencil,
-        for p in range(psm_config.MaxParticlesPerCell):
+                                                                        psm_config.object_force_field, p)
-            for i in range(stencil.D):
-                collision_assignments.append(
+    for i, main in enumerate(collision_rule.main_assignments):
-                    Assignment(
+        rhs = main.rhs
-                        psm_config.particle_force_field.center(p * stencil.D + i),
+        for p in range(psm_config.max_particles_per_cell):
-                        forces_rhs[p * stencil.D + i],
+            rhs += solid_collisions[p][i]
-                    )
+        collision_assignments.append(Assignment(main.lhs, rhs))
-                )
    collision_assignments = AssignmentCollection(collision_assignments)
-    ac = LbmCollisionRule(method, collision_assignments, collision_rule.subexpressions,
+    ac = LbmCollisionRule(method, replace_fraction_symbol_with_field(collision_assignments,
+                          psm_config.fraction_field_symbol, psm_config.fraction_field),
+                          replace_fraction_symbol_with_field(collision_rule.subexpressions,
+                          psm_config.fraction_field_symbol, psm_config.fraction_field),
                          collision_rule.simplification_hints)
    ac.topological_sort()
    return ac
No results found