lbmpy/fieldaccess.py → src/lbmpy/fieldaccess.py

@@ -4,13 +4,19 @@ 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',
            'PeriodicTwoFieldsAccessor', 'StreamPushTwoFieldsAccessor',
            'EsoTwistEvenTimeStepAccessor', 'EsoTwistOddTimeStepAccessor',
+           'EsoPullEvenTimeStepAccessor', 'EsoPullOddTimeStepAccessor',
+           'EsoPushEvenTimeStepAccessor', 'EsoPushOddTimeStepAccessor',
            'visualize_pdf_field_accessor', 'visualize_field_mapping']
 
 
@@ -51,11 +57,11 @@ class CollideOnlyInplaceAccessor(PdfFieldAccessor):
 
     @staticmethod
     def read(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
     @staticmethod
     def write(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
 
 class StreamPullTwoFieldsAccessor(PdfFieldAccessor):
@@ -67,7 +73,7 @@ class StreamPullTwoFieldsAccessor(PdfFieldAccessor):
 
     @staticmethod
     def write(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
 
 class StreamPushTwoFieldsAccessor(PdfFieldAccessor):
@@ -75,7 +81,7 @@ class StreamPushTwoFieldsAccessor(PdfFieldAccessor):
 
     @staticmethod
     def read(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
     @staticmethod
     def write(field, stencil):
@@ -109,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:
@@ -125,7 +131,7 @@ class PeriodicTwoFieldsAccessor(PdfFieldAccessor):
 
     @staticmethod
     def write(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
 
 class AAEvenTimeStepAccessor(PdfFieldAccessor):
@@ -133,11 +139,11 @@ class AAEvenTimeStepAccessor(PdfFieldAccessor):
 
     @staticmethod
     def read(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
     @staticmethod
     def write(field, stencil):
-        return [field(stencil.index(inverse_direction(d))) for d in stencil]
+        return [field(stencil.inverse_index(d)) for d in stencil]
 
 
 class AAOddTimeStepAccessor(PdfFieldAccessor):
@@ -145,57 +151,169 @@ class AAOddTimeStepAccessor(PdfFieldAccessor):
 
     @staticmethod
     def read(field, stencil):
-        res = []
-        for i, d in enumerate(stencil):
-            inv_dir = inverse_direction(d)
-            field_access = field[inv_dir](stencil.index(inv_dir))
-            res.append(field_access)
-        return res
+        return [field[inverse_direction(d)](stencil.inverse_index(d)) for i, d in enumerate(stencil)]
 
     @staticmethod
     def write(field, stencil):
         return [field[d](i) for i, d in enumerate(stencil)]
 
 
+class EsoTwistEvenTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
+    @staticmethod
+    def read(field, stencil):
+        return [field[tuple(max(-e, 0) for e in d)](i) for i, d in enumerate(stencil)]
+
+    @staticmethod
+    def write(field, stencil):
+        return [field[tuple(max(e, 0) for e in d)](stencil.inverse_index(d)) for d in stencil]
+
+
 class EsoTwistOddTimeStepAccessor(PdfFieldAccessor):
     is_inplace = True
 
     @staticmethod
     def read(field, stencil):
-        result = []
-        for direction in stencil:
-            inv_dir = inverse_direction(direction)
-            spatial_offset = tuple(max(e, 0) for e in inv_dir)
-            result.append(field[spatial_offset](stencil.index(inv_dir)))
+        return [field[tuple(max(e, 0) for e in inverse_direction(d))](stencil.inverse_index(d)) for d in stencil]
+
+    @staticmethod
+    def write(field, stencil):
+        return [field[tuple(max(e, 0) for e in d)](i) for i, d in enumerate(stencil)]
+
+
+class EsoPullEvenTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
+    @staticmethod
+    def read(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[inverse_direction(d)](i))
+            else:
+                result.append(field[center_cell](i))
         return result
 
     @staticmethod
     def write(field, stencil):
-        result = []
-        for i, direction in enumerate(stencil):
-            spatial_offset = tuple(max(e, 0) for e in direction)
-            result.append(field[spatial_offset](i))
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[center_cell](stencil.inverse_index(d)))
+            else:
+                result.append(field[d](stencil.inverse_index(d)))
+
         return result
 
 
-class EsoTwistEvenTimeStepAccessor(PdfFieldAccessor):
+class EsoPullOddTimeStepAccessor(PdfFieldAccessor):
     is_inplace = True
 
     @staticmethod
     def read(field, stencil):
-        result = []
-        for i, direction in enumerate(stencil):
-            spatial_offset = tuple(max(-e, 0) for e in direction)
-            result.append(field[spatial_offset](i))
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[inverse_direction(d)](stencil.inverse_index(d)))
+            else:
+                result.append(field[center_cell](stencil.inverse_index(d)))
         return result
 
     @staticmethod
     def write(field, stencil):
-        result = []
-        for direction in stencil:
-            inv_dir = inverse_direction(direction)
-            spatial_offset = tuple(max(e, 0) for e in direction)
-            result.append(field[spatial_offset](stencil.index(inv_dir)))
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[center_cell](i))
+            else:
+                result.append(field[d](i))
+
+        return result
+
+
+class EsoPushEvenTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
+    @staticmethod
+    def read(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[center_cell](stencil.inverse_index(d)))
+            else:
+                result.append(field[inverse_direction(d)](stencil.inverse_index(d)))
+
+        return result
+
+    @staticmethod
+    def write(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[d](i))
+            else:
+                result.append(field[center_cell](i))
+
+        return result
+
+
+class EsoPushOddTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
+    @staticmethod
+    def read(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            inv_dir = inverse_direction(d)
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[center_cell](i))
+            else:
+                result.append(field[inv_dir](i))
+
+        return result
+
+    @staticmethod
+    def write(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[d](stencil.inverse_index(d)))
+            else:
+                result.append(field[center_cell](stencil.inverse_index(d)))
+
         return result
 
 
@@ -214,15 +332,18 @@ def visualize_field_mapping(axes, stencil, field_mapping, inverted=False, color=
     grid.draw(axes)
 
 
-def visualize_pdf_field_accessor(pdf_field_accessor, title=True, read_plot_params={}, write_plot_params={},
+def visualize_pdf_field_accessor(pdf_field_accessor, title=True, read_plot_params=None, write_plot_params=None,
                                  figure=None):
-    from lbmpy.stencils import get_stencil
 
+    if write_plot_params is None:
+        write_plot_params = {}
+    if read_plot_params is None:
+        read_plot_params = {}
     if figure is None:
         import matplotlib.pyplot as plt
         figure = plt.gcf()
 
-    stencil = get_stencil('D2Q9')
+    stencil = LBStencil(Stencil.D2Q9)
 
     figure.patch.set_facecolor('white')
 
@@ -244,3 +365,25 @@ def visualize_pdf_field_accessor(pdf_field_accessor, title=True, read_plot_param
     if title:
         ax_left.set_title("Read")
         ax_right.set_title("Write")
+
+# -------------------------------------------- Helpers -----------------------------------------------------------
+
+
+def _get_lehmann_stencil(stencil):
+    """
+    EsoPull and EsoPush streaming is only simple to implement with a specific stencil ordering, that comes from
+    "High Performance Free Surface LBM on GPUs" by moritz lehmann
+
+    Args:
+        stencil: lattice Boltzmann stencil
+    """
+    if stencil.Q == 9:
+        return LBStencil(Stencil.D2Q9, ordering="lehmann")
+    elif stencil.Q == 15:
+        return LBStencil(Stencil.D3Q15, ordering="lehmann")
+    elif stencil.Q == 19:
+        return LBStencil(Stencil.D3Q19, ordering="lehmann")
+    elif stencil.Q == 27:
+        return LBStencil(Stencil.D3Q27, ordering="lehmann")
+    else:
+        ValueError("EsoPull or EsoPush is only available for D2Q9, D3Q15, D3Q19 and D3Q27 stencil")

src/lbmpy/flow_statistics.py

+import sympy as sp
+import pystencils as ps
+
+from pystencils.field import Field
+
+
+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 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 / 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 ::
+        \sigma_n^2 = \frac{M_{2,n}}{n}
+        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):
+        dim = field.values_per_cell()
+        welford_field = field.center
+    elif isinstance(field, Field.Access):
+        dim = field.field.values_per_cell()
+        welford_field = field
+    else:
+        raise ValueError("Vector field has to be a pystencils Field or Field.Access")
+
+    if isinstance(mean_field, Field):
+        welford_mean_field = mean_field.center
+    elif isinstance(mean_field, Field.Access):
+        welford_mean_field = mean_field
+    else:
+        raise ValueError("Mean vector field has to be a pystencils Field or Field.Access")
+
+    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_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("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}")
+
+    main_assignments = list()
+
+    for i in range(dim):
+        main_assignments.append(ps.Assignment(delta[i], welford_field.at_index(i) - welford_mean_field.at_index(i)))
+        main_assignments.append(ps.Assignment(welford_mean_field.at_index(i),
+                                              welford_mean_field.at_index(i) + delta[i] / counter))
+
+    if sum_of_squares_field is not None:
+        delta2 = sp.symbols(f"delta2_:{dim}")
+        for i in range(dim):
+            main_assignments.append(
+                ps.Assignment(delta2[i], welford_field.at_index(i) - welford_mean_field.at_index(i)))
+        for i in range(dim):
+            for j in range(dim):
+                idx = i * dim + j
+                main_assignments.append(ps.Assignment(
+                    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

lbmpy/fluctuatinglb.py → src/lbmpy/fluctuatinglb.py

@@ -3,33 +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
+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=(),
+                                       *,
+                                       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=(),
-                                       block_offsets=(0, 0, 0), seed=TypedSymbol("seed", np.uint32),
-                                       rng_node=PhiloxFourFloats, c_s_sq=sp.Rational(1, 3)):
-    """"""
+                                       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'.")
-
+    
     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):
@@ -41,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)]
@@ -60,3 +109,22 @@ def fluctuation_correction(method, rng_generator, amplitudes=SymbolGen("phi")):
 
     # corrections are applied in real space hence we need to convert to real space here
     return method.moment_matrix.inv() * amplitude_matrix * random_matrix
+
+
+class ThermalizedEquilibrium(ContinuousHydrodynamicMaxwellian):
+    """TODO: Remove Again! 
+        
+    This class is currently only used in the tutorial notebook `demo_thermalized_lbm.ipynb`
+    and has been added only temporarily, until the thermalized LBM is updated to our new
+    equilibrium framework."""
+    def __init__(self, random_number_symbols, **kwargs):
+        super().__init__(**kwargs)
+        self.random_number_symbols = random_number_symbols
+
+    def moment(self, exponent_tuple_or_polynomial):
+        value = super().moment(exponent_tuple_or_polynomial)
+        if is_shear_moment(exponent_tuple_or_polynomial, dim=self.dim):
+            value += self.random_number_symbols[0] * 0.001
+        elif get_order(exponent_tuple_or_polynomial) > 2:
+            value += self.random_number_symbols[1] * 0.001
+        return value

src/lbmpy/forcemodels.py

+r"""
+.. module:: forcemodels
+    :synopsis: Collection of forcing terms for hydrodynamic LBM simulations
+
+
+Get started:
+------------
+
+This module offers different models to introduce a body force in the lattice Boltzmann scheme.
+If you don't know which model to choose, use :class:`lbmpy.forcemodels.Guo`.
+
+
+Detailed information:
+---------------------
+
+Force models add a term :math:`C_F` to the collision equation:
+
+.. math ::
+
+    f(\mathbf{x} + c_q \Delta t, t + \Delta t) - f(\mathbf{x},t) = \Omega(f, f^{(\mathrm{eq})})
+                                                            + \underbrace{F_q}_{\mbox{forcing term}}
+
+The form of this term depends on the concrete force model: the first moment of this forcing term is equal
+to the acceleration :math:`\mathbf{a}` for all force models. Here :math:`\mathbf{F}` is the D dimensional force vector,
+which defines the force for each spatial dircetion.
+
+.. math ::
+
+    \sum_q \mathbf{c}_q \mathbf{F} = \mathbf{a}
+
+
+The second order moment is different for the forcing models - if it is zero the model is suited for
+incompressible flows. For weakly compressible collision operators a force model with a corrected second order moment
+should be chosen.
+
+.. math ::
+
+    \sum_q c_{qi} c_{qj} f_q &= F_i u_j + F_j u_i  &\qquad \mbox{for Guo, Luo models}
+    
+    \sum_q c_{qi} c_{qj} f_q &= 0  &\qquad \mbox{for Simple, Buick}
+    
+Models with zero second order moment have:
+
+.. math ::
+    
+    F_q = \frac{w_q}{c_s^2} c_{qi} \; a_i
+
+Models with nonzero second moment have:
+
+.. math ::
+    
+    F_q = \frac{w_q}{c_s^2} c_{qi} \; a_i + \frac{w_q}{c_s^4} (c_{qi} c_{qj} - c_s^2 \delta_{ij} ) u_j \, a_i
+
+
+For all force models the computation of the macroscopic velocity has to be adapted (shifted) by adding a term
+:math:`S_{macro}` that we call "macroscopic velocity shift"
+    
+    .. math ::
+    
+        \mathbf{u} &= \sum_q \mathbf{c}_q f_q + S_{\mathrm{macro}}
+        
+        S_{\mathrm{macro}} &= \frac{\Delta t}{2 \cdot \rho} \sum_q F_q
+
+
+Some models also shift the velocity entering the equilibrium distribution.
+
+Comparison
+----------
+ 
+Force models can be distinguished by 2 options:
+
+**Option 1**:
+    :math:`C_F = 1` and equilibrium velocity is not shifted, or :math:`C_F=(1 - \frac{\omega}{2})`
+    and equilibrum is shifted.
+ 
+**Option 2**: 
+    second velocity moment is zero or :math:`F_i u_j + F_j u_i`
+
+
+=====================  ====================  =================
+ Option2 \\ Option1    no equilibrium shift  equilibrium shift
+=====================  ====================  =================
+second moment zero      :class:`Simple`      :class:`Buick`
+second moment nonzero   :class:`Luo`         :class:`Guo`         
+=====================  ====================  =================
+  
+"""
+from warnings import warn
+import abc
+import sympy as sp
+
+from pystencils.sympyextensions import scalar_product
+
+from lbmpy.maxwellian_equilibrium import (
+    discrete_maxwellian_equilibrium, get_equilibrium_values_of_maxwell_boltzmann_function
+)
+from lbmpy.moments import (
+    MOMENT_SYMBOLS, exponent_tuple_sort_key, exponents_to_polynomial_representations,
+    extract_monomials, moment_sort_key, moment_matrix,
+    monomial_to_polynomial_transformation_matrix, set_up_shift_matrix)
+
+FORCE_SYMBOLS = sp.symbols("F_x, F_y, F_z")
+
+
+class AbstractForceModel(abc.ABC):
+    r"""
+    Abstract base class for all force models. All force models have to implement the __call__, which should return a
+    q dimensional vector added to the PDFs in the population space. If an MRT method is used, it is also possible
+    to apply the force directly in the moment space. This is done by additionally providing the function
+    moment_space_forcing. The MRT method will check if it is available and apply the force directly in the moment
+    space. For MRT methods in the central moment space the central_moment_space_forcing function can be provided,
+    which shifts the force vector to the central moment space. Applying the force in the collision space has the
+    advantage of saving FLOPs. Furthermore, it is sometimes easier to apply the correct force vector in the
+    collision space because often, the relaxation matrix has to be taken into account.
+
+    Args:
+        force: force vector of size dim which contains the force for each spatial dimension.
+    """
+
+    def __init__(self, force):
+        self._force = force
+
+        # All force models work internaly with a pure symbolic list of the forcing vector.
+        # Each entry of the original force vector which is not a symbol is mapped to a symbol and a subs dict is
+        # created. The subs dict should be used inside the LB method for the creation of the collision rule.
+        self._symbolic_force = [x if isinstance(x, sp.Symbol) else y for x, y in zip(force, FORCE_SYMBOLS)]
+        self._subs_dict_force = {x: y for (x, y) in zip(self._symbolic_force, force) if x != y}
+
+        # The following booleans should indicate if a force model is has the function moment_space_forcing which
+        # transfers the forcing terms to the moment space or central_moment_space_forcing which transfers them to the
+        # central moment space
+        self.has_moment_space_forcing = hasattr(self, "moment_space_forcing")
+        self.has_central_moment_space_forcing = hasattr(self, "central_moment_space_forcing")
+        self.has_symmetric_central_moment_forcing = hasattr(self, "symmetric_central_moment_forcing")
+
+    def __call__(self, lb_method):
+        r"""
+        This function returns a vector of size q which is added to the PDFs in the PDF space. It has to be implemented
+        by all forcing models and returns a sympy Matrix containing the q dimensional force vector.
+
+        Args:
+            lb_method: LB method, see lbmpy.creationfunctions.create_lb_method
+        """
+        raise NotImplementedError("Force model class has to overwrite __call__")
+
+    def macroscopic_velocity_shift(self, density):
+        r"""
+        macroscopic velocity shift by :math:`\frac{\Delta t}{2 \cdot \rho}`
+
+        Args:
+            density: Density symbol which is needed for the shift
+        """
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+    def macroscopic_momentum_density_shift(self, *_):
+        r"""
+        macroscopic momentum density shift by :math:`\frac{\Delta t}{2}`
+        """
+        return default_momentum_density_shift(self.symbolic_force_vector)
+
+    def equilibrium_velocity_shift(self, density):
+        r"""
+        Some models also shift the velocity entering the equilibrium distribution. By default the shift is zero
+
+        Args:
+            density: Density symbol which is needed for the shift
+        """
+        return [0] * len(self.symbolic_force_vector)
+
+    @property
+    def symbolic_force_vector(self):
+        return self._symbolic_force
+
+    @property
+    def subs_dict_force(self):
+        return self._subs_dict_force
+
+
+class Simple(AbstractForceModel):
+    r"""
+    A simple force model which introduces the following additional force term in the
+    collision process :math:`\frac{w_q}{c_s^2} \mathbf{c_{q}} \cdot \mathbf{F}` (often: force = rho * acceleration
+    where rho is the zeroth moment to be consistent with the above definition)
+    Should only be used with constant forces!
+    Shifts the macroscopic velocity by :math:`\frac{\mathbf{F}}{2}`, but does not change the equilibrium velocity.
+    """
+
+    def __call__(self, lb_method):
+        force = self.symbolic_force_vector
+        assert len(force) == lb_method.dim, "Force vectore must match with the dimensions of the lb method"
+        cs_sq = sp.Rational(1, 3)  # squared speed of sound
+
+        result = [(w_i / cs_sq) * scalar_product(force, direction)
+                  for direction, w_i in zip(lb_method.stencil, lb_method.weights)]
+
+        return sp.Matrix(result)
+
+    def moment_space_forcing(self, lb_method):
+        return (lb_method.moment_matrix * self(lb_method)).expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        moments = (lb_method.moment_matrix * sp.Matrix(self(lb_method))).expand()
+        return lb_method.shift_matrix * moments
+
+    def symmetric_central_moment_forcing(self, lb_method, central_moments):
+        u = lb_method.first_order_equilibrium_moment_symbols
+        cm_matrix = moment_matrix(central_moments, lb_method.stencil, shift_velocity=u)
+        before = sp.Matrix([0] * lb_method.stencil.Q)
+        after = cm_matrix @ sp.Matrix(self(lb_method))
+        return before, after
+
+
+class CentralMoment(AbstractForceModel):
+    r"""
+    A force model that only shifts the macroscopic and equilibrium velocities and does not introduce a forcing term to
+    the collision process. Forcing is then applied through relaxation of the first central moments in the shifted
+    frame of reference (cf. https://doi.org/10.1016/j.camwa.2015.05.001).
+    """
+    def __call__(self, lb_method):
+        raise ValueError("This force model can only be combined with the Cumulant collision model")
+
+    def symmetric_central_moment_forcing(self, lb_method, *_):
+        before = sp.Matrix([0, ] * lb_method.stencil.Q)
+        after = sp.Matrix([0, ] * lb_method.stencil.Q)
+        for i, sf in enumerate(self.symbolic_force_vector):
+            before[i + 1] = sp.Rational(1, 2) * sf
+            after[i + 1] = sp.Rational(1, 2) * sf
+
+        return before, after
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+class Luo(AbstractForceModel):
+    r"""Force model by Luo :cite:`luo1993lattice`.
+
+    Shifts the macroscopic velocity by :math:`\frac{\mathbf{F}}{2}`, but does not change the equilibrium velocity.
+    """
+
+    def __call__(self, lb_method):
+        u = sp.Matrix(lb_method.first_order_equilibrium_moment_symbols)
+        force = sp.Matrix(self.symbolic_force_vector)
+
+        cs_sq = sp.Rational(1, 3)  # squared speed of sound
+
+        result = []
+        for direction, w_i in zip(lb_method.stencil, lb_method.weights):
+            direction = sp.Matrix(direction)
+            first_summand = (direction - u) / cs_sq
+            second_summand = (direction * direction.dot(u)) / cs_sq ** 2
+
+            fq = w_i * force.dot(first_summand + second_summand)
+
+            result.append(fq.simplify())
+        return sp.Matrix(result)
+
+    def moment_space_forcing(self, lb_method):
+        return (lb_method.moment_matrix * self(lb_method)).expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        moments = lb_method.moment_matrix * self(lb_method)
+        return (lb_method.shift_matrix * moments).expand()
+
+    def symmetric_central_moment_forcing(self, lb_method, central_moments):
+        u = lb_method.first_order_equilibrium_moment_symbols
+        cm_matrix = moment_matrix(central_moments, lb_method.stencil, shift_velocity=u)
+        before = sp.Matrix([0] * lb_method.stencil.Q)
+        after = (cm_matrix @ sp.Matrix(self(lb_method))).expand()
+        return before, after
+
+
+class Guo(AbstractForceModel):
+    r"""
+    Force model by Guo  :cite:`guo2002discrete`, generalized to MRT,
+    which makes it equivalent to :cite:`schiller2008thermal`, equation 4.67
+    Adapts the calculation of the macroscopic velocity as well as the equilibrium velocity
+    (both shifted by :math:`\frac{\mathbf{F}}{2}`)!
+    """
+
+    def __call__(self, lb_method):
+        if len(set(lb_method.relaxation_rates)) == 1:
+            #   It's an SRT method!
+            rr = lb_method.symbolic_relaxation_matrix[0]
+            force_terms = Luo(self.symbolic_force_vector)(lb_method)
+            correction_factor = (1 - rr / 2)
+            result = correction_factor * force_terms
+        else:
+            force_terms = self.moment_space_forcing(lb_method)
+            result = (lb_method.moment_matrix.inv() * force_terms).expand()
+        return result
+
+    def moment_space_forcing(self, lb_method):
+        luo = Luo(self.symbolic_force_vector)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = correction_factor * (lb_method.moment_matrix * sp.Matrix(luo(lb_method))).expand()
+        return moments
+
+    def central_moment_space_forcing(self, lb_method):
+        luo = Luo(self.symbolic_force_vector)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = (lb_method.moment_matrix * sp.Matrix(luo(lb_method)))
+        central_moments = correction_factor * (lb_method.shift_matrix * moments).expand()
+
+        return central_moments
+
+    def symmetric_central_moment_forcing(self, lb_method, central_moments):
+        luo = Luo(self.symbolic_force_vector)
+        _, force_cms = luo.symmetric_central_moment_forcing(lb_method, central_moments)
+        force_cms = sp.Rational(1, 2) * force_cms
+        return force_cms, force_cms
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+class He(AbstractForceModel):
+    r"""
+    Force model by He  :cite:`HeForce`
+    Adapts the calculation of the macroscopic velocity as well as the equilibrium velocity
+    (both shifted by :math:`\frac{\mathbf{F}}{2}`)!
+
+    Force moments are derived from the continuous maxwellian equilibrium. From the
+    moment integrals of the continuous force term
+
+    .. math::
+
+        F (\mathbf{c})
+        = \frac{1}{\rho c_s^2}
+          \mathbf{F} \cdot ( \mathbf{c} - \mathbf{u} )
+          f^{\mathrm{eq}} (\mathbf{c})
+
+    the following analytical expresson for the monomial raw moments of the force is found:
+
+    .. math::
+
+        m_{\alpha\beta\gamma}^{F, \mathrm{He}}
+            = \frac{1}{\rho c_s^2} \left(
+                F_x m^{\mathrm{eq}}_{\alpha+1,\beta,\gamma}
+                + F_y m^{\mathrm{eq}}_{\alpha,\beta+1,\gamma}
+                + F_z m^{\mathrm{eq}}_{\alpha,\beta,\gamma+1}
+                - m^{eq}_{\alpha\beta\gamma} ( \mathbf{F} \cdot \mathbf{u} )
+            \right)
+    """
+
+    def __init__(self, force):
+        super(He, self).__init__(force)
+
+    def forcing_terms(self, lb_method):
+        u = sp.Matrix(lb_method.first_order_equilibrium_moment_symbols)
+        force = sp.Matrix(self.symbolic_force_vector)
+
+        cs_sq = sp.Rational(1, 3)  # squared speed of sound
+        # eq. 6.31 in the LB book by Krüger et al. shows that the equilibrium terms are devided by rho.
+        # This is done here with subs({rho: 1}) to be consistent with compressible and incompressible force models
+        cqc = lb_method.conserved_quantity_computation
+        eq_terms = discrete_maxwellian_equilibrium(lb_method.stencil, rho=sp.Integer(1),
+                                                   u=cqc.velocity_symbols, c_s_sq=sp.Rational(1, 3))
+
+        result = []
+        for direction, eq in zip(lb_method.stencil, eq_terms):
+            direction = sp.Matrix(direction)
+            eu_dot_f = (direction - u).dot(force)
+            result.append(eq * eu_dot_f / cs_sq)
+
+        return sp.Matrix(result)
+
+    def continuous_force_raw_moments(self, lb_method, moments=None):
+        rho = lb_method.zeroth_order_equilibrium_moment_symbol
+        u = lb_method.first_order_equilibrium_moment_symbols
+        dim = lb_method.dim
+        c_s_sq = sp.Rational(1, 3)
+        force = sp.Matrix(self.symbolic_force_vector)
+
+        moment_polynomials = lb_method.moments if moments is None else moments
+        moment_exponents = sorted(extract_monomials(moment_polynomials), key=exponent_tuple_sort_key)
+        moment_monomials = exponents_to_polynomial_representations(moment_exponents)
+        extended_monomials = set()
+        for m in moment_monomials:
+            extended_monomials |= {m} | {m * x for x in MOMENT_SYMBOLS[:dim]}
+
+        extended_monomials = sorted(extended_monomials, key=moment_sort_key)
+        moment_eq_values = get_equilibrium_values_of_maxwell_boltzmann_function(extended_monomials, dim, rho=rho,
+                                                                                u=u, c_s_sq=c_s_sq)
+        moment_to_eq_dict = {m: v for m, v in zip(extended_monomials, moment_eq_values)}
+
+        monomial_force_moments = []
+        for moment in moment_monomials:
+            m_base = moment_to_eq_dict[moment]
+            m_shifted = sp.Matrix([moment_to_eq_dict[moment * x] for x in MOMENT_SYMBOLS[:dim]])
+            m_force = (c_s_sq * rho)**(-1) * (force.dot(m_shifted) - m_base * force.dot(u))
+            monomial_force_moments.append(m_force.expand())
+
+        mono_to_poly_matrix = monomial_to_polynomial_transformation_matrix(moment_exponents, moment_polynomials)
+        polynomial_force_moments = mono_to_poly_matrix * sp.Matrix(monomial_force_moments)
+        return polynomial_force_moments
+
+    def continuous_force_central_moments(self, lb_method, moments=None):
+        if moments is None:
+            moments = lb_method.moments
+        raw_moments = self.continuous_force_raw_moments(lb_method, moments=moments)
+        u = lb_method.first_order_equilibrium_moment_symbols
+        shift_matrix = set_up_shift_matrix(moments, lb_method.stencil, velocity_symbols=u)
+        return (shift_matrix * raw_moments).expand()
+
+    def __call__(self, lb_method):
+        if len(set(lb_method.relaxation_rates)) == 1:
+            #   It's an SRT method!
+            rr = lb_method.symbolic_relaxation_matrix[0]
+            force_terms = self.forcing_terms(lb_method)
+            correction_factor = (1 - rr / 2)
+            result = correction_factor * force_terms
+        else:
+            force_terms = self.moment_space_forcing(lb_method)
+            result = (lb_method.moment_matrix.inv() * force_terms).expand()
+        return result
+
+    def moment_space_forcing(self, lb_method):
+        correction_factor = sp.eye(len(lb_method.stencil)) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = self.continuous_force_raw_moments(lb_method)
+        moments = (correction_factor * moments).expand()
+        return moments
+
+    def central_moment_space_forcing(self, lb_method):
+        correction_factor = sp.eye(len(lb_method.stencil)) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        central_moments = self.continuous_force_central_moments(lb_method)
+        central_moments = (correction_factor * central_moments).expand()
+        return central_moments
+
+    def symmetric_central_moment_forcing(self, lb_method, central_moments):
+        central_moments = exponents_to_polynomial_representations(central_moments)
+        force_cms = sp.Rational(1, 2) * self.continuous_force_central_moments(lb_method, moments=central_moments)
+        return force_cms, force_cms
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+class Schiller(Guo):
+    r"""
+    Force model by Schiller  :cite:`schiller2008thermal`, equation 4.67
+    Equivalent to the generalized Guo model.
+    """
+
+    def __init__(self, force):
+        warn("The Schiller force model is deprecated, please use the Guo model, which is equivalent",
+             DeprecationWarning)
+        super(Schiller, self).__init__(force)
+
+
+class Buick(AbstractForceModel):
+    r"""
+    This force model :cite:`buick2000gravity` has a force term with zero second moment. 
+    It is suited for incompressible lattice models. However it should be used with care because such a LB body form
+    model is only consistent when applied to the solution of steady - state hydrodynamic problems. More information
+    on an analysis of the Buick force model can be found in :cite:`silva2010` and in :cite:`silva2020`
+    """
+
+    def __call__(self, lb_method, **kwargs):
+        if len(set(lb_method.relaxation_rates)) == 1:
+            #   It's an SRT method!
+            rr = lb_method.symbolic_relaxation_matrix[0]
+            force_terms = Simple(self.symbolic_force_vector)(lb_method)
+            correction_factor = (1 - rr / 2)
+            result = correction_factor * force_terms
+        else:
+            force_terms = self.moment_space_forcing(lb_method)
+            result = (lb_method.moment_matrix.inv() * force_terms).expand()
+        return result
+
+    def moment_space_forcing(self, lb_method):
+        simple = Simple(self.symbolic_force_vector)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = correction_factor * (lb_method.moment_matrix * sp.Matrix(simple(lb_method)))
+        return moments.expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        simple = Simple(self.symbolic_force_vector)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = (lb_method.moment_matrix * sp.Matrix(simple(lb_method)))
+        central_moments = correction_factor * (lb_method.shift_matrix * moments)
+
+        return central_moments.expand()
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+class EDM(AbstractForceModel):
+    r"""Exact differencing force model as shown in :cite:`lbm_book` in eq. 6.32"""
+
+    def __call__(self, lb_method):
+        cqc = lb_method.conserved_quantity_computation
+        reference_density = cqc.density_symbol if cqc.compressible else 1
+        rho = cqc.density_symbol
+        delta_rho = cqc.density_deviation_symbol
+        rho_0 = cqc.background_density
+        u = cqc.velocity_symbols
+
+        equilibrium_terms = lb_method.get_equilibrium_terms()
+        equilibrium_terms = equilibrium_terms.subs({delta_rho: rho - rho_0})
+
+        shifted_u = (u_i + f_i / reference_density for u_i, f_i in zip(u, self._force))
+        shifted_eq = equilibrium_terms.subs({u_i: su_i for u_i, su_i in zip(u, shifted_u)})
+        return shifted_eq - equilibrium_terms
+
+    def moment_space_forcing(self, lb_method):
+        moments = lb_method.moment_matrix * self(lb_method)
+        return moments.expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        moments = lb_method.moment_matrix * self(lb_method)
+        central_moments = lb_method.shift_matrix * moments.expand()
+        return central_moments.expand()
+
+
+class ShanChen(AbstractForceModel):
+    r"""Shan and Chen force model. The implementation is done according to :cite:`silva2020`.
+        For reference compare table 1 which is the Shan and Chen model for an SRT collision operator. These terms are
+        transfered to the moment space and then all representations for the different collision operators are derived
+        from that.
+    """
+
+    def forcing_terms(self, lb_method):
+        q = len(lb_method.stencil)
+        cqc = lb_method.conserved_quantity_computation
+        rho = cqc.density_symbol if cqc.compressible else 1
+        u = cqc.velocity_symbols
+
+        F = sp.Matrix(self.symbolic_force_vector)
+        uf = sp.Matrix(u).dot(F)
+        F2 = sp.Matrix(F).dot(sp.Matrix(F))
+        Fq = sp.zeros(q, 1)
+        uq = sp.zeros(q, 1)
+        for i, cq in enumerate(lb_method.stencil):
+            Fq[i] = sp.Matrix(cq).dot(sp.Matrix(F))
+            uq[i] = sp.Matrix(cq).dot(u)
+
+        linear_term = sp.zeros(q, 1)
+        non_linear_term = sp.zeros(q, 1)
+        for i, w_i in enumerate(lb_method.weights):
+            linear_term[i] = w_i * (Fq[i] + 3 * uq[i] * Fq[i] - uf)
+            non_linear_term[i] = ((w_i / (2 * rho)) * (3 * Fq[i] ** 2 - F2))
+
+        return linear_term, non_linear_term
+
+    def __call__(self, lb_method):
+        force_terms = self.moment_space_forcing(lb_method)
+        result = lb_method.moment_matrix.inv() * force_terms
+        return result.expand()
+
+    def moment_space_forcing(self, lb_method):
+        linear_term, non_linear_term = self.forcing_terms(lb_method)
+        q = len(lb_method.stencil)
+
+        rel = lb_method.symbolic_relaxation_matrix
+        cs_sq = sp.Rational(1, 3)
+
+        correction_factor = 1 / cs_sq * (sp.eye(q) - sp.Rational(1, 2) * rel)
+        M = lb_method.moment_matrix
+        moments = correction_factor * (M * linear_term) + correction_factor ** 2 * (M * non_linear_term)
+        return moments.expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        linear_term, non_linear_term = self.forcing_terms(lb_method)
+        q = len(lb_method.stencil)
+
+        rel = lb_method.symbolic_relaxation_matrix
+        cs_sq = sp.Rational(1, 3)
+
+        correction_factor = 1 / cs_sq * (sp.eye(q) - sp.Rational(1, 2) * rel)
+        M = lb_method.moment_matrix
+        N = lb_method.shift_matrix
+        moments_linear_term = (M * linear_term)
+        moments_non_linear_term = (M * non_linear_term)
+        central_moments_linear_term = correction_factor * (N * moments_linear_term)
+        central_moments_non_linear_term = correction_factor ** 2 * (N * moments_non_linear_term)
+        central_moments = central_moments_linear_term + central_moments_non_linear_term
+        return central_moments.expand()
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+# --------------------------------  Helper functions  ------------------------------------------------------------------
+
+
+def default_velocity_shift(density, force):
+    return [f_i / (2 * density) for f_i in force]
+
+
+def default_momentum_density_shift(force):
+    return [f_i / 2 for f_i in force]

lbmpy/lbstep.py → src/lbmpy/lbstep.py

 from types import MappingProxyType
+from dataclasses import replace
 
 import numpy as np
 
 from lbmpy.boundaries.boundaryhandling import LatticeBoltzmannBoundaryHandling
-from lbmpy.creationfunctions import (
-    create_lb_function, switch_to_symbolic_relaxation_rates_for_omega_adapting_methods,
-    update_with_default_parameters)
+from lbmpy.creationfunctions import (create_lb_function, update_with_default_parameters)
+from lbmpy.enums import Stencil
 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 get_stencil
-from pystencils import create_data_handling, create_kernel, make_slice
+from lbmpy.stencils import LBStencil
+
+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,
-                 kernel_params=MappingProxyType({}), data_handling=None, name="lbm", optimization={},
+                 kernel_params=MappingProxyType({}), data_handling=None, name="lbm", optimization=None,
                  velocity_data_name=None, density_data_name=None, density_data_index=None,
                  compute_velocity_in_every_step=False, compute_density_in_every_step=False,
                  velocity_input_array_name=None, time_step_order='stream_collide', flag_interface=None,
-                 alignment_if_vectorized=64, fixed_loop_sizes=True, fixed_relaxation_rates=True,
-                 timeloop_creation_function=TimeLoop, **method_parameters):
-
+                 alignment_if_vectorized=64, fixed_loop_sizes=True,
+                 timeloop_creation_function=TimeLoop,
+                 lbm_config=None, lbm_optimisation=None,
+                 config: CreateKernelConfig | None = None,
+                 **method_parameters):
+
+        if optimization is None:
+            optimization = {}
         self._timeloop_creation_function = timeloop_creation_function
 
         # --- Parameter normalization  ---
@@ -32,7 +44,14 @@ class LatticeBoltzmannStep:
             if domain_size is not None:
                 raise ValueError("When passing a data_handling, the domain_size parameter can not be specified")
 
-        target = optimization.get('target', 'cpu')
+        if config is not None:
+            if IS_PYSTENCILS_2:
+                target = config.get_target()
+            else:
+                target = config.target
+        else:
+            target = optimization.get('target', Target.CPU)
+
         if data_handling is None:
             if domain_size is None:
                 raise ValueError("Specify either domain_size or data_handling")
@@ -42,18 +61,27 @@ class LatticeBoltzmannStep:
                                                  default_target=target,
                                                  parallel=False)
 
