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]

src/lbmpy/creationfunctions.py

@@ -58,16 +58,16 @@ from dataclasses import dataclass, field, replace
 from typing import Union, List, Tuple, Any, Type, Iterable
 from warnings import warn, filterwarnings
 
-import lbmpy.moment_transforms
-import pystencils.astnodes
+from ._compat import IS_PYSTENCILS_2
+
 import sympy as sp
-import sympy.core.numbers
 
 from lbmpy.enums import Stencil, Method, ForceModel, CollisionSpace, SubgridScaleModel
 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)
@@ -81,17 +81,22 @@ from lbmpy.stencils import LBStencil
 from lbmpy.turbulence_models import add_sgs_model
 from lbmpy.updatekernels import create_lbm_kernel, create_stream_pull_with_output_kernel
 from lbmpy.advanced_streaming.utility import Timestep, get_accessor
+from .forcemodels import AbstractForceModel
+
+import pystencils
 from pystencils import CreateKernelConfig, create_kernel
-from pystencils.astnodes import Conditional, Block
 from pystencils.cache import disk_cache_no_fallback
-from pystencils.node_collection import NodeCollection
-from pystencils.typing import collate_types
 from pystencils.field import Field
 from pystencils.simp import sympy_cse, SimplificationStrategy
 # needed for the docstring
 from lbmpy.methods.abstractlbmethod import LbmCollisionRule, AbstractLbMethod
 from lbmpy.methods.cumulantbased import CumulantBasedLbMethod
 
+if IS_PYSTENCILS_2:
+    from pystencils import Kernel as KernelFunction
+else:
+    from pystencils.astnodes import KernelFunction
+
 # Filter out JobLib warnings. They are not useful for use:
 # https://github.com/joblib/joblib/issues/683
 filterwarnings("ignore", message="Persisting input arguments took")
@@ -102,7 +107,7 @@ class LBMConfig:
     """
     **Below all parameters for the LBMConfig are explained**
     """
-    stencil: lbmpy.stencils.LBStencil = LBStencil(Stencil.D2Q9)
+    stencil: LBStencil = LBStencil(Stencil.D2Q9)
     """
     All stencils are defined in :class:`lbmpy.enums.Stencil`. From that :class:`lbmpy.stencils.LBStencil` 
     class will be created
@@ -161,7 +166,7 @@ class LBMConfig:
     truncated. Order 2 is sufficient to approximate Navier-Stokes. This parameter has no effect on cumulant-based
     methods, whose equilibrium terms have no contributions above order one.
     """
-    c_s_sq: sympy.Rational = sp.Rational(1, 3)
+    c_s_sq: sp.Expr = sp.Rational(1, 3)
     """
     The squared lattice speed of sound used to derive the LB method. It is very uncommon to use a value different 
     to 1 / 3.
@@ -178,7 +183,7 @@ class LBMConfig:
     If this argument is not provided, Gram-Schmidt orthogonalisation of the default modes is performed.
     """
 
-    force_model: Union[lbmpy.forcemodels.AbstractForceModel, ForceModel] = None
+    force_model: Union[AbstractForceModel, ForceModel] = None
     """
     Force model to determine how forcing terms enter the collision rule.
     Possibilities are defined in :class: `lbmpy.enums.ForceModel`
@@ -338,9 +343,9 @@ class LBMConfig:
     Instance of :class:`lbmpy.methods.LbmCollisionRule`. If this parameter is `None`,
     the update rule is derived via *create_lb_update_rule*.
     """
-    ast: pystencils.astnodes.KernelFunction = None
+    ast: KernelFunction = None
     """
-    Instance of *pystencils.astnodes.KernelFunction*. If this parameter is `None`,
+    Instance of *pystencils.KernelFunction*. If this parameter is `None`,
     the ast is derived via `create_lb_ast`.
     """
 