+        if lbm_config:
+            method_parameters['stencil'] = lbm_config.stencil
+
         if 'stencil' not in method_parameters:
-            method_parameters['stencil'] = 'D2Q9' if data_handling.dim == 2 else 'D3Q27'
+            method_parameters['stencil'] = LBStencil(Stencil.D2Q9) \
+                if data_handling.dim == 2 else LBStencil(Stencil.D3Q27)
 
-        method_parameters, optimization = update_with_default_parameters(method_parameters, optimization)
-        field_dtype = np.float64 if optimization['double_precision'] else np.float32
+        lbm_config, lbm_optimisation, config = update_with_default_parameters(method_parameters, optimization,
+                                                                              lbm_config, lbm_optimisation, config)
 
-        del method_parameters['kernel_type']
+        # the parallel datahandling understands only numpy datatypes. Strings lead to an errors
+        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 = len(lbm_kernel.method.stencil)
+            q = lbm_kernel.method.stencil.Q
         else:
-            q = len(get_stencil(method_parameters['stencil']))
+            q = lbm_config.stencil.Q
 
         self.name = name
         self._data_handling = data_handling
@@ -62,14 +90,22 @@ class LatticeBoltzmannStep:
         self.velocity_data_name = name + "_velocity" if velocity_data_name is None else velocity_data_name
         self.density_data_name = name + "_density" if density_data_name is None else density_data_name
         self.density_data_index = density_data_index
-        self._optimization = optimization
 
-        self._gpu = target == 'gpu' or target == 'opencl'
-        layout = optimization['field_layout']
+        if IS_PYSTENCILS_2:
+            self._gpu = target.is_gpu()
+        else:
+            self._gpu = target == Target.GPU
+
+        layout = lbm_optimisation.field_layout
 
         alignment = False
-        if optimization['target'] == 'cpu' and optimization['vectorization']:
-            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)
@@ -86,58 +122,71 @@ class LatticeBoltzmannStep:
                                           layout=layout, latex_name='ρ', dtype=field_dtype, alignment=alignment)
 
         if compute_velocity_in_every_step:
-            method_parameters['output']['velocity'] = self._data_handling.fields[self.velocity_data_name]
+            lbm_config.output['velocity'] = self._data_handling.fields[self.velocity_data_name]
         if compute_density_in_every_step:
             density_field = self._data_handling.fields[self.density_data_name]
             if self.density_data_index is not None:
                 density_field = density_field(density_data_index)
-            method_parameters['output']['density'] = density_field
+            lbm_config.output['density'] = density_field
         if velocity_input_array_name is not None:
-            method_parameters['velocity_input'] = self._data_handling.fields[velocity_input_array_name]
-        if method_parameters['omega_output_field'] and isinstance(method_parameters['omega_output_field'], str):
-            method_parameters['omega_output_field'] = data_handling.add_array(method_parameters['omega_output_field'],
-                                                                              dtype=field_dtype, alignment=alignment)
+            lbm_config = replace(lbm_config, velocity_input=self._data_handling.fields[velocity_input_array_name])
+        if isinstance(lbm_config.omega_output_field, str):
+            lbm_config = replace(lbm_config, omega_output_field=data_handling.add_array(lbm_config.omega_output_field,
+                                                                                        dtype=field_dtype,
+                                                                                        alignment=alignment,
+                                                                                        values_per_cell=1))
 
         self.kernel_params = kernel_params.copy()
 
         # --- Kernel creation ---
         if lbm_kernel is None:
-            switch_to_symbolic_relaxation_rates_for_omega_adapting_methods(method_parameters, self.kernel_params,
-                                                                           force=not fixed_relaxation_rates)
+
             if fixed_loop_sizes:
-                optimization['symbolic_field'] = data_handling.fields[self._pdf_arr_name]
-            method_parameters['field_name'] = self._pdf_arr_name
-            method_parameters['temporary_field_name'] = self._tmp_arr_name
+                lbm_optimisation = replace(lbm_optimisation, symbolic_field=data_handling.fields[self._pdf_arr_name])
+            lbm_config = replace(lbm_config, field_name=self._pdf_arr_name)
+            lbm_config = replace(lbm_config, temporary_field_name=self._tmp_arr_name)
+
             if time_step_order == 'stream_collide':
-                self._lbmKernels = [create_lb_function(optimization=optimization,
-                                                       **method_parameters)]
+                self._lbmKernels = [create_lb_function(lbm_config=lbm_config,
+                                                       lbm_optimisation=lbm_optimisation,
+                                                       config=config)]
             elif time_step_order == 'collide_stream':
-                self._lbmKernels = [create_lb_function(optimization=optimization,
-                                                       kernel_type='collide_only',
-                                                       **method_parameters),
-                                    create_lb_function(optimization=optimization,
-                                                       kernel_type='stream_pull_only',
-                                                       ** method_parameters)]
+                self._lbmKernels = [create_lb_function(lbm_config=lbm_config,
+                                                       lbm_optimisation=lbm_optimisation,
+                                                       config=config,
+                                                       kernel_type='collide_only'),
+                                    create_lb_function(lbm_config=lbm_config,
+                                                       lbm_optimisation=lbm_optimisation,
+                                                       config=config,
+                                                       kernel_type='stream_pull_only')]
 
         else:
             assert self._data_handling.dim == lbm_kernel.method.dim, \
-                "Error: %dD Kernel for %d dimensional domain" % (lbm_kernel.method.dim, self._data_handling.dim)
+                f"Error: {lbm_kernel.method.dim}D Kernel for {self._data_handling.dim} dimensional domain"
             self._lbmKernels = [lbm_kernel]
 
         self.method = self._lbmKernels[0].method
         self.ast = self._lbmKernels[0].ast
 
         # -- Boundary Handling  & Synchronization ---
-        stencil_name = method_parameters['stencil']
+        stencil_name = lbm_config.stencil.name
         self._sync_src = data_handling.synchronization_function([self._pdf_arr_name], stencil_name, target,
                                                                 stencil_restricted=True)
         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=optimization['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
+        self._config = config
 
         # -- Macroscopic Value Kernels
         self._getterKernel, self._setterKernel = self._compile_macroscopic_setter_and_getter()
@@ -190,6 +239,21 @@ class LatticeBoltzmannStep:
     def pdf_array_name(self):
         return self._pdf_arr_name
 
+    @property
+    def lbm_config(self):
+        """LBM configuration of the scenario"""
+        return self._lbm_config
+
+    @property
+    def lbm_optimisation(self):
+        """LBM optimisation parameters"""
+        return self._lbm_optimisation
+
+    @property
+    def config(self):
+        """Configutation of pystencils parameters"""
+        return self._config
+
     def _get_slice(self, data_name, slice_obj, masked):
         if slice_obj is None:
             slice_obj = make_slice[:, :] if self.dim == 2 else make_slice[:, :, 0.5]
@@ -343,11 +407,11 @@ class LatticeBoltzmannStep:
 
             collision_rule = create_advanced_velocity_setter_collision_rule(self.method, vel_backup_field,
                                                                             velocity_relaxation_rate)
-            optimization = self._optimization.copy()
-            optimization['symbolic_field'] = dh.fields[self._pdf_arr_name]
+            self._lbm_optimisation.symbolic_field = dh.fields[self._pdf_arr_name]
 
             kernel = create_lb_function(collision_rule=collision_rule, field_name=self._pdf_arr_name,
-                                        temporary_field_name=self._tmp_arr_name, optimization=optimization)
+                                        temporary_field_name=self._tmp_arr_name,
+                                        lbm_optimisation=self._lbm_optimisation)
             self._velocity_init_kernel = kernel
 
         def make_velocity_backup():
@@ -383,7 +447,7 @@ class LatticeBoltzmannStep:
             self._data_handling.all_to_cpu()
             self._data_handling.run_kernel(self._getterKernel, **self.kernel_params)
             global_residuum = compute_residuum()
-            print("Initialization iteration {}, residuum {}".format(steps_run, global_residuum))
+            print(f"Initialization iteration {steps_run}, residuum {global_residuum}")
             if np.isnan(global_residuum) or global_residuum < convergence_threshold:
                 break
 
@@ -391,8 +455,8 @@ class LatticeBoltzmannStep:
         converged = global_residuum < convergence_threshold
         if not converged:
             restore_velocity_backup()
-            raise ValueError("Iterative initialization did not converge after %d steps.\n"
-                             "Current residuum is %s" % (steps_run, global_residuum))
+            raise ValueError(f"Iterative initialization did not converge after {steps_run} steps.\n"
+                             f"Current residuum is {global_residuum}")
 
         return global_residuum, steps_run
 
@@ -404,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='cpu', cpu_openmp=self._optimization['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='cpu', cpu_openmp=self._optimization['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

lbmpy/macroscopic_value_kernels.py → 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, Assignment
 from pystencils.field import Field, get_layout_of_array
+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):
-    """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):
-        previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
-        field_accesses = previous_step_accessor.write(pdfs, lb_method.stencil)
-    elif streaming_pattern == 'pull':
+        accessor = get_accessor(streaming_pattern, previous_timestep)
+        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 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
 
     cqc = lb_method.conserved_quantity_computation
-    inp_eqs = cqc.equilibrium_input_equations_from_init_values(density, velocity)
+    inp_eqs = cqc.equilibrium_input_equations_from_init_values(density, velocity, force_substitution=False)
     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,
-                              streaming_pattern='pull', previous_timestep=Timestep.BOTH):
-    if isinstance(pdfs, Field):
-        previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
-        field_accesses = previous_step_accessor.write(pdfs, lb_method.stencil)
-    elif streaming_pattern == 'pull':
-        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}.")
+def macroscopic_values_getter(lb_method, density, velocity, pdfs, psm_config=None,
+                              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)
+
     cqc = lb_method.conserved_quantity_computation
     assert not (velocity is None and density is None)
     output_spec = {}
@@ -43,15 +90,57 @@ 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='cpu',
+                                      field_layout='numpy', target=Target.CPU,
                                       streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     """
     Create kernel to compute macroscopic value(s) from a pdf field (e.g. density or velocity)
@@ -65,7 +154,7 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
                       is determined automatically and assumed to be equal for all dimensions.        
         iteration_slice: if not None, iteration is done only over this slice of the field
         field_layout: layout for output field, also used for pdf field if pdf_arr is not given
-        target: 'cpu' or 'gpu'
+        target: `Target.CPU` or `Target.GPU`
         previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
 
     Returns:
@@ -116,16 +205,8 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
 
     eqs = cqc.output_equations_from_pdfs(pdf_symbols, output_mapping).all_assignments
 
-    if target == 'cpu':
-        import pystencils.cpu as cpu
-        kernel = cpu.make_python_function(cpu.create_kernel(
-            eqs, ghost_layers=ghost_layers, iteration_slice=iteration_slice))
-    elif target == 'gpu':
-        import pystencils.gpucuda as gpu
-        kernel = gpu.make_python_function(gpu.create_cuda_kernel(
-            eqs, ghost_layers=ghost_layers, iteration_slice=iteration_slice))
-    else:
-        raise ValueError("Unknown target '%s'. Possible targets are 'cpu' and 'gpu'" % (target,))
+    config = CreateKernelConfig(target=target, ghost_layers=ghost_layers, iteration_slice=iteration_slice)
+    kernel = create_kernel(eqs, config=config).compile()
 
     def getter(pdfs, **kwargs):
         if pdf_arr is not None:
@@ -141,7 +222,7 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
 
 def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None,
                                       ghost_layers=1, iteration_slice=None,
-                                      field_layout='numpy', target='cpu',
+                                      field_layout='numpy', target=Target.CPU,
                                       streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     """
     Creates a function that sets a pdf field to specified macroscopic quantities
@@ -156,7 +237,7 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
                       is determined automatically and assumed to be equal for all dimensions.        
         iteration_slice: if not None, iteration is done only over this slice of the field
         field_layout: layout of the pdf field if pdf_arr was not given
-        target: 'cpu' or 'gpu'
+        target: `Target.CPU` or `Target.GPU`
         previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
 
     Returns:
@@ -186,7 +267,7 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
 
         value_map[quantity_name] = value
 
-    cq_equations = cqc.equilibrium_input_equations_from_init_values(**value_map)
+    cq_equations = cqc.equilibrium_input_equations_from_init_values(**value_map, force_substitution=False)
 
     eq = lb_method.get_equilibrium(conserved_quantity_equations=cq_equations)
     if at_least_one_field_input:
@@ -201,16 +282,9 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
     substitutions = {sym: write_accesses[i] for i, sym in enumerate(lb_method.post_collision_pdf_symbols)}
     eq = eq.new_with_substitutions(substitutions).all_assignments
 
-    if target == 'cpu':
-        import pystencils.cpu as cpu
-        kernel = cpu.make_python_function(cpu.create_kernel(eq))
-        kernel = functools.partial(kernel, **fixed_kernel_parameters)
-    elif target == 'gpu':
-        import pystencils.gpucuda as gpu
-        kernel = gpu.make_python_function(gpu.create_cuda_kernel(eq))
-        kernel = functools.partial(kernel, **fixed_kernel_parameters)
-    else:
-        raise ValueError("Unknown target '%s'. Possible targets are 'cpu' and 'gpu'" % (target,))
+    config = CreateKernelConfig(target=target)
+    kernel = create_kernel(eq, config=config).compile()
+    kernel = functools.partial(kernel, **fixed_kernel_parameters)
 
     def setter(pdfs, **kwargs):
         if pdf_arr is not None:
@@ -236,15 +310,17 @@ def create_advanced_velocity_setter_collision_rule(method, velocity_field: Field
         LB collision rule
     """
     cqc = method.conserved_quantity_computation
-    density_symbol = cqc.defined_symbols(order=0)[1]
-    velocity_symbols = cqc.defined_symbols(order=1)[1]
+    density_symbol = cqc.density_symbol
+    velocity_symbols = cqc.velocity_symbols
 
     # density is computed from pdfs
-    eq_input_from_pdfs = cqc.equilibrium_input_equations_from_pdfs(method.pre_collision_pdf_symbols)
+    eq_input_from_pdfs = cqc.equilibrium_input_equations_from_pdfs(
+        method.pre_collision_pdf_symbols, force_substitution=False)
     eq_input_from_pdfs = eq_input_from_pdfs.new_filtered([density_symbol])
     # velocity is read from input field
     vel_symbols = [velocity_field(i) for i in range(method.dim)]
-    eq_input_from_field = cqc.equilibrium_input_equations_from_init_values(velocity=vel_symbols)
+    eq_input_from_field = cqc.equilibrium_input_equations_from_init_values(
+        velocity=vel_symbols, force_substitution=False)
     eq_input_from_field = eq_input_from_field.new_filtered(velocity_symbols)
     # then both are merged together
     eq_input = eq_input_from_pdfs.new_merged(eq_input_from_field)

lbmpy/max_domain_size_info.py → src/lbmpy/max_domain_size_info.py

@@ -27,16 +27,16 @@ import warnings
 
 import numpy as np
 
-# Optional packages cpuinfo, pycuda and psutil for hardware queries
+# Optional packages cpuinfo, cupy and psutil for hardware queries
 try:
     from cpuinfo import get_cpu_info
 except ImportError:
     get_cpu_info = None
 
 try:
-    from pycuda.autoinit import device
+    import cupy
 except ImportError:
-    device = None
+    cupy = None
 
 try:
     from psutil import virtual_memory
@@ -113,9 +113,11 @@ def memory_sizes_of_current_machine():
         if 'l3_cache_size' in cpu_info:
             result['L3'] = cpu_info['l3_cache_size']
 
-    if device:
-        size = device.total_memory() / (1024 * 1024)
-        result['GPU'] = "{0:.0f} MB".format(size)
+    if cupy:
+        for i in range(cupy.cuda.runtime.getDeviceCount()):
+            device = cupy.cuda.Device(i)
+            size = device.mem_info[1] / (1024 * 1024)
+            result[f'GPU{i}'] = "{0:.0f} MB".format(size)
 
     if virtual_memory:
         mem = virtual_memory()
@@ -124,7 +126,7 @@ def memory_sizes_of_current_machine():
 
     if not result:
         warnings.warn("Couldn't query for any local memory size."
-                      "Install py-cpuinfo to get cache sizes, psutil for RAM size and pycuda for GPU memory size.")
+                      "Install py-cpuinfo to get cache sizes, psutil for RAM size and cupy for GPU memory size.")
 
     return result
 

lbmpy/maxwellian_equilibrium.py → src/lbmpy/maxwellian_equilibrium.py

 """
-This module contains the continuous Maxwell-Boltzmann equilibrium and its discrete polynomial approximation, often
-used to formulate lattice-Boltzmann methods for hydrodynamics.
+This module contains a functional interface to the continuous Maxwell-Boltzmann equilibrium and its discrete 
+polynomial approximation, often used to formulate lattice-Boltzmann methods for hydrodynamics.
 Additionally functions are provided to compute moments and cumulants of these distributions.
+
+The functionality of this module has mostly been replaced by the :mod:`lbmpy.equilibrium` module.
+In particular, the continuous and discrete Maxwellians are now represented by 
+:class:`lbmpy.equilibrium.ContinuousHydrodynamicMaxwellian` and
+:class:`lbmpy.equilibrium.DiscreteHydrodynamicMaxwellian`.
 """
 
 import warnings
 
 import sympy as sp
-from sympy import Rational as R
+from sympy.core.numbers import Zero
 
 from pystencils.cache import disk_cache
+from pystencils.sympyextensions import remove_higher_order_terms
 
+from lbmpy.moments import MOMENT_SYMBOLS
+from lbmpy.continuous_distribution_measures import continuous_moment, continuous_central_moment, continuous_cumulant
 
-def get_weights(stencil, c_s_sq=sp.Rational(1, 3)):
-    q = len(stencil)
 
+def get_weights(stencil, c_s_sq=sp.Rational(1, 3)):
     if c_s_sq != sp.Rational(1, 3) and c_s_sq != sp.Symbol("c_s") ** 2:
         warnings.warn("Weights of discrete equilibrium are only valid if c_s^2 = 1/3")
 
     def weight_for_direction(direction):
         abs_sum = sum([abs(d) for d in direction])
-        return get_weights.weights[q][abs_sum]
+        return get_weights.weights[stencil.Q][abs_sum]
     return [weight_for_direction(d) for d in stencil]
 
 
 get_weights.weights = {
     9: {
-        0: R(4, 9),
-        1: R(1, 9),
-        2: R(1, 36),
+        0: sp.Rational(4, 9),
+        1: sp.Rational(1, 9),
+        2: sp.Rational(1, 36),
+    },
+    7: {
+        0: Zero(),
+        1: sp.Rational(1, 6),
     },
     15: {
-        0: R(2, 9),
-        1: R(1, 9),
-        3: R(1, 72),
+        0: sp.Rational(2, 9),
+        1: sp.Rational(1, 9),
+        3: sp.Rational(1, 72),
     },
     19: {
-        0: R(1, 3),
-        1: R(1, 18),
-        2: R(1, 36),
+        0: sp.Rational(1, 3),
+        1: sp.Rational(1, 18),
+        2: sp.Rational(1, 36),
     },
     27: {
-        0: R(8, 27),
-        1: R(2, 27),
-        2: R(1, 54),
-        3: R(1, 216),
+        0: sp.Rational(8, 27),
+        1: sp.Rational(2, 27),
+        2: sp.Rational(1, 54),
+        3: sp.Rational(1, 216),
     }
 }
 
 
 @disk_cache
-def discrete_maxwellian_equilibrium(stencil, rho=sp.Symbol("rho"), u=tuple(sp.symbols("u_0 u_1 u_2")), order=2,
+def discrete_maxwellian_equilibrium(stencil, rho=sp.Symbol("rho"), u=sp.symbols("u_:3"), order=2,
                                     c_s_sq=sp.Symbol("c_s") ** 2, compressible=True):
     """
     Returns the common discrete LBM equilibrium as a list of sympy expressions
@@ -64,65 +75,74 @@ def discrete_maxwellian_equilibrium(stencil, rho=sp.Symbol("rho"), u=tuple(sp.sy
         compressible: compressibility
     """
     weights = get_weights(stencil, c_s_sq)
-    assert len(stencil) == len(weights)
+    assert stencil.Q == len(weights)
+    u = u[:stencil.D]
 
-    dim = len(stencil[0])
-    u = u[:dim]
+    res = [discrete_equilibrium(e_q, u, rho, w_q, order, c_s_sq, compressible) for w_q, e_q in zip(weights, stencil)]
+    return tuple(res)
+
+
+def discrete_equilibrium(v=sp.symbols("v_:3"), u=sp.symbols("u_:3"), rho=sp.Symbol("rho"), weight=sp.Symbol("w"),
+                         order=2, c_s_sq=sp.Symbol("c_s") ** 2, compressible=True):
+    """
+    Returns the common discrete LBM equilibrium depending on the mesoscopic velocity and the directional lattice weight
 
+    Args:
+        v: symbols for mesoscopic velocity
+        u: symbols for macroscopic velocity
+        rho: sympy symbol for the density
+        weight: symbol for stencil weights
+        order: highest order of velocity terms (for hydrodynamics order 2 is sufficient)
+        c_s_sq: square of speed of sound
+        compressible: compressibility
+    """
     rho_outside = rho if compressible else sp.Rational(1, 1)
     rho_inside = rho if not compressible else sp.Rational(1, 1)
 
-    res = []
-    for w_q, e_q in zip(weights, stencil):
-        e_times_u = 0
-        for c_q_alpha, u_alpha in zip(e_q, u):
-            e_times_u += c_q_alpha * u_alpha
+    e_times_u = 0
+    for c_q_alpha, u_alpha in zip(v, u):
+        e_times_u += c_q_alpha * u_alpha
 
-        fq = rho_inside + e_times_u / c_s_sq
+    fq = rho_inside + e_times_u / c_s_sq
 
-        if order <= 1:
-            res.append(fq * rho_outside * w_q)
-            continue
+    if order <= 1:
+        return fq * rho_outside * weight
 
-        u_times_u = 0
-        for u_alpha in u:
-            u_times_u += u_alpha * u_alpha
-        fq += sp.Rational(1, 2) / c_s_sq**2 * e_times_u ** 2 - sp.Rational(1, 2) / c_s_sq * u_times_u
+    u_times_u = 0
+    for u_alpha in u:
+        u_times_u += u_alpha * u_alpha
+    fq += sp.Rational(1, 2) / c_s_sq**2 * e_times_u ** 2 - sp.Rational(1, 2) / c_s_sq * u_times_u
 
-        if order <= 2:
-            res.append(fq * rho_outside * w_q)
-            continue
+    if order <= 2:
+        return fq * rho_outside * weight
 
-        fq += sp.Rational(1, 6) / c_s_sq**3 * e_times_u**3 - sp.Rational(1, 2) / c_s_sq**2 * u_times_u * e_times_u
+    fq += sp.Rational(1, 6) / c_s_sq**3 * e_times_u**3 - sp.Rational(1, 2) / c_s_sq**2 * u_times_u * e_times_u
 
-        res.append(sp.expand(fq * rho_outside * w_q))
-
-    return tuple(res)
+    return sp.expand(fq * rho_outside * weight)
 
 
 @disk_cache
-def generate_equilibrium_by_matching_moments(stencil, moments, rho=sp.Symbol("rho"), u=tuple(sp.symbols("u_0 u_1 u_2")),
+def generate_equilibrium_by_matching_moments(stencil, moments, rho=sp.Symbol("rho"), u=sp.symbols("u_:3"),
                                              c_s_sq=sp.Symbol("c_s") ** 2, order=None):
     """
     Computes discrete equilibrium, by setting the discrete moments to values taken from the continuous Maxwellian.
     The number of moments has to match the number of directions in the stencil. For documentation of other parameters
-    see :func:`get_moments_of_continuous_maxwellian_equilibrium`
+    see :func:`get_equilibrium_values_of_maxwell_boltzmann_function`
     """
     from lbmpy.moments import moment_matrix
-    dim = len(stencil[0])
-    Q = len(stencil)
-    assert len(moments) == Q, "Moment count(%d) does not match stencil size(%d)" % (len(moments), Q)
-    continuous_moments_vector = get_moments_of_continuous_maxwellian_equilibrium(moments, dim, rho, u, c_s_sq, order)
+    assert len(moments) == stencil.Q, f"Moment count({len(moments)}) does not match stencil size({stencil.Q})"
+    continuous_moments_vector = get_equilibrium_values_of_maxwell_boltzmann_function(moments, stencil.D, rho, u, c_s_sq,
+                                                                                     order, space="moment")
     continuous_moments_vector = sp.Matrix(continuous_moments_vector)
     M = moment_matrix(moments, stencil)
-    assert M.rank() == Q, "Rank of moment matrix (%d) does not match stencil size (%d)" % (M.rank(), Q)
+    assert M.rank() == stencil.Q, f"Rank of moment matrix ({M.rank()}) does not match stencil size ({stencil.Q})"
     return M.inv() * continuous_moments_vector
 
 
 @disk_cache
 def continuous_maxwellian_equilibrium(dim=3, rho=sp.Symbol("rho"),
-                                      u=tuple(sp.symbols("u_0 u_1 u_2")),
-                                      v=tuple(sp.symbols("v_0 v_1 v_2")),
+                                      u=sp.symbols("u_:3"),
+                                      v=sp.symbols("v_:3"),
                                       c_s_sq=sp.Symbol("c_s") ** 2):
     """
     Returns sympy expression of Maxwell Boltzmann distribution
@@ -141,15 +161,14 @@ def continuous_maxwellian_equilibrium(dim=3, rho=sp.Symbol("rho"),
     return rho / (2 * sp.pi * c_s_sq) ** (sp.Rational(dim, 2)) * sp.exp(- vel_term / (2 * c_s_sq))
 
 
-# -------------------------------- Equilibrium moments/cumulants  ------------------------------------------------------
-
-
+# -------------------------------- Equilibrium moments  ----------------------------------------------------------------
 @disk_cache
-def get_moments_of_continuous_maxwellian_equilibrium(moments, dim, rho=sp.Symbol("rho"),
-                                                     u=tuple(sp.symbols("u_0 u_1 u_2")),
-                                                     c_s_sq=sp.Symbol("c_s") ** 2, order=None):
+def get_equilibrium_values_of_maxwell_boltzmann_function(moments, dim, rho=sp.Symbol("rho"),
+                                                         u=sp.symbols("u_:3"),
+                                                         c_s_sq=sp.Symbol("c_s") ** 2, order=None,
+                                                         space="moment"):
     """
-    Computes moments of the continuous Maxwell Boltzmann equilibrium distribution
+    Computes equilibrium values from the continuous Maxwell Boltzmann equilibrium.
 
     Args:
         moments: moments to compute, either in polynomial or exponent-tuple form
@@ -159,19 +178,27 @@ def get_moments_of_continuous_maxwellian_equilibrium(moments, dim, rho=sp.Symbol
         c_s_sq: symbol for speed of sound squared, defaults to symbol c_s**2
         order: if this parameter is not None, terms that have a higher polynomial order in the macroscopic velocity
                are removed
+        space: function space of the equilibrium values. Either moment, central moment or cumulant space are supported.
 
-    >>> get_moments_of_continuous_maxwellian_equilibrium( ( (0,0,0), (1,0,0), (0,1,0), (0,0,1), (2,0,0) ), dim=3 )
+    >>> get_equilibrium_values_of_maxwell_boltzmann_function( ( (0,0,0), (1,0,0), (0,1,0), (0,0,1), (2,0,0) ), dim=3 )
     [rho, rho*u_0, rho*u_1, rho*u_2, rho*(c_s**2 + u_0**2)]
     """
-    from pystencils.sympyextensions import remove_higher_order_terms
-    from lbmpy.moments import MOMENT_SYMBOLS
-    from lbmpy.continuous_distribution_measures import continuous_moment
-
     # trick to speed up sympy integration (otherwise it takes multiple minutes, or aborts):
     # use a positive, real symbol to represent c_s_sq -> then replace this symbol afterwards with the real c_s_sq
     c_s_sq_helper = sp.Symbol("csq_helper", positive=True, real=True)
     mb = continuous_maxwellian_equilibrium(dim, rho, u, MOMENT_SYMBOLS[:dim], c_s_sq_helper)
-    result = [continuous_moment(mb, moment, MOMENT_SYMBOLS[:dim]).subs(c_s_sq_helper, c_s_sq) for moment in moments]
+    if space == "moment":
+        result = [continuous_moment(mb, moment, MOMENT_SYMBOLS[:dim]).subs(c_s_sq_helper, c_s_sq)
+                  for moment in moments]
+    elif space == "central moment":
+        result = [continuous_central_moment(mb, moment, MOMENT_SYMBOLS[:dim], velocity=u).subs(c_s_sq_helper, c_s_sq)
+                  for moment in moments]
+    elif space == "cumulant":
+        result = [continuous_cumulant(mb, moment, MOMENT_SYMBOLS[:dim]).subs(c_s_sq_helper, c_s_sq)
+                  for moment in moments]
+    else:
+        raise ValueError("Only moment, central moment or cumulant space are supported")
+
     if order is not None:
         result = [remove_higher_order_terms(r, order=order, symbols=u) for r in result]
 
@@ -180,7 +207,7 @@ def get_moments_of_continuous_maxwellian_equilibrium(moments, dim, rho=sp.Symbol
 
 @disk_cache
 def get_moments_of_discrete_maxwellian_equilibrium(stencil, moments,
-                                                   rho=sp.Symbol("rho"), u=tuple(sp.symbols("u_0 u_1 u_2")),
+                                                   rho=sp.Symbol("rho"), u=sp.symbols("u_:3"),
                                                    c_s_sq=sp.Symbol("c_s") ** 2, order=None, compressible=True):
     """Compute moments of discrete maxwellian equilibrium.
 
@@ -201,7 +228,7 @@ def get_moments_of_discrete_maxwellian_equilibrium(stencil, moments,
 
 
 def compressible_to_incompressible_moment_value(term, rho, u):
-    """Compressible to incompressible equilibrium moments
+    """**[DEPRECATED]** Compressible to incompressible equilibrium moments
 
     Transforms so-called compressible equilibrium moments (as obtained from the continuous Maxwellian) by
     removing the density factor in all monomials where velocity components are multiplied to the density.
@@ -219,6 +246,8 @@ def compressible_to_incompressible_moment_value(term, rho, u):
     Returns:
         incompressible equilibrium value
     """
+    warnings.warn("Usage of `compressible_to_incompressible_moment_value` is deprecated,"
+                  " and the method will be removed soon.")
     term = sp.sympify(term)
     term = term.expand()
     if term.func != sp.Add:
@@ -235,32 +264,12 @@ def compressible_to_incompressible_moment_value(term, rho, u):
             res += t
     return res
 
-
-# -------------------------------- Equilibrium moments -----------------------------------------------------------------
-
-
-def get_cumulants_of_continuous_maxwellian_equilibrium(cumulants, dim, rho=sp.Symbol("rho"),
-                                                       u=tuple(sp.symbols("u_0 u_1 u_2")), c_s_sq=sp.Symbol("c_s") ** 2,
-                                                       order=None):
-    from lbmpy.moments import MOMENT_SYMBOLS
-    from lbmpy.continuous_distribution_measures import continuous_cumulant
-    from pystencils.sympyextensions import remove_higher_order_terms
-
-    # trick to speed up sympy integration (otherwise it takes multiple minutes, or aborts):
-    # use a positive, real symbol to represent c_s_sq -> then replace this symbol afterwards with the real c_s_sq
-    c_s_sq_helper = sp.Symbol("csq_helper", positive=True, real=True)
-    mb = continuous_maxwellian_equilibrium(dim, rho, u, MOMENT_SYMBOLS[:dim], c_s_sq_helper)
-    result = [continuous_cumulant(mb, cumulant, MOMENT_SYMBOLS[:dim]).subs(c_s_sq_helper, c_s_sq)
-              for cumulant in cumulants]
-    if order is not None:
-        result = [remove_higher_order_terms(r, order=order, symbols=u) for r in result]
-
-    return result
+# -------------------------------- Equilibrium cumulants ---------------------------------------------------------------
 
 
 @disk_cache
 def get_cumulants_of_discrete_maxwellian_equilibrium(stencil, cumulants,
-                                                     rho=sp.Symbol("rho"), u=tuple(sp.symbols("u_0 u_1 u_2")),
+                                                     rho=sp.Symbol("rho"), u=sp.symbols("u_:3"),
                                                      c_s_sq=sp.Symbol("c_s") ** 2, order=None, compressible=True):
     from lbmpy.cumulants import discrete_cumulant
     if order is None:

src/lbmpy/methods/__init__.py

+from .creationfunctions import (
+    CollisionSpaceInfo,
+    create_mrt_orthogonal, create_mrt_raw, create_central_moment, create_srt, create_trt, create_trt_kbc,
+    create_trt_with_magic_number, create_with_continuous_maxwellian_equilibrium,
+    create_with_discrete_maxwellian_equilibrium, create_from_equilibrium,
+    create_cumulant, create_with_default_polynomial_cumulants, create_with_monomial_cumulants)
+
+from .default_moment_sets import mrt_orthogonal_modes_literature, cascaded_moment_sets_literature
+
+from .abstractlbmethod import LbmCollisionRule, AbstractLbMethod, RelaxationInfo
+from .conservedquantitycomputation import AbstractConservedQuantityComputation, DensityVelocityComputation
+
+
+__all__ = ['CollisionSpaceInfo', 'RelaxationInfo', 
+           'AbstractLbMethod', 'LbmCollisionRule',
+           'AbstractConservedQuantityComputation', 'DensityVelocityComputation',
+           'create_srt', 'create_trt', 'create_trt_with_magic_number', 'create_trt_kbc',
+           'create_mrt_orthogonal', 'create_mrt_raw', 'create_central_moment',
+           'create_with_continuous_maxwellian_equilibrium', 'create_with_discrete_maxwellian_equilibrium',
+           'create_from_equilibrium',
+           'mrt_orthogonal_modes_literature', 'cascaded_moment_sets_literature',
+           'create_cumulant', 'create_with_default_polynomial_cumulants',
+           'create_with_monomial_cumulants']

lbmpy/methods/abstractlbmethod.py → src/lbmpy/methods/abstractlbmethod.py

@@ -2,13 +2,18 @@ import abc
 from collections import namedtuple
 
 import sympy as sp
+from sympy.core.numbers import Zero
 
-from pystencils import AssignmentCollection
+from pystencils import Assignment, AssignmentCollection
+from lbmpy.stencils import LBStencil
 
 RelaxationInfo = namedtuple('RelaxationInfo', ['equilibrium_value', 'relaxation_rate'])
 
 
 class LbmCollisionRule(AssignmentCollection):
+    """
+    A pystencils AssignmentCollection that additionally holds an `AbstractLbMethod`
+    """
     def __init__(self, lb_method, *args, **kwargs):
         super(LbmCollisionRule, self).__init__(*args, **kwargs)
         self.method = lb_method
@@ -17,34 +22,66 @@ class LbmCollisionRule(AssignmentCollection):
 class AbstractLbMethod(abc.ABC):
     """Abstract base class for all LBM methods."""
 
-    def __init__(self, stencil):
+    def __init__(self, stencil: LBStencil):
         self._stencil = stencil
 
     @property
     def stencil(self):
-        """Discrete set of velocities, represented as nested tuple"""
+        """Discrete set of velocities, represented by :class:`lbmpy.stencils.LBStencil`"""
         return self._stencil
 
     @property
     def dim(self):
-        return len(self.stencil[0])
+        """The method's spatial dimensionality"""
+        return self._stencil.D
 
     @property
     def pre_collision_pdf_symbols(self):
         """Tuple of symbols representing the pdf values before collision"""
-        return sp.symbols("f_:%d" % (len(self.stencil),))
+        return sp.symbols(f"f_:{self._stencil.Q}")
 
     @property
     def post_collision_pdf_symbols(self):
         """Tuple of symbols representing the pdf values after collision"""
-        return sp.symbols("d_:%d" % (len(self.stencil),))
+        return sp.symbols(f"d_:{self._stencil.Q}")
+
+    @property
+    @abc.abstractmethod
+    def relaxation_rates(self):
+        """Tuple containing the relaxation rates of each moment"""
+
+    @property
+    def relaxation_matrix(self):
+        """Returns a qxq diagonal matrix which contains the relaxation rate for each moment on the diagonal"""
+        d = sp.zeros(len(self.relaxation_rates))
+        for i, w in enumerate(self.relaxation_rates):
+            # note that 0.0 is converted to sp.Zero here. It is not possible to prevent this.
+            d[i, i] = w
+        return d
+
+    @property
+    def symbolic_relaxation_matrix(self):
+        """Returns a qxq diagonal matrix which contains the relaxation rate for each moment on the diagonal.
+           In contrast to the normal relaxation matrix all numeric values are replaced by sympy symbols"""
+        _, d = self._generate_symbolic_relaxation_matrix()
+        return d
+
+    @property
+    def subs_dict_relaxation_rate(self):
+        """returns a dictionary which maps the replaced numerical relaxation rates to its original numerical value"""
+        result = dict()
+        for i in range(self._stencil.Q):
+            result[self.symbolic_relaxation_matrix[i, i]] = self.relaxation_matrix[i, i]
+        return result
 
     # ------------------------- Abstract Methods & Properties ----------------------------------------------------------
 
+    @property
     @abc.abstractmethod
     def conserved_quantity_computation(self):
         """Returns an instance of class :class:`lbmpy.methods.AbstractConservedQuantityComputation`"""
 
+    @property
     @abc.abstractmethod
     def weights(self):
         """Returns a sequence of weights, one for each lattice direction"""
@@ -59,3 +96,36 @@ class AbstractLbMethod(abc.ABC):
     def get_collision_rule(self):
         """Returns an LbmCollisionRule i.e. an equation collection with a reference to the method.
          This collision rule defines the collision operator."""
+
+    # -------------------------------- Helper Functions ----------------------------------------------------------------
+
+    def _generate_symbolic_relaxation_matrix(self, relaxation_rates=None, relaxation_rates_modifier=None):
+        """
+        This function replaces the numbers in the relaxation matrix with symbols in this case, and returns also
+        the subexpressions, that assign the number to the newly introduced symbol
+        """
+        rr = relaxation_rates if relaxation_rates is not None else self.relaxation_rates
+        subexpressions = {}
+        symbolic_relaxation_rates = list()
+        for relaxation_rate in rr:
+            relaxation_rate = sp.sympify(relaxation_rate)
+            if isinstance(relaxation_rate, sp.Symbol):
+                symbolic_relaxation_rate = relaxation_rate
+            else:
+                if isinstance(relaxation_rate, Zero):
+                    relaxation_rate = 0.0
+                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)
+
+        substitutions = [Assignment(e[1], e[0]) for e in subexpressions.items()]
+        if relaxation_rates_modifier is not None:
+            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(*symbolic_relaxation_rates)

lbmpy/methods/conservedquantitycomputation.py → src/lbmpy/methods/conservedquantitycomputation.py

 import abc
-from collections import OrderedDict
 
+from warnings import warn
 import sympy as sp
 
 from pystencils import Assignment, AssignmentCollection, Field
@@ -48,7 +48,7 @@ class AbstractConservedQuantityComputation(abc.ABC):
         """
 
     @abc.abstractmethod
-    def equilibrium_input_equations_from_pdfs(self, pdfs):
+    def equilibrium_input_equations_from_pdfs(self, pdfs, **kwargs):
         """
         Returns an equation collection that defines all necessary quantities to compute the equilibrium as functions
         of the pdfs.
@@ -59,7 +59,7 @@ class AbstractConservedQuantityComputation(abc.ABC):
         """
 
     @abc.abstractmethod
-    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols):
+    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols, **kwargs):
         """
         Returns an equation collection that defines conserved quantities for output. These conserved quantities might
         be slightly different that the ones used as input for the equilibrium e.g. due to a force model.
@@ -82,119 +82,259 @@ class AbstractConservedQuantityComputation(abc.ABC):
 
 
 class DensityVelocityComputation(AbstractConservedQuantityComputation):
-    def __init__(self, stencil, compressible, force_model=None,
-                 zeroth_order_moment_symbol=sp.Symbol("rho"),
-                 first_order_moment_symbols=sp.symbols("u_:3")):
-        dim = len(stencil[0])
+    r"""
+    This class emits equations for in- and output of the conserved quantities of regular
+    hydrodynamic lattice Boltzmann methods, which are density and velocity. The following symbolic
+    quantities are manages by this class:
+
+    - Density :math:`\rho`, background density :math:`\rho_0` (typically set to :math:`1`) and the
+      density deviation :math:`\delta \rho`. They are connected through :math:`\rho = \rho_0 + \delta\rho`.
+    - Velocity :math:`\mathbf{u} = (u_0, u_1, u_2)`.
+
+    Furthermore, this class provides output functionality for the second-order moment tensor :math:`p`.
+
+    Parameters:
+        stencil: see :class:`lbmpy.stencils.LBStencil`
+        compressible: `True` indicates the usage of a compressible equilibrium 
+                        (see :class:`lbmpy.equilibrium.ContinuousHydrodynamicMaxwellian`),
+                        and sets the reference density to :math:`\rho`.
+                        `False` indicates an incompressible equilibrium, using :math:`\rho` as
+                        reference density.
+        zero_centered: Indicates whether or not PDFs are stored in regular or zero-centered format.
+        force_model: Employed force model. See :mod:`lbmpy.forcemodels`.
+    """
+    
+    def __init__(self, stencil, compressible, zero_centered, force_model=None,
+                 background_density=sp.Integer(1),
+                 density_symbol=sp.Symbol("rho"),
+                 density_deviation_symbol=sp.Symbol("delta_rho"),
+                 velocity_symbols=sp.symbols("u_:3"),
+                 c_s_sq=sp.Symbol("c_s")**2,
+                 second_order_moment_symbols=sp.symbols("p_:9"),
+                 zeroth_order_moment_symbol=None,
+                 first_order_moment_symbols=None):
+        if zeroth_order_moment_symbol is not None:
+            warn("Usage of parameter 'zeroth_order_moment_symbol' is deprecated."
+                 " Use 'density_symbol' or 'density_deviation_symbol' instead.")
+            density_symbol = zeroth_order_moment_symbol
+
+        if first_order_moment_symbols is not None:
+            warn("Usage of parameter 'first_order_moment_symbols' is deprecated."
+                 " Use 'velocity_symbols' instead.")
+            velocity_symbols = first_order_moment_symbols
+
+        dim = stencil.D
         self._stencil = stencil
         self._compressible = compressible
-        self._forceModel = force_model
-        self._symbolOrder0 = zeroth_order_moment_symbol
-        self._symbolsOrder1 = first_order_moment_symbols[:dim]
+        self._zero_centered = zero_centered
+        self._force_model = force_model
+        self._background_density = background_density
+        self._density_symbol = density_symbol
+        self._density_deviation_symbol = density_deviation_symbol
+        self._velocity_symbols = velocity_symbols[:dim]
+        self._second_order_moment_symbols = second_order_moment_symbols[:(dim * dim)]
+        self._c_s_sq = c_s_sq
 
     @property
     def conserved_quantities(self):
         return {'density': 1,
-                'velocity': len(self._stencil[0])}
+                'density_deviation': 1,
+                'velocity': self._stencil.D}
 
     @property
     def compressible(self):
+        """Indicates whether a compressible or incompressible equilibrium is used."""
         return self._compressible
 
     def defined_symbols(self, order='all'):
         if order == 'all':
-            return {'density': self._symbolOrder0,
-                    'velocity': self._symbolsOrder1}
+            return {'density': self._density_symbol,
+                    'density_deviation': self._density_deviation_symbol,
+                    'velocity': self._velocity_symbols}
         elif order == 0:
-            return 'density', self._symbolOrder0
+            return {'density': self._density_symbol,
+                    'density_deviation': self._density_deviation_symbol, }
         elif order == 1:
-            return 'velocity', self._symbolsOrder1
+            return {'velocity': self._velocity_symbols}
         else:
-            return None
+            return dict()
 
     @property
     def zero_centered_pdfs(self):
-        return not self._compressible
+        """Whether regular or zero-centered storage is employed."""
+        return self._zero_centered
 
     @property
     def zeroth_order_moment_symbol(self):
-        return self._symbolOrder0
+        """Symbol corresponding to the zeroth-order moment of the stored PDFs, 
+        i.e. `density_deviation_symbol` if zero-centered, else `density_symbol`."""
+        warn("Usage of 'zeroth_order_moment_symbol' is deprecated."
+             "Use 'density_symbol' or 'density_deviation_symbol' instead.")
+        # to be consistent with previous behaviour, this method returns `density_symbol` for non-zero-centered methods,
+        # and `density_deviation_symbol` for zero-centered methods
+        return self._density_deviation_symbol if self.zero_centered_pdfs else self._density_symbol
 
     @property
     def first_order_moment_symbols(self):
-        return self._symbolsOrder1
+        """Symbol corresponding to the first-order moment vector of the stored PDFs, 
+        divided by the reference density."""
+        warn("Usage of 'first_order_moment_symbols' is deprecated. Use 'velocity_symbols' instead.")
+        return self._velocity_symbols
+
+    @property
+    def density_symbol(self):
+        """Symbol for the density."""
+        return self._density_symbol
+
+    @property
+    def density_deviation_symbol(self):
+        """Symbol for the density deviation."""
+        return self._density_deviation_symbol
+
+    @property
+    def background_density(self):
+        """Symbol or value of the background density."""
+        return self._background_density
+
+    @property
+    def velocity_symbols(self):
+        """Symbols for the velocity."""
+        return self._velocity_symbols
+
+    @property
+    def force_model(self):
+        return self._force_model
+
+    def set_force_model(self, force_model):
+        self._force_model = force_model
 
     @property
     def default_values(self):
-        result = {self._symbolOrder0: 1}
-        for s in self._symbolsOrder1:
+        result = {self._density_symbol: 1, self._density_deviation_symbol: 0}
+        for s in self._velocity_symbols:
             result[s] = 0
         return result
 
-    def equilibrium_input_equations_from_pdfs(self, pdfs):
-        dim = len(self._stencil[0])
-        eq_coll = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs, self._symbolOrder0,
-                                                                  self._symbolsOrder1[:dim])
-        if self._compressible:
-            eq_coll = divide_first_order_moments_by_rho(eq_coll, dim)
-
-        eq_coll = apply_force_model_shift('equilibrium_velocity_shift', dim, eq_coll,
-                                          self._forceModel, self._compressible)
-        return eq_coll
-
-    def equilibrium_input_equations_from_init_values(self, density=1, velocity=(0, 0, 0)):
-        dim = len(self._stencil[0])
-        zeroth_order_moment = density
-        first_order_moments = velocity[:dim]
-        vel_offset = [0] * dim
+    def equilibrium_input_equations_from_pdfs(self, pdfs, force_substitution=True):
+        #   Compute moments from stored PDFs.
+        #   If storage is zero_centered, this yields the density_deviation, else density as zeroth moment.
+        zeroth_moment_symbol = self._density_deviation_symbol if self.zero_centered_pdfs else self._density_symbol
+        reference_density = self._density_symbol if self.compressible else self._background_density
+        ac = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs, zeroth_moment_symbol,
+                                                             self._velocity_symbols)
+
+        main_assignments_dict = ac.main_assignments_dict
 