@@ -376,8 +381,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):
-            raise ValueError("Entropic and fluctuating methods can only be created with `zero_centered=False`.")
+        if self.zero_centered and self.entropic:
+            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 +473,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()]
@@ -555,7 +560,6 @@ def create_lb_function(ast=None, lbm_config=None, lbm_optimisation=None, config=
 
     res.method = ast.method
     res.update_rule = ast.update_rule
-    res.ast = ast
     return res
 
 
@@ -571,9 +575,7 @@ def create_lb_ast(update_rule=None, lbm_config=None, lbm_optimisation=None, conf
         update_rule = create_lb_update_rule(lbm_config.collision_rule, lbm_config=lbm_config,
                                             lbm_optimisation=lbm_optimisation, config=config)
 
-    field_types = set(fa.field.dtype for fa in update_rule.defined_symbols if isinstance(fa, Field.Access))
-
-    config = replace(config, data_type=collate_types(field_types), ghost_layers=1)
+    config = replace(config, ghost_layers=1)
     ast = create_kernel(update_rule, config=config)
 
     ast.method = update_rule.method
@@ -599,7 +601,11 @@ def create_lb_update_rule(collision_rule=None, lbm_config=None, lbm_optimisation
 
     lb_method = collision_rule.method
 
-    field_data_type = config.data_type[lbm_config.field_name].numpy_dtype
+    if IS_PYSTENCILS_2:
+        fallback_field_data_type = config.get_option("default_dtype")
+    else:
+        fallback_field_data_type = config.data_type[lbm_config.field_name].numpy_dtype
+
     q = collision_rule.method.stencil.Q
 
     if lbm_optimisation.symbolic_field is not None:
@@ -607,16 +613,17 @@ def create_lb_update_rule(collision_rule=None, lbm_config=None, lbm_optimisation
     elif lbm_optimisation.field_size:
         field_size = tuple([s + 2 for s in lbm_optimisation.field_size] + [q])
         src_field = Field.create_fixed_size(lbm_config.field_name, field_size, index_dimensions=1,
-                                            layout=lbm_optimisation.field_layout, dtype=field_data_type)
+                                            layout=lbm_optimisation.field_layout, dtype=fallback_field_data_type)
     else:
         src_field = Field.create_generic(lbm_config.field_name, spatial_dimensions=collision_rule.method.dim,
-                                         index_shape=(q,), layout=lbm_optimisation.field_layout, dtype=field_data_type)
+                                         index_shape=(q,), layout=lbm_optimisation.field_layout,
+                                         dtype=fallback_field_data_type)
 
     if lbm_optimisation.symbolic_temporary_field is not None:
         dst_field = lbm_optimisation.symbolic_temporary_field
     else:
         dst_field = src_field.new_field_with_different_name(lbm_config.temporary_field_name)
-
+   
     kernel_type = lbm_config.kernel_type
     if kernel_type == 'stream_pull_only':
         update_rule = create_stream_pull_with_output_kernel(lb_method, src_field, dst_field, lbm_config.output)
@@ -684,11 +691,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 +708,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 +790,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 +876,45 @@ def create_lb_method(lbm_config=None, **params):
 
 
 def create_psm_update_rule(lbm_config, lbm_optimisation):
-    node_collection = []
+    if IS_PYSTENCILS_2:
+        raise NotImplementedError(
+            "`create_psm_update_rule` is not yet available when using pystencils 2.0. "
+            "To instead derive a (potentially less efficient) PSM kernel without branches, "
+            "use `create_lb_update_rule` with a `PsmConfig` object instead."
+        )
+    
+    from pystencils.astnodes import Conditional, Block
+    from pystencils.node_collection import NodeCollection
+
+    if lbm_config.psm_config is None:
+        raise ValueError("Specify a PSM Config in the LBM Config, when creating a psm update rule")
+
+    config_without_particles = copy.deepcopy(lbm_config)
+    config_without_particles.psm_config.max_particles_per_cell = 0
 
-    # Use regular lb update rule for no overlapping particles
-    config_without_psm = copy.deepcopy(lbm_config)
-    config_without_psm.psm_config = None
-    # TODO: the force is still multiplied by (1.0 - self.psm_config.fraction_field.center)
-    #  (should not harm if memory bound since self.psm_config.fraction_field.center should always be 0.0)
     lb_update_rule = create_lb_update_rule(
-        lbm_config=config_without_psm, lbm_optimisation=lbm_optimisation
-    )
-    node_collection.append(
-        Conditional(
-            lbm_config.psm_config.fraction_field.center(0) <= 0.0,
-            Block(lb_update_rule.all_assignments),
-        )
-    )
+        lbm_config=config_without_particles, lbm_optimisation=lbm_optimisation)
+
+    node_collection = 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)
 

src/lbmpy/custom_code_nodes.py

 import numpy as np
 import sympy as sp
 
+from ._compat import IS_PYSTENCILS_2
+
+if IS_PYSTENCILS_2:
+    raise ImportError("`lbmpy.custom_code_nodes` is only available when running with pystencils 1.x")
+
 from pystencils.typing import TypedSymbol, create_type
 from pystencils.backends.cbackend import CustomCodeNode
 
@@ -44,14 +49,14 @@ class MirroredStencilDirections(CustomCodeNode):
         return tuple(direction)
 
     @staticmethod
-    def _mirrored_symbol(mirror_axis):
+    def _mirrored_symbol(mirror_axis, _stencil):
         axis = ['x', 'y', 'z']
         return TypedSymbol(f"{axis[mirror_axis]}_axis_mirrored_stencil_dir", create_type('int32'))
 
     def __init__(self, stencil, mirror_axis, dtype=np.int32):
         offsets_dtype = create_type(dtype)
 
-        mirrored_stencil_symbol = MirroredStencilDirections._mirrored_symbol(mirror_axis)
+        mirrored_stencil_symbol = MirroredStencilDirections._mirrored_symbol(mirror_axis, stencil)
         mirrored_directions = [stencil.index(MirroredStencilDirections.mirror_stencil(direction, mirror_axis))
                                for direction in stencil]
         code = "\n"
@@ -68,7 +73,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):

src/lbmpy/fieldaccess.py

@@ -4,11 +4,12 @@ import sympy as sp
 
 from pystencils import Field
 # ------------------------------------------------ Interface -----------------------------------------------------------
-from pystencils.astnodes import LoopOverCoordinate
+
 from pystencils.stencil import inverse_direction
 
 from lbmpy.enums import Stencil
 from lbmpy.stencils import LBStencil