-        if self._compressible:
-            if self._forceModel and hasattr(self._forceModel, 'macroscopic_velocity_shift'):
-                vel_offset = self._forceModel.macroscopic_velocity_shift(zeroth_order_moment)
+        #   If zero_centered, rho must still be computed, else delta_rho
+        if self.zero_centered_pdfs:
+            main_assignments_dict[self._density_symbol] = self._background_density + self._density_deviation_symbol
         else:
-            if self._forceModel and hasattr(self._forceModel, 'macroscopic_velocity_shift'):
-                vel_offset = self._forceModel.macroscopic_velocity_shift(sp.Rational(1, 1))
-            zeroth_order_moment -= sp.Rational(1, 1)
-        eqs = [Assignment(self._symbolOrder0, zeroth_order_moment)]
+            main_assignments_dict[self._density_deviation_symbol] = self._density_symbol - self._background_density
 
-        first_order_moments = [a - b for a, b in zip(first_order_moments, vel_offset)]
-        eqs += [Assignment(l, r) for l, r in zip(self._symbolsOrder1, first_order_moments)]
+        #   If compressible, velocities must be divided by rho
+        for u in self._velocity_symbols:
+            main_assignments_dict[u] /= reference_density
 
-        return AssignmentCollection(eqs, [])
+        #   If force model is present, apply force model shift
+        if self._force_model is not None:
+            velocity_shifts = self._force_model.equilibrium_velocity_shift(reference_density)
+            for u, s in zip(self._velocity_symbols, velocity_shifts):
+                main_assignments_dict[u] += s
 
-    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols):
-        dim = len(self._stencil[0])
+            if force_substitution:
+                ac = add_symbolic_force_substitutions(ac, self._force_model._subs_dict_force)
 
-        ac = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
-                                                             self._symbolOrder0, self._symbolsOrder1)
+        ac.set_main_assignments_from_dict(main_assignments_dict)
+        ac.topological_sort()
+        return ac
+
+    def equilibrium_input_equations_from_init_values(self, density=1, velocity=(0, 0, 0), force_substitution=True):
+        dim = self._stencil.D
+        density_deviation = self._density_symbol - self._background_density
+        velocity = velocity[:dim]
+        vel_offset = [0] * dim
+        reference_density = self._density_symbol if self.compressible else self._background_density
+
+        if self._force_model is not None:
+            vel_offset = self._force_model.macroscopic_velocity_shift(reference_density)
+
+        eqs = [Assignment(self._density_symbol, density),
+               Assignment(self._density_deviation_symbol, density_deviation)]
+
+        velocity_eqs = [u - o for u, o in zip(velocity, vel_offset)]
+        eqs += [Assignment(l, r) for l, r in zip(self._velocity_symbols, velocity_eqs)]
+
+        result = AssignmentCollection(eqs, [])
+
+        if self._force_model is not None and force_substitution:
+            result = add_symbolic_force_substitutions(result, self._force_model._subs_dict_force)
+
+        return result
 
-        if self._compressible:
-            ac = divide_first_order_moments_by_rho(ac, dim)
+    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols, force_substitution=True):
+        dim = self._stencil.D
+
+        zeroth_moment_symbol = self._density_deviation_symbol if self.zero_centered_pdfs else self._density_symbol
+        first_moment_symbols = sp.symbols(f'momdensity_:{dim}')
+        reference_density = self._density_symbol if self.compressible else self._background_density
+        ac = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
+                                                             zeroth_moment_symbol, first_moment_symbols,
+                                                             self._second_order_moment_symbols)
+
+        eqs = {**ac.main_assignments_dict, **ac.subexpressions_dict}
+
+        #   If zero_centered, rho must still be computed, else delta_rho
+        if self.zero_centered_pdfs:
+            eqs[self._density_symbol] = self._background_density + self._density_deviation_symbol
+            dim = self._stencil.D
+            for j in range(dim):
+                #   Diagonal entries are shifted by c_s^2
+                eqs[self._second_order_moment_symbols[dim * j + j]] += self._c_s_sq
         else:
-            ac = add_density_offset(ac)
+            eqs[self._density_deviation_symbol] = self._density_symbol - self._background_density
+
+        #   If compressible, velocities must be divided by rho
+        for m, u in zip(first_moment_symbols, self._velocity_symbols):
+            eqs[u] = m / reference_density
+
+        #   If force model is present, apply force model shift
+        mom_density_shifts = [0] * dim
+        if self._force_model is not None:
+            velocity_shifts = self._force_model.macroscopic_velocity_shift(reference_density)
+            for u, s in zip(self._velocity_symbols, velocity_shifts):
+                eqs[u] += s
 
-        ac = apply_force_model_shift('macroscopic_velocity_shift', dim, ac, self._forceModel, self._compressible)
+            mom_density_shifts = self._force_model.macroscopic_momentum_density_shift()
 
         main_assignments = []
-        eqs = OrderedDict([(eq.lhs, eq.rhs) for eq in ac.all_assignments])
 
         if 'density' in output_quantity_names_to_symbols:
             density_output_symbol = output_quantity_names_to_symbols['density']
             if isinstance(density_output_symbol, Field):
                 density_output_symbol = density_output_symbol.center
-            if density_output_symbol != self._symbolOrder0:
-                main_assignments.append(Assignment(density_output_symbol, self._symbolOrder0))
+            if density_output_symbol != self._density_symbol:
+                main_assignments.append(Assignment(density_output_symbol, self._density_symbol))
             else:
-                main_assignments.append(Assignment(self._symbolOrder0, eqs[self._symbolOrder0]))
-                del eqs[self._symbolOrder0]
+                main_assignments.append(Assignment(self._density_symbol, eqs[self._density_symbol]))
+                del eqs[self._density_symbol]
+        if 'density_deviation' in output_quantity_names_to_symbols:
+            density_deviation_output_symbol = output_quantity_names_to_symbols['density_deviation']
+            if isinstance(density_deviation_output_symbol, Field):
+                density_deviation_output_symbol = density_deviation_output_symbol.center
+            if density_deviation_output_symbol != self._density_deviation_symbol:
+                main_assignments.append(Assignment(density_deviation_output_symbol, self._density_deviation_symbol))
+            else:
+                main_assignments.append(Assignment(self._density_deviation_symbol,
+                                        eqs[self._density_deviation_symbol]))
+                del eqs[self._density_deviation_symbol]
         if 'velocity' in output_quantity_names_to_symbols:
             vel_output_symbols = output_quantity_names_to_symbols['velocity']
             if isinstance(vel_output_symbols, Field):
                 vel_output_symbols = vel_output_symbols.center_vector
-            if tuple(vel_output_symbols) != tuple(self._symbolsOrder1):
-                main_assignments += [Assignment(a, b) for a, b in zip(vel_output_symbols, self._symbolsOrder1)]
+            if tuple(vel_output_symbols) != tuple(self._velocity_symbols):
+                main_assignments += [Assignment(a, b) for a, b in zip(vel_output_symbols, self._velocity_symbols)]
             else:
-                for u_i in self._symbolsOrder1:
+                for u_i in self._velocity_symbols:
                     main_assignments.append(Assignment(u_i, eqs[u_i]))
                     del eqs[u_i]
         if 'momentum_density' in output_quantity_names_to_symbols:
@@ -204,13 +344,8 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
             # Is not optimal when velocity and momentum_density are calculated together,
             # but this is usually not the case
             momentum_density_output_symbols = output_quantity_names_to_symbols['momentum_density']
-            mom_density_eq_coll = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
-                                                                                  self._symbolOrder0,
-                                                                                  self._symbolsOrder1)
-            mom_density_eq_coll = apply_force_model_shift('macroscopic_momentum_density_shift', dim, 
-                                                          mom_density_eq_coll, self._forceModel, self._compressible)
-            for sym, val in zip(momentum_density_output_symbols, mom_density_eq_coll.main_assignments[1:]):
-                main_assignments.append(Assignment(sym, val.rhs))
+            for sym, m, s in zip(momentum_density_output_symbols, first_moment_symbols, mom_density_shifts):
+                main_assignments.append(Assignment(sym, m + s))
         if 'moment0' in output_quantity_names_to_symbols:
             moment0_output_symbol = output_quantity_names_to_symbols['moment0']
             if isinstance(moment0_output_symbol, Field):
@@ -222,19 +357,31 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
                 moment1_output_symbol = moment1_output_symbol.center_vector
             main_assignments.extend([Assignment(lhs, sum(d[i] * pdf for d, pdf in zip(self._stencil, pdfs)))
                                      for i, lhs in enumerate(moment1_output_symbol)])
+        if 'moment2' in output_quantity_names_to_symbols:
+            moment2_output_symbol = output_quantity_names_to_symbols['moment2']
+            if isinstance(moment2_output_symbol, Field):
+                moment2_output_symbol = moment2_output_symbol.center_vector
+            for i, p in enumerate(moment2_output_symbol):
+                main_assignments.append(Assignment(p, eqs[self._second_order_moment_symbols[i]]))
+                del eqs[self._second_order_moment_symbols[i]]
+
+        ac = AssignmentCollection(main_assignments, eqs)
+        if self._force_model is not None and force_substitution:
+            ac = add_symbolic_force_substitutions(ac, self._force_model._subs_dict_force)
+
+        ac.topological_sort()
 
-        ac = ac.copy(main_assignments, [Assignment(a, b) for a, b in eqs.items()])
         return ac.new_without_unused_subexpressions()
 
     def __repr__(self):
-        return "ConservedValueComputation for %s" % (", " .join(self.conserved_quantities.keys()),)
+        return "ConservedValueComputation for %s" % (", ".join(self.conserved_quantities.keys()),)
 
 
 # -----------------------------------------  Helper functions ----------------------------------------------------------
 
 
-def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
-                                                    symbolic_zeroth_moment, symbolic_first_moments):
+def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs, symbolic_zeroth_moment,
+                                                    symbolic_first_moments, symbolic_second_moments=None):
     r"""
     Returns an equation system that computes the zeroth and first order moments with the least amount of operations
 
@@ -244,13 +391,16 @@ def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
 
         \rho = \sum_{d \in S} f_d
         u_j = \sum_{d \in S} f_d u_jd
+        p_j = \sum_{d \in S} {d \in S} f_d u_jd
 
     Args:
         stencil: called :math:`S` above
         symbolic_pdfs: called :math:`f` above
         symbolic_zeroth_moment:  called :math:`\rho` above
         symbolic_first_moments: called :math:`u` above
+        symbolic_second_moments: called :math:`p` above
     """
+
     def filter_out_plus_terms(expr):
         result = 0
         for term in expr.args:
@@ -258,32 +408,46 @@ def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
                 result += term
         return result
 
-    dim = len(stencil[0])
+    dim = stencil.D
 
-    subexpressions = []
+    subexpressions = list()
     pdf_sum = sum(symbolic_pdfs)
     u = [0] * dim
     for f, offset in zip(symbolic_pdfs, stencil):
         for i in range(dim):
             u[i] += f * int(offset[i])
 
+    p = [0] * dim * dim
+    for f, offset in zip(symbolic_pdfs, stencil):
+        for i in range(dim):
+            for j in range(dim):
+                p[dim * i + j] += f * int(offset[i]) * int(offset[j])
+
     plus_terms = [set(filter_out_plus_terms(u_i).args) for u_i in u]
+
+    velo_terms = sp.symbols(f"vel:{dim}Term")
     for i in range(dim):
         rhs = plus_terms[i]
         for j in range(i):
             rhs -= plus_terms[j]
-        eq = Assignment(sp.Symbol("vel%dTerm" % (i,)), sum(rhs))
+        eq = Assignment(velo_terms[i], sum(rhs))
         subexpressions.append(eq)
+        if len(rhs) == 0:  # if one of the substitutions is not found the simplification can not be applied
+            subexpressions = []
+            break
 
     for subexpression in subexpressions:
         pdf_sum = pdf_sum.subs(subexpression.rhs, subexpression.lhs)
 
-    for i in range(dim):
-        u[i] = u[i].subs(subexpressions[i].rhs, subexpressions[i].lhs)
+    if len(subexpressions) > 0:
+        for i in range(dim):
+            u[i] = u[i].subs(subexpressions[i].rhs, subexpressions[i].lhs)
 
     equations = []
     equations += [Assignment(symbolic_zeroth_moment, pdf_sum)]
     equations += [Assignment(u_i_sym, u_i) for u_i_sym, u_i in zip(symbolic_first_moments, u)]
+    if symbolic_second_moments:
+        equations += [Assignment(symbolic_second_moments[i], p[i]) for i in range(dim ** 2)]
 
     return AssignmentCollection(equations, subexpressions)
 
@@ -313,23 +477,45 @@ def add_density_offset(assignment_collection, offset=sp.Rational(1, 1)):
     return assignment_collection.copy([new_density] + old_eqs[1:])
 
 
-def apply_force_model_shift(shift_member_name, dim, assignment_collection, force_model, compressible, reverse=False):
+def apply_force_model_shift(shift_func, dim, assignment_collection, compressible, reverse=False):
     """
     Modifies the first order moment equations in assignment collection according to the force model shift.
     It is applied if force model has a method named shift_member_name. The equations 1: dim+1 of the passed
     equation collection are assumed to be the velocity equations.
+
+    Args:
+        shift_func: shift function which is applied. See lbmpy.forcemodels.AbstractForceModel for details
+        dim: number of spatial dimensions
+        assignment_collection: assignment collection containing the conserved quantity computation
+        compressible: True if a compressible LB method is used. Otherwise the Helmholtz decomposition was applied
+                      for rho
+        reverse: If True the sign of the shift is flipped
+    """
+    old_eqs = assignment_collection.main_assignments
+    density = old_eqs[0].lhs if compressible else sp.Rational(1, 1)
+    old_vel_eqs = old_eqs[1:dim + 1]
+    vel_offsets = shift_func(density)
+    if reverse:
+        vel_offsets = [-v for v in vel_offsets]
+    shifted_velocity_eqs = [Assignment(old_eq.lhs, old_eq.rhs + offset)
+                            for old_eq, offset in zip(old_vel_eqs, vel_offsets)]
+    new_eqs = [old_eqs[0]] + shifted_velocity_eqs + old_eqs[dim + 1:]
+    return assignment_collection.copy(new_eqs)
+
+
+def add_symbolic_force_substitutions(assignment_collection, subs_dict):
+    """
+    Every force model holds a symbolic representation of the forcing terms internally. This function adds the
+    equations for the D-dimensional force vector to the symbolic replacements
+
+    Args:
+        assignment_collection: assignment collection which will be modified
+        subs_dict: substitution dict which can be obtained from the force model
     """
-    if force_model is not None and hasattr(force_model, shift_member_name):
-        old_eqs = assignment_collection.main_assignments
-        density = old_eqs[0].lhs if compressible else sp.Rational(1, 1)
-        old_vel_eqs = old_eqs[1:dim + 1]
-        shift_func = getattr(force_model, shift_member_name)
-        vel_offsets = shift_func(density)
-        if reverse:
-            vel_offsets = [-v for v in vel_offsets]
-        shifted_velocity_eqs = [Assignment(old_eq.lhs, old_eq.rhs + offset)
-                                for old_eq, offset in zip(old_vel_eqs, vel_offsets)]
-        new_eqs = [old_eqs[0]] + shifted_velocity_eqs + old_eqs[dim + 1:]
-        return assignment_collection.copy(new_eqs)
-    else:
-        return assignment_collection
+    old_eqs = assignment_collection.subexpressions
+    subs_equations = []
+    for key, value in zip(subs_dict.keys(), subs_dict.values()):
+        subs_equations.append(Assignment(key, value))
+
+    new_eqs = subs_equations + old_eqs
+    return assignment_collection.copy(main_assignments=None, subexpressions=new_eqs)

lbmpy/methods/creationfunctions.py → src/lbmpy/methods/creationfunctions.py

+from warnings import warn
+from dataclasses import dataclass
+from typing import Type
+
 import itertools
 import operator
 from collections import OrderedDict
@@ -5,193 +9,252 @@ from functools import reduce
 
 import sympy as sp
 
-from lbmpy.maxwellian_equilibrium import (
-    compressible_to_incompressible_moment_value, get_cumulants_of_continuous_maxwellian_equilibrium,
-    get_moments_of_continuous_maxwellian_equilibrium,
-    get_moments_of_discrete_maxwellian_equilibrium, get_weights)
+from lbmpy.maxwellian_equilibrium import get_weights
+from lbmpy.equilibrium import ContinuousHydrodynamicMaxwellian, DiscreteHydrodynamicMaxwellian
 
-from lbmpy.methods.abstractlbmethod import RelaxationInfo
+from lbmpy.methods.default_moment_sets import cascaded_moment_sets_literature
 
-from lbmpy.methods.centeredcumulant import CenteredCumulantBasedLbMethod
-from lbmpy.methods.centeredcumulant.centered_cumulants import get_default_polynomial_cumulants_for_stencil
-from lbmpy.methods.centeredcumulant.cumulant_transform import CentralMomentsToCumulantsByGeneratingFunc
+from lbmpy.moment_transforms import CentralMomentsToCumulantsByGeneratingFunc
 
-from lbmpy.methods.conservedquantitycomputation import DensityVelocityComputation
+from .conservedquantitycomputation import DensityVelocityComputation
 
-from lbmpy.methods.momentbased.momentbasedmethod import MomentBasedLbMethod
-from lbmpy.methods.momentbased.moment_transforms import PdfsToCentralMomentsByShiftMatrix
+from .momentbased.momentbasedmethod import MomentBasedLbMethod
+from .momentbased.centralmomentbasedmethod import CentralMomentBasedLbMethod
+from .cumulantbased import CumulantBasedLbMethod
+from lbmpy.moment_transforms import (
+    AbstractMomentTransform, BinomialChimeraTransform, PdfsToMomentsByChimeraTransform)
+from lbmpy.moment_transforms.rawmomenttransforms import AbstractRawMomentTransform
+from lbmpy.moment_transforms.centralmomenttransforms import AbstractCentralMomentTransform
 
 from lbmpy.moments import (
-    MOMENT_SYMBOLS, discrete_moment, exponents_to_polynomial_representations,
+    MOMENT_SYMBOLS, exponents_to_polynomial_representations,
     get_default_moment_set_for_stencil, gram_schmidt, is_even, moments_of_order,
     moments_up_to_component_order, sort_moments_into_groups_of_same_order,
     is_bulk_moment, is_shear_moment, get_order)
 
 from lbmpy.relaxationrates import relaxation_rate_from_magic_number
-from lbmpy.stencils import get_stencil
-from pystencils.stencil import have_same_entries
+from lbmpy.enums import Stencil, CollisionSpace
+from lbmpy.stencils import LBStencil
 from pystencils.sympyextensions import common_denominator
 
 
-def create_with_discrete_maxwellian_eq_moments(stencil, moment_to_relaxation_rate_dict, compressible=False,
-                                               force_model=None, equilibrium_order=2, c_s_sq=sp.Rational(1, 3)):
-    r"""Creates a moment-based LBM by taking a list of moments with corresponding relaxation rate.
+@dataclass
+class CollisionSpaceInfo:
+    """
+    This class encapsulates necessary details about a method's collision space.
+    """
+    collision_space: CollisionSpace
+    """
+    The method's collision space.
+    """
+    raw_moment_transform_class: Type[AbstractRawMomentTransform] = None
+    """
+    Python class that determines how PDFs are transformed to raw moment space. If left as 'None', this parameter
+    will be inferred from `collision_space`, defaulting to 
+    :class:`lbmpy.moment_transforms.PdfsToMomentsByChimeraTransform`
+    if `CollisionSpace.RAW_MOMENTS` is given, or `None` otherwise.
+    """
+    central_moment_transform_class: Type[AbstractCentralMomentTransform] = None
+    """
+    Python class that determines how PDFs are transformed to central moment space. If left as 'None', this parameter
+    will be inferred from `collision_space`, defaulting to 
+    :class:`lbmpy.moment_transforms.BinomialChimeraTransform`
+    if `CollisionSpace.CENTRAL_MOMENTS` or `CollisionSpace.CUMULANTS` is given, or `None` otherwise.
+    """
+    cumulant_transform_class: Type[AbstractMomentTransform] = None
+    """
+    Python class that determines how central moments are transformed to cumulant space. If left as 'None', this 
+    parameter will be inferred from `collision_space`, defaulting to 
+    :class:`lbmpy.moment_transforms.CentralMomentsToCumulantsByGeneratingFunc`
+    if `CollisionSpace.CUMULANTS` is given, or `None` otherwise.
+    """
+
+    def __post_init__(self):
+        if self.collision_space == CollisionSpace.RAW_MOMENTS and self.raw_moment_transform_class is None:
+            self.raw_moment_transform_class = PdfsToMomentsByChimeraTransform
+        if self.collision_space in (CollisionSpace.CENTRAL_MOMENTS, CollisionSpace.CUMULANTS) \
+                and self.central_moment_transform_class is None:
+            self.central_moment_transform_class = BinomialChimeraTransform
+        if self.collision_space == CollisionSpace.CUMULANTS and self.cumulant_transform_class is None:
+            self.cumulant_transform_class = CentralMomentsToCumulantsByGeneratingFunc
+
+
+def create_with_discrete_maxwellian_equilibrium(stencil, moment_to_relaxation_rate_dict,
+                                                compressible=False, zero_centered=False, delta_equilibrium=False,
+                                                force_model=None, equilibrium_order=2, c_s_sq=sp.Rational(1, 3),
+                                                **kwargs):
+    r"""Creates a moment-based LBM by taking a dictionary of moments with corresponding relaxation rates.
 
-    These moments are relaxed against the moments of the discrete Maxwellian distribution.
+    These moments are relaxed against the moments of the discrete Maxwellian distribution
+    (see :class:`lbmpy.equilibrium.DiscreteHydrodynamicMaxwellian`).
+
+    Internally, this method calls :func:`lbmpy.methods.create_from_equilibrium`.
 
     Args:
-        stencil: nested tuple defining the discrete velocity space. See `get_stencil`
+        stencil: instance of :class:`lbmpy.stencils.LBStenil`
         moment_to_relaxation_rate_dict: dict that has as many entries as the stencil. Each moment, which can be
                                         represented by an exponent tuple or in polynomial form
                                         (see `lbmpy.moments`), is mapped to a relaxation rate.
-        compressible: incompressible LBM methods split the density into :math:`\rho = \rho_0 + \Delta \rho`
-                      where :math:`\rho_0` is chosen as one, and the first moment of the pdfs is :math:`\Delta \rho` .
-                      This approximates the incompressible Navier-Stokes equations better than the standard
-                      compressible model.
-        force_model: force model instance, or None if no external forces
+        compressible: If `False`, the incompressible equilibrium formulation is used to better approximate the
+                      incompressible Navier-Stokes equations. Otherwise, the default textbook equilibrium is used.
+        zero_centered: If `True`, the zero-centered storage format for PDFs is used, storing only their deviation from
+                       the background distribution (given by the lattice weights).
+        delta_equilibrium: Takes effect only if zero-centered storage is used. If `True`, the equilibrium distribution
+                           is also given only as its deviation from the background distribution.
+        force_model: instance of :class:`lbmpy.forcemodels.AbstractForceModel`, or None if no external forces are 
+                     present.
         equilibrium_order: approximation order of macroscopic velocity :math:`\mathbf{u}` in the equilibrium
         c_s_sq: Speed of sound squared
+        kwargs: See :func:`lbmpy.methods.create_from_equilibrium`
 
     Returns:
-        `lbmpy.methods.momentbased.MomentBasedLbMethod` instance
+        Instance of a subclass of :class:`lbmpy.methods.AbstractLbMethod`.
     """
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
-    mom_to_rr_dict = OrderedDict(moment_to_relaxation_rate_dict)
-    assert len(mom_to_rr_dict) == len(stencil), \
-        "The number of moments has to be the same as the number of stencil entries"
-
-    density_velocity_computation = DensityVelocityComputation(stencil, compressible, force_model)
-    eq_values = get_moments_of_discrete_maxwellian_equilibrium(stencil, tuple(mom_to_rr_dict.keys()),
-                                                               c_s_sq=c_s_sq, compressible=compressible,
-                                                               order=equilibrium_order)
-
-    rr_dict = OrderedDict([(mom, RelaxationInfo(eq_mom, rr))
-                           for mom, rr, eq_mom in zip(mom_to_rr_dict.keys(), mom_to_rr_dict.values(), eq_values)])
-
-    return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model)
-
-
-def create_with_continuous_maxwellian_eq_moments(stencil, moment_to_relaxation_rate_dict, compressible=False,
-                                                 force_model=None, equilibrium_order=2, c_s_sq=sp.Rational(1, 3)):
+    cqc = DensityVelocityComputation(stencil, compressible, zero_centered, force_model=force_model, c_s_sq=c_s_sq)
+    equilibrium = DiscreteHydrodynamicMaxwellian(stencil, compressible=compressible,
+                                                 deviation_only=delta_equilibrium,
+                                                 order=equilibrium_order,
+                                                 c_s_sq=c_s_sq)
+    return create_from_equilibrium(stencil, equilibrium, cqc, moment_to_relaxation_rate_dict,
+                                   zero_centered=zero_centered, force_model=force_model, **kwargs)
+
+
+def create_with_continuous_maxwellian_equilibrium(stencil, moment_to_relaxation_rate_dict,
+                                                  compressible=False, zero_centered=False, delta_equilibrium=False,
+                                                  force_model=None, equilibrium_order=2, c_s_sq=sp.Rational(1, 3),
+                                                  **kwargs):
     r"""
-    Creates a moment-based LBM by taking a list of moments with corresponding relaxation rate. These moments are
-    relaxed against the moments of the continuous Maxwellian distribution.
-    For parameter description see :func:`lbmpy.methods.create_with_discrete_maxwellian_eq_moments`.
-    By using the continuous Maxwellian we automatically get a compressible model.
+    Creates a moment-based LBM by taking a dictionary of moments with corresponding relaxation rates. 
+    These moments are relaxed against the moments of the continuous Maxwellian distribution
+    (see :class:`lbmpy.equilibrium.ContinuousHydrodynamicMaxwellian`).
+
+    Internally, this method calls :func:`lbmpy.methods.create_from_equilibrium`.
 
     Args:
-        stencil: nested tuple defining the discrete velocity space. See `get_stencil`
+        stencil: instance of :class:`lbmpy.stencils.LBStenil`
         moment_to_relaxation_rate_dict: dict that has as many entries as the stencil. Each moment, which can be
                                         represented by an exponent tuple or in polynomial form
                                         (see `lbmpy.moments`), is mapped to a relaxation rate.
-        compressible: incompressible LBM methods split the density into :math:`\rho = \rho_0 + \Delta \rho`
-                      where :math:`\rho_0` is chosen as one, and the first moment of the pdfs is :math:`\Delta \rho` .
-                      This approximates the incompressible Navier-Stokes equations better than the standard
-                      compressible model.
-        force_model: force model instance, or None if no external forces
+        compressible: If `False`, the incompressible equilibrium formulation is used to better approximate the
+                      incompressible Navier-Stokes equations. Otherwise, the default textbook equilibrium is used.
+        zero_centered: If `True`, the zero-centered storage format for PDFs is used, storing only their deviation from 
+                       the background distribution (given by the lattice weights).
+        delta_equilibrium: Takes effect only if zero-centered storage is used. If `True`, the equilibrium 
+                           distribution is also given only as its deviation from the background distribution.
+        force_model: Instance of :class:`lbmpy.forcemodels.AbstractForceModel`, or None if no external forces 
+                     are present.
         equilibrium_order: approximation order of macroscopic velocity :math:`\mathbf{u}` in the equilibrium
         c_s_sq: Speed of sound squared
+        kwargs: See :func:`lbmpy.methods.create_from_equilibrium`
 
     Returns:
-        `lbmpy.methods.momentbased.MomentBasedLbMethod` instance
+        Instance of a subclass of :class:`lbmpy.methods.AbstractLbMethod`.
     """
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
-    mom_to_rr_dict = OrderedDict(moment_to_relaxation_rate_dict)
-    assert len(mom_to_rr_dict) == len(stencil), "The number of moments has to be equal to the number of stencil entries"
-    dim = len(stencil[0])
-    density_velocity_computation = DensityVelocityComputation(stencil, compressible, force_model)
-    eq_values = get_moments_of_continuous_maxwellian_equilibrium(tuple(mom_to_rr_dict.keys()), dim, c_s_sq=c_s_sq,
-                                                                 order=equilibrium_order)
-
-    if not compressible:
-        rho = density_velocity_computation.defined_symbols(order=0)[1]
-        u = density_velocity_computation.defined_symbols(order=1)[1]
-        eq_values = [compressible_to_incompressible_moment_value(em, rho, u) for em in eq_values]
-
-    rr_dict = OrderedDict([(mom, RelaxationInfo(eq_mom, rr))
-                           for mom, rr, eq_mom in zip(mom_to_rr_dict.keys(), mom_to_rr_dict.values(), eq_values)])
-
-    return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model)
-
-
-def create_generic_mrt(stencil, moment_eq_value_relaxation_rate_tuples, compressible=False,
-                       force_model=None):
+    cqc = DensityVelocityComputation(stencil, compressible, zero_centered, force_model=force_model, c_s_sq=c_s_sq)
+    equilibrium = ContinuousHydrodynamicMaxwellian(dim=stencil.D, compressible=compressible,
+                                                   deviation_only=delta_equilibrium,
+                                                   order=equilibrium_order,
+                                                   c_s_sq=c_s_sq)
+    return create_from_equilibrium(stencil, equilibrium, cqc, moment_to_relaxation_rate_dict,
+                                   zero_centered=zero_centered, force_model=force_model, **kwargs)
+
+
+def create_from_equilibrium(stencil, equilibrium, conserved_quantity_computation,
+                            moment_to_relaxation_rate_dict,
+                            collision_space_info=CollisionSpaceInfo(CollisionSpace.POPULATIONS),
+                            zero_centered=False, force_model=None, fraction_field=None):
     r"""
-    Creates a generic moment-based LB method.
-
-    Args:
-        stencil: sequence of lattice velocities
-        moment_eq_value_relaxation_rate_tuples: sequence of tuples containing (moment, equilibrium value, relax. rate)
-        compressible: compressibility, determines calculation of velocity for force models
-        force_model: see create_with_discrete_maxwellian_eq_moments
-    """
-    density_velocity_computation = DensityVelocityComputation(stencil, compressible, force_model)
+    Creates a lattice Boltzmann method in either population, moment, or central moment space, from a given
+    discrete velocity set and equilibrium distribution. 
 
-    rr_dict = OrderedDict()
-    for moment, eq_value, rr in moment_eq_value_relaxation_rate_tuples:
-        moment = sp.sympify(moment)
-        rr_dict[moment] = RelaxationInfo(eq_value, rr)
-    return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model)
-
-
-def create_from_equilibrium(stencil, equilibrium, moment_to_relaxation_rate_dict,
-                            compressible=False, force_model=None):
-    r"""
-    Creates a moment-based LB method using a given equilibrium distribution function
+    This function takes a stencil, an equilibrium distribution, an appropriate conserved quantity computation
+    instance, a dictionary mapping moment polynomials to their relaxation rates, and a collision space info
+    determining the desired collision space. It returns a method instance relaxing the given moments against
+    their equilibrium values computed from the given distribution, in the given collision space.
 
     Args:
-        stencil: see create_with_discrete_maxwellian_eq_moments
-        equilibrium: list of equilibrium terms, dependent on rho and u, one for each stencil direction
-        moment_to_relaxation_rate_dict: relaxation rate for each moment, or a symbol/float if all should relaxed with
-                                        the same rate
-        compressible: see create_with_discrete_maxwellian_eq_moments
-        force_model: see create_with_discrete_maxwellian_eq_moments
+        stencil: Instance of :class:`lbmpy.stencils.LBStencil`
+        equilibrium: Instance of a subclass of :class:`lbmpy.equilibrium.AbstractEquilibrium`.
+        conserved_quantity_computation: Instance of a subclass of 
+                                        :class:`lbmpy.methods.AbstractConservedQuantityComputation`,
+                                        which must provide equations for the conserved quantities used in
+                                        the equilibrium.
+        moment_to_relaxation_rate_dict: Dictionary mapping moment polynomials to relaxation rates.
+        collision_space_info: Instance of :class:`CollisionSpaceInfo`, defining the method's desired collision space
+                              and the manner of transformation to this space. Cumulant-based methods are not supported
+                              yet.
+        zero_centered: Whether or not the zero-centered storage format should be used. If `True`, the given equilibrium
+                       must either be a deviation-only formulation, or must provide a background distribution for PDFs
+                       to be centered around.
+        force_model: Instance of :class:`lbmpy.forcemodels.AbstractForceModel`, or None if no external forces are
+                     present.
     """
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
     if any(isinstance(moment_to_relaxation_rate_dict, t) for t in (sp.Symbol, float, int)):
         moment_to_relaxation_rate_dict = {m: moment_to_relaxation_rate_dict
                                           for m in get_default_moment_set_for_stencil(stencil)}
 
     mom_to_rr_dict = OrderedDict(moment_to_relaxation_rate_dict)
-    assert len(mom_to_rr_dict) == len(stencil), "The number of moments has to be equal to the number of stencil entries"
-    density_velocity_computation = DensityVelocityComputation(stencil, compressible, force_model)
-
-    rr_dict = OrderedDict([(mom, RelaxationInfo(discrete_moment(equilibrium, mom, stencil).expand(), rr))
-                           for mom, rr in zip(mom_to_rr_dict.keys(), mom_to_rr_dict.values())])
-    return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model)
+    assert len(mom_to_rr_dict) == stencil.Q, "The number of moments has to be equal to the number of stencil entries"
+
+    cqc = conserved_quantity_computation
+    cspace = collision_space_info
+
+    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,
+                                   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,
+                                   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)
 
 
 # ------------------------------------ SRT / TRT/ MRT Creators ---------------------------------------------------------
 
 
-def create_srt(stencil, relaxation_rate, maxwellian_moments=False, **kwargs):
+def create_srt(stencil, relaxation_rate, continuous_equilibrium=True, **kwargs):
     r"""Creates a single relaxation time (SRT) lattice Boltzmann model also known as BGK model.
 
+    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
+    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.
+
+    If not specified otherwise, collision equations will be derived in population space.
+
     Args:
-        stencil: nested tuple defining the discrete velocity space. See :func:`lbmpy.stencils.get_stencil`
+        stencil: instance of :class:`lbmpy.stencils.LBStencil`
         relaxation_rate: relaxation rate (inverse of the relaxation time)
                         usually called :math:`\omega` in LBM literature
-        maxwellian_moments: determines if the discrete or continuous maxwellian equilibrium is
+        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
                         used to compute the equilibrium moments
 
     Returns:
         :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
     """
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
+    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
+    check_and_set_mrt_space(CollisionSpace.POPULATIONS)
     moments = get_default_moment_set_for_stencil(stencil)
     rr_dict = OrderedDict([(m, relaxation_rate) for m in moments])
-    if maxwellian_moments:
-        return create_with_continuous_maxwellian_eq_moments(stencil, rr_dict, **kwargs)
+    if continuous_equilibrium:
+        return create_with_continuous_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
     else:
-        return create_with_discrete_maxwellian_eq_moments(stencil, rr_dict, **kwargs)
+        return create_with_discrete_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
 
 
 def create_trt(stencil, relaxation_rate_even_moments, relaxation_rate_odd_moments,
-               maxwellian_moments=False, **kwargs):
+               continuous_equilibrium=True, **kwargs):
     """
     Creates a two relaxation time (TRT) lattice Boltzmann model, where even and odd moments are relaxed differently.
     In the SRT model the exact wall position of no-slip boundaries depends on the viscosity, the TRT method does not