+from ._compat import get_loop_counter_symbol
 
 __all__ = ['PdfFieldAccessor', 'CollideOnlyInplaceAccessor', 'StreamPullTwoFieldsAccessor',
            'AAEvenTimeStepAccessor', 'AAOddTimeStepAccessor',
@@ -114,7 +115,7 @@ class PeriodicTwoFieldsAccessor(PdfFieldAccessor):
                 lower_limit = self._ghostLayers
                 upper_limit = field.spatial_shape[coord_id] - 1 - self._ghostLayers
                 limit_diff = upper_limit - lower_limit
-                loop_counter = LoopOverCoordinate.get_loop_counter_symbol(coord_id)
+                loop_counter = get_loop_counter_symbol(coord_id)
                 if dir_element == 0:
                     periodic_pull_direction.append(0)
                 elif dir_element == 1:

src/lbmpy/flow_statistics.py

@@ -4,17 +4,18 @@ import pystencils as ps
 from pystencils.field import Field
 
 
-def welford_assignments(field, mean_field, sum_of_products_field=None):
+def welford_assignments(field, mean_field, sum_of_squares_field=None, sum_of_cubes_field=None):
     r"""
     Creates the assignments for the kernel creation of Welford's algorithm
     (https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm).
     This algorithm is an online algorithm for the mean and variance calculation of sample data, here implemented for
-    the execution on vector fields, e.g., velocity.
-    The calculation of the variance is optional and only performed if a field for the sum of products is given.
+    the execution on scalar or vector fields, e.g., velocity.
+    The calculation of the variance / the third-order central moments is optional and only performed if a field for
+    the sum of squares / sum of cubes is given.
 
     The mean value is directly updated in the mean vector field.
-    The variance must be retrieved in a post-processing step. Let :math `M_{2,n}` denote the value of the sum of
-    products after the first :math `n` samples. According to Welford the biased sample variance
+    The variance / covariance must be retrieved in a post-processing step. Let :math `M_{2,n}` denote the value of the
+    sum of squares after the first :math `n` samples. According to Welford the biased sample variance
     :math `\sigma_n^2` and the unbiased sample variance :math `s_n^2` are given by
 
     .. math ::
@@ -22,6 +23,9 @@ def welford_assignments(field, mean_field, sum_of_products_field=None):
         s_n^2 = \frac{M_{2,n}}{n-1},
 
     respectively.
+
+    Liekwise, to the 3rd-order central moment(s) of the (vector) field, the sum of cubes field must be divided
+    by :math `n` in a post-processing step.
     """
 
     if isinstance(field, Field):
@@ -40,22 +44,41 @@ def welford_assignments(field, mean_field, sum_of_products_field=None):
     else:
         raise ValueError("Mean vector field has to be a pystencils Field or Field.Access")
 
-    if sum_of_products_field is None:
-        # sum of products will not be calculated, i.e., the variance is not retrievable
-        welford_sum_of_products_field = None
+    if sum_of_squares_field is None:
+        # sum of products will not be calculated, i.e., the covariance matrix is not retrievable
+        welford_sum_of_squares_field = None
+    else:
+        if isinstance(sum_of_squares_field, Field):
+            welford_sum_of_squares_field = sum_of_squares_field.center
+        elif isinstance(sum_of_squares_field, Field.Access):
+            welford_sum_of_squares_field = sum_of_squares_field
+        else:
+            raise ValueError("Sum of squares field has to be a pystencils Field or Field.Access")
+
+        assert welford_sum_of_squares_field.field.values_per_cell() == dim ** 2, \
+            (f"Sum of squares field does not have the right layout. "
+             f"Index dimension: {welford_sum_of_squares_field.field.values_per_cell()}, expected: {dim ** 2}")
+
+    if sum_of_cubes_field is None:
+        # sum of cubes will not be calculated, i.e., the 3rd-order central moments are not retrievable
+        welford_sum_of_cubes_field = None
     else:
-        if isinstance(sum_of_products_field, Field):
-            welford_sum_of_products_field = sum_of_products_field.center
-            assert sum_of_products_field.values_per_cell() == dim**2, \
-                (f"Sum of products field does not have the right layout. "
-                 f"Index dimension: {sum_of_products_field.values_per_cell()}, expected: {dim**2}")
-        elif isinstance(sum_of_products_field, Field.Access):
-            welford_sum_of_products_field = sum_of_products_field
-            assert sum_of_products_field.field.values_per_cell() == dim**2, \
-                (f"Sum of products field does not have the right layout. "
-                 f"Index dimension: {sum_of_products_field.field.values_per_cell()}, expected: {dim**2}")
+        if isinstance(sum_of_cubes_field, Field):
+            welford_sum_of_cubes_field = sum_of_cubes_field.center
+        elif isinstance(sum_of_cubes_field, Field.Access):
+            welford_sum_of_cubes_field = sum_of_cubes_field
         else:
-            raise ValueError("Variance vector field has to be a pystencils Field or Field.Access")
+            raise ValueError("Sum of cubes field has to be a pystencils Field or Field.Access")
+
+        assert welford_sum_of_cubes_field.field.values_per_cell() == dim ** 3, \
+            (f"Sum of cubes field does not have the right layout. "
+             f"Index dimension: {welford_sum_of_cubes_field.field.values_per_cell()}, expected: {dim ** 3}")
+
+    # for the calculation of the thrid-order moments, the variance must also be calculated
+    if welford_sum_of_cubes_field is not None:
+        assert welford_sum_of_squares_field is not None
+
+    # actual assignments
 
     counter = sp.Symbol('counter')
     delta = sp.symbols(f"delta_:{dim}")
@@ -67,7 +90,7 @@ def welford_assignments(field, mean_field, sum_of_products_field=None):
         main_assignments.append(ps.Assignment(welford_mean_field.at_index(i),
                                               welford_mean_field.at_index(i) + delta[i] / counter))
 
-    if sum_of_products_field is not None:
+    if sum_of_squares_field is not None:
         delta2 = sp.symbols(f"delta2_:{dim}")
         for i in range(dim):
             main_assignments.append(
@@ -76,7 +99,21 @@ def welford_assignments(field, mean_field, sum_of_products_field=None):
             for j in range(dim):
                 idx = i * dim + j
                 main_assignments.append(ps.Assignment(
-                    welford_sum_of_products_field.at_index(idx),
-                    welford_sum_of_products_field.at_index(idx) + delta[i] * delta2[j]))
+                    welford_sum_of_squares_field.at_index(idx),
+                    welford_sum_of_squares_field.at_index(idx) + delta[i] * delta2[j]))
+
+        if sum_of_cubes_field is not None:
+            for i in range(dim):
+                for j in range(dim):
+                    for k in range(dim):
+                        idx = (i * dim + j) * dim + k
+                        main_assignments.append(ps.Assignment(
+                            welford_sum_of_cubes_field.at_index(idx),
+                            welford_sum_of_cubes_field.at_index(idx)
+                            - delta[k] / counter * welford_sum_of_squares_field(i * dim + j)
+                            - delta[j] / counter * welford_sum_of_squares_field(k * dim + i)
+                            - delta[i] / counter * welford_sum_of_squares_field(j * dim + k)
+                            + delta2[i] * (2 * delta[j] - delta2[j]) * delta[k]
+                        ))
 
     return main_assignments