@@ -200,16 +263,24 @@ def create_trt(stencil, relaxation_rate_even_moments, relaxation_rate_odd_moment
     Parameters are similar to :func:`lbmpy.methods.create_srt`, but instead of one relaxation rate there are
     two relaxation rates: one for even moments (determines viscosity) and one for odd moments.
     If unsure how to choose the odd relaxation rate, use the function :func:`lbmpy.methods.create_trt_with_magic_number`
+
+    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
+    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.
+
+    If not specified otherwise, collision equations will be derived in population space.
+
+    Returns:
+        :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
     """
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
+    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
+    check_and_set_mrt_space(CollisionSpace.POPULATIONS)
     moments = get_default_moment_set_for_stencil(stencil)
     rr_dict = OrderedDict([(m, relaxation_rate_even_moments if is_even(m) else relaxation_rate_odd_moments)
                            for m in moments])
-    if maxwellian_moments:
-        return create_with_continuous_maxwellian_eq_moments(stencil, rr_dict, **kwargs)
+    if continuous_equilibrium:
+        return create_with_continuous_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
     else:
-        return create_with_discrete_maxwellian_eq_moments(stencil, rr_dict, **kwargs)
+        return create_with_discrete_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
 
 
 def create_trt_with_magic_number(stencil, relaxation_rate, magic_number=sp.Rational(3, 16), **kwargs):
@@ -218,8 +289,13 @@ def create_trt_with_magic_number(stencil, relaxation_rate, magic_number=sp.Ratio
     determines from the even moment relaxation time and a "magic number".
     For possible parameters see :func:`lbmpy.methods.create_trt`
 
+    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
+    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.
+
+    If not specified otherwise, collision equations will be derived in population space.
+
     Args:
-        stencil: nested tuple defining the discrete velocity space. See :func:`lbmpy.stencils.get_stencil`
+        stencil: instance of :class:`lbmpy.stencils.LBStencil`
         relaxation_rate: relaxation rate (inverse of the relaxation time)
                         usually called :math:`\omega` in LBM literature
         magic_number: magic number which is used to calculate the relxation rate for the odd moments.
@@ -232,44 +308,107 @@ def create_trt_with_magic_number(stencil, relaxation_rate, magic_number=sp.Ratio
                       relaxation_rate_odd_moments=rr_odd, **kwargs)
 
 
-def create_mrt_raw(stencil, relaxation_rates, maxwellian_moments=False, **kwargs):
+def create_mrt_raw(stencil, relaxation_rates, continuous_equilibrium=True, conserved_moments=True, **kwargs):
     r"""
     Creates a MRT method using non-orthogonalized moments.
 
+    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
+    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.
+
+    If not specified otherwise, collision equations will be derived in raw moment space.
+
     Args:
-        stencil: nested tuple defining the discrete velocity space. See :func:`lbmpy.stencils.get_stencil`
+        stencil: instance of :class:`lbmpy.stencils.LBStencil`
         relaxation_rates: relaxation rates (inverse of the relaxation times) for each moment
-        maxwellian_moments: determines if the discrete or continuous maxwellian equilibrium is
+        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.
     Returns:
         :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
     """
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
+    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
+    check_and_set_mrt_space(CollisionSpace.RAW_MOMENTS)
     moments = get_default_moment_set_for_stencil(stencil)
-    rr_dict = OrderedDict(zip(moments, relaxation_rates))
-    if maxwellian_moments:
-        return create_with_continuous_maxwellian_eq_moments(stencil, rr_dict, **kwargs)
+    nested_moments = [(c,) for c in moments]
+    rr_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments, stencil.D, conserved_moments)
+    if continuous_equilibrium:
+        return create_with_continuous_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
+    else:
+        return create_with_discrete_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
+
+
+def create_central_moment(stencil, relaxation_rates, nested_moments=None,
+                          continuous_equilibrium=True, conserved_moments=True, **kwargs):
+    r"""
+    Creates moment based LB method where the collision takes place in the central moment space.
+
+    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
+    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.
+
+    Args:
+        stencil: instance of :class:`lbmpy.stencils.LBStencil`
+        relaxation_rates: relaxation rates (inverse of the relaxation times) for each moment
+        nested_moments: a list of lists of modes, grouped by common relaxation times.
+        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.
+    Returns:
+        :class:`lbmpy.methods.momentbased.CentralMomentBasedLbMethod` instance
+    """
+    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
+
+    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.CENTRAL_MOMENTS))
+    if kwargs['collision_space_info'].collision_space != CollisionSpace.CENTRAL_MOMENTS:
+        raise ValueError("Central moment-based methods can only be derived in central moment space.")
+
+    if nested_moments and not isinstance(nested_moments[0], list):
+        nested_moments = list(sort_moments_into_groups_of_same_order(nested_moments).values())
+        second_order_moments = nested_moments[2]
+        bulk_moment = [m for m in second_order_moments if is_bulk_moment(m, stencil.D)]
+        shear_moments = [m for m in second_order_moments if is_shear_moment(m, stencil.D)]
+        assert len(shear_moments) + len(bulk_moment) == len(second_order_moments)
+        nested_moments[2] = shear_moments
+        nested_moments.insert(3, bulk_moment)
+
+    if not nested_moments:
+        nested_moments = cascaded_moment_sets_literature(stencil)
+
+    rr_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments, stencil.D, conserved_moments)
+    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)
+
+    if continuous_equilibrium:
+        return create_with_continuous_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
     else:
-        return create_with_discrete_maxwellian_eq_moments(stencil, rr_dict, **kwargs)
+        return create_with_discrete_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
 
 
 def create_trt_kbc(dim, shear_relaxation_rate, higher_order_relaxation_rate, method_name='KBC-N4',
-                   maxwellian_moments=False, **kwargs):
+                   continuous_equilibrium=True, **kwargs):
     """
     Creates a method with two relaxation rates, one for lower order moments which determines the viscosity and
     one for higher order moments. In entropic models this second relaxation rate is chosen subject to an entropy
     condition. Which moments are relaxed by which rate is determined by the method_name
 
+    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
+    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.
+
+    If not specified otherwise, collision equations will be derived in population space.
+
     Args:
         dim: 2 or 3, leads to stencil D2Q9 or D3Q27
         shear_relaxation_rate: relaxation rate that determines viscosity
         higher_order_relaxation_rate: relaxation rate for higher order moments
         method_name: string 'KBC-Nx' where x can be an number from 1 to 4, for details see
                     "Karlin 2015: Entropic multi relaxation lattice Boltzmann models for turbulent flows"
-        maxwellian_moments: determines if the discrete or continuous maxwellian equilibrium is
+        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
                         used to compute the equilibrium moments.
     """
+    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
+    check_and_set_mrt_space(CollisionSpace.POPULATIONS)
+
     def product(iterable):
         return reduce(operator.mul, iterable, 1)
 
@@ -290,7 +429,7 @@ def create_trt_kbc(dim, shear_relaxation_rate, higher_order_relaxation_rate, met
     explicitly_defined = set(rho + velocity + shear_tensor_off_diagonal
                              + shear_tensor_diagonal + energy_transport_tensor)
     rest = list(set(poly_repr(moments_up_to_component_order(2, dim))) - explicitly_defined)
-    assert len(rest) + len(explicitly_defined) == 3**dim
+    assert len(rest) + len(explicitly_defined) == 3 ** dim
 
     # naming according to paper Karlin 2015: Entropic multi relaxation lattice Boltzmann models for turbulent flows
     d = shear_tensor_off_diagonal + shear_tensor_trace_free_diagonal[:-1]
@@ -315,235 +454,87 @@ def create_trt_kbc(dim, shear_relaxation_rate, higher_order_relaxation_rate, met
                        [omega_s] * len(shear_part) + \
                        [omega_h] * len(rest_part)
 
-    stencil = get_stencil("D2Q9") if dim == 2 else get_stencil("D3Q27")
+    stencil = LBStencil(Stencil.D2Q9) if dim == 2 else LBStencil(Stencil.D3Q27)
     all_moments = rho + velocity + shear_part + rest_part
     moment_to_rr = OrderedDict((m, rr) for m, rr in zip(all_moments, relaxation_rates))
 
-    if maxwellian_moments:
-        return create_with_continuous_maxwellian_eq_moments(stencil, moment_to_rr, **kwargs)
+    if continuous_equilibrium:
+        return create_with_continuous_maxwellian_equilibrium(stencil, moment_to_rr, **kwargs)
     else:
-        return create_with_discrete_maxwellian_eq_moments(stencil, moment_to_rr, **kwargs)
+        return create_with_discrete_maxwellian_equilibrium(stencil, moment_to_rr, **kwargs)
 
 
-def create_mrt_orthogonal(stencil, relaxation_rate_getter, maxwellian_moments=False, weighted=None,
-                          nested_moments=None, **kwargs):
+def create_mrt_orthogonal(stencil, relaxation_rates, continuous_equilibrium=True, weighted=None,
+                          nested_moments=None, conserved_moments=True, **kwargs):
     r"""
     Returns an orthogonal multi-relaxation time model for the stencils D2Q9, D3Q15, D3Q19 and D3Q27.
     These MRT methods are just one specific version - there are many MRT methods possible for all these stencils
     which differ by the linear combination of moments and the grouping into equal relaxation times.
-    To create a generic MRT method use `create_with_discrete_maxwellian_eq_moments`
+    To create a generic MRT method use `create_with_discrete_maxwellian_equilibrium`
+
+    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
+    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.
+
+    If not specified otherwise, collision equations will be derived in raw moment space.
 
     Args:
-        stencil: nested tuple defining the discrete velocity space. See `get_stencil`
-        relaxation_rate_getter: function getting a list of moments as argument, returning the associated relaxation
-        maxwellian_moments: determines if the discrete or continuous maxwellian equilibrium is
+        stencil: instance of :class:`lbmpy.stencils.LBStencil`
+        relaxation_rates: relaxation rates for the moments
+        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
                                                used to compute the equilibrium moments
         weighted: whether to use weighted or unweighted orthogonality
-        nested_moments: a list of lists of modes, grouped by common relaxation times. This is usually used in
-                        conjunction with `mrt_orthogonal_modes_literature`. If this argument is not provided,
-                        Gram-Schmidt orthogonalization of the default modes is performed.
+        nested_moments: a list of lists of modes, grouped by common relaxation times. If this argument is not provided,
+                        Gram-Schmidt orthogonalization of the default modes is performed. The default modes equal the
+                        raw moments except for the separation of the shear and bulk viscosity.
+        conserved_moments: If lower order moments are conserved or not.
     """
-    dim = len(stencil[0])
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
+    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
+    check_and_set_mrt_space(CollisionSpace.RAW_MOMENTS)
 
-    if weighted is None and not nested_moments:
-        raise ValueError("Please specify whether you want weighted or unweighted orthogonality using 'weighted='")
-    elif weighted:
+    if weighted:
         weights = get_weights(stencil, sp.Rational(1, 3))
     else:
         weights = None
 
-    moment_to_relaxation_rate_dict = OrderedDict()
     if not nested_moments:
         moments = get_default_moment_set_for_stencil(stencil)
         x, y, z = MOMENT_SYMBOLS
-        if dim == 2:
+        if stencil.D == 2:
             diagonal_viscous_moments = [x ** 2 + y ** 2, x ** 2]
         else:
             diagonal_viscous_moments = [x ** 2 + y ** 2 + z ** 2, x ** 2, y ** 2 - z ** 2]
-        for i, d in enumerate(MOMENT_SYMBOLS[:dim]):
-            moments[moments.index(d**2)] = diagonal_viscous_moments[i]
+
+        for i, d in enumerate(MOMENT_SYMBOLS[:stencil.D]):
+            if d ** 2 in moments:
+                moments[moments.index(d ** 2)] = diagonal_viscous_moments[i]
         orthogonal_moments = gram_schmidt(moments, stencil, weights)
         orthogonal_moments_scaled = [e * common_denominator(e) for e in orthogonal_moments]
         nested_moments = list(sort_moments_into_groups_of_same_order(orthogonal_moments_scaled).values())
         # second order moments: separate bulk from shear
         second_order_moments = nested_moments[2]
-        bulk_moment = [m for m in second_order_moments if is_bulk_moment(m, dim)]
-        shear_moments = [m for m in second_order_moments if is_shear_moment(m, dim)]
+        bulk_moment = [m for m in second_order_moments if is_bulk_moment(m, stencil.D)]
+        shear_moments = [m for m in second_order_moments if is_shear_moment(m, stencil.D)]
         assert len(shear_moments) + len(bulk_moment) == len(second_order_moments)
         nested_moments[2] = shear_moments
         nested_moments.insert(3, bulk_moment)
-    for moment_list in nested_moments:
-        rr = relaxation_rate_getter(moment_list)
-        for m in moment_list:
-            moment_to_relaxation_rate_dict[m] = rr
-
-    if maxwellian_moments:
-        return create_with_continuous_maxwellian_eq_moments(stencil, moment_to_relaxation_rate_dict, **kwargs)
-    else:
-        return create_with_discrete_maxwellian_eq_moments(stencil, moment_to_relaxation_rate_dict, **kwargs)
-
 
-def mrt_orthogonal_modes_literature(stencil, is_weighted):
-    """
-    Returns a list of lists of modes, grouped by common relaxation times.
-    This is for commonly used MRT models found in literature.
+    moment_to_relaxation_rate_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments,
+                                                               stencil.D, conserved_moments)
 
-    Args:
-        stencil: nested tuple defining the discrete velocity space. See `get_stencil`
-        is_weighted: whether to use weighted or unweighted orthogonality
-
-    MRT schemes as described in the following references are used
-    """
-    x, y, z = MOMENT_SYMBOLS
-    one = sp.Rational(1, 1)
-
-    if have_same_entries(stencil, get_stencil("D2Q9")) and is_weighted:
-        # Reference:
-        # Duenweg, B., Schiller, U. D., & Ladd, A. J. (2007). Statistical mechanics of the fluctuating
-        # lattice Boltzmann equation. Physical Review E, 76(3)
-        sq = x ** 2 + y ** 2
-        all_moments = [one, x, y, 3 * sq - 2, 2 * x ** 2 - sq, x * y,
-                       (3 * sq - 4) * x, (3 * sq - 4) * y, 9 * sq ** 2 - 15 * sq + 2]
-        nested_moments = list(sort_moments_into_groups_of_same_order(all_moments).values())
-        return nested_moments
-    elif have_same_entries(stencil, get_stencil("D3Q15")) and is_weighted:
-        sq = x ** 2 + y ** 2 + z ** 2
-        nested_moments = [
-            [one, x, y, z],  # [0, 3, 5, 7]
-            [sq - 1],  # [1]
-            [3 * sq ** 2 - 9 * sq + 4],  # [2]
-            [(3 * sq - 5) * x, (3 * sq - 5) * y, (3 * sq - 5) * z],  # [4, 6, 8]
-            [3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z],  # [9, 10, 11, 12, 13]
-            [x * y * z]
-        ]
-    elif have_same_entries(stencil, get_stencil("D3Q19")) and is_weighted:
-        # This MRT variant mentioned in the dissertation of Ulf Schiller
-        # "Thermal fluctuations and boundary conditions in the lattice Boltzmann method" (2008), p. 24ff
-        # There are some typos in the moment matrix on p.27
-        # The here implemented ordering of the moments is however different from that reference (Eq. 2.61-2.63)
-        # The moments are weighted-orthogonal (Eq. 2.58)
-
-        # Further references:
-        # Duenweg, B., Schiller, U. D., & Ladd, A. J. (2007). Statistical mechanics of the fluctuating
-        # lattice Boltzmann equation. Physical Review E, 76(3)
-        # Chun, B., & Ladd, A. J. (2007). Interpolated boundary condition for lattice Boltzmann simulations of
-        # flows in narrow gaps. Physical review E, 75(6)
-        sq = x ** 2 + y ** 2 + z ** 2
-        nested_moments = [
-            [one, x, y, z],  # [0, 3, 5, 7]
-            [sq - 1],  # [1]
-            [3 * sq ** 2 - 6 * sq + 1],  # [2]
-            [(3 * sq - 5) * x, (3 * sq - 5) * y, (3 * sq - 5) * z],  # [4, 6, 8]
-            [3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z],  # [9, 11, 13, 14, 15]
-            [(2 * sq - 3) * (3 * x ** 2 - sq), (2 * sq - 3) * (y ** 2 - z ** 2)],  # [10, 12]
-            [(y ** 2 - z ** 2) * x, (z ** 2 - x ** 2) * y, (x ** 2 - y ** 2) * z]  # [16, 17, 18]
-        ]
-    elif have_same_entries(stencil, get_stencil("D3Q27")) and not is_weighted:
-        xsq, ysq, zsq = x ** 2, y ** 2, z ** 2
-        all_moments = [
-            sp.Rational(1, 1),  # 0
-            x, y, z,  # 1, 2, 3
-            x * y, x * z, y * z,  # 4, 5, 6
-            xsq - ysq,  # 7
-            (xsq + ysq + zsq) - 3 * zsq,  # 8
-            (xsq + ysq + zsq) - 2,  # 9
-            3 * (x * ysq + x * zsq) - 4 * x,  # 10
-            3 * (xsq * y + y * zsq) - 4 * y,  # 11
-            3 * (xsq * z + ysq * z) - 4 * z,  # 12
-            x * ysq - x * zsq,  # 13
-            xsq * y - y * zsq,  # 14
-            xsq * z - ysq * z,  # 15
-            x * y * z,  # 16
-            3 * (xsq * ysq + xsq * zsq + ysq * zsq) - 4 * (xsq + ysq + zsq) + 4,  # 17
-            3 * (xsq * ysq + xsq * zsq - 2 * ysq * zsq) - 2 * (2 * xsq - ysq - zsq),  # 18
-            3 * (xsq * ysq - xsq * zsq) - 2 * (ysq - zsq),  # 19
-            3 * (xsq * y * z) - 2 * (y * z),  # 20
-            3 * (x * ysq * z) - 2 * (x * z),  # 21
-            3 * (x * y * zsq) - 2 * (x * y),  # 22
-            9 * (x * ysq * zsq) - 6 * (x * ysq + x * zsq) + 4 * x,  # 23
-            9 * (xsq * y * zsq) - 6 * (xsq * y + y * zsq) + 4 * y,  # 24
-            9 * (xsq * ysq * z) - 6 * (xsq * z + ysq * z) + 4 * z,  # 25
-            27 * (xsq * ysq * zsq) - 18 * (xsq * ysq + xsq * zsq + ysq * zsq) + 12 * (xsq + ysq + zsq) - 8,  # 26
-        ]
-        nested_moments = list(sort_moments_into_groups_of_same_order(all_moments).values())
+    if continuous_equilibrium:
+        return create_with_continuous_maxwellian_equilibrium(stencil,
+                                                             moment_to_relaxation_rate_dict, **kwargs)
     else:
-        raise NotImplementedError("No MRT model is available (yet) for this stencil. "
-                                  "Create a custom MRT using 'create_with_discrete_maxwellian_eq_moments'")
+        return create_with_discrete_maxwellian_equilibrium(stencil,
+                                                           moment_to_relaxation_rate_dict, **kwargs)
 
-    return nested_moments
 
 # ----------------------------------------- Cumulant method creators ---------------------------------------------------
-
-
-def create_centered_cumulant_model(stencil, cumulant_to_rr_dict, force_model=None,
-                                   equilibrium_order=None, c_s_sq=sp.Rational(1, 3),
-                                   galilean_correction=False,
-                                   central_moment_transform_class=PdfsToCentralMomentsByShiftMatrix,
-                                   cumulant_transform_class=CentralMomentsToCumulantsByGeneratingFunc):
-    r"""Creates a cumulant lattice Boltzmann model.
-
-    Args:
-        stencil: nested tuple defining the discrete velocity space. See :func:`lbmpy.stencils.get_stencil`
-        cumulant_to_rr_dict: dict that has as many entries as the stencil. Each cumulant, which can be
-                             represented by an exponent tuple or in polynomial form is mapped to a relaxation rate.
-                             See `get_default_polynomial_cumulants_for_stencil`
-        force_model: force model used for the collision. For cumulant LB method a good choice is
-                     `lbmpy.methods.centeredcumulant.CenteredCumulantForceModel`
-        equilibrium_order: approximation order of macroscopic velocity :math:`\mathbf{u}` in the equilibrium
-        c_s_sq: Speed of sound squared
-        galilean_correction: special correction for D3Q27 cumulant collisions. See Appendix H in
-                             :cite:`geier2015`. Implemented in :mod:`lbmpy.methods.centeredcumulant.galilean_correction`
-        central_moment_transform_class: Class which defines the transformation to the central moment space
-                                        (see :mod:`lbmpy.methods.momentbased.moment_transforms`)
-        cumulant_transform_class: Class which defines the transformation from the central moment space to the
-                                  cumulant space (see :mod:`lbmpy.methods.centeredcumulant.cumulant_transform`)
-
-    Returns:
-        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
-    """
-
-    one = sp.Integer(1)
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
-
-    assert len(cumulant_to_rr_dict) == len(
-        stencil), "The number of moments has to be equal to the number of stencil entries"
-    dim = len(stencil[0])
-
-    # CQC
-    cqc = DensityVelocityComputation(stencil, True, force_model=force_model)
-    density_symbol = cqc.zeroth_order_moment_symbol
-    velocity_symbols = cqc.first_order_moment_symbols
-
-    #   Equilibrium Values
-    higher_order_polynomials = list(filter(lambda x: get_order(x) > 1, cumulant_to_rr_dict.keys()))
-
-    #   Lower Order Equilibria
-    cumulants_to_relaxation_info_dict = {one: RelaxationInfo(density_symbol, cumulant_to_rr_dict[one])}
-    for d in MOMENT_SYMBOLS[:dim]:
-        cumulants_to_relaxation_info_dict[d] = RelaxationInfo(0, cumulant_to_rr_dict[d])
-
-    #   Polynomial Cumulant Equilibria
-    polynomial_equilibria = get_cumulants_of_continuous_maxwellian_equilibrium(
-        higher_order_polynomials, dim, rho=density_symbol, u=velocity_symbols, c_s_sq=c_s_sq, order=equilibrium_order)
-    polynomial_equilibria = [density_symbol * v for v in polynomial_equilibria]
-
-    for i, c in enumerate(higher_order_polynomials):
-        cumulants_to_relaxation_info_dict[c] = RelaxationInfo(polynomial_equilibria[i], cumulant_to_rr_dict[c])
-
-    return CenteredCumulantBasedLbMethod(stencil, cumulants_to_relaxation_info_dict, cqc, force_model,
-                                         galilean_correction=galilean_correction,
-                                         central_moment_transform_class=central_moment_transform_class,
-                                         cumulant_transform_class=cumulant_transform_class)
-
-
-def create_with_polynomial_cumulants(stencil, relaxation_rates, cumulant_groups, **kwargs):
-    r"""Creates a cumulant lattice Boltzmann model based on a default polynomial set.
+def create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments=True, **kwargs):
+    r"""Creates a cumulant-based lattice Boltzmann method.
 
     Args:
-        stencil: nested tuple defining the discrete velocity space. See :func:`lbmpy.stencils.get_stencil`
+        stencil: instance of :class:`lbmpy.stencils.LBStencil`
         relaxation_rates: relaxation rates for each cumulant group. If None are provided a list of symbolic relaxation
                           rates is created and used. If only a list with one entry is provided this relaxation rate is
                           used for determine the viscosity of the simulation. All other cumulants are relaxed with one.
@@ -551,102 +542,178 @@ def create_with_polynomial_cumulants(stencil, relaxation_rates, cumulant_groups,
                           that the force is applied correctly to the momentum groups
         cumulant_groups: Nested sequence of polynomial expressions defining the cumulants to be relaxed. All cumulants 
                          within one group are relaxed with the same relaxation rate.
-        kwargs: See :func:`create_centered_cumulant_model`
+        conserved_moments: If lower order moments are conserved or not.
+        kwargs: See :func:`create_with_continuous_maxwellian_equilibrium`
 
     Returns:
-        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
+        :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
     """
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
-
-    cumulant_to_rr_dict = OrderedDict()
-    rr_iter = iter(relaxation_rates)
-    for group in cumulant_groups:
-        if all(get_order(c) <= 1 for c in group):
-            for cumulant in group:
-                cumulant_to_rr_dict[cumulant] = 0
-        else:
-            try:
-                rr = next(rr_iter)
-                for cumulant in group:
-                    cumulant_to_rr_dict[cumulant] = rr
-            except StopIteration:
-                raise ValueError('Not enough relaxation rates specified.')
+    cumulant_to_rr_dict = _get_relaxation_info_dict(relaxation_rates, cumulant_groups, stencil.D, conserved_moments)
 
-    return create_centered_cumulant_model(stencil, cumulant_to_rr_dict, **kwargs)
+    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)
 
+    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.CUMULANTS))
 
-def create_with_monomial_cumulants(stencil, relaxation_rates, **kwargs):
-    r"""Creates a cumulant lattice Boltzmann model based on a default polinomial set.
+    if kwargs['collision_space_info'].collision_space != CollisionSpace.CUMULANTS:
+        raise ValueError("Cumulant-based methods can only be derived in cumulant space.")
+
+    return create_with_continuous_maxwellian_equilibrium(stencil, cumulant_to_rr_dict, 
+                                                         compressible=True, delta_equilibrium=False,
+                                                         **kwargs)
+
+
+def create_with_monomial_cumulants(stencil, relaxation_rates, conserved_moments=True, **kwargs):
+    r"""Creates a cumulant lattice Boltzmann model using the given stencil's canonical monomial cumulants.
 
     Args:
-        stencil: nested tuple defining the discrete velocity space. See :func:`lbmpy.stencils.get_stencil`
+        stencil: instance of :class:`lbmpy.stencils.LBStencil`
         relaxation_rates: relaxation rates for each cumulant group. If None are provided a list of symbolic relaxation
                           rates is created and used. If only a list with one entry is provided this relaxation rate is
                           used for determine the viscosity of the simulation. All other cumulants are relaxed with one.
                           If a cumulant force model is provided the first order cumulants are relaxed with two to ensure
                           that the force is applied correctly to the momentum groups
-        kwargs: See :func:`create_centered_cumulant_model`
+        conserved_moments: If lower order moments are conserved or not.
+        kwargs: See :func:`create_cumulant`
 
     Returns:
-        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
+        :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
     """
-
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
-
-    dim = len(stencil[0])
-
     # Get monomial moments
     cumulants = get_default_moment_set_for_stencil(stencil)
-
-    if len(relaxation_rates) == 1:
-        r_rates = []
-        for c in cumulants:
-            order = get_order(c)
-            if order <= 1:
-                #   Conserved moments
-                continue
-            elif is_shear_moment(c, dim):
-                #   Viscosity-governing moments
-                r_rates.append(relaxation_rates[0])
-            else:
-                #   The rest
-                r_rates.append(1)
-    else:
-        r_rates = relaxation_rates
-
     cumulant_groups = [(c,) for c in cumulants]
+    return create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments, **kwargs)
 
-    return create_with_polynomial_cumulants(stencil, r_rates, cumulant_groups, **kwargs)
 
+def create_with_default_polynomial_cumulants(stencil, relaxation_rates, conserved_moments=True, **kwargs):
+    r"""Creates a cumulant lattice Boltzmann model based on the default polynomial set of :cite:`geier2015`.
 
-def create_with_default_polynomial_cumulants(stencil, relaxation_rates, **kwargs):
-    r"""Creates a cumulant lattice Boltzmann model based on a default polynomial set.
-
-    Args: See :func:`create_with_polynomial_cumulants`.
+    Args: See :func:`create_cumulant`.
 
     Returns:
-        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
+        :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
     """
-
-    if isinstance(stencil, str):
-        stencil = get_stencil(stencil)
-
     # Get polynomial groups
-    cumulant_groups = get_default_polynomial_cumulants_for_stencil(stencil)
+    cumulant_groups = cascaded_moment_sets_literature(stencil)
+    return create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments, **kwargs)
 
-    if len(relaxation_rates) == 1:
-        r_rates = [relaxation_rates[0]]     # For correct viscosity
-        r_rates += [1] * (len(cumulant_groups) - 3)
-    else:
-        assert len(relaxation_rates) >= len(cumulant_groups) - 2, \
-            f"Number of relaxation rates must at least match the number of non-conserved cumulant groups. " \
-            f"For this stencil we have {len(cumulant_groups) - 2} such cumulant groups"
-        r_rates = relaxation_rates
 
-    return create_with_polynomial_cumulants(stencil, r_rates, cumulant_groups, **kwargs)
+def _get_relaxation_info_dict(relaxation_rates, nested_moments, dim,
+                              conserved_moments=True, relaxation_rates_modifier=None):
+    r"""Creates a dictionary where each moment is mapped to a relaxation rate.
 
+    Args:
+        relaxation_rates: list of relaxation rates which should be used. This can also be a function which
+                          takes a moment group in the list of nested moments and returns a list of relaxation rates.
+                          This list has to have the length of the moment group and is then used for the moments
+                          in the moment group.
+        nested_moments: list of lists containing the moments.
+        dim: dimension
+        conserved_moments: If lower order moments are conserved or not.
+    """
+    result = OrderedDict()
+
+    if callable(relaxation_rates):
+        for group in nested_moments:
+            rr = iter(relaxation_rates(group))
+            for moment in group:
+                result[moment] = next(rr)
+        if relaxation_rates_modifier is not None:
+            for key in result:
+                result[key] *= relaxation_rates_modifier
+        return result
+
+    number_of_moments = 0
+    shear_moments = 0
+    bulk_moments = 0
+
+    for group in nested_moments:
+        for moment in group:
+            number_of_moments += 1
+            if is_shear_moment(moment, dim):
+                shear_moments += 1
+            if is_bulk_moment(moment, dim):
+                bulk_moments += 1
+
+    # if only one relaxation rate is specified it is used as the shear relaxation rate
+    if len(relaxation_rates) == 1:
+        for group in nested_moments:
+            for moment in group:
+                if conserved_moments:
+                    if get_order(moment) <= 1:
+                        result[moment] = 0.0
+                    elif is_shear_moment(moment, dim):
+                        result[moment] = relaxation_rates[0]
+                    else:
+                        result[moment] = 1.0
+                else:
+                    if is_shear_moment(moment, dim) or get_order(moment) <= 1:
+                        result[moment] = relaxation_rates[0]
+                    else:
+                        result[moment] = 1.0
+
+    # if relaxation rate for each moment is specified they are all used
+    if len(relaxation_rates) == number_of_moments:
+        rr_iter = iter(relaxation_rates)
+        for group in nested_moments:
+            for moment in group:
+                rr = next(rr_iter)
+                result[moment] = rr
+
+    # Fallback case, relaxes each group with the same relaxation rate and separates shear and bulk moments
+    next_rr = True
+    if len(relaxation_rates) != 1 and len(relaxation_rates) != number_of_moments:
+        try:
+            rr_iter = iter(relaxation_rates)
+            if shear_moments > 0:
+                shear_rr = next(rr_iter)
+            if bulk_moments > 0:
+                bulk_rr = next(rr_iter)
+            for group in nested_moments:
+                if next_rr:
+                    rr = next(rr_iter)
+                next_rr = False
+                for moment in group:
+                    if conserved_moments:
+                        if get_order(moment) <= 1:
+                            result[moment] = 0.0
+                        elif is_shear_moment(moment, dim):
+                            result[moment] = shear_rr
+                        elif is_bulk_moment(moment, dim):
+                            result[moment] = bulk_rr
+                        else:
+                            next_rr = True
+                            result[moment] = rr
+                    else:
+                        if is_shear_moment(moment, dim) or get_order(moment) <= 1:
+                            result[moment] = shear_rr
+                        elif is_bulk_moment(moment, dim):
+                            result[moment] = bulk_rr
+                        else:
+                            next_rr = True
+                            result[moment] = rr
+
+        except StopIteration:
+            raise ValueError("Not enough relaxation rates are specified. You can either specify one relaxation rate, "
+                             "which is used as the shear relaxation rate. In this case, conserved moments are "
+                             "relaxed with 0, and higher-order moments are relaxed with 1. Another "
+                             "possibility is to specify a relaxation rate for shear, bulk and one for each order "
+                             "moment. In this case, conserved moments are also "
+                             "relaxed with 0. The last possibility is to specify a relaxation rate for each moment, "
+                             "including conserved moments")
+    if relaxation_rates_modifier is not None:
+        for key in result:
+            result[key] *= relaxation_rates_modifier
+    return result
+
+
+def check_and_set_mrt_space(default, **kwargs):
+    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(default))
+
+    if kwargs['collision_space_info'].collision_space not in (CollisionSpace.RAW_MOMENTS, CollisionSpace.POPULATIONS):
+        raise ValueError("Raw moment-based methods can only be derived in population or raw moment space.")
 
 # ----------------------------------------- Comparison view for notebooks ----------------------------------------------
 