src/lbmpy/fluctuatinglb.py

@@ -3,36 +3,84 @@
 to generate a fluctuating LBM the equilibrium moment values have to be scaled and an additive (random)
 correction term is added to the collision rule
 """
+from typing import overload
+
+from ._compat import IS_PYSTENCILS_2
+
 import numpy as np
 import sympy as sp
 
 from lbmpy.moments import MOMENT_SYMBOLS, is_shear_moment, get_order
 from lbmpy.equilibrium import ContinuousHydrodynamicMaxwellian
 from pystencils import Assignment, TypedSymbol
-from pystencils.rng import PhiloxFourFloats, random_symbol
 from pystencils.simp.assignment_collection import SymbolGen
 
+if IS_PYSTENCILS_2:
+    from pystencils.sympyextensions.random import RngBase, Philox
+    from pystencils.sympyextensions import tcast
+else:
+    from pystencils.rng import PhiloxFourFloats, random_symbol
+
+
+@overload
+def add_fluctuations_to_collision_rule(collision_rule, temperature=None, amplitudes=(),
+                                       *,
+                                       block_offsets, seed, rng_node, c_s_sq):
+    """Fluctuating LB implementation for pystencils 1.3"""
+
 
+@overload
 def add_fluctuations_to_collision_rule(collision_rule, temperature=None, amplitudes=(),
-                                       block_offsets=(0, 0, 0), seed=TypedSymbol("seed", np.uint32),
-                                       rng_node=PhiloxFourFloats, c_s_sq=sp.Rational(1, 3)):
-    """"""
+                                       *,
+                                       rng: 'RngBase | None' = None, c_s_sq):
+    """Fluctuating LB implementation for pystencils 2.0
+    
+    Args:
+        collision_rule: The base collision rule
+        temperature: Expression representing the fluid temperature
+        amplitudes: If ``temperature`` was not specified, expression representing the fluctuation amplitude
+        rng: Random number generator instance used to compute the fluctuations.
+            If `None`, the float32 Philox RNG will be used.
+    """
+
+
+def add_fluctuations_to_collision_rule(collision_rule, temperature=None, amplitudes=(),
+                                       c_s_sq=sp.Rational(1, 3), **kwargs):
     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)
-    if block_offsets == 'walberla':
-        block_offsets = tuple(TypedSymbol("block_offset_{}".format(i), np.uint32) for i in range(3))
 
     if not method.is_weighted_orthogonal:
         raise ValueError("Fluctuations can only be added to weighted-orthogonal methods")
 
-    rng_symbol_gen = random_symbol(collision_rule.subexpressions, seed=seed,
-                                   rng_node=rng_node, dim=method.dim, offsets=block_offsets)
+    if IS_PYSTENCILS_2:
+        rng: RngBase = kwargs.get("rng", Philox("fluctuation_rng", np.float32, TypedSymbol("seed", np.uint32)))
+        ts = TypedSymbol("time_step", np.uint32)
+
+        def _rng_symbol_gen():
+            while True:
+                rx, rasm = rng.get_random_vector(ts)
+                collision_rule.subexpressions.insert(0, rasm)
+                for i in range(rng.vector_size):
+                    yield tcast.as_numeric(rx[i])
+
+        rng_symbol_gen = _rng_symbol_gen()
+    else:
+        block_offsets = kwargs.get("block_offsets", (0, 0, 0))
+        rng_node = kwargs.get("rng_node", PhiloxFourFloats)
+        seed = kwargs.get("seed", TypedSymbol("seed", np.uint32))
+
+        if block_offsets == 'walberla':
+            block_offsets = tuple(TypedSymbol("block_offset_{}".format(i), np.uint32) for i in range(3))
+
+        rng_symbol_gen = random_symbol(
+            collision_rule.subexpressions, seed=seed,
+            rng_node=rng_node, dim=method.dim, offsets=block_offsets
+        )
+
     correction = fluctuation_correction(method, rng_symbol_gen, amplitudes)
 
     for i, corr in enumerate(correction):
@@ -44,9 +92,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
-    if method.conserved_quantity_computation.zero_centered_pdfs:
-        density += 1
+    density = method._cqc.density_symbol
     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)]

src/lbmpy/lbstep.py

@@ -10,11 +10,18 @@ from lbmpy.macroscopic_value_kernels import (
     create_advanced_velocity_setter_collision_rule, pdf_initialization_assignments)
 from lbmpy.simplificationfactory import create_simplification_strategy
 from lbmpy.stencils import LBStencil
-from pystencils import create_data_handling, create_kernel, make_slice, Target, Backend
+
+from pystencils import CreateKernelConfig
+from pystencils import create_data_handling, create_kernel, make_slice, Target
 from pystencils.slicing import SlicedGetter
 from pystencils.timeloop import TimeLoop
 
 
+from ._compat import IS_PYSTENCILS_2
+if not IS_PYSTENCILS_2:
+    from pystencils import Backend
+
+
 class LatticeBoltzmannStep:
 
     def __init__(self, domain_size=None, lbm_kernel=None, periodicity=False,
@@ -24,7 +31,9 @@ class LatticeBoltzmannStep:
                  velocity_input_array_name=None, time_step_order='stream_collide', flag_interface=None,
                  alignment_if_vectorized=64, fixed_loop_sizes=True,
                  timeloop_creation_function=TimeLoop,
-                 lbm_config=None, lbm_optimisation=None, config=None, **method_parameters):
+                 lbm_config=None, lbm_optimisation=None,
+                 config: CreateKernelConfig | None = None,
+                 **method_parameters):
 
         if optimization is None:
             optimization = {}
@@ -36,7 +45,10 @@ class LatticeBoltzmannStep:
                 raise ValueError("When passing a data_handling, the domain_size parameter can not be specified")
 
         if config is not None:
-            target = config.target
+            if IS_PYSTENCILS_2:
+                target = config.get_target()
+            else:
+                target = config.target
         else:
             target = optimization.get('target', Target.CPU)
 
@@ -60,7 +72,11 @@ class LatticeBoltzmannStep:
                                                                               lbm_config, lbm_optimisation, config)
 
         # the parallel datahandling understands only numpy datatypes. Strings lead to an errors
-        field_dtype = config.data_type.default_factory().numpy_dtype
+        if IS_PYSTENCILS_2:
+            from pystencils import create_type
+            field_dtype = create_type(config.get_option("default_dtype")).numpy_dtype
+        else:
+            field_dtype = config.data_type.default_factory().numpy_dtype
 
         if lbm_kernel:
             q = lbm_kernel.method.stencil.Q
@@ -75,12 +91,21 @@ class LatticeBoltzmannStep:
         self.density_data_name = name + "_density" if density_data_name is None else density_data_name
         self.density_data_index = density_data_index
 
-        self._gpu = target == Target.GPU
+        if IS_PYSTENCILS_2:
+            self._gpu = target.is_gpu()
+        else:
+            self._gpu = target == Target.GPU
+
         layout = lbm_optimisation.field_layout
 
         alignment = False
-        if config.backend == Backend.C and config.cpu_vectorize_info:
-            alignment = alignment_if_vectorized
+
+        if IS_PYSTENCILS_2:
+            if config.get_target().is_vector_cpu() and config.cpu.vectorize.enable:
+                alignment = alignment_if_vectorized
+        else:
+            if config.backend == Backend.C and config.cpu_vectorize_info:
+                alignment = alignment_if_vectorized
 
         self._data_handling.add_array(self._pdf_arr_name, values_per_cell=q, gpu=self._gpu, layout=layout,
                                       latex_name='src', dtype=field_dtype, alignment=alignment)
@@ -150,10 +175,14 @@ class LatticeBoltzmannStep:
         self._sync_tmp = data_handling.synchronization_function([self._tmp_arr_name], stencil_name, target,
                                                                 stencil_restricted=True)
 
-        self._boundary_handling = LatticeBoltzmannBoundaryHandling(self.method, self._data_handling, self._pdf_arr_name,
-                                                                   name=name + "_boundary_handling",
-                                                                   flag_interface=flag_interface,
-                                                                   target=target, openmp=config.cpu_openmp)
+        self._boundary_handling = LatticeBoltzmannBoundaryHandling(
+            self.method, self._data_handling, self._pdf_arr_name,
+            name=name + "_boundary_handling",
+            flag_interface=flag_interface,
+            target=target,
+            openmp=config.cpu_openmp,
+            **({"default_dtype": field_dtype} if IS_PYSTENCILS_2 else dict())
+        )
 
         self._lbm_config = lbm_config
         self._lbm_optimisation = lbm_optimisation
@@ -223,7 +252,7 @@ class LatticeBoltzmannStep:
     @property
     def config(self):
         """Configutation of pystencils parameters"""
-        return self.config
+        return self._config
 
     def _get_slice(self, data_name, slice_obj, masked):
         if slice_obj is None:
@@ -439,12 +468,19 @@ class LatticeBoltzmannStep:
         rho_field = rho_field.center if self.density_data_index is None else rho_field(self.density_data_index)
         vel_field = self._data_handling.fields[self.velocity_data_name]
 
+        if IS_PYSTENCILS_2:
+            gen_config = CreateKernelConfig(target=Target.CPU)
+            gen_config.cpu.openmp.enable = self._config.cpu.openmp.get_option("enable")
+            gen_config.default_dtype = self._config.get_option("default_dtype")
+        else:
+            gen_config = CreateKernelConfig(target=Target.CPU, cpu_openmp=self._config.cpu_openmp)
+
         getter_eqs = cqc.output_equations_from_pdfs(pdf_field.center_vector,
                                                     {'density': rho_field, 'velocity': vel_field})
-        getter_kernel = create_kernel(getter_eqs, target=Target.CPU, cpu_openmp=self._config.cpu_openmp).compile()
+        getter_kernel = create_kernel(getter_eqs, config=gen_config).compile()
 
         setter_eqs = pdf_initialization_assignments(lb_method, rho_field,
                                                     vel_field.center_vector, pdf_field.center_vector)
         setter_eqs = create_simplification_strategy(lb_method)(setter_eqs)
-        setter_kernel = create_kernel(setter_eqs, target=Target.CPU, cpu_openmp=self._config.cpu_openmp).compile()
+        setter_kernel = create_kernel(setter_eqs, config=gen_config).compile()
         return getter_kernel, setter_kernel

src/lbmpy/lookup_tables.py

+from typing import Sequence, Any
+from abc import ABC, abstractmethod
+import numpy as np
+import sympy as sp
+
+from ._compat import IS_PYSTENCILS_2
+
+if not IS_PYSTENCILS_2:
+    raise ImportError("`lbmpy.lookup_tables` is only available when running with pystencils 2.x")
+
+from pystencils import Assignment
+from pystencils.sympyextensions import TypedSymbol
+from pystencils.types.quick import Arr
+from pystencils.types import UserTypeSpec, create_type
+
+
+class LookupTables(ABC):
+    @abstractmethod
+    def get_array_declarations(self) -> list[Assignment]:
+        pass
+
+
+class NeighbourOffsetArrays(LookupTables):
+
+    @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(stencil)
+                ]
+            )
+
+    @staticmethod
+    def _offset_symbols(stencil):
+        q = len(stencil)
+        dim = len(stencil[0])
+        return [
+            TypedSymbol(f"neighbour_offset_{d}", Arr(create_type("int32"), q))
+            for d in ["x", "y", "z"][:dim]
+        ]
+
+    def __init__(self, stencil, offsets_dtype: UserTypeSpec = np.int32):
+        self._offsets_dtype = create_type(
+            offsets_dtype
+        )  # TODO: Currently, this has no effect
+        self._stencil = stencil
+        self._dim = len(stencil[0])
+
+    def get_array_declarations(self) -> list[Assignment]:
+        array_symbols = NeighbourOffsetArrays._offset_symbols(self._stencil)
+        return [
+            Assignment(arrsymb, tuple((d[i] for d in self._stencil)))
+            for i, arrsymb in enumerate(array_symbols)
+        ]
+
+
+class MirroredStencilDirections(LookupTables):
+
+    @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, stencil):
+        axis = ["x", "y", "z"]
+        q = len(stencil)
+        return TypedSymbol(
+            f"{axis[mirror_axis]}_axis_mirrored_stencil_dir", Arr(create_type("int32"), q)
+        )
+
+    def __init__(self, stencil, mirror_axis, dtype=np.int32):
+        self._offsets_dtype = create_type(dtype)  # TODO: Currently, this has no effect
+
+        self._mirrored_stencil_symbol = MirroredStencilDirections._mirrored_symbol(
+            mirror_axis, stencil
+        )
+        self._mirrored_directions = tuple(
+            stencil.index(
+                MirroredStencilDirections.mirror_stencil(direction, mirror_axis)
+            )
+            for direction in stencil
+        )
+
+    def get_array_declarations(self) -> list[Assignment]:
+        return [Assignment(self._mirrored_stencil_symbol, self._mirrored_directions)]
+
+
+class LbmWeightInfo(LookupTables):
+    def __init__(self, lb_method, data_type="double"):
+        self._weights = lb_method.weights
+        self._weights_array = TypedSymbol("weights", Arr(create_type(data_type), len(self._weights)))
+
+    def weight_of_direction(self, dir_idx, lb_method=None):
+        if isinstance(sp.sympify(dir_idx), sp.Integer):
+            assert lb_method is not None
+            return lb_method.weights[dir_idx].evalf(17)
+        else:
+            return sp.IndexedBase(self._weights_array, shape=(1,))[dir_idx]
+
+    def get_array_declarations(self) -> list[Assignment]:
+        return [Assignment(self._weights_array, tuple(self._weights))]
+
+
+class TranslationArraysNode(LookupTables):
+
+    def __init__(self, array_content: Sequence[tuple[TypedSymbol, Sequence[Any]]]):
+        self._decls = [
+            Assignment(symb, tuple(content)) for symb, content in array_content
+        ]
+
+    def __str__(self):
+        return "Variable PDF Access Translation Arrays"
+
+    def __repr__(self):
+        return "Variable PDF Access Translation Arrays"
+
+    def get_array_declarations(self) -> list[Assignment]:
+        return self._decls

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 pystencils 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,
-                                   streaming_pattern='pull', previous_timestep=Timestep.BOTH,
-                                   set_pre_collision_pdfs=False):
-    """Assignments to initialize the pdf field with equilibrium"""
+def get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, pre_collision_pdfs):
     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)
-        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}.")
+
+    field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, use_pre_collision_pdfs)
+
     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,

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:
-                # special treatment for zero, sp.Zero would be an integer ..
+            if isinstance(relaxation_rate, sp.Symbol):
+                symbolic_relaxation_rate = relaxation_rate
+            else:
                 if isinstance(relaxation_rate, Zero):
                     relaxation_rate = 0.0
-                if not isinstance(relaxation_rate, sp.Symbol):
-                    rt_symbol = sp.Symbol(f"rr_{len(subexpressions)}")
-                    subexpressions[relaxation_rate] = rt_symbol
-            unique_relaxation_rates.add(relaxation_rate)
+                if relaxation_rate in subexpressions:
+                    symbolic_relaxation_rate = subexpressions[relaxation_rate]
+                else:
+                    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)

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:
-        relaxation_rates_modifier = (1.0 - fraction_field.center)
+    if 'fraction_field' in kwargs and kwargs['fraction_field'] is not None:
+        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:
-        relaxation_rates_modifier = (1.0 - fraction_field.center)
+    if 'fraction_field' in kwargs and kwargs['fraction_field'] is not None:
+        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)
 

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.

src/lbmpy/methods/cumulantbased/fourth_order_correction.py

@@ -3,8 +3,7 @@ import sympy as sp
 from lbmpy.moment_transforms import PRE_COLLISION_MONOMIAL_CUMULANT, POST_COLLISION_CUMULANT
 from lbmpy.moments import MOMENT_SYMBOLS, statistical_quantity_symbol
 from lbmpy.stencils import Stencil, LBStencil
-from pystencils import Assignment
-from pystencils.simp.assignment_collection import AssignmentCollection
+from pystencils import Assignment, AssignmentCollection
 from .cumulantbasedmethod import CumulantBasedLbMethod
 
 X, Y, Z = MOMENT_SYMBOLS

src/lbmpy/methods/cumulantbased/galilean_correction.py

 import sympy as sp
 
-from pystencils.simp.assignment_collection import AssignmentCollection
-from pystencils import Assignment
+from pystencils import Assignment, AssignmentCollection
 
 from lbmpy.stencils import Stencil, LBStencil
 from lbmpy.moments import MOMENT_SYMBOLS, statistical_quantity_symbol

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.

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:
-            relaxation_rates_modifier = (1.0 - self.fraction_field.center)
+        if self._fraction_field is not None:
+            relaxation_rates_modifier = (1.0 - self._fraction_field)
             rr_sub_expressions, d = self._generate_symbolic_relaxation_matrix(
                 relaxation_rates_modifier=relaxation_rates_modifier)
         else:

src/lbmpy/moment_transforms/abstractmomenttransform.py

 from abc import abstractmethod
 import sympy as sp
 
-from pystencils.simp import (SimplificationStrategy, sympy_cse)
+from pystencils.simp import SimplificationStrategy, sympy_cse
 
 from lbmpy.moments import (
     exponents_to_polynomial_representations, extract_monomials, exponent_tuple_sort_key)

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,

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_solid = []
-        for eq in equilibrium_fluid:
-            eq_sol = eq
-            for i in range(stencil.D):
-                eq_sol = eq_sol.subs(sp.Symbol("u_" + str(i)),
-                                     psm_config.object_velocity_field.center(p * stencil.D + i), )
-            equilibrium_solid.append(eq_sol)
-
-        # Build solid collision
-        for i, (eqFluid, eqSolid, f, offset) in enumerate(
-                zip(equilibrium_fluid, equilibrium_solid, pre_collision_pdf_symbols, stencil)):
-            inverse_direction_index = stencil.stencil_entries.index(stencil.inverse_stencil_entries[i])
-            if psm_config.SC == 1:
-                solid_collision = psm_config.individual_fraction_field.center(p) * (
-                    (
-                        pre_collision_pdf_symbols[inverse_direction_index]
-                        - equilibrium_fluid[inverse_direction_index]
-                    )
-                    - (f - eqSolid)
-                )
-            elif psm_config.SC == 2:
-                # TODO get relaxation rate vector from method and use the right relaxation rate [i]
-                solid_collision = psm_config.individual_fraction_field.center(p) * (
-                    (eqSolid - f) + (1.0 - method.relaxation_rates[0]) * (f - eqFluid)
-                )
-            elif psm_config.SC == 3:
-                solid_collision = psm_config.individual_fraction_field.center(p) * (
-                    (
-                        pre_collision_pdf_symbols[inverse_direction_index]
-                        - equilibrium_solid[inverse_direction_index]
-                    )
-                    - (f - eqSolid)
-                )
-            else:
-                raise ValueError("Only SC=1, SC=2 and SC=3 are supported.")
-            solid_collisions[i] += solid_collision
-            for j in range(stencil.D):
-                forces_rhs[p * stencil.D + j] -= solid_collision * int(offset[j])
-
-    # Add solid collision to main assignments of collision rule
+    equilibrium_fluid = method.get_equilibrium_terms()
+    equilibrium_solid = []
+
+    # get equilibrium form object velocity
+    for eq in equilibrium_fluid:
+        eq_sol = eq
+        for i in range(stencil.D):
+            eq_sol = eq_sol.subs(sp.Symbol("u_" + str(i)),
+                                 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(
+            zip(equilibrium_fluid, equilibrium_solid, pre_collision_pdf_symbols, stencil)):
+        inverse_direction_index = stencil.stencil_entries.index(stencil.inverse_stencil_entries[i])
+        if psm_config.solid_collision == 1:
+            solid_collision = (fraction_field.center(particle_per_cell_counter)
+                               * ((pre_collision_pdf_symbols[inverse_direction_index]
+                                   - equilibrium_fluid[inverse_direction_index]) - (f - eqSolid)))
+        elif psm_config.solid_collision == 2:
+            # TODO get relaxation rate vector from method and use the right relaxation rate [i]
+            solid_collision = (fraction_field.center(particle_per_cell_counter)
+                               * ((eqSolid - f) + (1.0 - method.relaxation_rates[0]) * (f - eqFluid)))
+        elif psm_config.solid_collision == 3:
+            solid_collision = (fraction_field.center(particle_per_cell_counter)
+                               * ((pre_collision_pdf_symbols[inverse_direction_index]
+                                   - equilibrium_solid[inverse_direction_index]) - (f - eqSolid)))
+        else:
+            raise ValueError("Only sc=1, sc=2 and sc=3 are supported.")
+
+        solid_collisions[i] += solid_collision
+
+    return solid_collisions
+
+
+def get_psm_force_from_solid_collision(solid_collisions, stencil, object_force_field, particle_per_cell_counter):
+    force_assignments = []
+    for d in range(stencil.D):
+        forces_rhs = 0
+        for sc, offset in zip(solid_collisions, stencil):
+            forces_rhs -= sc * int(offset[d])
+
+        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):
-        collision_assignments.append(Assignment(main.lhs, main.rhs + sc))
-
-    # Add hydrodynamic force calculations to collision assignments if two-way coupling is used
-    # (i.e., force field is not None)
-    if psm_config.particle_force_field is not None:
-        for p in range(psm_config.MaxParticlesPerCell):
-            for i in range(stencil.D):
-                collision_assignments.append(
-                    Assignment(
-                        psm_config.particle_force_field.center(p * stencil.D + i),
-                        forces_rhs[p * stencil.D + i],
-                    )
-                )
+    solid_collisions = [0] * psm_config.max_particles_per_cell
+    for p in range(psm_config.max_particles_per_cell):
+        solid_collisions[p] = get_psm_solid_collision_term(collision_rule, psm_config, p)
+
+        if psm_config.object_force_field is not None:
+            collision_assignments += get_psm_force_from_solid_collision(solid_collisions[p], method.stencil,
+                                                                        psm_config.object_force_field, p)
+
+    for i, main in enumerate(collision_rule.main_assignments):
+        rhs = main.rhs
+        for p in range(psm_config.max_particles_per_cell):
+            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

src/lbmpy/phasefield/contact_angle_circle_fitting.py

@@ -39,8 +39,8 @@ def circle_intersections(midpoint0, midpoint1, radius0, radius1):
 def interface_region(concentration_arr, phase0, phase1, area=3):
     import scipy.ndimage as sc_image
 
-    area_phase0 = sc_image.morphology.binary_dilation(concentration_arr[..., phase0] > 0.5, iterations=area)
-    area_phase1 = sc_image.morphology.binary_dilation(concentration_arr[..., phase1] > 0.5, iterations=area)
+    area_phase0 = sc_image.binary_dilation(concentration_arr[..., phase0] > 0.5, iterations=area)
+    area_phase1 = sc_image.binary_dilation(concentration_arr[..., phase1] > 0.5, iterations=area)
     return np.logical_and(area_phase0, area_phase1)