@@ -666,10 +733,10 @@ def compare_moment_based_lb_methods(reference, other, show_deviations_only=False
         if show_deviations_only and diff == 0:
             pass
         else:
-            table.append(["$%s$" % (sp.latex(moment), ),
-                          "$%s$" % (sp.latex(reference_value), ),
-                          "$%s$" % (sp.latex(other_value), ),
-                          "$%s$" % (sp.latex(diff),)])
+            table.append([f"${sp.latex(moment)}$",
+                          f"${sp.latex(reference_value)}$",
+                          f"${sp.latex(other_value)}$",
+                          f"${sp.latex(diff)}$"])
 
     only_in_ref = reference_moments - other_moments
     if only_in_ref:
@@ -677,8 +744,8 @@ def compare_moment_based_lb_methods(reference, other, show_deviations_only=False
         table.append(['Only in Ref', 'value', '', ''])
         for moment in only_in_ref:
             val = reference.relaxation_info_dict[moment].equilibrium_value
-            table.append(["$%s$" % (sp.latex(moment),),
-                          "$%s$" % (sp.latex(val),),
+            table.append([f"${sp.latex(moment)}$",
+                          f"${sp.latex(val)}$",
                           " ", " "])
 
     only_in_other = other_moments - reference_moments
@@ -687,9 +754,9 @@ def compare_moment_based_lb_methods(reference, other, show_deviations_only=False
         table.append(['Only in Other', '', 'value', ''])
         for moment in only_in_other:
             val = other.relaxation_info_dict[moment].equilibrium_value
-            table.append(["$%s$" % (sp.latex(moment),),
+            table.append([f"${sp.latex(moment)}$",
                           " ",
-                          "$%s$" % (sp.latex(val),),
+                          f"${sp.latex(val)}$",
                           " "])
 
     table_display = ipy_table.make_table(table)
@@ -697,3 +764,15 @@ def compare_moment_based_lb_methods(reference, other, show_deviations_only=False
         for col in range(4):
             ipy_table.set_cell_style(row_idx, col, color='#bbbbbb')
     return table_display
+
+# ----------------------------------------- Deprecation of Maxwellian Moments -----------------------------------------
+
+
+def _deprecate_maxwellian_moments(continuous_equilibrium, kwargs):
+    if 'maxwellian_moments' in kwargs:
+        warn("Argument 'maxwellian_moments' is deprecated and will be removed by version 0.5."
+             "Use `continuous_equilibrium` instead.",
+             stacklevel=2)
+        return kwargs['maxwellian_moments']
+    else:
+        return continuous_equilibrium

src/lbmpy/methods/cumulantbased/__init__.py

+from .cumulantbasedmethod import CumulantBasedLbMethod
+from .galilean_correction import add_galilean_correction
+from .fourth_order_correction import add_fourth_order_correction
+
+__all__ = ['CumulantBasedLbMethod', 'add_galilean_correction', 'add_fourth_order_correction']

src/lbmpy/methods/cumulantbased/cumulant_simplifications.py

+import sympy as sp
+from pystencils.simp.subexpression_insertion import insert_subexpressions
+
+from warnings import warn
+
+
+def insert_logs(ac, **kwargs):
+    def callback(exp):
+        rhs = exp.rhs
+        logs = rhs.atoms(sp.log)
+        return len(logs) > 0
+
+    return insert_subexpressions(ac, callback, **kwargs)
+
+
+def insert_log_products(ac, **kwargs):
+    def callback(asm):
+        rhs = asm.rhs
+        if rhs.find(sp.log):
+            return True
+        return False
+
+    return insert_subexpressions(ac, callback, **kwargs)
+
+
+def expand_post_collision_central_moments(ac):
+    if 'post_collision_monomial_central_moments' in ac.simplification_hints:
+        subexpr_dict = ac.subexpressions_dict
+        cm_symbols = ac.simplification_hints['post_collision_monomial_central_moments']
+        for s in cm_symbols:
+            if s in subexpr_dict:
+                subexpr_dict[s] = subexpr_dict[s].expand()
+        ac = ac.copy()
+        ac.set_sub_expressions_from_dict(subexpr_dict)
+    return ac
+
+
+def check_for_logarithms(ac):
+    logs = ac.atoms(sp.log)
+    if len(logs) > 0:
+        warn("""There are logarithms remaining in your cumulant-based collision operator!
+                This will let your kernel's performance and numerical accuracy deterioate severly.
+                Either you have disabled simplification, or it unexpectedly failed.
+                If the presence of logarithms is intended, please inspect the kernel to make sure
+                if this warning can be ignored.
+                Otherwise, if setting `simplification='auto'` in your optimization config does not resolve
+                the problem, try a different parametrization, or contact the developers.""")

src/lbmpy/methods/cumulantbased/cumulantbasedmethod.py

+import sympy as sp
+from collections import OrderedDict
+from typing import Set
+from warnings import filterwarnings
+
+from pystencils import Assignment, AssignmentCollection
+from pystencils.sympyextensions import is_constant
+from pystencils.simp import apply_to_all_assignments
+
+from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
+from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
+
+from lbmpy.moments import moment_matrix, MOMENT_SYMBOLS, statistical_quantity_symbol
+
+from lbmpy.moment_transforms import (
+    PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT, POST_COLLISION_MONOMIAL_CENTRAL_MOMENT,
+    CentralMomentsToCumulantsByGeneratingFunc,
+    BinomialChimeraTransform)
+
+
+class CumulantBasedLbMethod(AbstractLbMethod):
+    """
+    This class implements cumulant-based lattice boltzmann methods which relax all the non-conserved quantities
+    as either monomial or polynomial cumulants. It is mostly inspired by the work presented in :cite:`geier2015`.
+
+    This method is implemented modularly as the transformation from populations to central moments to cumulants
+    is governed by subclasses of :class:`lbmpy.moment_transforms.AbstractMomentTransform`
+    which can be specified by constructor argument. This allows the selection of the most efficient transformation
+    for a given setup.
+
+    Parameters:
+        stencil: see :class:`lbmpy.stencils.LBStencil`
+        equilibrium: Instance of :class:`lbmpy.equilibrium.AbstractEquilibrium`, defining the equilibrium distribution
+                     used by this method.
+        relaxation_dict: a dictionary mapping cumulants in either tuple or polynomial formulation
+                         to their relaxation rate.
+        conserved_quantity_computation: instance of :class:`lbmpy.methods.AbstractConservedQuantityComputation`.
+                                        This determines how conserved quantities are computed, and defines
+                                        the symbols used in the equilibrium moments like e.g. density and velocity.
+        force_model: Instance of :class:`lbmpy.forcemodels.AbstractForceModel`, or None if no forcing terms are required
+        zero_centered: Determines the PDF storage format, regular or centered around the equilibrium's
+                       background distribution.
+        central_moment_transform_class: transformation class to transform PDFs to central moment space (subclass of 
+                                        :class:`lbmpy.moment_transforms.AbstractCentralMomentTransform`)
+        cumulant_transform_class: transform class to get from the central moment space to the cumulant space
+    """
+
+    def __init__(self, stencil, equilibrium, relaxation_dict,
+                 conserved_quantity_computation=None,
+                 force_model=None, zero_centered=False, fraction_field=None,
+                 central_moment_transform_class=BinomialChimeraTransform,
+                 cumulant_transform_class=CentralMomentsToCumulantsByGeneratingFunc):
+        assert isinstance(conserved_quantity_computation,
+                          AbstractConservedQuantityComputation)
+        super(CumulantBasedLbMethod, self).__init__(stencil)
+
+        if force_model is not None:
+            if not force_model.has_symmetric_central_moment_forcing:
+                raise ValueError(f"Force model {force_model} does not offer symmetric central moment forcing.")
+
+        self._equilibrium = equilibrium
+        self._relaxation_dict = OrderedDict(relaxation_dict)
+        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
+
+    @property
+    def force_model(self):
+        """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.
+        Beware: Changes to this dictionary are not reflected in the method. For changing relaxation rates,
+        use `relaxation_rate_dict` instead."""
+        return OrderedDict({m: RelaxationInfo(v, rr)
+                            for (m, rr), v in zip(self._relaxation_dict.items(), self.cumulant_equilibrium_values)})
+
+    @property
+    def conserved_quantity_computation(self):
+        """Returns an instance of class :class:`lbmpy.methods.AbstractConservedQuantityComputation`"""
+        return self._cqc
+
+    @property
+    def equilibrium_distribution(self):
+        """Returns this method's equilibrium distribution (see :class:`lbmpy.equilibrium.AbstractEquilibrium`"""
+        return self._equilibrium
+
+    @property
+    def central_moment_transform_class(self):
+        """The transform class (subclass of :class:`lbmpy.moment_transforms.AbstractCentralMomentTransform` defining the
+        transformation of populations to central moment space."""
+        return self._central_moment_transform_class
+
+    @property
+    def cumulant_transform_class(self):
+        """The transform class defining the transform from central moment to cumulant space."""
+        return self._cumulant_transform_class
+
+    @property
+    def cumulants(self):
+        """Cumulants relaxed by this method."""
+        return tuple(self._relaxation_dict.keys())
+
+    @property
+    def cumulant_equilibrium_values(self):
+        """Equilibrium values of this method's :attr:`cumulants`."""
+        return self._equilibrium.cumulants(self.cumulants, rescale=True)
+
+    @property
+    def relaxation_rates(self):
+        """Relaxation rates for this method's :attr:`cumulants`."""
+        return tuple(self._relaxation_dict.values())
+
+    @property
+    def relaxation_rate_dict(self):
+        """Dictionary mapping cumulants to relaxation rates. Changes are reflected by the method."""
+        return self._relaxation_dict
+
+    @property
+    def zeroth_order_equilibrium_moment_symbol(self, ):
+        """Returns a symbol referring to the zeroth-order moment of this method's equilibrium distribution,
+        which is the area under its curve
+        (see :attr:`lbmpy.equilibrium.AbstractEquilibrium.zeroth_order_moment_symbol`)."""
+        return self._equilibrium.zeroth_order_moment_symbol
+
+    @property
+    def first_order_equilibrium_moment_symbols(self, ):
+        """Returns a vector of symbols referring to the first-order moment of this method's equilibrium distribution,
+        which is its mean value. (see :attr:`lbmpy.equilibrium.AbstractEquilibrium.first_order_moment_symbols`)."""
+        return self._equilibrium.first_order_moment_symbols
+
+    def set_zeroth_moment_relaxation_rate(self, relaxation_rate):
+        e = sp.Rational(1, 1)
+        self._relaxation_dict[e] = relaxation_rate
+
+    def set_first_moment_relaxation_rate(self, relaxation_rate):
+        for e in MOMENT_SYMBOLS[:self.dim]:
+            assert e in self._relaxation_dict, \
+                "First cumulants are not relaxed separately by this method"
+        for e in MOMENT_SYMBOLS[:self.dim]:
+            self._relaxation_dict[e] = relaxation_rate
+
+    def set_conserved_moments_relaxation_rate(self, relaxation_rate):
+        self.set_zeroth_moment_relaxation_rate(relaxation_rate)
+        self.set_first_moment_relaxation_rate(relaxation_rate)
+
+    def set_force_model(self, force_model):
+        if not force_model.has_symmetric_central_moment_forcing:
+            raise ValueError("Given force model does not support symmetric central moment forcing.")
+        self._force_model = force_model
+
+    def _repr_html_(self):
+        def stylized_bool(b):
+            return "✓" if b else "✗"
+
+        html = f"""
+        <table style="border:none; width: 100%">
+            <tr>
+                <th colspan="3" style="text-align: left">
+                    Cumulant-Based Method
+                </th>
+                <td>Stencil: {self.stencil.name}</td>
+                <td>Zero-Centered Storage: {stylized_bool(self._zero_centered)}</td>
+                <td>Force Model: {"None" if self._force_model is None else type(self._force_model).__name__}</td>
+            </tr>
+        </table>
+        """
+
+        html += self._equilibrium._repr_html_()
+
+        html += """
+        <table style="border:none; width: 100%">
+            <tr> <th colspan="3" style="text-align: left"> Relaxation Info </th> </tr>
+            <tr>
+                <th>Cumulant</th>
+                <th>Eq. Value </th>
+                <th>Relaxation Rate</th>
+            </tr>
+        """
+
+        for cumulant, (eq_value, rr) in self.relaxation_info_dict.items():
+            vals = {
+                'rr': sp.latex(rr),
+                'cumulant': sp.latex(cumulant),
+                'eq_value': sp.latex(eq_value),
+                'nb': 'style="border:none"',
+            }
+            html += """<tr {nb}>
+                            <td {nb}>${cumulant}$</td>
+                            <td {nb}>${eq_value}$</td>
+                            <td {nb}>${rr}$</td>
+                         </tr>\n""".format(**vals)
+
+        html += "</table>"
+        return html
+
+    #   ----------------------- Overridden Abstract Members --------------------------
+
+    @property
+    def weights(self):
+        """Returns a sequence of weights, one for each lattice direction"""
+        if self._weights is None:
+            self._weights = self._compute_weights()
+        return self._weights
+
+    def override_weights(self, weights):
+        assert len(weights) == len(self.stencil)
+        self._weights = weights
+
+    def get_equilibrium(self, conserved_quantity_equations: AssignmentCollection = None, subexpressions: bool = False,
+                        pre_simplification: bool = False, keep_cqc_subexpressions: bool = True,
+                        include_force_terms: bool = False) -> LbmCollisionRule:
+        """Returns equation collection, to compute equilibrium values.
+        The equations have the post collision symbols as left-hand sides and are
+        functions of the conserved quantities
+
+        Args:
+            conserved_quantity_equations: equations to compute conserved quantities.
+            subexpressions: if set to false all subexpressions of the equilibrium assignments are plugged
+                            into the main assignments
+            pre_simplification: with or without pre_simplifications for the calculation of the collision
+            keep_cqc_subexpressions: if equilibrium is returned without subexpressions keep_cqc_subexpressions
+                                     determines if also subexpressions to calculate conserved quantities should be
+                                     plugged into the main assignments
+            include_force_terms: if set to True the equilibrium is shifted by forcing terms coming from the force model
+                                 of the method.
+        """
+        r_info_dict = OrderedDict({c: RelaxationInfo(info.equilibrium_value, sp.Integer(1))
+                                   for c, info in self.relaxation_info_dict.items()})
+        ac = self._centered_cumulant_collision_rule(cumulant_to_relaxation_info_dict=r_info_dict,
+                                                    conserved_quantity_equations=conserved_quantity_equations,
+                                                    pre_simplification=pre_simplification,
+                                                    include_force_terms=include_force_terms,
+                                                    symbolic_relaxation_rates=False)
+
+        expand_all_assignments = apply_to_all_assignments(sp.expand)
+
+        if not subexpressions:
+            if keep_cqc_subexpressions:
+                bs = self._bound_symbols_cqc(conserved_quantity_equations)
+                ac = expand_all_assignments(ac.new_without_subexpressions(subexpressions_to_keep=bs))
+                return ac.new_without_unused_subexpressions()
+            else:
+                ac = expand_all_assignments(ac.new_without_subexpressions())
+                return ac.new_without_unused_subexpressions()
+        else:
+            return ac.new_without_unused_subexpressions()
+
+    def get_equilibrium_terms(self) -> sp.Matrix:
+        equilibrium = self.get_equilibrium()
+        return sp.Matrix([eq.rhs for eq in equilibrium.main_assignments])
+
+    def get_collision_rule(self, conserved_quantity_equations: AssignmentCollection = None,
+                           pre_simplification: bool = False) -> AssignmentCollection:
+        """Returns an LbmCollisionRule i.e. an equation collection with a reference to the method.
+        This collision rule defines the collision operator."""
+        return self._centered_cumulant_collision_rule(cumulant_to_relaxation_info_dict=self.relaxation_info_dict,
+                                                      conserved_quantity_equations=conserved_quantity_equations,
+                                                      pre_simplification=pre_simplification,
+                                                      include_force_terms=True, symbolic_relaxation_rates=True)
+
+    #   ------------------------------- Internals --------------------------------------------
+
+    def _bound_symbols_cqc(self, conserved_quantity_equations: AssignmentCollection = None) -> Set[sp.Symbol]:
+        f = self.pre_collision_pdf_symbols
+        cqe = conserved_quantity_equations
+
+        if cqe is None:
+            cqe = self._cqc.equilibrium_input_equations_from_pdfs(f, False)
+
+        return cqe.bound_symbols
+
+    def _compute_weights(self):
+        bg = self.equilibrium_distribution.background_distribution
+        assert bg is not None, "Could not compute weights, since no background distribution is given."
+        if bg.discrete_populations is not None:
+            #   Compute lattice weights as the discrete populations of the background distribution ...
+            weights = bg.discrete_populations
+        else:
+            #   or, if those are not available, by moment matching.
+            moments = self.cumulants
+            mm_inv = moment_matrix(moments, self.stencil).inv()
+            bg_moments = bg.moments(moments)
+            weights = (mm_inv * sp.Matrix(bg_moments)).expand()
+
+        for w in weights:
+            assert is_constant(w)
+
+        return [w for w in weights]
+
+    def _centered_cumulant_collision_rule(self, cumulant_to_relaxation_info_dict: OrderedDict,
+                                          conserved_quantity_equations: AssignmentCollection = None,
+                                          pre_simplification: bool = False,
+                                          include_force_terms: bool = False,
+                                          symbolic_relaxation_rates: bool = False) -> LbmCollisionRule:
+
+        # Filter out JobLib warnings. They are not usefull for use:
+        # https://github.com/joblib/joblib/issues/683
+        filterwarnings("ignore", message="Persisting input arguments took")
+
+        stencil = self.stencil
+        f = self.pre_collision_pdf_symbols
+        density = self.zeroth_order_equilibrium_moment_symbol
+        velocity = self.first_order_equilibrium_moment_symbols
+        cqe = conserved_quantity_equations
+
+        polynomial_cumulants = self.cumulants
+
+        rrs = [cumulant_to_relaxation_info_dict[c].relaxation_rate for c in polynomial_cumulants]
+        if symbolic_relaxation_rates:
+            subexpressions_relaxation_rates, d = self._generate_symbolic_relaxation_matrix(relaxation_rates=rrs)
+        else:
+            subexpressions_relaxation_rates = []
+            d = sp.zeros(len(rrs))
+            for i, w in enumerate(rrs):
+                # note that 0.0 is converted to sp.Zero here. It is not possible to prevent this.
+                d[i, i] = w
+
+        if cqe is None:
+            cqe = self._cqc.equilibrium_input_equations_from_pdfs(f, False)
+
+        forcing_subexpressions = AssignmentCollection([])
+        if self._force_model is not None:
+            forcing_subexpressions = AssignmentCollection(self._force_model.subs_dict_force)
+        else:
+            include_force_terms = False
+
+        #   See if a background shift is necessary
+        if self._zero_centered:
+            assert not self._equilibrium.deviation_only
+            background_distribution = self._equilibrium.background_distribution
+            assert background_distribution is not None
+        else:
+            background_distribution = None
+
+        #   1) Get Forward and Backward Transformations between central moment and cumulant space,
+        #      and find required central moments
+        k_to_c_transform = self._cumulant_transform_class(stencil, polynomial_cumulants, density, velocity)
+        k_to_c_eqs = k_to_c_transform.forward_transform(simplification=pre_simplification)
+        c_post_to_k_post_eqs = k_to_c_transform.backward_transform(simplification=pre_simplification)
+
+        C_pre = k_to_c_transform.pre_collision_symbols
+        C_post = k_to_c_transform.post_collision_symbols
+        central_moments = k_to_c_transform.required_central_moments
+
+        #   2) Get Forward Transformation from PDFs to central moments
+        pdfs_to_k_transform = self._central_moment_transform_class(
+            stencil, None, density, velocity, moment_exponents=central_moments, conserved_quantity_equations=cqe,
+            background_distribution=background_distribution)
+        pdfs_to_k_eqs = pdfs_to_k_transform.forward_transform(
+            f, simplification=pre_simplification, return_monomials=True)
+
+        #   3) Symmetric forcing
+        if include_force_terms:
+            force_before, force_after = self._force_model.symmetric_central_moment_forcing(self, central_moments)
+            k_asms_dict = pdfs_to_k_eqs.main_assignments_dict
+            for cm_exp, kappa_f in zip(central_moments, force_before):
+                cm_symb = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT, cm_exp)
+                k_asms_dict[cm_symb] += kappa_f
+            pdfs_to_k_eqs.set_main_assignments_from_dict(k_asms_dict)
+
+            k_post_asms_dict = c_post_to_k_post_eqs.main_assignments_dict
+            for cm_exp, kappa_f in zip(central_moments, force_after):
+                cm_symb = statistical_quantity_symbol(POST_COLLISION_MONOMIAL_CENTRAL_MOMENT, cm_exp)
+                k_post_asms_dict[cm_symb] += kappa_f
+            c_post_to_k_post_eqs.set_main_assignments_from_dict(k_post_asms_dict)
+
+        #   4) Add relaxation rules for polynomial cumulants
+        C_eq = sp.Matrix(self.cumulant_equilibrium_values)
+
+        C_pre_vec = sp.Matrix(C_pre)
+        collision_rule = C_pre_vec + d @ (C_eq - C_pre_vec)
+        cumulant_collision_eqs = [Assignment(lhs, rhs) for lhs, rhs in zip(C_post, collision_rule)]
+        cumulant_collision_eqs = AssignmentCollection(cumulant_collision_eqs)
+
+        #   5) Get backward transformation from central moments to PDFs
+        d = self.post_collision_pdf_symbols
+        k_post_to_pdfs_eqs = pdfs_to_k_transform.backward_transform(
+            d, simplification=pre_simplification, start_from_monomials=True)
+
+        #   6) That's all. Now, put it all together.
+        all_acs = [] if pdfs_to_k_transform.absorbs_conserved_quantity_equations else [cqe]
+        subexpressions_relaxation_rates = AssignmentCollection(subexpressions_relaxation_rates)
+        all_acs += [subexpressions_relaxation_rates, forcing_subexpressions, pdfs_to_k_eqs, k_to_c_eqs,
+                    cumulant_collision_eqs, c_post_to_k_post_eqs]
+        subexpressions = [ac.all_assignments for ac in all_acs]
+        subexpressions += k_post_to_pdfs_eqs.subexpressions
+        main_assignments = k_post_to_pdfs_eqs.main_assignments
+
+        simplification_hints = cqe.simplification_hints.copy()
+        simplification_hints.update(self._cqc.defined_symbols())
+        simplification_hints['relaxation_rates'] = [rr for rr in self.relaxation_rates]
+        simplification_hints['post_collision_monomial_central_moments'] = \
+            pdfs_to_k_transform.post_collision_monomial_symbols
+
+        #   Aaaaaand we're done.
+        return LbmCollisionRule(self, main_assignments, subexpressions, simplification_hints)

src/lbmpy/methods/cumulantbased/fourth_order_correction.py

+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, AssignmentCollection
+from .cumulantbasedmethod import CumulantBasedLbMethod
+
+X, Y, Z = MOMENT_SYMBOLS
+CORRECTED_FOURTH_ORDER_POLYNOMIALS = [X ** 2 * Y ** 2 - 2 * X ** 2 * Z ** 2 + Y ** 2 * Z ** 2,
+                                      X ** 2 * Y ** 2 + X ** 2 * Z ** 2 - 2 * Y ** 2 * Z ** 2,
+                                      X ** 2 * Y ** 2 + X ** 2 * Z ** 2 + Y ** 2 * Z ** 2,
+                                      X ** 2 * Y * Z,
+                                      X * Y ** 2 * Z,
+                                      X * Y * Z ** 2]
+FOURTH_ORDER_CORRECTION_SYMBOLS = sp.symbols("corr_fourth_:6")
+FOURTH_ORDER_RELAXATION_RATE_SYMBOLS = sp.symbols("corr_rr_:10")
+
+
+def add_fourth_order_correction(collision_rule, shear_relaxation_rate, bulk_relaxation_rate, limiter):
+    """Adds the fourth order correction terms (:cite:`geier2017`, eq. 44-49) to a given polynomial D3Q27
+    cumulant collision rule."""
+    method = collision_rule.method
+
+    if not isinstance(method, CumulantBasedLbMethod) or method.stencil != LBStencil(Stencil.D3Q27):
+        raise ValueError("Fourth-order correction is only defined for D3Q27 cumulant methods.")
+
+    polynomials = method.cumulants
+    rho = method.zeroth_order_equilibrium_moment_symbol
+
+    if not (set(CORRECTED_FOURTH_ORDER_POLYNOMIALS) < set(polynomials)):
+        raise ValueError("Fourth order correction requires polynomial cumulants "
+                         f"{', '.join(CORRECTED_FOURTH_ORDER_POLYNOMIALS)} to be present")
+
+    #   Call
+    #   PC1 = (x^2 * y^2 - 2 * x^2 * z^2 + y^2 * z^2),
+    #   PC2 = (x^2 * y^2 + x^2 * z^2 - 2 * y^2 * z^2),
+    #   PC3 = (x^2 * y^2 + x^2 * z^2 + y^2 * z^2),
+    #   PC4 = (x^2 * y * z),
+    #   PC5 = (x * y^2 * z),
+    #   PC6 = (x * y * z^2)
+    poly_symbols = [sp.Symbol(f'{POST_COLLISION_CUMULANT}_{polynomials.index(poly)}')
+                    for poly in CORRECTED_FOURTH_ORDER_POLYNOMIALS]
+
+    a_symb, b_symb = sp.symbols("a_corr, b_corr")
+    a, b = get_optimal_additional_parameters(shear_relaxation_rate, bulk_relaxation_rate)
+    correction_terms = get_fourth_order_correction_terms(method.relaxation_rate_dict, rho, a_symb, b_symb)
+    optimal_parametrisation = get_optimal_parametrisation_with_limiters(collision_rule, shear_relaxation_rate,
+                                                                        bulk_relaxation_rate, limiter)
+
+    subexp_dict = collision_rule.subexpressions_dict
+    subexp_dict = {**subexp_dict,
+                   a_symb: a,
+                   b_symb: b,
+                   **correction_terms.subexpressions_dict,
+                   **optimal_parametrisation.subexpressions_dict,
+                   **correction_terms.main_assignments_dict,
+                   **optimal_parametrisation.main_assignments_dict}
+    for sym, cor in zip(poly_symbols, FOURTH_ORDER_CORRECTION_SYMBOLS):
+        subexp_dict[sym] += cor
+
+    collision_rule.set_sub_expressions_from_dict(subexp_dict)
+    collision_rule.topological_sort()
+
+    return collision_rule
+
+
+def get_optimal_additional_parameters(shear_relaxation_rate, bulk_relaxation_rate):
+    """
+    Calculates the optimal parametrization for the additional parameters provided in :cite:`geier2017`
+    Equations 115-116.
+    """
+
+    omega_1 = shear_relaxation_rate
+    omega_2 = bulk_relaxation_rate
+
+    omega_11 = omega_1 * omega_1
+    omega_12 = omega_1 * omega_2
+    omega_22 = omega_2 * omega_2
+
+    two = sp.Float(2)
+    three = sp.Float(3)
+    four = sp.Float(4)
+    nine = sp.Float(9)
+
+    denom_ab = (omega_1 - omega_2) * (omega_2 * (two + three * omega_1) - sp.Float(8) * omega_1)
+
+    a = (four * omega_11 + two * omega_12 * (omega_1 - sp.Float(6)) + omega_22 * (
+        omega_1 * (sp.Float(10) - three * omega_1) - four)) / denom_ab
+    b = (four * omega_12 * (nine * omega_1 - sp.Float(16)) - four * omega_11 - two * omega_22 * (
+        two + nine * omega_1 * (omega_1 - two))) / (three * denom_ab)
+
+    return a, b
+
+
+def get_fourth_order_correction_terms(rrate_dict, rho, a, b):
+    pc1, pc2, pc3, pc4, pc5, pc6 = CORRECTED_FOURTH_ORDER_POLYNOMIALS
+
+    omega_1 = rrate_dict[X * Y]
+    omega_2 = rrate_dict[X ** 2 + Y ** 2 + Z ** 2]
+
+    try:
+        omega_6 = rrate_dict[pc1]
+        assert omega_6 == rrate_dict[pc2], \
+            "Cumulants (x^2 - y^2) and (x^2 - z^2) must have the same relaxation rate"
+        omega_7 = rrate_dict[pc3]
+        omega_8 = rrate_dict[pc4]
+        assert omega_8 == rrate_dict[pc5] == rrate_dict[pc6]
+    except IndexError:
+        raise ValueError("For the fourth order correction, all six polynomial cumulants"
+                         "(x^2 * y^2 - 2 * x^2 * z^2 + y^2 * z^2),"
+                         "(x^2 * y^2 + x^2 * z^2 - 2 * y^2 * z^2),"
+                         "(x^2 * y^2 + x^2 * z^2 + y^2 * z^2),"
+                         "(x^2 * y * z), (x * y^2 * z) and (x * y * z^2) must be present!")
+
+    dxu, dyv, dzw, dxvdyu, dxwdzu, dywdzv = sp.symbols('Dx, Dy, Dz, DxvDyu, DxwDzu, DywDzv')
+    c_xy = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (1, 1, 0))
+    c_xz = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (1, 0, 1))
+    c_yz = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (0, 1, 1))
+    c_xx = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (2, 0, 0))
+    c_yy = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (0, 2, 0))
+    c_zz = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (0, 0, 2))
+
+    cor1, cor2, cor3, cor4, cor5, cor6 = FOURTH_ORDER_CORRECTION_SYMBOLS
+
+    one = sp.Float(1)
+    two = sp.Float(2)
+    three = sp.Float(3)
+
+    #   Derivative Approximations
+    subexpressions = [
+        Assignment(dxu, - omega_1 / (two * rho) * (two * c_xx - c_yy - c_zz)
+                   - omega_2 / (two * rho) * (c_xx + c_yy + c_zz - rho)),
+        Assignment(dyv, dxu + (three * omega_1) / (two * rho) * (c_xx - c_yy)),
+        Assignment(dzw, dxu + (three * omega_1) / (two * rho) * (c_xx - c_zz)),
+        Assignment(dxvdyu, - three * omega_1 / rho * c_xy),
+        Assignment(dxwdzu, - three * omega_1 / rho * c_xz),
+        Assignment(dywdzv, - three * omega_1 / rho * c_yz)]
+
+    one_half = sp.Rational(1, 2)
+    one_over_three = sp.Rational(1, 3)
+    two_over_three = sp.Rational(2, 3)
+    four_over_three = sp.Rational(4, 3)
+
+    #   Correction Terms
+    main_assignments = [
+        Assignment(cor1, two_over_three * (one / omega_1 - one_half) * omega_6 * a * rho * (dxu - two * dyv + dzw)),
+        Assignment(cor2, two_over_three * (one / omega_1 - one_half) * omega_6 * a * rho * (dxu + dyv - two * dzw)),
+        Assignment(cor3, - four_over_three * (one / omega_1 - one_half) * omega_7 * a * rho * (dxu + dyv + dzw)),
+        Assignment(cor4, - one_over_three * (one / omega_1 - one_half) * omega_8 * b * rho * dywdzv),
+        Assignment(cor5, - one_over_three * (one / omega_1 - one_half) * omega_8 * b * rho * dxwdzu),
+        Assignment(cor6, - one_over_three * (one / omega_1 - one_half) * omega_8 * b * rho * dxvdyu)]
+
+    return AssignmentCollection(main_assignments=main_assignments, subexpressions=subexpressions)
+
+
+def get_optimal_parametrisation_with_limiters(collision_rule, shear_relaxation_rate, bulk_relaxation_rate, limiter):
+    """
+    Calculates the optimal parametrization for the relaxation rates provided in :cite:`geier2017`
+    Equations 112-114.
+    """
+
+    # if omega numbers
+    # assert omega_1 != omega_2, "Relaxation rates associated with shear and bulk viscosity must not be identical."
+    # assert omega_1 > 7/4
+
+    omega_1 = FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[0]
+    omega_2 = FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[1]
+
+    non_limited_omegas = sp.symbols("non_limited_omega_3:6")
+    limited_omegas = sp.symbols("limited_omega_3:6")
+
+    omega_11 = omega_1 * omega_1
+    omega_12 = omega_1 * omega_2
+    omega_22 = omega_2 * omega_2
+
+    one = sp.Float(1)
+    two = sp.Float(2)
+    three = sp.Float(3)
+    five = sp.Float(5)
+    six = sp.Float(6)
+    seven = sp.Float(7)
+    eight = sp.Float(8)
+    nine = sp.Float(9)
+
+    omega_3 = (eight * (omega_1 - two) * (omega_2 * (three * omega_1 - one) - five * omega_1)) / (
+        eight * (five - two * omega_1) * omega_1 + omega_2 * (eight + omega_1 * (nine * omega_1 - sp.Float(26))))
+
+    omega_4 = (eight * (omega_1 - two) * (omega_1 + omega_2 * (three * omega_1 - seven))) / (
+        omega_2 * (sp.Float(56) - sp.Float(42) * omega_1 + nine * omega_11) - eight * omega_1)
+
+    omega_5 = (sp.Float(24) * (omega_1 - two) * (sp.Float(4) * omega_11 + omega_12 * (
+        sp.Float(18) - sp.Float(13) * omega_1) + omega_22 * (two + omega_1 * (
+            six * omega_1 - sp.Float(11))))) / (sp.Float(16) * omega_11 * (omega_1 - six) - two * omega_12 * (
+                sp.Float(216) + five * omega_1 * (nine * omega_1 - sp.Float(46))) + omega_22 * (omega_1 * (
+                    three * omega_1 - sp.Float(10)) * (sp.Float(15) * omega_1 - sp.Float(28)) - sp.Float(48)))
+
+    rho = collision_rule.method.zeroth_order_equilibrium_moment_symbol
+
+    # add limiters to improve stability
+    c_xyy = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (1, 2, 0))
+    c_xzz = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (1, 0, 2))
+    c_xyz = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (1, 1, 1))
+    abs_c_xyy_plus_xzz = sp.Abs(c_xyy + c_xzz)
+    abs_c_xyy_minus_xzz = sp.Abs(c_xyy - c_xzz)
+    abs_c_xyz = sp.Abs(c_xyz)
+
+    limited_omega_3 = non_limited_omegas[0] + ((one - non_limited_omegas[0]) * abs_c_xyy_plus_xzz) / \
+        (rho * limiter + abs_c_xyy_plus_xzz)
+    limited_omega_4 = non_limited_omegas[1] + ((one - non_limited_omegas[1]) * abs_c_xyy_minus_xzz) / \
+        (rho * limiter + abs_c_xyy_minus_xzz)
+    limited_omega_5 = non_limited_omegas[2] + ((one - non_limited_omegas[2]) * abs_c_xyz) / (rho * limiter + abs_c_xyz)
+
+    subexpressions = [
+        Assignment(non_limited_omegas[0], omega_3),
+        Assignment(non_limited_omegas[1], omega_4),
+        Assignment(non_limited_omegas[2], omega_5),
+        Assignment(limited_omegas[0], limited_omega_3),
+        Assignment(limited_omegas[1], limited_omega_4),
+        Assignment(limited_omegas[2], limited_omega_5)]
+
+    #   Correction Terms
+    main_assignments = [
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[0], shear_relaxation_rate),
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[1], bulk_relaxation_rate),
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[2], limited_omegas[0]),
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[3], limited_omegas[1]),
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[4], limited_omegas[2]),
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[5], one),
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[6], one),
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[7], one),
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[8], one),
+        Assignment(FOURTH_ORDER_RELAXATION_RATE_SYMBOLS[9], one),
+    ]
+
+    return AssignmentCollection(main_assignments=main_assignments, subexpressions=subexpressions)

lbmpy/methods/centeredcumulant/galilean_correction.py → src/lbmpy/methods/cumulantbased/galilean_correction.py

-from pystencils.simp.assignment_collection import AssignmentCollection
 import sympy as sp
-from pystencils import Assignment
 
-from lbmpy.moments import MOMENT_SYMBOLS
-from lbmpy.methods.centeredcumulant.centered_cumulants import statistical_quantity_symbol
-from lbmpy.methods.centeredcumulant.cumulant_transform import PRE_COLLISION_CUMULANT
+from pystencils import Assignment, AssignmentCollection
 
-x, y, z = MOMENT_SYMBOLS
-corrected_polynomials = [x**2 - y**2, x**2 - z**2, x**2 + y**2 + z**2]
+from lbmpy.stencils import Stencil, LBStencil
+from lbmpy.moments import MOMENT_SYMBOLS, statistical_quantity_symbol
+from lbmpy.moment_transforms import PRE_COLLISION_MONOMIAL_CUMULANT, POST_COLLISION_CUMULANT
+
+from .cumulantbasedmethod import CumulantBasedLbMethod
+
+X, Y, Z = MOMENT_SYMBOLS
+CORRECTED_POLYNOMIALS = [X**2 - Y**2, X**2 - Z**2, X**2 + Y**2 + Z**2]
+CORRECTION_SYMBOLS = sp.symbols("corr_:3")
 
 
 def contains_corrected_polynomials(polynomials):
-    return all(cp in polynomials for cp in corrected_polynomials)
+    return all(cp in polynomials for cp in CORRECTED_POLYNOMIALS)
+
+
+def add_galilean_correction(collision_rule):
+    """Adds the galilean correction terms (:cite:`geier2015`, eq. 58-63) to a given polynomial D3Q27
+    cumulant collision rule."""
+    method = collision_rule.method
+
+    if not isinstance(method, CumulantBasedLbMethod) or method.stencil != LBStencil(Stencil.D3Q27):
+        raise ValueError("Galilean correction is only defined for D3Q27 cumulant methods.")
 
+    polynomials = method.cumulants
+    rho = method.zeroth_order_equilibrium_moment_symbol
+    u = method.first_order_equilibrium_moment_symbols
+
+    if not (set(CORRECTED_POLYNOMIALS) < set(polynomials)):
+        raise ValueError("Galilean correction requires polynomial cumulants "
+                         f"{', '.join(CORRECTED_POLYNOMIALS)} to be present")
 
-def add_galilean_correction(poly_relaxation_eqs, polynomials, correction_terms):
     #   Call PC1 = (x^2 - y^2), PC2 = (x^2 - z^2), PC3 = (x^2 + y^2 + z^2)
-    try:
-        index_pc1 = polynomials.index(corrected_polynomials[0])
-        index_pc2 = polynomials.index(corrected_polynomials[1])
-        index_pc3 = polynomials.index(corrected_polynomials[2])
-    except ValueError:
-        raise ValueError("For the galilean correction, all three polynomial cumulants"
-                         + "(x^2 - y^2), (x^2 - z^2) and (x^2 + y^2 + z^2) need to be present!")
+    poly_symbols = [sp.Symbol(f'{POST_COLLISION_CUMULANT}_{polynomials.index(poly)}')
+                    for poly in CORRECTED_POLYNOMIALS]
+
+    correction_terms = get_galilean_correction_terms(method.relaxation_rate_dict, rho, u)
 
-    cor1 = correction_terms.main_assignments[0].lhs
-    cor2 = correction_terms.main_assignments[1].lhs
-    cor3 = correction_terms.main_assignments[2].lhs
+    subexp_dict = collision_rule.subexpressions_dict
+    subexp_dict = {**subexp_dict,
+                   **correction_terms.subexpressions_dict,
+                   **correction_terms.main_assignments_dict}
+    for sym, cor in zip(poly_symbols, CORRECTION_SYMBOLS):
+        subexp_dict[sym] += cor
 
-    poly_relaxation_eqs[index_pc1] += cor1
-    poly_relaxation_eqs[index_pc2] += cor2
-    poly_relaxation_eqs[index_pc3] += cor3
+    collision_rule.set_sub_expressions_from_dict(subexp_dict)
+    collision_rule.topological_sort()
 
-    return poly_relaxation_eqs
+    return collision_rule
 
 
-def get_galilean_correction_terms(cumulant_to_relaxation_info_dict, rho, u_vector,
-                                  pre_collision_cumulant_base=PRE_COLLISION_CUMULANT):
+def get_galilean_correction_terms(rrate_dict, rho, u_vector):
 
-    pc1 = corrected_polynomials[0]
-    pc2 = corrected_polynomials[1]
-    pc3 = corrected_polynomials[2]
+    pc1 = CORRECTED_POLYNOMIALS[0]
+    pc2 = CORRECTED_POLYNOMIALS[1]
+    pc3 = CORRECTED_POLYNOMIALS[2]
 
     try:
-        omega_1 = cumulant_to_relaxation_info_dict[pc1].relaxation_rate
-        assert omega_1 == cumulant_to_relaxation_info_dict[pc2].relaxation_rate, \
+        omega_1 = rrate_dict[pc1]
+        assert omega_1 == rrate_dict[pc2], \
             "Cumulants (x^2 - y^2) and (x^2 - z^2) must have the same relaxation rate"
-        omega_2 = cumulant_to_relaxation_info_dict[pc3].relaxation_rate
+        omega_2 = rrate_dict[pc3]
     except IndexError:
         raise ValueError("For the galilean correction, all three polynomial cumulants"
                          + "(x^2 - y^2), (x^2 - z^2) and (x^2 + y^2 + z^2) must be present!")
 
     dx, dy, dz = sp.symbols('Dx, Dy, Dz')
-    c_xx = statistical_quantity_symbol(pre_collision_cumulant_base, (2, 0, 0))
-    c_yy = statistical_quantity_symbol(pre_collision_cumulant_base, (0, 2, 0))
-    c_zz = statistical_quantity_symbol(pre_collision_cumulant_base, (0, 0, 2))
+    c_xx = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (2, 0, 0))
+    c_yy = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (0, 2, 0))
+    c_zz = statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CUMULANT, (0, 0, 2))
 
-    cor1, cor2, cor3 = sp.symbols("corr_:3")
+    cor1, cor2, cor3 = CORRECTION_SYMBOLS
 
     #   Derivative Approximations
     subexpressions = [