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
+from functools import reduce
+
+import sympy as sp
+
+from lbmpy.maxwellian_equilibrium import get_weights
+from lbmpy.equilibrium import ContinuousHydrodynamicMaxwellian, DiscreteHydrodynamicMaxwellian
+
+from lbmpy.methods.default_moment_sets import cascaded_moment_sets_literature
+
+from lbmpy.moment_transforms import CentralMomentsToCumulantsByGeneratingFunc
+
+from .conservedquantitycomputation import DensityVelocityComputation
+
+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, 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.enums import Stencil, CollisionSpace
+from lbmpy.stencils import LBStencil
+from pystencils.sympyextensions import common_denominator
+
+
+@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
+    (see :class:`lbmpy.equilibrium.DiscreteHydrodynamicMaxwellian`).
+
+    Internally, this method calls :func:`lbmpy.methods.create_from_equilibrium`.
+
+    Args:
+        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: 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:
+        Instance of a subclass of :class:`lbmpy.methods.AbstractLbMethod`.
+    """
+    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 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: 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: 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:
+        Instance of a subclass of :class:`lbmpy.methods.AbstractLbMethod`.
+    """
+    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 lattice Boltzmann method in either population, moment, or central moment space, from a given
+    discrete velocity set and equilibrium distribution. 
+
+    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: 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 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) == 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, 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: instance of :class:`lbmpy.stencils.LBStencil`
+        relaxation_rate: relaxation rate (inverse of the relaxation time)
+                        usually called :math:`\omega` in LBM literature
+        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
+                        used to compute the equilibrium moments
+
+    Returns:
+        :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
+    """
+    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 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_trt(stencil, relaxation_rate_even_moments, relaxation_rate_odd_moments,
+               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
+    have this problem.
+
+    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
+    """
+    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 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_trt_with_magic_number(stencil, relaxation_rate, magic_number=sp.Rational(3, 16), **kwargs):
+    r"""
+    Creates a two relaxation time (TRT) lattice Boltzmann method, where the relaxation time for odd moments is
+    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: 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.
+
+    Returns:
+        :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
+    """
+    rr_odd = relaxation_rate_from_magic_number(relaxation_rate, magic_number)
+    return create_trt(stencil, relaxation_rate_even_moments=relaxation_rate,
+                      relaxation_rate_odd_moments=rr_odd, **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: instance of :class:`lbmpy.stencils.LBStencil`
+        relaxation_rates: relaxation rates (inverse of the relaxation times) for each moment
+        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
+    """
+    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
+    check_and_set_mrt_space(CollisionSpace.RAW_MOMENTS)
+    moments = get_default_moment_set_for_stencil(stencil)
+    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_equilibrium(stencil, rr_dict, **kwargs)
+
+
+def create_trt_kbc(dim, shear_relaxation_rate, higher_order_relaxation_rate, method_name='KBC-N4',
+                   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"
+        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)
+
+    the_moment = MOMENT_SYMBOLS[:dim]
+
+    rho = [sp.Rational(1, 1)]
+    velocity = list(the_moment)
+
+    shear_tensor_off_diagonal = [product(t) for t in itertools.combinations(the_moment, 2)]
+    shear_tensor_diagonal = [m_i * m_i for m_i in the_moment]
+    shear_tensor_trace = sum(shear_tensor_diagonal)
+    shear_tensor_trace_free_diagonal = [dim * d - shear_tensor_trace for d in shear_tensor_diagonal]
+
+    poly_repr = exponents_to_polynomial_representations
+    energy_transport_tensor = list(poly_repr([a for a in moments_of_order(3, dim, True)
+                                              if 3 not in a]))
+
+    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
+
+    # 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]
+    t = [shear_tensor_trace]
+    q = energy_transport_tensor
+    if method_name == 'KBC-N1':
+        decomposition = [d, t + q + rest]
+    elif method_name == 'KBC-N2':
+        decomposition = [d + t, q + rest]
+    elif method_name == 'KBC-N3':
+        decomposition = [d + q, t + rest]
+    elif method_name == 'KBC-N4':
+        decomposition = [d + t + q, rest]
+    else:
+        raise ValueError("Unknown model. Supported models KBC-Nx where x in (1,2,3,4)")
+
+    omega_s, omega_h = shear_relaxation_rate, higher_order_relaxation_rate
+    shear_part, rest_part = decomposition
+
+    relaxation_rates = [omega_s] + \
+                       [omega_s] * len(velocity) + \
+                       [omega_s] * len(shear_part) + \
+                       [omega_h] * len(rest_part)
+
+    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 continuous_equilibrium:
+        return create_with_continuous_maxwellian_equilibrium(stencil, moment_to_rr, **kwargs)
+    else:
+        return create_with_discrete_maxwellian_equilibrium(stencil, moment_to_rr, **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_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: 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. 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.
+    """
+    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
+    check_and_set_mrt_space(CollisionSpace.RAW_MOMENTS)
+
+    if weighted:
+        weights = get_weights(stencil, sp.Rational(1, 3))
+    else:
+        weights = None
+
+    if not nested_moments:
+        moments = get_default_moment_set_for_stencil(stencil)
+        x, y, z = MOMENT_SYMBOLS
+        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[: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, 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)
+
+    moment_to_relaxation_rate_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments,
+                                                               stencil.D, conserved_moments)
+
+    if continuous_equilibrium:
+        return create_with_continuous_maxwellian_equilibrium(stencil,
+                                                             moment_to_relaxation_rate_dict, **kwargs)
+    else:
+        return create_with_discrete_maxwellian_equilibrium(stencil,
+                                                           moment_to_relaxation_rate_dict, **kwargs)
+
+
+# ----------------------------------------- Cumulant method creators ---------------------------------------------------
+def create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments=True, **kwargs):
+    r"""Creates a cumulant-based lattice Boltzmann method.
+
+    Args:
+        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
+        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.
+        conserved_moments: If lower order moments are conserved or not.
+        kwargs: See :func:`create_with_continuous_maxwellian_equilibrium`
+
+    Returns:
+        :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
+    """
+    cumulant_to_rr_dict = _get_relaxation_info_dict(relaxation_rates, cumulant_groups, stencil.D, conserved_moments)
+
+    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))
+
+    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: 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
+        conserved_moments: If lower order moments are conserved or not.
+        kwargs: See :func:`create_cumulant`
+
+    Returns:
+        :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
+    """
+    # Get monomial moments
+    cumulants = get_default_moment_set_for_stencil(stencil)
+    cumulant_groups = [(c,) for c in cumulants]
+    return create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments, **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`.
+
+    Args: See :func:`create_cumulant`.
+
+    Returns:
+        :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
+    """
+    # Get polynomial groups
+    cumulant_groups = cascaded_moment_sets_literature(stencil)
+    return create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments, **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 ----------------------------------------------
+
+
+def compare_moment_based_lb_methods(reference, other, show_deviations_only=False):
+    import ipy_table
+    table = []
+    caption_rows = [len(table)]
+    table.append(['Shared Moment', 'ref', 'other', 'difference'])
+
+    reference_moments = set(reference.moments)
+    other_moments = set(other.moments)
+    for moment in reference_moments.intersection(other_moments):
+        reference_value = reference.relaxation_info_dict[moment].equilibrium_value
+        other_value = other.relaxation_info_dict[moment].equilibrium_value
+        diff = sp.simplify(reference_value - other_value)
+        if show_deviations_only and diff == 0:
+            pass
+        else:
+            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:
+        caption_rows.append(len(table))
+        table.append(['Only in Ref', 'value', '', ''])
+        for moment in only_in_ref:
+            val = reference.relaxation_info_dict[moment].equilibrium_value
+            table.append([f"${sp.latex(moment)}$",
+                          f"${sp.latex(val)}$",
+                          " ", " "])
+
+    only_in_other = other_moments - reference_moments
+    if only_in_other:
+        caption_rows.append(len(table))
+        table.append(['Only in Other', '', 'value', ''])
+        for moment in only_in_other:
+            val = other.relaxation_info_dict[moment].equilibrium_value
+            table.append([f"${sp.latex(moment)}$",
+                          " ",
+                          f"${sp.latex(val)}$",
+                          " "])
+
+    table_display = ipy_table.make_table(table)
+    for row_idx in caption_rows:
+        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
+from pystencils.simp.assignment_collection import 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)

src/lbmpy/methods/cumulantbased/galilean_correction.py

+import sympy as sp
+
+from pystencils.simp.assignment_collection import AssignmentCollection
+from pystencils import Assignment
+
+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)
+
+
+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")
+
+    #   Call PC1 = (x^2 - y^2), PC2 = (x^2 - z^2), PC3 = (x^2 + y^2 + z^2)
+    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)
+
+    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
+
+    collision_rule.set_sub_expressions_from_dict(subexp_dict)
+    collision_rule.topological_sort()
+
+    return collision_rule
+
+
+def get_galilean_correction_terms(rrate_dict, rho, u_vector):
+
+    pc1 = CORRECTED_POLYNOMIALS[0]
+    pc2 = CORRECTED_POLYNOMIALS[1]
+    pc3 = CORRECTED_POLYNOMIALS[2]
+
+    try:
+        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 = 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_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 = CORRECTION_SYMBOLS
+
+    #   Derivative Approximations
+    subexpressions = [
+        Assignment(dx, - omega_1 / (2 * rho) * (2 * c_xx - c_yy - c_zz)
+                   - omega_2 / (2 * rho) * (c_xx + c_yy + c_zz - rho)),
+        Assignment(dy, dx + (3 * omega_1) / (2 * rho) * (c_xx - c_yy)),
+        Assignment(dz, dx + (3 * omega_1) / (2 * rho) * (c_xx - c_zz))]
+
+    #   Correction Terms
+    main_assignments = [
+        Assignment(cor1, - 3 * rho * (1 - omega_1 / 2) * (u_vector[0]**2 * dx - u_vector[1]**2 * dy)),
+        Assignment(cor2, - 3 * rho * (1 - omega_1 / 2) * (u_vector[0]**2 * dx - u_vector[2]**2 * dz)),
+        Assignment(cor3, - 3 * rho * (1 - omega_2 / 2)
+                   * (u_vector[0]**2 * dx + u_vector[1]**2 * dy + u_vector[2]**2 * dz))]
+    return AssignmentCollection(main_assignments=main_assignments, subexpressions=subexpressions)

src/lbmpy/methods/default_moment_sets.py

+import sympy as sp
+
+from lbmpy.enums import Stencil
+from lbmpy.moments import MOMENT_SYMBOLS, sort_moments_into_groups_of_same_order
+from lbmpy.stencils import LBStencil
+from pystencils.stencil import have_same_entries
+
+
+def cascaded_moment_sets_literature(stencil):
+    """
+    Returns default groups of central moments or cumulants to be relaxed with common relaxation rates 
+    as stated in literature.
+    Groups are ordered like this:
+
+    - First group is density
+    - Second group are the momentum modes
+    - Third group are the shear modes
+    - Fourth group is the bulk mode
+    - Remaining groups do not govern hydrodynamic properties
+
+    Args:
+        stencil: instance of :class:`lbmpy.stencils.LBStencil`. Can be D2Q9, D3Q7, D3Q15, D3Q19 or D3Q27
+    """
+    x, y, z = MOMENT_SYMBOLS
+    if have_same_entries(stencil, LBStencil(Stencil.D2Q9)):
+        #   Cumulants of the D2Q9 stencil up to third order are equal to
+        #   the central moments; only the fourth-order cumulant x**2 * y**2
+        #   has a more complicated form. They can be arranged into groups
+        #   for the preservation of rotational invariance as described by
+        #   Martin Geier in his dissertation.
+        #
+        #   Reference: Martin Geier. Ab inito derivation of the cascaded Lattice Boltzmann
+        #   Automaton. Dissertation. University of Freiburg. 2006.
+        return [
+            [sp.sympify(1)],  # density is conserved
+            [x, y],  # momentum is relaxed for cumulant forcing
+
+            [x * y, x ** 2 - y ** 2],  # shear
+
+            [x ** 2 + y ** 2],  # bulk
+
+            [x ** 2 * y, x * y ** 2],
+            [x ** 2 * y ** 2]
+        ]
+
+    elif have_same_entries(stencil, LBStencil(Stencil.D3Q7)):
+        # D3Q7 moments: https://arxiv.org/ftp/arxiv/papers/1611/1611.03329.pdf
+        return [
+            [sp.sympify(1)],  # density is conserved
+            [x, y, z],  # momentum might be affected by forcing
+
+            [x ** 2 - y ** 2,
+             x ** 2 - z ** 2],  # shear
+
+            [x ** 2 + y ** 2 + z ** 2],  # bulk
+        ]
+
+    elif have_same_entries(stencil, LBStencil(Stencil.D3Q15)):
+        #   D3Q15 central moments by Premnath et al. https://arxiv.org/pdf/1202.6081.pdf.
+        return [
+            [sp.sympify(1)],  # density is conserved
+            [x, y, z],  # momentum might be affected by forcing
+
+            [x * y,
+             x * z,
+             y * z,
+             x ** 2 - y ** 2,
+             x ** 2 - z ** 2],  # shear
+
+            [x ** 2 + y ** 2 + z ** 2],  # bulk
+
+            [x * (x ** 2 + y ** 2 + z ** 2),
+             y * (x ** 2 + y ** 2 + z ** 2),
+             z * (x ** 2 + y ** 2 + z ** 2)],
+
+            [x * y * z],
+
+            [x ** 2 * y ** 2 + x ** 2 * z ** 2 + y ** 2 * z ** 2]
+        ]
+
+    elif have_same_entries(stencil, LBStencil(Stencil.D3Q19)):
+        #   D3Q19 cumulants are obtained by pruning the D3Q27 cumulant set as
+        #   described by Coreixas, 2019.
+        return [
+            [sp.sympify(1)],  # density is conserved
+            [x, y, z],  # momentum might be affected by forcing
+
+            [x * y,
+             x * z,
+             y * z,
+             x ** 2 - y ** 2,
+             x ** 2 - z ** 2],  # shear
+
+            [x ** 2 + y ** 2 + z ** 2],  # bulk
+
+            [x * y ** 2 + x * z ** 2,
+             x ** 2 * y + y * z ** 2,
+             x ** 2 * z + y ** 2 * z],
+
+            [x * y ** 2 - x * z ** 2,
+             x ** 2 * y - y * z ** 2,
+             x ** 2 * z - y ** 2 * z],
+
+            [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]
+        ]
+
+    elif have_same_entries(stencil, LBStencil(Stencil.D3Q27)):
+        #   Cumulants grouped to preserve rotational invariance as described by Geier et al, 2015
+        return [
+            [sp.sympify(1)],  # density is conserved
+            [x, y, z],  # momentum might be affected by forcing
+
+            [x * y,
+             x * z,
+             y * z,
+             x ** 2 - y ** 2,
+             x ** 2 - z ** 2],  # shear
+
+            [x ** 2 + y ** 2 + z ** 2],  # bulk
+
+            [x * y ** 2 + x * z ** 2,
+             x ** 2 * y + y * z ** 2,
+             x ** 2 * z + y ** 2 * z],
+
+            [x * y ** 2 - x * z ** 2,
+             x ** 2 * y - y * z ** 2,
+             x ** 2 * z - y ** 2 * z],
+
+            [x * y * z],
+
+            [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],
+
+            [x ** 2 * y ** 2 * z,
+             x ** 2 * y * z ** 2,
+             x * y ** 2 * z ** 2],
+
+            [x ** 2 * y ** 2 * z ** 2]
+        ]
+    else:
+        raise ValueError("No default set of cascaded moments is available for this stencil. "
+                         "Please specify your own set of cascaded moments.")
+
+
+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.
+
+    Args:
+        stencil: instance of :class:`lbmpy.stencils.LBStencil`. Can be D2Q9, D3Q15, D3Q19 or D3Q27
+        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, LBStencil(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, LBStencil(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, LBStencil(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, LBStencil(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())
+    else:
+        raise NotImplementedError("No MRT model is available (yet) for this stencil. "
+                                  "Create a custom MRT using 'create_with_discrete_maxwellian_equilibrium'")
+
+    return nested_moments

src/lbmpy/methods/momentbased/__init__.py

+from lbmpy.methods.momentbased.momentbasedmethod import MomentBasedLbMethod
+from .centralmomentbasedmethod import CentralMomentBasedLbMethod
+
+__all__ = ["MomentBasedLbMethod", "CentralMomentBasedLbMethod"]

src/lbmpy/methods/momentbased/centralmomentbasedmethod.py

+import sympy as sp
+from collections import OrderedDict
+from typing import Set
+
+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.moment_transforms import BinomialChimeraTransform
+
+from lbmpy.moments import MOMENT_SYMBOLS, moment_matrix, set_up_shift_matrix
+
+
+def relax_central_moments(pre_collision_symbols, post_collision_symbols,
+                          relaxation_rates, equilibrium_values,
+                          force_terms):
+    equilibrium_vec = sp.Matrix(equilibrium_values)
+    moment_vec = sp.Matrix(pre_collision_symbols)
+    relaxation_matrix = sp.diag(*relaxation_rates)
+    moment_vec = moment_vec + relaxation_matrix * (equilibrium_vec - moment_vec) + force_terms
+    main_assignments = [Assignment(s, eq) for s, eq in zip(post_collision_symbols, moment_vec)]
+
+    return AssignmentCollection(main_assignments)
+
+
+#   =============================== LB Method Implementation ===========================================================
+
+
+class CentralMomentBasedLbMethod(AbstractLbMethod):
+    """
+    Central Moment based LBM is a class to represent the single (SRT), two (TRT) and multi relaxation time (MRT)
+    methods, where the collision is performed in the central moment space.
+    These methods work by transforming the pdfs into moment space using a linear transformation and then shiftig
+    them into the central moment space. In the central moment space each component (moment) is relaxed to an
+    equilibrium moment by a certain relaxation rate. These equilibrium moments can e.g. be determined by taking the
+    equilibrium moments of the continuous Maxwellian.
+
+    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 moments 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`)
+    """
+
+    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):
+        assert isinstance(conserved_quantity_computation, AbstractConservedQuantityComputation)
+        super(CentralMomentBasedLbMethod, self).__init__(stencil)
+
+        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._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 moments 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.moment_equilibrium_values)})
+    
+    @property
+    def conserved_quantity_computation(self):
+        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 moments(self):
+        """Central moments relaxed by this method."""
+        return tuple(self._relaxation_dict.keys())
+
+    @property
+    def moment_equilibrium_values(self):
+        """Equilibrium values of this method's :attr:`moments`."""
+        return self._equilibrium.central_moments(self.moments, self.first_order_equilibrium_moment_symbols)
+
+    @property
+    def relaxation_rates(self):
+        """Relaxation rates for this method's :attr:`moments`."""
+        return tuple(self._relaxation_dict.values())
+        
+    @property
+    def relaxation_rate_dict(self):
+        """Dictionary mapping moments 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 it's 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
+
+    @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 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 moments 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):
+        self._force_model = force_model
+
+    @property
+    def moment_matrix(self) -> sp.Matrix:
+        return moment_matrix(self.moments, self.stencil)
+
+    @property
+    def shift_matrix(self) -> sp.Matrix:
+        return set_up_shift_matrix(self.moments, self.stencil)
+
+    @property
+    def relaxation_matrix(self) -> sp.Matrix:
+        d = sp.zeros(len(self.relaxation_rates))
+        for i in range(0, len(self.relaxation_rates)):
+            d[i, i] = self.relaxation_rates[i]
+        return d
+
+    def __getstate__(self):
+        # Workaround for a bug in joblib
+        self._moment_to_relaxation_info_dict_to_pickle = [i for i in self._relaxation_dict.items()]
+        return self.__dict__
+
+    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">
+                    Central-Moment-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>Central Moment</th>
+                <th>Eq. Value </th>
+                <th>Relaxation Rate</th>
+            </tr>
+        """
+
+        for moment, (eq_value, rr) in self.relaxation_info_dict.items():
+            vals = {
+                'rr': sp.latex(rr),
+                'moment': sp.latex(moment),
+                'eq_value': sp.latex(eq_value),
+                'nb': 'style="border:none"',
+            }
+            html += """<tr {nb}>
+                            <td {nb}>${moment}$</td>
+                            <td {nb}>${eq_value}$</td>
+                            <td {nb}>${rr}$</td>
+                         </tr>\n""".format(**vals)
+
+        html += "</table>"
+        return html
+
+    #   ----------------------- Overridden Abstract Members --------------------------
+
+    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.
+        """
+        _, d = self._generate_symbolic_relaxation_matrix()
+        rr_sub_expressions = set([Assignment(d[i, i], sp.Integer(1)) for i in range(len(self.relaxation_rates))])
+        r_info_dict = OrderedDict({c: RelaxationInfo(info.equilibrium_value, sp.Integer(1))
+                                   for c, info in self.relaxation_info_dict.items()})
+        ac = self._central_moment_collision_rule(moment_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)
+        ac.subexpressions = list(rr_sub_expressions) + ac.subexpressions
+        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) -> LbmCollisionRule:
+        """Returns an LbmCollisionRule i.e. an equation collection with a reference to the method.
+        This collision rule defines the collision operator."""
+        return self._central_moment_collision_rule(moment_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.
+            mm_inv = self.moment_matrix.inv()
+            bg_moments = bg.moments(self.moments)
+            weights = (mm_inv * sp.Matrix(bg_moments)).expand()
+        
+        for w in weights:
+            assert is_constant(w)
+
+        return [w for w in weights]
+
+    def _central_moment_collision_rule(self, moment_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:
+        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
+
+        relaxation_info_dict = dict()
+        subexpressions_relaxation_rates = []
+        if symbolic_relaxation_rates:
+            subexpressions_relaxation_rates, sd = self._generate_symbolic_relaxation_matrix()
+            for i, moment in enumerate(moment_to_relaxation_info_dict):
+                relaxation_info_dict[moment] = RelaxationInfo(moment_to_relaxation_info_dict[moment][0], sd[i, i])
+        else:
+            relaxation_info_dict = moment_to_relaxation_info_dict
+
+        if cqe is None:
+            cqe = self._cqc.equilibrium_input_equations_from_pdfs(f, False)
+
+        forcing_subexpressions = AssignmentCollection([])
+        moment_space_forcing = False
+        if self._force_model is not None:
+            if include_force_terms:
+                moment_space_forcing = self.force_model.has_central_moment_space_forcing
+            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 and not self._equilibrium.deviation_only:
+            background_distribution = self._equilibrium.background_distribution
+            assert background_distribution is not None
+        else:
+            background_distribution = None
+
+        #   1) Get Forward Transformation from PDFs to central moments
+        pdfs_to_c_transform = self.central_moment_transform_class(
+            stencil, self.moments, density, velocity, conserved_quantity_equations=cqe,
+            background_distribution=background_distribution)
+        pdfs_to_c_eqs = pdfs_to_c_transform.forward_transform(f, simplification=pre_simplification)
+
+        #   2) Collision
+        k_pre = pdfs_to_c_transform.pre_collision_symbols
+        k_post = pdfs_to_c_transform.post_collision_symbols
+
+        relaxation_infos = [relaxation_info_dict[m] for m in self.moments]
+        relaxation_rates = [info.relaxation_rate for info in relaxation_infos]
+        equilibrium_value = [info.equilibrium_value for info in relaxation_infos]
+
+        if moment_space_forcing:
+            force_model_terms = self._force_model.central_moment_space_forcing(self)
+        else:
+            force_model_terms = sp.Matrix([0] * stencil.Q)
+
+        collision_eqs = relax_central_moments(k_pre, k_post, tuple(relaxation_rates),
+                                              tuple(equilibrium_value), force_terms=force_model_terms)
+
+        #   3) Get backward transformation from central moments to PDFs
+        post_collision_values = self.post_collision_pdf_symbols
+        c_post_to_pdfs_eqs = pdfs_to_c_transform.backward_transform(post_collision_values,
+                                                                    simplification=pre_simplification)
+
+        #   4) Now, put it all together.
+        all_acs = [] if pdfs_to_c_transform.absorbs_conserved_quantity_equations else [cqe]
+        subexpressions_relaxation_rates = AssignmentCollection(subexpressions_relaxation_rates)
+        all_acs += [subexpressions_relaxation_rates, forcing_subexpressions, pdfs_to_c_eqs, collision_eqs]
+        subexpressions = [ac.all_assignments for ac in all_acs]
+        subexpressions += c_post_to_pdfs_eqs.subexpressions
+        main_assignments = c_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]
+
+        #   5) Maybe add forcing terms.
+        if include_force_terms and not moment_space_forcing:
+            force_model_terms = self._force_model(self)
+            force_term_symbols = sp.symbols(f"forceTerm_:{len(force_model_terms)}")
+            force_subexpressions = [Assignment(sym, force_model_term)
+                                    for sym, force_model_term in zip(force_term_symbols, force_model_terms)]
+            subexpressions += force_subexpressions
+            main_assignments = [Assignment(eq.lhs, eq.rhs + force_term_symbol)
+                                for eq, force_term_symbol in zip(main_assignments, force_term_symbols)]
+            simplification_hints['force_terms'] = force_term_symbols
+
+        return LbmCollisionRule(self, main_assignments, subexpressions)

lbmpy/methods/entropic.py → src/lbmpy/methods/momentbased/entropic.py

@@ -43,12 +43,11 @@ def add_entropy_condition(collision_rule, omega_output_field=None):
         ds.append(entry[0])
 
     stencil = collision_rule.method.stencil
-    q = len(stencil)
     f_symbols = collision_rule.method.pre_collision_pdf_symbols
 
-    ds_symbols = [sp.Symbol("entropicDs_%d" % (i,)) for i in range(q)]
-    dh_symbols = [sp.Symbol("entropicDh_%d" % (i,)) for i in range(q)]
-    feq_symbols = [sp.Symbol("entropicFeq_%d" % (i,)) for i in range(q)]
+    ds_symbols = [sp.Symbol(f"entropicDs_{i}") for i in range(stencil.Q)]
+    dh_symbols = [sp.Symbol(f"entropicDh_{i}") for i in range(stencil.Q)]
+    feq_symbols = [sp.Symbol(f"entropicFeq_{i}") for i in range(stencil.Q)]
 
     subexpressions = [Assignment(a, b) for a, b in zip(ds_symbols, ds)] + \
                      [Assignment(a, b) for a, b in zip(dh_symbols, dh)] + \
@@ -74,6 +73,17 @@ def add_entropy_condition(collision_rule, omega_output_field=None):
     if omega_output_field:
         new_collision_rule.main_assignments.append(Assignment(omega_output_field.center, omega_h))
 
+    try:
+        new_collision_rule.topological_sort()
+    except ValueError as e:
+        print("After adding the entropic condition, a cyclic dependency has been detected. This problem occurred most "
+              "likely due to the use of a force model combined with the entropic method. As described by Silva et al. "
+              "(https://doi.org/10.1103/PhysRevE.102.063307), most force schemes for the TRT collision operator depend "
+              "on both relaxation times. However, the force is also needed to calculate the free relaxation parameter "
+              "in the first place for entropic methods. Thus a cyclic dependency appears. The problem does not appear "
+              "with the SIMPLE, LUO or EDM force model.")
+        raise e
+
     return new_collision_rule
 
 
@@ -125,19 +135,19 @@ def add_iterative_entropy_condition(collision_rule, free_omega=None, newton_iter
 
     # 2) get equilibrium from method and define subexpressions for it
     eq_terms = [eq.rhs for eq in collision_rule.method.get_equilibrium().main_assignments]
-    eq_symbols = sp.symbols("entropicFeq_:%d" % (len(eq_terms, )))
+    eq_symbols = sp.symbols(f"entropicFeq_:{len(eq_terms)}")
     eq_subexpressions = [Assignment(a, b) for a, b in zip(eq_symbols, eq_terms)]
 
     # 3) find coefficients of entropy derivatives
     entropy_diff = sp.diff(discrete_approx_entropy(rr_polynomials, eq_symbols), free_omega)
     coefficients_first_diff = [c.expand() for c in reversed(sp.poly(entropy_diff, free_omega).all_coeffs())]
-    sym_coeff_diff1 = sp.symbols("entropicDiffCoeff_:%d" % (len(coefficients_first_diff, )))
+    sym_coeff_diff1 = sp.symbols(f"entropicDiffCoeff_:{len(coefficients_first_diff)}")
     coefficient_eqs = [Assignment(a, b) for a, b in zip(sym_coeff_diff1, coefficients_first_diff)]
     sym_coeff_diff2 = [(i + 1) * coeff for i, coeff in enumerate(sym_coeff_diff1[1:])]
 
     # 4) define Newtons method update iterations
     newton_iteration_equations = []
-    intermediate_omegas = [sp.Symbol("omega_iter_%i" % (i,)) for i in range(newton_iterations + 1)]
+    intermediate_omegas = [sp.Symbol(f"omega_iter_{i}") for i in range(newton_iterations + 1)]
     intermediate_omegas[0] = initial_value
     intermediate_omegas[-1] = free_omega
     for omega_idx in range(len(intermediate_omegas) - 1):
@@ -158,6 +168,17 @@ def add_iterative_entropy_condition(collision_rule, free_omega=None, newton_iter
         from lbmpy.updatekernels import write_quantities_to_field
         new_collision_rule = write_quantities_to_field(new_collision_rule, free_omega, omega_output_field)
 
+    try:
+        new_collision_rule.topological_sort()
+    except ValueError as e:
+        print("After adding the entropic condition, a cyclic dependency has been detected. This problem occurred most "
+              "likely due to the use of a force model combined with the entropic method. As described by Silva et al. "
+              "(https://doi.org/10.1103/PhysRevE.102.063307), most force schemes for the TRT collision operator depend "
+              "on both relaxation times. However, the force is also needed to calculate the free relaxation parameter "
+              "in the first place for entropic methods. Thus a cyclic dependency appears. The problem does not appear "
+              "with the SIMPLE, LUO or EDM force model.")
+        raise e
+
     return new_collision_rule
 
 
@@ -195,22 +216,18 @@ def _get_entropy_maximizing_omega(omega_s, f_eq, ds, dh):
     return 1 - ((omega_s - 1) * ds_dh / dh_dh)
 
 
-class RelaxationRatePolynomialDecomposition(object):
+class RelaxationRatePolynomialDecomposition:
 
     def __init__(self, collision_rule, free_relaxation_rates, fixed_relaxation_rates):
         self._collisionRule = collision_rule
         self._free_relaxation_rates = free_relaxation_rates
         self._fixed_relaxation_rates = fixed_relaxation_rates
         self._all_relaxation_rates = fixed_relaxation_rates + free_relaxation_rates
-        for se in collision_rule.subexpressions:
-            for rr in free_relaxation_rates:
-                assert rr not in se.rhs.atoms(sp.Symbol), \
-                    "Decomposition not possible since free relaxation rates are already in subexpressions"
 
     def symbolic_relaxation_rate_factors(self, relaxation_rate, power):
         q = len(self._collisionRule.method.stencil)
         omega_idx = self._all_relaxation_rates.index(relaxation_rate)
-        return [sp.Symbol("entFacOmega_%d_%d_%d" % (i, omega_idx, power)) for i in range(q)]
+        return [sp.Symbol(f"entFacOmega_{i}_{omega_idx}_{power}") for i in range(q)]
 
     def relaxation_rate_factors(self, relaxation_rate):
         update_equations = self._collisionRule.main_assignments
@@ -236,7 +253,7 @@ class RelaxationRatePolynomialDecomposition(object):
 
     @staticmethod
     def symbolic_constant_expr(i):
-        return sp.Symbol("entOffset_%d" % (i,))
+        return sp.Symbol(f"entOffset_{i}")
 
     def constant_exprs(self):
         subs_dict = {rr: 0 for rr in self._free_relaxation_rates}
@@ -252,7 +269,7 @@ class RelaxationRatePolynomialDecomposition(object):
 
     def symbolic_equilibrium(self):
         q = len(self._collisionRule.method.stencil)
-        return [sp.Symbol("entFeq_%d" % (i,)) for i in range(q)]
+        return [sp.Symbol(f"entFeq_{i}") for i in range(q)]
 
 
 def _get_relaxation_rates(collision_rule):
@@ -267,6 +284,12 @@ def _get_relaxation_rates(collision_rule):
 
     method = collision_rule.method
     omega_s = get_shear_relaxation_rate(method)
+
+    # if the shear relaxation rate is not specified as a symbol look for its symbolic counter part in the subs dict
+    for symbolic_rr, rr in method.subs_dict_relaxation_rate.items():
+        if omega_s == rr:
+            omega_s = symbolic_rr
+
     assert omega_s in relaxation_rates
 
     relaxation_rates_without_omega_s = relaxation_rates - {omega_s}

src/lbmpy/methods/momentbased/momentbasedmethod.py

+from collections import OrderedDict
+from typing import Iterable, Set
+
+import sympy as sp
+import numpy as np
+
+from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
+from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation, DensityVelocityComputation
+from lbmpy.moments import MOMENT_SYMBOLS, moment_matrix
+from pystencils.sympyextensions import is_constant
+from pystencils import Assignment, AssignmentCollection
+
+from lbmpy.moment_transforms import PdfsToMomentsByChimeraTransform
+
+
+class MomentBasedLbMethod(AbstractLbMethod):
+    """
+    Moment based LBM is a class to represent the single (SRT), two (TRT) and multi relaxation time (MRT) methods.
+    These methods work by transforming the pdfs into moment space using a linear transformation. In the moment-space
+    each component (moment) is relaxed to an equilibrium moment by a certain relaxation rate. These
+    equilibrium moments can e.g. be determined by taking the equilibrium moments of the continuous Maxwellian.
+
+    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 moments 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.
+        moment_transform_class: transformation class to transform PDFs to the moment space (subclass of 
+                                :class:`lbmpy.moment_transforms.AbstractRawMomentTransform`), or `None`
+                                if equations are to be derived in population space. 
+    """
+
+    def __init__(self, stencil, equilibrium, relaxation_dict,
+                 conserved_quantity_computation=None, force_model=None, zero_centered=False,
+                 fraction_field=None, moment_transform_class=PdfsToMomentsByChimeraTransform):
+        assert isinstance(conserved_quantity_computation, AbstractConservedQuantityComputation)
+        super(MomentBasedLbMethod, self).__init__(stencil)
+
+        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._moment_transform_class = 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 moments 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.moment_equilibrium_values)})
+
+    @property
+    def conserved_quantity_computation(self):
+        return self._cqc
+
+    @property
+    def equilibrium_distribution(self):
+        """Returns this method's equilibrium distribution (see :class:`lbmpy.equilibrium.AbstractEquilibrium`"""
+        return self._equilibrium
+
+    @property
+    def moment_transform_class(self):
+        """The transform class (subclass of :class:`lbmpy.moment_transforms.AbstractRawMomentTransform` defining the
+        transformation of populations to moment space."""
+        return self._moment_transform_class
+
+    @property
+    def moment_space_collision(self):
+        """Returns whether collision is derived in terms of moments or in terms of populations only."""
+        return self._moment_transform_class is not None
+
+    @property
+    def moments(self):
+        """Moments relaxed by this method."""
+        return tuple(self._relaxation_dict.keys())
+
+    @property
+    def moment_equilibrium_values(self):
+        """Equilibrium values of this method's :attr:`moments`."""
+        return self._equilibrium.moments(self.moments)
+
+    @property
+    def relaxation_rates(self):
+        """Relaxation rates for this method's :attr:`moments`."""
+        return tuple(self._relaxation_dict.values())
+
+    @property
+    def relaxation_rate_dict(self):
+        """Dictionary mapping moments 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 it's 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
+
+    @property
+    def weights(self):
+        if self._weights is None:
+            self._weights = self._compute_weights()
+        return self._weights
+
+    def override_weights(self, weights):
+        """Manually set this method's lattice weights."""
+        assert len(weights) == len(self.stencil)
+        self._weights = weights
+
+    def get_equilibrium(self, conserved_quantity_equations: AssignmentCollection = None,
+                        include_force_terms: bool = False, pre_simplification: bool = False,
+                        subexpressions: bool = False, keep_cqc_subexpressions: bool = True) -> 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.
+            include_force_terms: if set to True the equilibrium is shifted by forcing terms coming from the force model
+                                 of the method.
+            pre_simplification: with or without pre-simplifications for the calculation of the collision
+            subexpressions: if set to false all subexpressions of the equilibrium assignments are plugged
+                            into the main assignments
+            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
+
+        """
+        _, d = self._generate_symbolic_relaxation_matrix()
+        rr_sub_expressions = set([Assignment(d[i, i], sp.Integer(1)) for i in range(len(self.relaxation_rates))])
+
+        ac = self._collision_rule_with_relaxation_matrix(d=d,
+                                                         additional_subexpressions=rr_sub_expressions,
+                                                         include_force_terms=include_force_terms,
+                                                         conserved_quantity_equations=conserved_quantity_equations,
+                                                         pre_simplification=pre_simplification)
+
+        if not subexpressions:
+            if keep_cqc_subexpressions:
+                bs = self._bound_symbols_cqc(conserved_quantity_equations)
+                ac = ac.new_without_subexpressions(subexpressions_to_keep=bs)
+                return ac.new_without_unused_subexpressions()
+            else:
+                ac = ac.new_without_subexpressions()
+                return ac.new_without_unused_subexpressions()
+        else:
+            return ac.new_without_unused_subexpressions()
+
+    def get_equilibrium_terms(self):
+        """Returns this method's equilibrium populations as a vector."""
+        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 = True) -> LbmCollisionRule:
+
+        if self._fraction_field is not None:
+            relaxation_rates_modifier = (1.0 - self._fraction_field)
+            rr_sub_expressions, d = self._generate_symbolic_relaxation_matrix(
+                relaxation_rates_modifier=relaxation_rates_modifier)
+        else:
+            rr_sub_expressions, d = self._generate_symbolic_relaxation_matrix()
+
+        ac = self._collision_rule_with_relaxation_matrix(d=d,
+                                                         additional_subexpressions=rr_sub_expressions,
+                                                         include_force_terms=True,
+                                                         conserved_quantity_equations=conserved_quantity_equations,
+                                                         pre_simplification=pre_simplification)
+        return ac
+
+    def set_zeroth_moment_relaxation_rate(self, relaxation_rate):
+        """Alters the relaxation rate of the zeroth-order moment."""
+        one = sp.Rational(1, 1)
+        self._relaxation_dict[one] = relaxation_rate
+
+    def set_first_moment_relaxation_rate(self, relaxation_rate):
+        """Alters the relaxation rates of the first-order moments."""
+        for e in MOMENT_SYMBOLS[:self.dim]:
+            assert e in self._relaxation_dict, "First moments 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):
+        """Alters the relaxation rates of the zeroth- and first-order moments."""
+        self.set_zeroth_moment_relaxation_rate(relaxation_rate)
+        self.set_first_moment_relaxation_rate(relaxation_rate)
+
+    def set_force_model(self, force_model):
+        """Updates this method's force model."""
+        self._force_model = force_model
+        if isinstance(self._cqc, DensityVelocityComputation):
+            self._cqc.set_force_model(force_model)
+
+    @property
+    def collision_matrix(self) -> sp.Matrix:
+        pdfs_to_moments = self.moment_matrix
+        d = self.relaxation_matrix
+        return pdfs_to_moments.inv() * d * pdfs_to_moments
+
+    @property
+    def inverse_collision_matrix(self) -> sp.Matrix:
+        pdfs_to_moments = self.moment_matrix
+        inverse_relaxation_matrix = self.relaxation_matrix.inv()
+        return pdfs_to_moments.inv() * inverse_relaxation_matrix * pdfs_to_moments
+
+    @property
+    def moment_matrix(self) -> sp.Matrix:
+        return moment_matrix(self.moments, self.stencil)
+
+    @property
+    def is_orthogonal(self) -> bool:
+        return (self.moment_matrix * self.moment_matrix.T).is_diagonal()
+
+    @property
+    def is_weighted_orthogonal(self) -> bool:
+        weights_matrix = sp.Matrix([self.weights] * len(self.weights))
+        moment_matrix_times_weights = sp.Matrix(np.multiply(self.moment_matrix, weights_matrix))
+
+        return (moment_matrix_times_weights * self.moment_matrix.T).is_diagonal()
+
+    def __getstate__(self):
+        # Workaround for a bug in joblib
+        self._momentToRelaxationInfoDictToPickle = [i for i in self._relaxation_dict.items()]
+        return self.__dict__
+
+    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">
+                    Moment-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>Moment</th>
+                <th>Eq. Value </th>
+                <th>Relaxation Rate</th>
+            </tr>
+        """
+
+        for moment, (eq_value, rr) in self.relaxation_info_dict.items():
+            vals = {
+                'rr': sp.latex(rr),
+                'moment': sp.latex(moment),
+                'eq_value': sp.latex(eq_value),
+                'nb': 'style="border:none"',
+            }
+            html += """<tr {nb}>
+                            <td {nb}>${moment}$</td>
+                            <td {nb}>${eq_value}$</td>
+                            <td {nb}>${rr}$</td>
+                         </tr>\n""".format(**vals)
+
+        html += "</table>"
+        return html
+
+    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.
+            mm_inv = self.moment_matrix.inv()
+            bg_moments = bg.moments(self.moments)
+            weights = (mm_inv * sp.Matrix(bg_moments)).expand()
+
+        for w in weights:
+            assert is_constant(w)
+
+        return [w for w in weights]
+
+    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 _collision_rule_with_relaxation_matrix(self, d: sp.Matrix,
+                                               additional_subexpressions: Iterable[Assignment] = None,
+                                               include_force_terms: bool = True,
+                                               conserved_quantity_equations: AssignmentCollection = None,
+                                               pre_simplification: bool = False) -> LbmCollisionRule:
+        if additional_subexpressions is None:
+            additional_subexpressions = list()
+        f = sp.Matrix(self.pre_collision_pdf_symbols)
+        moment_polynomials = list(self.moments)
+
+        cqe = conserved_quantity_equations
+        if cqe is None:
+            cqe = self._cqc.equilibrium_input_equations_from_pdfs(f, False)
+
+        if self._force_model is None:
+            include_force_terms = False
+
+        moment_space_forcing = False
+
+        if include_force_terms and self._moment_transform_class:
+            if self._force_model is not None:
+                moment_space_forcing = self._force_model.has_moment_space_forcing
+
+        forcing_subexpressions = []
+        if self._force_model is not None:
+            forcing_subexpressions = AssignmentCollection(self._force_model.subs_dict_force).all_assignments
+
+        rho = self.zeroth_order_equilibrium_moment_symbol
+        u = self.first_order_equilibrium_moment_symbols
+
+        #   See if a background shift is necessary
+        if self._zero_centered and not self._equilibrium.deviation_only:
+            background_distribution = self._equilibrium.background_distribution
+            assert background_distribution is not None
+        else:
+            background_distribution = None
+
+        m_eq = sp.Matrix(self.moment_equilibrium_values)
+
+        if self._moment_transform_class:
+            # Derive equations in moment space if a transform is given
+            pdf_to_m_transform = self._moment_transform_class(self.stencil, moment_polynomials, rho, u,
+                                                              conserved_quantity_equations=cqe,
+                                                              background_distribution=background_distribution)
+
+            m_pre = pdf_to_m_transform.pre_collision_symbols
+            m_post = pdf_to_m_transform.post_collision_symbols
+
+            pdf_to_m_eqs = pdf_to_m_transform.forward_transform(self.pre_collision_pdf_symbols,
+                                                                simplification=pre_simplification)
+            m_post_to_f_post_eqs = pdf_to_m_transform.backward_transform(self.post_collision_pdf_symbols,
+                                                                         simplification=pre_simplification)
+
+            m_pre_vec = sp.Matrix(m_pre)
+            collision_rule = m_pre_vec + d * (m_eq - m_pre_vec)
+
+            if include_force_terms and moment_space_forcing:
+                collision_rule += self._force_model.moment_space_forcing(self)
+
+            collision_eqs = [Assignment(lhs, rhs) for lhs, rhs in zip(m_post, collision_rule)]
+            collision_eqs = AssignmentCollection(collision_eqs)
+
+            all_acs = [] if pdf_to_m_transform.absorbs_conserved_quantity_equations else [cqe]
+            all_acs += [pdf_to_m_eqs, collision_eqs]
+            subexpressions = list(additional_subexpressions) + forcing_subexpressions + [ac.all_assignments for ac in
+                                                                                         all_acs]
+            subexpressions += m_post_to_f_post_eqs.subexpressions
+            main_assignments = m_post_to_f_post_eqs.main_assignments
+        else:
+            #   For SRT, TRT by default, and whenever customarily required, derive equations entirely in
+            #   population space
+            if self._zero_centered and not self._equilibrium.deviation_only:
+                raise Exception("Can only derive population-space equations for zero-centered storage"
+                                " if delta equilibrium is used.")
+            pdf_to_moment = self.moment_matrix
+            collision_rule = f + pdf_to_moment.inv() * d * (m_eq - pdf_to_moment * f)
+            collision_eqs = [Assignment(lhs, rhs) for lhs, rhs in zip(self.post_collision_pdf_symbols, collision_rule)]
+            subexpressions = list(additional_subexpressions) + forcing_subexpressions + cqe.all_assignments
+            main_assignments = collision_eqs
+
+        simplification_hints = cqe.simplification_hints.copy()
+        simplification_hints.update(self._cqc.defined_symbols())
+        simplification_hints['relaxation_rates'] = [d[i, i] for i in range(d.rows)]
+
+        if include_force_terms and not moment_space_forcing:
+            force_model_terms = self._force_model(self)
+            force_term_symbols = sp.symbols(f"forceTerm_:{len(force_model_terms)}")
+            force_subexpressions = [Assignment(sym, force_model_term)
+                                    for sym, force_model_term in zip(force_term_symbols, force_model_terms)]
+            subexpressions += force_subexpressions
+            main_assignments = [Assignment(eq.lhs, eq.rhs + force_term_symbol)
+                                for eq, force_term_symbol in zip(main_assignments, force_term_symbols)]
+            simplification_hints['force_terms'] = force_term_symbols
+
+        ac = LbmCollisionRule(self, main_assignments, subexpressions, simplification_hints)
+        ac.topological_sort()
+        return ac

lbmpy/methods/momentbasedsimplifications.py → src/lbmpy/methods/momentbased/momentbasedsimplifications.py

@@ -4,11 +4,15 @@ All of these transformations operate on :class:`pystencils.AssignmentCollection`
 simplification hints, which are set by the MomentBasedLbMethod.
 """
 import sympy as sp
+from itertools import product
 
 from lbmpy.methods.abstractlbmethod import LbmCollisionRule
-from pystencils import Assignment
-from pystencils.sympyextensions import (
-    extract_most_common_factor, replace_second_order_products, subs_additive)
+from pystencils import Assignment, AssignmentCollection
+from pystencils.stencil import inverse_direction
+from pystencils.simp.subexpression_insertion import insert_subexpressions, is_constant
+from pystencils.sympyextensions import extract_most_common_factor, replace_second_order_products, subs_additive
+
+from collections import defaultdict
 
 
 def replace_second_order_velocity_products(cr: LbmCollisionRule):
@@ -39,7 +43,7 @@ def factor_relaxation_rates(cr: LbmCollisionRule):
     """
     sh = cr.simplification_hints
     assert 'relaxation_rates' in sh, "Needs simplification hint 'relaxation_rates': Sequence of relaxation rates"
-    if len(sh['relaxation_rates']) > 19:  # heuristics, works well if there is a small number of relaxation rates
+    if len(set(sh['relaxation_rates'])) > 19:  # heuristics, works well if there is a small number of relaxation rates
         return cr
 
     relaxation_rates = sp.Matrix(sh['relaxation_rates']).atoms(sp.Symbol)
@@ -200,6 +204,7 @@ def cse_in_opposing_directions(cr: LbmCollisionRule):
 
                 found_subexpressions, new_terms = sp.cse(handled_terms, symbols=replacement_symbol_generator,
                                                          order='None', optimizations=[])
+
                 substitutions += [Assignment(f[0], f[1]) for f in found_subexpressions]
 
                 update_rules = [Assignment(ur.lhs, ur.rhs.subs(relaxation_rate * old_term, new_coefficient * new_term))
@@ -216,8 +221,147 @@ def cse_in_opposing_directions(cr: LbmCollisionRule):
     return res
 
 
+def substitute_moments_in_conserved_quantity_equations(ac: AssignmentCollection):
+    """
+    Applied on an assignment collection containing both equations for raw moments
+    and conserved quantities, this simplification attempts to express the conserved
+    quantities in terms of the zeroth- and first-order moments.
+
+    For example, :math:`rho =  f_0 + f_1 + ... + f_8` will be replaced by the zeroth-
+    order moment: :math:`rho = m_{00}`
+
+    Required simplification hints:
+        - cq_symbols_to_moments: A dictionary mapping the conserved quantity symbols
+          to their corresponding moment symbols (like `{rho : m_00, u_0 : m_10, u_1 : m_01}`).
+    """
+    sh = ac.simplification_hints
+    if 'cq_symbols_to_moments' not in sh:
+        raise ValueError("Simplification hint 'cq_symbols_to_moments' missing.")
+
+    cq_symbols_to_moments = sh['cq_symbols_to_moments']
+    if len(cq_symbols_to_moments) == 0:
+        return ac
+
+    required_symbols = list(cq_symbols_to_moments.keys()) + list(cq_symbols_to_moments.values())
+    reduced_ac = ac.new_filtered(required_symbols).new_without_subexpressions()
+    main_asm_dict = ac.main_assignments_dict
+    subexp_dict = ac.subexpressions_dict
+    reduced_assignments = reduced_ac.main_assignments_dict
+
+    for cq_sym, moment_sym in cq_symbols_to_moments.items():
+        moment_eq = reduced_assignments[moment_sym]
+        assert moment_eq.count(cq_sym) == 0, f"Expressing conserved quantity {cq_sym} using moment {moment_sym} " \
+                                             "would introduce a circular dependency."
+        cq_eq = subs_additive(reduced_assignments[cq_sym], moment_sym, moment_eq)
+        if cq_sym in main_asm_dict:
+            main_asm_dict[cq_sym] = cq_eq
+        else:
+            assert moment_sym in subexp_dict, f"Cannot express subexpression {cq_sym}" \
+                                              f" using main assignment {moment_sym}!"
+            subexp_dict[cq_sym] = cq_eq
+
+    main_assignments = [Assignment(lhs, rhs) for lhs, rhs in main_asm_dict.items()]
+    subexpressions = [Assignment(lhs, rhs) for lhs, rhs in subexp_dict.items()]
+    ac = ac.copy(main_assignments=main_assignments, subexpressions=subexpressions)
+    ac.topological_sort()
+    return ac.new_without_unused_subexpressions()
+
+
+def split_pdf_main_assignments_by_symmetry(ac: AssignmentCollection):
+    """
+    Splits assignments to post-collision PDFs streaming in opposite directions
+    into their symmetric and asymetric parts, which are extracted as subexpressions.
+    Useful especially when computing PDF values from post-collision raw moments, where
+    symmetric splitting can reduce the number of required additions by one half.
+
+    Required simplification hints:
+        - stencil: Velocity set of the LB method as a nested tuple of directions
+        - post_collision_pdf_symbols: Sequence of symbols corresponding to the stencil velocities
+    """
+    sh = ac.simplification_hints
+    if 'stencil' not in sh:
+        raise ValueError("Symmetric splitting requires the stencil as a simplification hint.")
+    if 'post_collision_pdf_symbols' not in sh:
+        raise ValueError("Symmetric splitting requires the post-collision pdf symbols as a simplification hint.")
+
+    stencil = sh['stencil']
+    pdf_symbols = sh['post_collision_pdf_symbols']
+
+    asm_dict = ac.main_assignments_dict
+    subexpressions = ac.subexpressions
+    done = set()
+    subexp_to_symbol_dict = defaultdict(lambda: next(ac.subexpression_symbol_generator))
+    half = sp.Rational(1, 2)
+    for i, f in enumerate(pdf_symbols):
+        if i in done:
+            continue
+        c = stencil[i]
+        if all(cj == 0 for cj in c):
+            continue
+        c_inv = inverse_direction(c)
+        i_inv = stencil.index(c_inv)
+        f_inv = pdf_symbols[i_inv]
+        done |= {i, i_inv}
+
+        f_eq = asm_dict[f]
+        f_inv_eq = asm_dict[f_inv]
+
+        symmetric_part = half * (f_eq + f_inv_eq)
+        asymmetric_part = half * (f_eq - f_inv_eq)
+
+        symmetric_symb = subexp_to_symbol_dict[symmetric_part]
+        asymmetric_symb = subexp_to_symbol_dict[asymmetric_part]
+
+        asm_dict[f] = symmetric_symb + asymmetric_symb
+        asm_dict[f_inv] = symmetric_symb - asymmetric_symb
+    for subexp, sym in subexp_to_symbol_dict.items():
+        subexpressions.append(Assignment(sym, subexp))
+    main_assignments = [Assignment(lhs, rhs) for lhs, rhs in asm_dict.items()]
+    return ac.copy(main_assignments=main_assignments, subexpressions=subexpressions)
+
+
+def insert_pure_products(ac, symbols, **kwargs):
+    """Inserts any subexpression whose RHS is a product containing exclusively factors
+    from the given sequence of symbols."""
+    def callback(exp):
+        rhs = exp.rhs
+        if isinstance(rhs, sp.Symbol) and rhs in symbols:
+            return True
+        elif isinstance(rhs, sp.Mul):
+            if all((is_constant(arg) or (arg in symbols)) for arg in rhs.args):
+                return True
+        return False
+
+    return insert_subexpressions(ac, callback, **kwargs)
+
+
+def insert_conserved_quantity_products(cr, **kwargs):
+    from lbmpy.moments import statistical_quantity_symbol as sq_sym
+    from lbmpy.moment_transforms import PRE_COLLISION_MONOMIAL_RAW_MOMENT as m
+    
+    rho = cr.method.zeroth_order_equilibrium_moment_symbol
+    u = cr.method.first_order_equilibrium_moment_symbols
+    m000 = sq_sym(m, (0,) * cr.method.dim)
+    symbols = (rho, m000) + u
+
+    return insert_pure_products(cr, symbols)
+
+
+def insert_half_force(cr, **kwargs):
+    fmodel = cr.method.force_model
+    if not fmodel:
+        return cr
+    force = fmodel.symbolic_force_vector
+    force_exprs = set(c * f / 2 for c, f in product((1, -1), force))
+
+    def callback(expr):
+        return expr.rhs in force_exprs
+
+    return insert_subexpressions(cr, callback, **kwargs)
+
 # -------------------------------------- Helper Functions --------------------------------------------------------------
 
+
 def __get_common_quadratic_and_constant_terms(cr: LbmCollisionRule):
     """Determines a common subexpression useful for most LBM model often called f_eq_common.
     It contains the quadratic and constant terms of the center update rule."""
@@ -241,10 +385,15 @@ def __get_common_quadratic_and_constant_terms(cr: LbmCollisionRule):
         t = t.subs({ft: 0 for ft in sh['force_terms']})
 
     weight = t
+    weight = weight.subs(sh['density_deviation'], 1)
+    weight = weight.subs(sh['density'], 1)
 
     for u in sh['velocity']:
         weight = weight.subs(u, 0)
-    weight = weight / sh['density']
+    # weight = weight / sh['density']
     if weight == 0:
         return None
+
+    # t = t.subs(sh['density_deviation'], 0)
+
     return t / weight

src/lbmpy/moment_transforms/__init__.py

+from .abstractmomenttransform import (
+    PRE_COLLISION_MONOMIAL_RAW_MOMENT, POST_COLLISION_MONOMIAL_RAW_MOMENT,
+    PRE_COLLISION_RAW_MOMENT, POST_COLLISION_RAW_MOMENT,
+    PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT, POST_COLLISION_MONOMIAL_CENTRAL_MOMENT,
+    PRE_COLLISION_CENTRAL_MOMENT, POST_COLLISION_CENTRAL_MOMENT,
+    PRE_COLLISION_CUMULANT, POST_COLLISION_CUMULANT,
+    PRE_COLLISION_MONOMIAL_CUMULANT, POST_COLLISION_MONOMIAL_CUMULANT
+)
+
+from .abstractmomenttransform import AbstractMomentTransform
+
+from .rawmomenttransforms import (
+    PdfsToMomentsByMatrixTransform, PdfsToMomentsByChimeraTransform
+)
+
+from .centralmomenttransforms import (
+    PdfsToCentralMomentsByMatrix,
+    BinomialChimeraTransform,
+    PdfsToCentralMomentsByShiftMatrix,
+    FastCentralMomentTransform
+)
+
+from .cumulanttransforms import CentralMomentsToCumulantsByGeneratingFunc
+
+__all__ = [
+    "AbstractMomentTransform",
+    "PdfsToMomentsByMatrixTransform", "PdfsToMomentsByChimeraTransform",
+    "PdfsToCentralMomentsByMatrix",
+    "BinomialChimeraTransform",
+    "PdfsToCentralMomentsByShiftMatrix",
+    "FastCentralMomentTransform",
+    "CentralMomentsToCumulantsByGeneratingFunc",
+    "PRE_COLLISION_MONOMIAL_RAW_MOMENT", "POST_COLLISION_MONOMIAL_RAW_MOMENT",
+    "PRE_COLLISION_RAW_MOMENT", "POST_COLLISION_RAW_MOMENT",
+    "PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT", "POST_COLLISION_MONOMIAL_CENTRAL_MOMENT",
+    "PRE_COLLISION_CENTRAL_MOMENT", "POST_COLLISION_CENTRAL_MOMENT",
+    "PRE_COLLISION_CUMULANT", "POST_COLLISION_CUMULANT",
+    "PRE_COLLISION_MONOMIAL_CUMULANT", "POST_COLLISION_MONOMIAL_CUMULANT"
+]

src/lbmpy/moment_transforms/abstractmomenttransform.py

+from abc import abstractmethod
+import sympy as sp
+
+from pystencils.simp import (SimplificationStrategy, sympy_cse)
+
+from lbmpy.moments import (
+    exponents_to_polynomial_representations, extract_monomials, exponent_tuple_sort_key)
+
+from lbmpy.moments import statistical_quantity_symbol as sq_sym
+
+
+PRE_COLLISION_MONOMIAL_RAW_MOMENT = 'm'
+POST_COLLISION_MONOMIAL_RAW_MOMENT = 'm_post'
+
+PRE_COLLISION_RAW_MOMENT = 'M'
+POST_COLLISION_RAW_MOMENT = 'M_post'
+
+PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT = 'kappa'
+POST_COLLISION_MONOMIAL_CENTRAL_MOMENT = 'kappa_post'
+
+PRE_COLLISION_CENTRAL_MOMENT = 'K'
+POST_COLLISION_CENTRAL_MOMENT = 'K_post'
+
+PRE_COLLISION_MONOMIAL_CUMULANT = 'c'
+POST_COLLISION_MONOMIAL_CUMULANT = 'c_post'
+
+PRE_COLLISION_CUMULANT = 'C'
+POST_COLLISION_CUMULANT = 'C_post'
+
+
+class AbstractMomentTransform:
+    r"""Abstract Base Class for classes providing transformations between moment spaces."""
+
+    def __init__(self, stencil,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 moment_exponents=None,
+                 moment_polynomials=None,
+                 conserved_quantity_equations=None,
+                 background_distribution=None,
+                 pre_collision_symbol_base=None,
+                 post_collision_symbol_base=None,
+                 pre_collision_monomial_symbol_base=None,
+                 post_collision_monomial_symbol_base=None):
+        """Abstract Base Class constructor.
+
+        Args:
+            stencil: Nested tuple defining the velocity set
+            equilibrium_density: Symbol of the equilibrium density used in the collision rule
+            equilibrium_velocity: Tuple of symbols of the equilibrium velocity used in the collision rule
+            moment_exponents=None: Exponent tuples of the monomial basis of the collision space
+            moment_polynomials=None: Polynomial basis of the collision space
+            conserved_quantity_equations: Optionally, an assignment collection computing the conserved quantities
+                                          (typically density and velocity) from pre-collision populations
+            background_distribution: If not `None`, zero-centered storage of the populations is assumed and the 
+                                     moments of the passed distribution (instance of 
+                                     `lbmpy.equilibrium.AbstractEquilibrium`) are included in the transform equations.
+
+        """
+        if moment_exponents is not None and moment_polynomials is not None:
+            raise ValueError("Both moment_exponents and moment_polynomials were given. Pass only one of them!")
+
+        self.stencil = stencil
+        self.dim = stencil.D
+        self.q = stencil.Q
+
+        if moment_exponents is not None:
+            self.moment_exponents = list(moment_exponents)
+            self.moment_polynomials = exponents_to_polynomial_representations(self.moment_exponents)
+        elif moment_polynomials is not None:
+            self.moment_polynomials = moment_polynomials
+            moment_exponents = list(extract_monomials(moment_polynomials, dim=self.dim))
+            self.moment_exponents = sorted(moment_exponents, key=exponent_tuple_sort_key)
+        else:
+            raise ValueError("You must provide either moment_exponents or moment_polynomials!")
+
+        self.cqe = conserved_quantity_equations
+        self.equilibrium_density = equilibrium_density
+        self.equilibrium_velocity = equilibrium_velocity
+
+        self.background_distribution = background_distribution
+
+        self.base_pre = pre_collision_symbol_base
+        self.base_post = post_collision_symbol_base
+        self.mono_base_pre = pre_collision_monomial_symbol_base
+        self.mono_base_post = post_collision_monomial_symbol_base
+
+    @property
+    def pre_collision_symbols(self):
+        """List of symbols corresponding to the pre-collision quantities
+        that will be the left-hand sides of assignments returned by :func:`forward_transform`."""
+        return sp.symbols(f'{self.base_pre}_:{self.q}')
+
+    @property
+    def post_collision_symbols(self):
+        """List of symbols corresponding to the post-collision quantities
+        that are input to the right-hand sides of assignments returned by:func:`backward_transform`."""
+        return sp.symbols(f'{self.base_post}_:{self.q}')
+
+    @property
+    def pre_collision_monomial_symbols(self):
+        """List of symbols corresponding to the pre-collision monomial quantities
+        that might exist as left-hand sides of subexpressions in the assignment collection 
+        returned by :func:`forward_transform`."""
+        return tuple(sq_sym(self.mono_base_pre, e) for e in self.moment_exponents)
+
+    @property
+    def post_collision_monomial_symbols(self):
+        """List of symbols corresponding to the post-collision monomial quantities
+        that might exist as left-hand sides of subexpressions in the assignment collection
+        returned by :func:`backward_transform`."""
+        return tuple(sq_sym(self.mono_base_post, e) for e in self.moment_exponents)
+
+    @abstractmethod
+    def forward_transform(self, *args, **kwargs):
+        """Implemented in a subclass, will return the forward transform equations."""
+        raise NotImplementedError("forward_transform must be implemented in a subclass")
+
+    @abstractmethod
+    def backward_transform(self, *args, **kwargs):
+        """Implemented in a subclass, will return the backward transform equations."""
+        raise NotImplementedError("backward_transform must be implemented in a subclass")
+
+    @property
+    def absorbs_conserved_quantity_equations(self):
+        """Whether or not the given conserved quantity equations will be included in
+        the assignment collection returned by :func:`forward_transform`, possibly in simplified
+        form."""
+        return False
+
+    @property
+    def _default_simplification(self):
+        return SimplificationStrategy()
+
+    def _get_simp_strategy(self, simplification, direction=None):
+        if isinstance(simplification, bool):
+            simplification = 'default' if simplification else 'none'
+
+        if simplification == 'default' or simplification == 'default_with_cse':
+            simp = self._default_simplification if direction is None else self._default_simplification[direction]
+            if simplification == 'default_with_cse':
+                simp.add(sympy_cse)
+            return simp
+        else:
+            return None

src/lbmpy/moment_transforms/centralmomenttransforms.py

+from functools import partial
+import sympy as sp
+
+from pystencils import Assignment, AssignmentCollection
+from pystencils.simp import (
+    SimplificationStrategy, add_subexpressions_for_constants)
+from pystencils.simp.assignment_collection import SymbolGen
+from pystencils.sympyextensions import subs_additive, fast_subs
+
+from lbmpy.moments import (
+    moment_matrix, monomial_to_polynomial_transformation_matrix,
+    set_up_shift_matrix, contained_moments, moments_up_to_order,
+    moments_of_order,
+    central_moment_reduced_monomial_to_polynomial_matrix)
+
+from lbmpy.moments import statistical_quantity_symbol as sq_sym
+
+from .abstractmomenttransform import (
+    AbstractMomentTransform,
+    PRE_COLLISION_CENTRAL_MOMENT, POST_COLLISION_CENTRAL_MOMENT,
+    PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT, POST_COLLISION_MONOMIAL_CENTRAL_MOMENT
+)
+
+from .rawmomenttransforms import PdfsToMomentsByChimeraTransform
+
+
+class AbstractCentralMomentTransform(AbstractMomentTransform):
+    """Abstract base class for all transformations between population space
+    and central-moment space."""
+
+    def __init__(self, stencil, moment_polynomials,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 pre_collision_symbol_base=PRE_COLLISION_CENTRAL_MOMENT,
+                 post_collision_symbol_base=POST_COLLISION_CENTRAL_MOMENT,
+                 pre_collision_monomial_symbol_base=PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT,
+                 post_collision_monomial_symbol_base=POST_COLLISION_MONOMIAL_CENTRAL_MOMENT,
+                 **kwargs):
+        super(AbstractCentralMomentTransform, self).__init__(
+            stencil, equilibrium_density, equilibrium_velocity,
+            moment_polynomials=moment_polynomials,
+            pre_collision_symbol_base=pre_collision_symbol_base,
+            post_collision_symbol_base=post_collision_symbol_base,
+            pre_collision_monomial_symbol_base=pre_collision_monomial_symbol_base,
+            post_collision_monomial_symbol_base=post_collision_monomial_symbol_base,
+            **kwargs
+        )
+
+        assert len(self.moment_polynomials) == self.q, 'Number of moments must match stencil'
+
+    def _cm_background_shift(self, central_moments):
+        if self.background_distribution is not None:
+            shift = self.background_distribution.central_moments(central_moments, self.equilibrium_velocity)
+        else:
+            shift = (0,) * self.q
+        return sp.Matrix(shift)
+# end class AbstractRawMomentTransform
+
+
+class PdfsToCentralMomentsByMatrix(AbstractCentralMomentTransform):
+    """Transform from populations to central moment space by matrix-vector multiplication."""
+
+    def __init__(self, stencil, moment_polynomials,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 **kwargs):
+        super(PdfsToCentralMomentsByMatrix, self).__init__(
+            stencil, moment_polynomials, equilibrium_density, equilibrium_velocity, **kwargs)
+
+        moment_matrix_without_shift = moment_matrix(self.moment_polynomials, self.stencil)
+        shift_matrix = set_up_shift_matrix(self.moment_polynomials, self.stencil, equilibrium_velocity)
+
+        self.forward_matrix = shift_matrix * moment_matrix_without_shift
+        self.backward_matrix = moment_matrix_without_shift.inv() * shift_matrix.inv()
+
+    def forward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_f_to_k',
+                          return_monomials=False):
+        r"""Returns an assignment collection containing equations for pre-collision polynomial
+        central moments, expressed in terms of the pre-collision populations by matrix-multiplication.
+
+        The central moment transformation matrix :math:`K` provided by :func:`lbmpy.moments.moment_matrix` 
+        is used to compute the pre-collision moments as :math:`\mathbf{K} = K \cdot \mathbf{f}`,
+        which are returned element-wise.
+
+        Args:
+            pdf_symbols: List of symbols that represent the pre-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            return_monomials: Return equations for monomial moments. Use only when specifying 
+                              ``moment_exponents`` in constructor!
+        """
+        simplification = self._get_simp_strategy(simplification)
+
+        if return_monomials:
+            assert len(self.moment_exponents) == self.q, "Could not derive invertible monomial transform." \
+                f"Expected {self.q} monomials, but got {len(self.moment_exponents)}."
+            km = moment_matrix(self.moment_exponents, self.stencil, shift_velocity=self.equilibrium_velocity)
+            background_shift = self._cm_background_shift(self.moment_exponents)
+            pre_collision_moments = self.pre_collision_monomial_symbols
+        else:
+            km = self.forward_matrix
+            background_shift = self._cm_background_shift(self.moment_polynomials)
+            pre_collision_moments = self.pre_collision_symbols
+
+        f_to_k_vec = km * sp.Matrix(pdf_symbols) + background_shift
+        main_assignments = [Assignment(k, eq) for k, eq in zip(pre_collision_moments, f_to_k_vec)]
+
+        symbol_gen = SymbolGen(subexpression_base)
+        ac = AssignmentCollection(main_assignments, subexpression_symbol_generator=symbol_gen)
+
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    def backward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_k_to_f',
+                           start_from_monomials=False):
+        r"""Returns an assignment collection containing equations for post-collision populations, 
+        expressed in terms of the post-collision polynomial central moments by matrix-multiplication.
+
+        The moment transformation matrix :math:`K` provided by :func:`lbmpy.moments.moment_matrix` is
+        inverted and used to compute the pre-collision moments as 
+        :math:`\mathbf{f}^{\ast} = K^{-1} \cdot \mathbf{K}_{\mathrm{post}}`, which is returned element-wise.
+
+        Args:
+            pdf_symbols: List of symbols that represent the post-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            start_from_monomials: Return equations for monomial moments. Use only when specifying 
+                                  ``moment_exponents`` in constructor!
+        """
+        simplification = self._get_simp_strategy(simplification)
+
+        if start_from_monomials:
+            assert len(self.moment_exponents) == self.q, "Could not derive invertible monomial transform." \
+                f"Expected {self.q} monomials, but got {len(self.moment_exponents)}."
+            mm_inv = moment_matrix(self.moment_exponents, self.stencil).inv()
+            shift_inv = set_up_shift_matrix(self.moment_exponents, self.stencil, self.equilibrium_velocity).inv()
+            km_inv = mm_inv * shift_inv
+            background_shift = self._cm_background_shift(self.moment_exponents)
+            post_collision_moments = self.post_collision_monomial_symbols
+        else:
+            km_inv = self.backward_matrix
+            background_shift = self._cm_background_shift(self.moment_polynomials)
+            post_collision_moments = self.post_collision_symbols
+
+        symbol_gen = SymbolGen(subexpression_base)
+
+        subexpressions = [Assignment(xi, m - s)
+                          for xi, m, s in zip(symbol_gen, post_collision_moments, background_shift)]
+
+        m_to_f_vec = km_inv * sp.Matrix([s.lhs for s in subexpressions])
+        main_assignments = [Assignment(f, eq) for f, eq in zip(pdf_symbols, m_to_f_vec)]
+
+        ac = AssignmentCollection(main_assignments, subexpressions=subexpressions,
+                                  subexpression_symbol_generator=symbol_gen)
+
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    @property
+    def _default_simplification(self):
+        simplification = SimplificationStrategy()
+        return simplification
+# end class PdfsToCentralMomentsByMatrix
+
+
+class BinomialChimeraTransform(AbstractCentralMomentTransform):
+    """Transform from populations to central moments using a chimera transform implementing the binomial expansion."""
+
+    def __init__(self, stencil, moment_polynomials,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 conserved_quantity_equations=None,
+                 **kwargs):
+        super(BinomialChimeraTransform, self).__init__(
+            stencil, moment_polynomials, equilibrium_density, equilibrium_velocity,
+            conserved_quantity_equations=conserved_quantity_equations, **kwargs)
+
+        #   Potentially, de-aliasing is required
+        if len(self.moment_exponents) != self.q:
+            P, m_reduced = central_moment_reduced_monomial_to_polynomial_matrix(self.moment_polynomials,
+                                                                                self.stencil,
+                                                                                velocity_symbols=equilibrium_velocity)
+            self.mono_to_poly_matrix = P
+            self.moment_exponents = m_reduced
+        else:
+            self.mono_to_poly_matrix = monomial_to_polynomial_transformation_matrix(self.moment_exponents,
+                                                                                    self.moment_polynomials)
+
+        if 'moment_exponents' in kwargs:
+            del kwargs['moment_exponents']
+
+        self.raw_moment_transform = PdfsToMomentsByChimeraTransform(
+            stencil, None, equilibrium_density, equilibrium_velocity,
+            conserved_quantity_equations=conserved_quantity_equations,
+            moment_exponents=self.moment_exponents,
+            **kwargs)
+
+        self.poly_to_mono_matrix = self.mono_to_poly_matrix.inv()
+
+    @property
+    def absorbs_conserved_quantity_equations(self):
+        return True
+
+    def forward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_f_to_k',
+                          return_monomials=False):
+        r"""Returns equations for polynomial central moments, computed from pre-collision populations
+        through a cascade of three steps.
+
+        First, the monomial raw moment vector :math:`\mathbf{m}` is computed using the raw-moment
+        chimera transform (see `lbmpy.moment_transforms.PdfsToMomentsByChimeraTransform`).
+
+        Second, we obtain monomial central moments from monomial raw moments using the binomial
+        chimera transform:
+
+        .. math::
+
+            \kappa_{ab|\gamma} &:= \sum_{c = 0}^{\gamma} \binom{\gamma}{c} v_z^{\gamma - c} m_{abc} \\
+            \kappa_{a|\beta\gamma} &:= \sum_{b = 0}^{\beta} \binom{\beta}{b} v_z^{\beta - b} \kappa_{ab|\gamma} \\
+            \kappa_{\alpha\beta\gamma} &:=
+                \sum_{a = 0}^{\alpha} \binom{\alpha}{a} v_z^{\alpha - a} \kappa_{a|\beta\gamma} \\
+
+        Lastly, the polynomial central moments are computed using the polynomialization matrix
+        as :math:`\mathbf{K} = P \mathbf{\kappa}`.
+
+        **Conserved Quantity Equations**
+
+        If given, this transform absorbs the conserved quantity equations and simplifies them
+        using the raw moment equations, if simplification is enabled.
+
+        **Simplification**
+
+        If simplification is enabled, the absorbed conserved quantity equations are - if possible - 
+        rewritten using the monomial symbols. If the conserved quantities originate somewhere else
+        than in the lower-order moments (like from an external field), they are not affected by this
+        simplification.
+
+        The raw moment chimera transform is simplified by propagation of aliases.
+
+        The equations of the binomial chimera transform are simplified by expressing conserved raw moments
+        in terms of the conserved quantities, and subsequent propagation of aliases, constants, and any
+        expressions that are purely products of conserved quantities.
+
+        **De-Aliasing**
+
+        If more than :math:`q` monomial moments are extracted from the polynomial set, they
+        are de-aliased and reduced to a set of only :math:`q` moments using the same rules
+        as for raw moments. For polynomialization, a special reduced matrix :math:`\tilde{P}`
+        is used, which is computed using `lbmpy.moments.central_moment_reduced_monomial_to_polynomial_matrix`.
+
+
+        Args:
+            pdf_symbols: List of symbols that represent the pre-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            return_monomials: Return equations for monomial moments. Use only when specifying 
+                              ``moment_exponents`` in constructor!
+
+        """
+        simplification = self._get_simp_strategy(simplification, 'forward')
+
+        mono_raw_moment_base = self.raw_moment_transform.mono_base_pre
+        mono_central_moment_base = self.mono_base_pre
+
+        mono_cm_symbols = self.pre_collision_monomial_symbols
+
+        rm_ac = self.raw_moment_transform.forward_transform(pdf_symbols, simplification=False, return_monomials=True)
+        cq_symbols_to_moments = self.raw_moment_transform.get_cq_to_moment_symbols_dict(mono_raw_moment_base)
+
+        chim = self.BinomialChimera(tuple(-u for u in self.equilibrium_velocity),
+                                    mono_raw_moment_base, mono_central_moment_base)
+        chim_ac = chim(self.moment_exponents)
+
+        cq_subs = dict()
+        if simplification:
+            from lbmpy.methods.momentbased.momentbasedsimplifications import (
+                substitute_moments_in_conserved_quantity_equations)
+            rm_ac = substitute_moments_in_conserved_quantity_equations(rm_ac)
+
+            #   Compute replacements for conserved moments in terms of the CQE
+            rm_asm_dict = rm_ac.main_assignments_dict
+            for cq_sym, moment_sym in cq_symbols_to_moments.items():
+                cq_eq = rm_asm_dict[cq_sym]
+                solutions = sp.solve(cq_eq - cq_sym, moment_sym)
+                if len(solutions) > 0:
+                    cq_subs[moment_sym] = solutions[0]
+
+            chim_ac = chim_ac.new_with_substitutions(cq_subs, substitute_on_lhs=False)
+
+            fo_kappas = [sq_sym(mono_central_moment_base, es) for es in moments_of_order(1, dim=self.stencil.D)]
+            ac_filtered = chim_ac.new_filtered(fo_kappas).new_without_subexpressions()
+            chim_asm_dict = chim_ac.main_assignments_dict
+            for asm in ac_filtered.main_assignments:
+                chim_asm_dict[asm.lhs] = asm.rhs
+            chim_ac.set_main_assignments_from_dict(chim_asm_dict)
+
+        subexpressions = rm_ac.all_assignments + chim_ac.subexpressions
+
+        if return_monomials:
+            main_assignments = chim_ac.main_assignments
+        else:
+            subexpressions += chim_ac.main_assignments
+            poly_eqs = self.mono_to_poly_matrix * sp.Matrix(mono_cm_symbols)
+            main_assignments = [Assignment(m, v) for m, v in zip(self.pre_collision_symbols, poly_eqs)]
+
+        symbol_gen = SymbolGen(subexpression_base)
+        ac = AssignmentCollection(main_assignments=main_assignments, subexpressions=subexpressions,
+                                  subexpression_symbol_generator=symbol_gen)
+
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    def backward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_k_to_f',
+                           start_from_monomials=False):
+        r"""Returns an assignment collection containing equations for post-collision populations, 
+        expressed in terms of the post-collision polynomial central moments by three steps.
+
+        The post-collision monomial central moments :math:`\mathbf{\kappa}_{\mathrm{post}}` are first 
+        obtained from the polynomials through multiplication with :math:`P^{-1}`.
+
+        Afterward, monomial post-collision raw moments are obtained from monomial central moments using the binomial
+        chimera transform:
+
+        .. math::
+
+            m^{\ast}_{ab|\gamma} &:= \sum_{c = 0}^{\gamma} \binom{\gamma}{c} v_z^{\gamma - c} \kappa^{\ast}_{abc} \\
+            m^{\ast}_{a|\beta\gamma} &:= \sum_{b = 0}^{\beta} \binom{\beta}{b} v_z^{\beta - b} m^{\ast}_{ab|\gamma} \\
+            m^{\ast}_{\alpha\beta\gamma} &:=
+                \sum_{a = 0}^{\alpha} \binom{\alpha}{a} v_z^{\alpha - a} m^{\ast}_{a|\beta\gamma} \\
+
+        Finally, the monomial raw moment transformation 
+        matrix :math:`M_r` provided by :func:`lbmpy.moments.moment_matrix` 
+        is inverted and used to compute the pre-collision moments as 
+        :math:`\mathbf{f}_{\mathrm{post}} = M_r^{-1} \cdot \mathbf{m}_{\mathrm{post}}`.
+
+        **De-Aliasing**: 
+
+        See `PdfsToCentralMomentsByShiftMatrix.forward_transform`.
+
+        **Simplifications**
+
+        If simplification is enabled, the inverse shift matrix equations are simplified by recursively 
+        inserting lower-order moments into equations for higher-order moments. To this end, these equations 
+        are factored recursively by the velocity symbols.
+
+        The equations of the binomial chimera transform are simplified by propagation of aliases.
+
+        Further, the equations for populations :math:`f_i` and :math:`f_{\bar{i}}` 
+        of opposite stencil directions :math:`\mathbf{c}_i` and :math:`\mathbf{c}_{\bar{i}} = - \mathbf{c}_i`
+        are split into their symmetric and antisymmetric parts :math:`f_i^{\mathrm{sym}}, f_i^{\mathrm{anti}}`, such
+        that
+
+        .. math::
+
+            f_i = f_i^{\mathrm{sym}} + f_i^{\mathrm{anti}}
+
+            f_{\bar{i}} = f_i^{\mathrm{sym}} - f_i^{\mathrm{anti}}
+
+
+        Args:
+            pdf_symbols: List of symbols that represent the post-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            start_from_monomials: Return equations for monomial moments. Use only when specifying 
+                                  ``moment_exponents`` in constructor!
+        """
+        simplification = self._get_simp_strategy(simplification, 'backward')
+
+        mono_cm_symbols = self.post_collision_monomial_symbols
+
+        subexpressions = []
+        if not start_from_monomials:
+            mono_eqs = self.poly_to_mono_matrix * sp.Matrix(self.post_collision_symbols)
+            subexpressions += [Assignment(cm, v) for cm, v in zip(mono_cm_symbols, mono_eqs)]
+
+        mono_raw_moment_base = self.raw_moment_transform.mono_base_post
+        mono_central_moment_base = self.mono_base_post
+
+        chim = self.BinomialChimera(self.equilibrium_velocity, mono_central_moment_base, mono_raw_moment_base)
+        chim_ac = chim(self.moment_exponents)
+
+        if simplification:
+            from pystencils.simp import insert_aliases
+            chim_ac = insert_aliases(chim_ac)
+
+        subexpressions += chim_ac.all_assignments
+
+        rm_ac = self.raw_moment_transform.backward_transform(
+            pdf_symbols, simplification=False, start_from_monomials=True)
+        subexpressions += rm_ac.subexpressions
+
+        ac = rm_ac.copy(subexpressions=subexpressions)
+        if simplification:
+            ac = simplification.apply(ac)
+
+        return ac
+
+    #   ----------------------------- Private Members -----------------------------
+
+    class BinomialChimera:
+        def __init__(self, v, from_base, to_base):
+            self._v = v
+            self._from_base = from_base
+            self._to_base = to_base
+            self._chim_dict = None
+
+        def _chimera_symbol(self, fixed_directions, remaining_exponents):
+            if not fixed_directions:
+                return None
+
+            fixed_str = '_'.join(str(direction) for direction in fixed_directions)
+            exp_str = '_'.join(str(exp) for exp in remaining_exponents)
+            return sp.Symbol(f"chimera_{self._to_base}_{fixed_str}_e_{exp_str}")
+
+        @property
+        def chimera_assignments(self):
+            assert self._chim_dict is not None
+            return [Assignment(lhs, rhs) for lhs, rhs in self._chim_dict.items()]
+
+        def _derive(self, exponents, depth):
+            if depth == len(exponents):
+                return sq_sym(self._from_base, exponents)
+
+            v = self._v
+
+            fixed = exponents[:depth]
+            remaining = exponents[depth:]
+            chim_symb = self._chimera_symbol(fixed, remaining)
+            if chim_symb in self._chim_dict:
+                return chim_symb
+
+            choose = sp.binomial
+
+            alpha = exponents[depth]
+            s = sp.Integer(0)
+            for a in range(alpha + 1):
+                rec_exps = list(exponents)
+                rec_exps[depth] = a
+                s += choose(alpha, a) * v[depth] ** (alpha - a) * self._derive(rec_exps, depth + 1)
+
+            if chim_symb is not None:
+                self._chim_dict[chim_symb] = s
+                return chim_symb
+            else:
+                return Assignment(sq_sym(self._to_base, exponents), s)
+
+        def __call__(self, monos):
+            self._chim_dict = dict()
+            ac = []
+            for m in monos:
+                ac.append(self._derive(m, 0))
+            return AssignmentCollection(ac, self._chim_dict)
+
+    @property
+    def _default_simplification(self):
+        from pystencils.simp import insert_aliases, insert_constants
+        from lbmpy.methods.momentbased.momentbasedsimplifications import insert_pure_products
+
+        cq = (self.equilibrium_density,) + self.equilibrium_velocity
+        fw_skip = cq + self.raw_moment_transform.pre_collision_monomial_symbols + self.pre_collision_monomial_symbols
+
+        forward_simp = SimplificationStrategy()
+        forward_simp.add(partial(insert_pure_products, symbols=cq, skip=fw_skip))
+        forward_simp.add(partial(insert_aliases, skip=fw_skip))
+        forward_simp.add(partial(insert_constants, skip=fw_skip))
+
+        from lbmpy.methods.momentbased.momentbasedsimplifications import split_pdf_main_assignments_by_symmetry
+
+        bw_skip = self.raw_moment_transform.post_collision_monomial_symbols + self.post_collision_monomial_symbols
+
+        backward_simp = SimplificationStrategy()
+        backward_simp.add(partial(insert_aliases, skip=bw_skip))
+        backward_simp.add(split_pdf_main_assignments_by_symmetry)
+        backward_simp.add(add_subexpressions_for_constants)
+
+        return {
+            'forward': forward_simp,
+            'backward': backward_simp
+        }
+
+# end class PdfsToCentralMomentsByShiftMatrix
+
+
+class FastCentralMomentTransform(AbstractCentralMomentTransform):
+    """Transform from populations to central moments, using the fast central-moment
+    transform equations introduced by :cite:`geier2015`.
+
+    **Attention:** The fast central moment transform has originally been designed for the
+    D3Q27 stencil, and is also tested and safely usable with D2Q9 and D3Q19. While the forward-
+    transform does not pose any problems, the backward equations may be inefficient, or
+    even not cleanly derivable for other stencils. Use with care!"""
+
+    def __init__(self, stencil,
+                 moment_polynomials,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 conserved_quantity_equations=None,
+                 **kwargs):
+        super(FastCentralMomentTransform, self).__init__(
+            stencil, moment_polynomials, equilibrium_density, equilibrium_velocity,
+            conserved_quantity_equations=conserved_quantity_equations, **kwargs)
+
+        #   Potentially, de-aliasing is required
+        if len(self.moment_exponents) != self.q:
+            P, m_reduced = central_moment_reduced_monomial_to_polynomial_matrix(self.moment_polynomials,
+                                                                                self.stencil,
+                                                                                velocity_symbols=equilibrium_velocity)
+            self.mono_to_poly_matrix = P
+            self.moment_exponents = m_reduced
+        else:
+            self.mono_to_poly_matrix = monomial_to_polynomial_transformation_matrix(self.moment_exponents,
+                                                                                    self.moment_polynomials)
+
+        self.poly_to_mono_matrix = self.mono_to_poly_matrix.inv()
+
+        moment_matrix_without_shift = moment_matrix(self.moment_exponents, self.stencil)
+        shift_matrix = set_up_shift_matrix(self.moment_exponents, self.stencil, equilibrium_velocity)
+        self.inv_monomial_matrix = moment_matrix_without_shift.inv() * shift_matrix.inv()
+
+    def forward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_f_to_k',
+                          return_monomials=False):
+        r"""Returns an assignment collection containing equations for pre-collision polynomial
+        central moments, expressed in terms of the pre-collision populations.
+
+        The monomial central moments are computed from populations through the central-moment
+        chimera transform:
+
+        .. math::
+
+            f_{xyz} &:= f_i \text{ such that } c_i = (x,y,z)^T \\
+            \kappa_{xy|\gamma} &:= \sum_{z \in \{-1, 0, 1\} } f_{xyz} \cdot (z - u_z)^{\gamma} \\
+            \kappa_{x|\beta \gamma} &:= \sum_{y \in \{-1, 0, 1\}} \kappa_{xy|\gamma} \cdot (y - u_y)^{\beta} \\
+            \kappa_{\alpha \beta \gamma} &:= \sum_{x \in \{-1, 0, 1\}} \kappa_{x|\beta \gamma} \cdot (x - u_x)^{\alpha}
+
+        The polynomial moments are afterward computed from the monomials by matrix-multiplication 
+        using the polynomialization matrix :math:`P`.
+
+        **De-Aliasing**
+
+        If more than :math:`q` monomial moments are extracted from the polynomial set, they
+        are de-aliased and reduced to a set of only :math:`q` moments using the same rules
+        as for raw moments. For polynomialization, a special reduced matrix :math:`\tilde{P}`
+        is used, which is computed using `lbmpy.moments.central_moment_reduced_monomial_to_polynomial_matrix`.
+
+
+        Args:
+            pdf_symbols: List of symbols that represent the pre-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            return_monomials: Return equations for monomial moments. Use only when specifying 
+                              ``moment_exponents`` in constructor!
+        """
+        simplification = self._get_simp_strategy(simplification, 'forward')
+        monomial_symbol_base = self.mono_base_pre
+
+        def _partial_kappa_symbol(fixed_directions, remaining_exponents):
+            fixed_str = '_'.join(str(direction) for direction in fixed_directions).replace('-', 'm')
+            exp_str = '_'.join(str(exp) for exp in remaining_exponents).replace('-', 'm')
+            return sp.Symbol(f"partial_{monomial_symbol_base}_{fixed_str}_e_{exp_str}")
+
+        partial_sums_dict = dict()
+        monomial_eqs = []
+
+        def collect_partial_sums(exponents, dimension=0, fixed_directions=tuple()):
+            if dimension == self.dim:
+                #   Base Case
+                if fixed_directions in self.stencil:
+                    return pdf_symbols[self.stencil.index(fixed_directions)]
+                else:
+                    return 0
+            else:
+                #   Recursive Case
+                summation = sp.sympify(0)
+                for d in [-1, 0, 1]:
+                    next_partial = collect_partial_sums(
+                        exponents, dimension=dimension + 1, fixed_directions=fixed_directions + (d,))
+                    summation += next_partial * (d - self.equilibrium_velocity[dimension])**exponents[dimension]
+
+                if dimension == 0:
+                    lhs_symbol = sq_sym(monomial_symbol_base, exponents)
+                    monomial_eqs.append(Assignment(lhs_symbol, summation))
+                else:
+                    lhs_symbol = _partial_kappa_symbol(fixed_directions, exponents[dimension:])
+                    partial_sums_dict[lhs_symbol] = summation
+                return lhs_symbol
+
+        for e in self.moment_exponents:
+            collect_partial_sums(e)
+
+        subexpressions = [Assignment(lhs, rhs) for lhs, rhs in partial_sums_dict.items()]
+
+        if return_monomials:
+            shift = self._cm_background_shift(self.moment_exponents)
+            main_assignments = [Assignment(a.lhs, a.rhs + s) for a, s in zip(monomial_eqs, shift)]
+        else:
+            subexpressions += monomial_eqs
+            moment_eqs = self.mono_to_poly_matrix * sp.Matrix(self.pre_collision_monomial_symbols)
+            moment_eqs += self._cm_background_shift(self.moment_polynomials)
+            main_assignments = [Assignment(m, v) for m, v in zip(self.pre_collision_symbols, moment_eqs)]
+
+        symbol_gen = SymbolGen(subexpression_base)
+        ac = AssignmentCollection(main_assignments, subexpressions=subexpressions,
+                                  subexpression_symbol_generator=symbol_gen)
+        if simplification:
+            ac = self._simplify_lower_order_moments(ac, monomial_symbol_base, return_monomials)
+            ac = simplification.apply(ac)
+        return ac
+
+    def backward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_k_to_f',
+                           start_from_monomials=False):
+        r"""Returns an assignment collection containing equations for post-collision populations, 
+        expressed in terms of the post-collision polynomial central moments using the backward
+        fast central moment transform.
+
+        First, monomial central moments are obtained from the polynomial moments by multiplication
+        with :math:`P^{-1}`. Then, the elementwise equations of the matrix 
+        multiplication :math:`K^{-1} \cdot \mathbf{K}` with the monomial central moment matrix 
+        (see `PdfsToCentralMomentsByMatrix`) are recursively simplified by extracting certain linear 
+        combinations of velocities, to obtain equations similar to the ones given in :cite:`geier2015`.
+
+        The backward transform is designed for D3Q27, inherently generalizes to D2Q9, and is tested
+        for D3Q19. It also returns correct equations for D3Q15, whose efficiency is however questionable.
+
+        **De-Aliasing**: 
+
+        See `FastCentralMomentTransform.forward_transform`.
+
+        Args:
+            pdf_symbols: List of symbols that represent the post-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            start_from_monomials: Return equations for monomial moments. Use only when specifying 
+                                  ``moment_exponents`` in constructor!
+        """
+        simplification = self._get_simp_strategy(simplification, 'backward')
+
+        post_collision_moments = self.post_collision_symbols
+        post_collision_monomial_moments = self.post_collision_monomial_symbols
+
+        symbol_gen = SymbolGen(subexpression_base)
+
+        subexpressions = []
+        if not start_from_monomials:
+            background_shift = self._cm_background_shift(self.moment_polynomials)
+            shift_equations = [Assignment(xi, m - s)
+                               for xi, m, s in zip(symbol_gen, post_collision_moments, background_shift)]
+            monomial_eqs = self.poly_to_mono_matrix * sp.Matrix([s.lhs for s in shift_equations])
+            subexpressions += shift_equations
+            subexpressions += [Assignment(m, v) for m, v in zip(post_collision_monomial_moments, monomial_eqs)]
+        else:
+            background_shift = self._cm_background_shift(self.moment_exponents)
+            shift_equations = [Assignment(xi, m - s)
+                               for xi, m, s in zip(symbol_gen, post_collision_monomial_moments, background_shift)]
+            subexpressions += shift_equations
+            post_collision_monomial_moments = [s.lhs for s in shift_equations]
+
+        raw_equations = self.inv_monomial_matrix * sp.Matrix(post_collision_monomial_moments)
+        raw_equations = [Assignment(f, eq) for f, eq in zip(pdf_symbols, raw_equations)]
+
+        ac = self._split_backward_equations(raw_equations, symbol_gen)
+        ac.subexpressions = subexpressions + ac.subexpressions
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    #   ----------------------------- Private Members -----------------------------
+
+    @property
+    def _default_simplification(self):
+        forward_simp = SimplificationStrategy()
+
+        backward_simp = SimplificationStrategy()
+        backward_simp.add(add_subexpressions_for_constants)
+
+        return {
+            'forward': forward_simp,
+            'backward': backward_simp
+        }
+
+    def _simplify_lower_order_moments(self, ac, moment_base, search_in_main_assignments):
+        if self.cqe is None:
+            return ac
+
+        moment_symbols = [sq_sym(moment_base, e) for e in moments_up_to_order(1, dim=self.dim)]
+
+        if search_in_main_assignments:
+            f_to_cm_dict = ac.main_assignments_dict
+            f_to_cm_dict_reduced = ac.new_without_subexpressions().main_assignments_dict
+        else:
+            f_to_cm_dict = ac.subexpressions_dict
+            f_to_cm_dict_reduced = ac.new_without_subexpressions(moment_symbols).subexpressions_dict
+
+        cqe_subs = self.cqe.new_without_subexpressions().main_assignments_dict
+        for m in moment_symbols:
+            m_eq = fast_subs(fast_subs(f_to_cm_dict_reduced[m], cqe_subs), cqe_subs)
+            m_eq = m_eq.expand().cancel()
+            for cqe_sym, cqe_exp in cqe_subs.items():
+                m_eq = subs_additive(m_eq, cqe_sym, cqe_exp)
+            f_to_cm_dict[m] = m_eq
+
+        if search_in_main_assignments:
+            main_assignments = [Assignment(lhs, rhs) for lhs, rhs in f_to_cm_dict.items()]
+            return ac.copy(main_assignments=main_assignments)
+        else:
+            subexpressions = [Assignment(lhs, rhs) for lhs, rhs in f_to_cm_dict.items()]
+            return ac.copy(subexpressions=subexpressions)
+
+    def _split_backward_equations_recursive(self, assignment, all_subexpressions,
+                                            stencil_direction, subexp_symgen, known_coeffs_dict,
+                                            step=0):
+        #   Base Cases
+        # if step == self.dim:
+        #     return assignment
+
+        #   Base Case
+        #   If there are no more velocity symbols in the subexpression,
+        #   don't split it up further
+        if assignment.rhs.atoms(sp.Symbol).isdisjoint(set(self.equilibrium_velocity)):
+            return assignment
+
+        #   Recursive Case
+
+        u = self.equilibrium_velocity[-1 - step]
+        d = stencil_direction[-1 - step]
+        one = sp.Integer(1)
+        two = sp.Integer(2)
+
+        #   Factors to group terms by
+        grouping_factors = {
+            -1: [one, 2 * u - 1, u**2 - u],
+            0: [-one, -2 * u, 1 - u**2],
+            1: [one, 2 * u + 1, u**2 + u]
+        }
+        factors = grouping_factors[d]
+
+        #   Common Integer factor to extract from all groups
+        common_factor = one if d == 0 else two
+
+        #   Proxy for factor grouping
+        v = sp.Symbol('v')
+        square_factor_eq = (factors[2] - v**2)
+        lin_factor_eq = (factors[1] - v)
+        sub_for_u_sq = sp.solve(square_factor_eq, u**2)[0]
+        sub_for_u = sp.solve(lin_factor_eq, u)[0]
+        subs = {u**2: sub_for_u_sq, u: sub_for_u}
+        rhs_grouped_by_v = assignment.rhs.subs(subs).expand().collect(v)
+
+        new_rhs = sp.Integer(0)
+        for k in range(3):
+            coeff = rhs_grouped_by_v.coeff(v, k)
+            coeff_subexp = common_factor * coeff
+
+            #   Explicitly divide out the constant factor in the zero case
+            if k == 0:
+                coeff_subexp = coeff_subexp / factors[0]
+
+            #   MEMOISATION:
+            #   The subexpression just generated might already have been found
+            #   If so, reuse the existing symbol and skip forward.
+            #   Otherwise, create it anew and continue recursively
+            coeff_symb = known_coeffs_dict.get(coeff_subexp, None)
+            if coeff_symb is None:
+                coeff_symb = next(subexp_symgen)
+                known_coeffs_dict[coeff_subexp] = coeff_symb
+                coeff_assignment = Assignment(coeff_symb, coeff_subexp)
+
+                #   Recursively split the coefficient term
+                coeff_assignment = self._split_backward_equations_recursive(
+                    coeff_assignment, all_subexpressions, stencil_direction, subexp_symgen,
+                    known_coeffs_dict, step=step + 1)
+                all_subexpressions.append(coeff_assignment)
+
+            new_rhs += factors[k] * coeff_symb
+
+        new_rhs = sp.Rational(1, common_factor) * new_rhs
+
+        return Assignment(assignment.lhs, new_rhs)
+
+    def _split_backward_equations(self, backward_assignments, subexp_symgen):
+        all_subexpressions = []
+        split_main_assignments = []
+        known_coeffs_dict = dict()
+        for asm, stencil_dir in zip(backward_assignments, self.stencil):
+            split_asm = self._split_backward_equations_recursive(
+                asm, all_subexpressions, stencil_dir, subexp_symgen, known_coeffs_dict)
+            split_main_assignments.append(split_asm)
+        ac = AssignmentCollection(split_main_assignments, subexpressions=all_subexpressions,
+                                  subexpression_symbol_generator=subexp_symgen)
+        ac.topological_sort(sort_main_assignments=False)
+        return ac
+
+# end class FastCentralMomentTransform
+
+
+class PdfsToCentralMomentsByShiftMatrix(AbstractCentralMomentTransform):
+    """Transform from populations to central moments using a shift matrix."""
+
+    def __init__(self, stencil, moment_polynomials,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 conserved_quantity_equations=None,
+                 **kwargs):
+        super(PdfsToCentralMomentsByShiftMatrix, self).__init__(
+            stencil, moment_polynomials, equilibrium_density, equilibrium_velocity,
+            conserved_quantity_equations=conserved_quantity_equations, **kwargs)
+
+        #   Potentially, de-aliasing is required
+        if len(self.moment_exponents) != self.q:
+            P, m_reduced = central_moment_reduced_monomial_to_polynomial_matrix(self.moment_polynomials,
+                                                                                self.stencil,
+                                                                                velocity_symbols=equilibrium_velocity)
+            self.mono_to_poly_matrix = P
+            self.moment_exponents = m_reduced
+        else:
+            self.mono_to_poly_matrix = monomial_to_polynomial_transformation_matrix(self.moment_exponents,
+                                                                                    self.moment_polynomials)
+
+        if 'moment_exponents' in kwargs:
+            del kwargs['moment_exponents']
+
+        self.raw_moment_transform = PdfsToMomentsByChimeraTransform(
+            stencil, None, equilibrium_density, equilibrium_velocity,
+            conserved_quantity_equations=conserved_quantity_equations,
+            moment_exponents=self.moment_exponents,
+            **kwargs)
+
+        self.shift_matrix = set_up_shift_matrix(self.moment_exponents, self.stencil, self.equilibrium_velocity)
+        self.inv_shift_matrix = self.shift_matrix.inv()
+        self.poly_to_mono_matrix = self.mono_to_poly_matrix.inv()
+
+    @property
+    def absorbs_conserved_quantity_equations(self):
+        return True
+
+    def forward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_f_to_k',
+                          return_monomials=False):
+        r"""Returns equations for polynomial central moments, computed from pre-collision populations
+        through a cascade of three steps.
+
+        First, the monomial raw moment vector :math:`\mathbf{m}` is computed using the raw-moment
+        chimera transform (see `lbmpy.moment_transforms.PdfsToMomentsByChimeraTransform`). Then, the
+        monomial shift matrix :math:`N` provided by `lbmpy.moments.set_up_shift_matrix` is used to compute
+        the monomial central moment vector as :math:`\mathbf{\kappa} = N \mathbf{m}`. Lastly, the polynomial
+        central moments are computed using the polynomialization matrix as :math:`\mathbf{K} = P \mathbf{\kappa}`.
+
+        **Conserved Quantity Equations**
+
+        If given, this transform absorbs the conserved quantity equations and simplifies them
+        using the raw moment equations, if simplification is enabled.
+
+
+        **Simplification**
+
+        If simplification is enabled, the absorbed conserved quantity equations are - if possible - 
+        rewritten using the monomial symbols. If the conserved quantities originate somewhere else
+        than in the lower-order moments (like from an external field), they are not affected by this
+        simplification.
+
+        The relations between conserved quantities and raw moments are used to simplify the equations
+        obtained from the shift matrix. Further, these equations are simplified by recursively inserting
+        lower-order moments into equations for higher-order moments.
+
+         **De-Aliasing**
+
+        If more than :math:`q` monomial moments are extracted from the polynomial set, they
+        are de-aliased and reduced to a set of only :math:`q` moments using the same rules
+        as for raw moments. For polynomialization, a special reduced matrix :math:`\tilde{P}`
+        is used, which is computed using `lbmpy.moments.central_moment_reduced_monomial_to_polynomial_matrix`.
+
+
+        Args:
+            pdf_symbols: List of symbols that represent the pre-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            return_monomials: Return equations for monomial moments. Use only when specifying 
+                              ``moment_exponents`` in constructor!
+
+        """
+        simplification = self._get_simp_strategy(simplification, 'forward')
+
+        raw_moment_base = self.raw_moment_transform.mono_base_pre
+        central_moment_base = self.mono_base_pre
+
+        mono_rm_symbols = self.raw_moment_transform.pre_collision_monomial_symbols
+        mono_cm_symbols = self.pre_collision_monomial_symbols
+
+        rm_ac = self.raw_moment_transform.forward_transform(pdf_symbols, simplification=False, return_monomials=True)
+        cq_symbols_to_moments = self.raw_moment_transform.get_cq_to_moment_symbols_dict(raw_moment_base)
+        rm_to_cm_vec = self.shift_matrix * sp.Matrix(mono_rm_symbols)
+
+        cq_subs = dict()
+        if simplification:
+            from lbmpy.methods.momentbased.momentbasedsimplifications import (
+                substitute_moments_in_conserved_quantity_equations)
+            rm_ac = substitute_moments_in_conserved_quantity_equations(rm_ac)
+
+            #   Compute replacements for conserved moments in terms of the CQE
+            rm_asm_dict = rm_ac.main_assignments_dict
+            for cq_sym, moment_sym in cq_symbols_to_moments.items():
+                cq_eq = rm_asm_dict[cq_sym]
+                solutions = sp.solve(cq_eq - cq_sym, moment_sym)
+                if len(solutions) > 0:
+                    cq_subs[moment_sym] = solutions[0]
+
+            rm_to_cm_vec = fast_subs(rm_to_cm_vec, cq_subs)
+
+        rm_to_cm_dict = {cm: rm for cm, rm in zip(mono_cm_symbols, rm_to_cm_vec)}
+
+        if simplification:
+            rm_to_cm_dict = self._simplify_raw_to_central_moments(
+                rm_to_cm_dict, self.moment_exponents, raw_moment_base, central_moment_base)
+            rm_to_cm_dict = self._undo_remaining_cq_subexpressions(rm_to_cm_dict, cq_subs)
+
+        subexpressions = rm_ac.all_assignments
+
+        if return_monomials:
+            main_assignments = [Assignment(lhs, rhs) for lhs, rhs in rm_to_cm_dict.items()]
+        else:
+            subexpressions += [Assignment(lhs, rhs) for lhs, rhs in rm_to_cm_dict.items()]
+            poly_eqs = self.mono_to_poly_matrix * sp.Matrix(mono_cm_symbols)
+            main_assignments = [Assignment(m, v) for m, v in zip(self.pre_collision_symbols, poly_eqs)]
+
+        symbol_gen = SymbolGen(subexpression_base)
+        ac = AssignmentCollection(main_assignments=main_assignments, subexpressions=subexpressions,
+                                  subexpression_symbol_generator=symbol_gen)
+
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    def backward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_k_to_f',
+                           start_from_monomials=False):
+        r"""Returns an assignment collection containing equations for post-collision populations, 
+        expressed in terms of the post-collision polynomial central moments by matrix-multiplication
+        including the shift matrix.
+
+        The post-collision monomial central moments :math:`\mathbf{\kappa}_{\mathrm{post}}` are first 
+        obtained from the polynomials through multiplication with :math:`P^{-1}`.
+        The shift-matrix is inverted as well, to obtain the monomial raw moments as 
+        :math:`\mathbf{m}_{post} = N^{-1} \mathbf{\kappa}_{post}`. Finally, the monomial raw moment transformation 
+        matrix :math:`M_r` provided by :func:`lbmpy.moments.moment_matrix` 
+        is inverted and used to compute the pre-collision moments as 
+        :math:`\mathbf{f}_{\mathrm{post}} = M_r^{-1} \cdot \mathbf{m}_{\mathrm{post}}`.
+
+         **De-Aliasing**: 
+
+        See `PdfsToCentralMomentsByShiftMatrix.forward_transform`.
+
+        **Simplifications**
+
+        If simplification is enabled, the inverse shift matrix equations are simplified by recursively 
+        inserting lower-order moments into equations for higher-order moments. To this end, these equations 
+        are factored recursively by the velocity symbols.
+
+        Further, the equations for populations :math:`f_i` and :math:`f_{\bar{i}}` 
+        of opposite stencil directions :math:`\mathbf{c}_i` and :math:`\mathbf{c}_{\bar{i}} = - \mathbf{c}_i`
+        are split into their symmetric and antisymmetric parts :math:`f_i^{\mathrm{sym}}, f_i^{\mathrm{anti}}`, such
+        that
+
+        .. math::
+
+            f_i = f_i^{\mathrm{sym}} + f_i^{\mathrm{anti}}
+
+            f_{\bar{i}} = f_i^{\mathrm{sym}} - f_i^{\mathrm{anti}}
+
+
+        Args:
+            pdf_symbols: List of symbols that represent the post-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            start_from_monomials: Return equations for monomial moments. Use only when specifying 
+                                  ``moment_exponents`` in constructor!
+        """
+        simplification = self._get_simp_strategy(simplification, 'backward')
+
+        mono_rm_symbols = self.raw_moment_transform.post_collision_monomial_symbols
+        mono_cm_symbols = self.post_collision_monomial_symbols
+
+        subexpressions = []
+        if not start_from_monomials:
+            mono_eqs = self.poly_to_mono_matrix * sp.Matrix(self.post_collision_symbols)
+            subexpressions += [Assignment(cm, v) for cm, v in zip(mono_cm_symbols, mono_eqs)]
+
+        cm_to_rm_vec = self.inv_shift_matrix * sp.Matrix(mono_cm_symbols)
+        cm_to_rm_dict = {rm: eq for rm, eq in zip(mono_rm_symbols, cm_to_rm_vec)}
+
+        if simplification:
+            cm_to_rm_dict = self._factor_backward_eqs_by_velocities(mono_rm_symbols, cm_to_rm_dict)
+
+        rm_ac = self.raw_moment_transform.backward_transform(
+            pdf_symbols, simplification=False, start_from_monomials=True)
+        subexpressions += [Assignment(lhs, rhs) for lhs, rhs in cm_to_rm_dict.items()]
+        subexpressions += rm_ac.subexpressions
+        ac = rm_ac.copy(subexpressions=subexpressions)
+        if simplification:
+            ac = simplification.apply(ac)
+
+        return ac
+
+    #   ----------------------------- Private Members -----------------------------
+
+    def _simplify_raw_to_central_moments(self, rm_to_cm_dict, moment_exponents, raw_moment_base, central_moment_base):
+        for cm in moment_exponents:
+            if sum(cm) < 2:
+                continue
+            cm_symb = sq_sym(central_moment_base, cm)
+            cm_asm_rhs = rm_to_cm_dict[cm_symb]
+            for m in contained_moments(cm, min_order=2)[::-1]:
+                contained_rm_symb = sq_sym(raw_moment_base, m)
+                contained_cm_symb = sq_sym(central_moment_base, m)
+                contained_cm_eq = rm_to_cm_dict[contained_cm_symb]
+                rm_in_terms_of_cm = sp.solve(contained_cm_eq - contained_cm_symb, contained_rm_symb)[0]
+                cm_asm_rhs = cm_asm_rhs.subs({contained_rm_symb: rm_in_terms_of_cm}).expand()
+            rm_to_cm_dict[cm_symb] = cm_asm_rhs
+        return rm_to_cm_dict
+
+    def _undo_remaining_cq_subexpressions(self, rm_to_cm_dict, cq_subs):
+        for cm, cm_eq in rm_to_cm_dict.items():
+            for rm, rm_subexp in cq_subs.items():
+                cm_eq = subs_additive(cm_eq, rm, rm_subexp)
+            rm_to_cm_dict[cm] = cm_eq
+        return rm_to_cm_dict
+
+    def _factor_backward_eqs_by_velocities(self, symbolic_rms, cm_to_rm_dict, required_match_replacement=0.75):
+        velocity_by_occurences = dict()
+        for rm, rm_eq in cm_to_rm_dict.items():
+            velocity_by_occurences[rm] = sorted(self.equilibrium_velocity, key=rm_eq.count, reverse=True)
+        for d in range(self.dim):
+            for rm, rm_eq in cm_to_rm_dict.items():
+                u_sorted = velocity_by_occurences[rm]
+                cm_to_rm_dict[rm] = rm_eq.expand().collect(u_sorted[d])
+
+            for i, rm1 in enumerate(symbolic_rms):
+                for _, rm2 in enumerate(symbolic_rms[i + 1:]):
+                    cm_to_rm_dict[rm2] = subs_additive(
+                        cm_to_rm_dict[rm2], rm1, cm_to_rm_dict[rm1],
+                        required_match_replacement=required_match_replacement)
+        return cm_to_rm_dict
+
+    @property
+    def _default_simplification(self):
+        forward_simp = SimplificationStrategy()
+
+        from lbmpy.methods.momentbased.momentbasedsimplifications import split_pdf_main_assignments_by_symmetry
+
+        backward_simp = SimplificationStrategy()
+        backward_simp.add(split_pdf_main_assignments_by_symmetry)
+        backward_simp.add(add_subexpressions_for_constants)
+
+        return {
+            'forward': forward_simp,
+            'backward': backward_simp
+        }
+
+# end class PdfsToCentralMomentsByShiftMatrix

src/lbmpy/moment_transforms/cumulanttransforms.py

+import numpy as np
+import sympy as sp
+
+from pystencils import Assignment, AssignmentCollection
+from pystencils.simp import SimplificationStrategy
+from pystencils.simp.assignment_collection import SymbolGen
+
+from lbmpy.moments import (
+    moments_up_to_order, statistical_quantity_symbol, exponent_tuple_sort_key,
+    monomial_to_polynomial_transformation_matrix
+)
+
+from itertools import product, chain
+
+from .abstractmomenttransform import (
+    AbstractMomentTransform,
+    PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT, POST_COLLISION_MONOMIAL_CENTRAL_MOMENT,
+    PRE_COLLISION_CUMULANT, POST_COLLISION_CUMULANT,
+    PRE_COLLISION_MONOMIAL_CUMULANT, POST_COLLISION_MONOMIAL_CUMULANT
+)
+
+#   ======================= Central Moments <-> Cumulants ==============================================================
+
+WAVE_NUMBER_SYMBOLS = sp.symbols('Xi_x, Xi_y, Xi_z')
+
+
+def moment_index_from_derivative(d, variables):
+    diffs = d.args[1:]
+    indices = [0] * len(variables)
+    for var, count in diffs:
+        indices[variables.index(var)] = count
+    return tuple(indices)
+
+
+def derivative_as_statistical_quantity(d, variables, quantity_name):
+    indices = moment_index_from_derivative(d, variables)
+    return statistical_quantity_symbol(quantity_name, indices)
+
+
+def count_derivatives(derivative):
+    return np.sum(np.fromiter((d[1] for d in derivative.args[1:]), int))
+
+
+#   ============= Transformation through cumulant-generating function =============
+
+class CentralMomentsToCumulantsByGeneratingFunc(AbstractMomentTransform):
+
+    def __init__(self, stencil, cumulant_polynomials,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 cumulant_exponents=None,
+                 pre_collision_symbol_base=PRE_COLLISION_CUMULANT,
+                 post_collision_symbol_base=POST_COLLISION_CUMULANT,
+                 pre_collision_monomial_symbol_base=PRE_COLLISION_MONOMIAL_CUMULANT,
+                 post_collision_monomial_symbol_base=POST_COLLISION_MONOMIAL_CUMULANT,
+                 **kwargs):
+        super(CentralMomentsToCumulantsByGeneratingFunc, self).__init__(
+            stencil, equilibrium_density, equilibrium_velocity,
+            moment_polynomials=cumulant_polynomials,
+            moment_exponents=cumulant_exponents,
+            pre_collision_symbol_base=pre_collision_symbol_base,
+            post_collision_symbol_base=post_collision_symbol_base,
+            pre_collision_monomial_symbol_base=pre_collision_monomial_symbol_base,
+            post_collision_monomial_symbol_base=post_collision_monomial_symbol_base,
+            **kwargs)
+
+        self.cumulant_exponents = self.moment_exponents
+        self.cumulant_polynomials = self.moment_polynomials
+
+        if len(self.cumulant_exponents) != stencil.Q:
+            raise ValueError("Number of cumulant exponent tuples must match stencil size.")
+
+        if len(self.cumulant_polynomials) != stencil.Q:
+            raise ValueError("Number of cumulant polynomials must match stencil size.")
+
+        self.central_moment_exponents = self.compute_required_central_moments()
+
+        self.mono_to_poly_matrix = monomial_to_polynomial_transformation_matrix(self.cumulant_exponents,
+                                                                                self.cumulant_polynomials)
+        self.poly_to_mono_matrix = self.mono_to_poly_matrix.inv()
+
+    @property
+    def required_central_moments(self):
+        """The required central moments as a sorted list of exponent tuples"""
+        return self.central_moment_exponents
+
+    def compute_required_central_moments(self):
+        def _contained_moments(m):
+            ranges = (range(i + 1) for i in m)
+            return product(*ranges)
+
+        #   Always require zeroth and first moments
+        required_moments = set(moments_up_to_order(1, dim=self.dim))
+        #   In differentiating the generating function, all derivatives contained in c will occur
+        #   --> all of these moments are required
+        for c in self.cumulant_exponents:
+            required_moments |= set(_contained_moments(c))
+
+        assert len(required_moments) == self.stencil.Q, 'Number of required central moments must match stencil size.'
+
+        return sorted(list(required_moments), key=exponent_tuple_sort_key)
+
+    def forward_transform(self,
+                          central_moment_base=PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT,
+                          simplification=True,
+                          subexpression_base='sub_k_to_C',
+                          return_monomials=False):
+        simplification = self._get_simp_strategy(simplification)
+
+        monomial_equations = []
+        for c_symbol, exp in zip(self.pre_collision_monomial_symbols, self.cumulant_exponents):
+            eq = self.cumulant_from_central_moments(exp, central_moment_base)
+            monomial_equations.append(Assignment(c_symbol, eq))
+
+        if return_monomials:
+            subexpressions = []
+            main_assignments = monomial_equations
+        else:
+            subexpressions = monomial_equations
+            poly_eqs = self.mono_to_poly_matrix @ sp.Matrix(self.pre_collision_monomial_symbols)
+            main_assignments = [Assignment(c, v) for c, v in zip(self.pre_collision_symbols, poly_eqs)]
+
+        symbol_gen = SymbolGen(subexpression_base)
+        ac = AssignmentCollection(main_assignments, subexpressions=subexpressions,
+                                  subexpression_symbol_generator=symbol_gen)
+
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    def backward_transform(self,
+                           central_moment_base=POST_COLLISION_MONOMIAL_CENTRAL_MOMENT,
+                           simplification=True,
+                           subexpression_base='sub_C_to_k',
+                           start_from_monomials=False):
+        simplification = self._get_simp_strategy(simplification)
+
+        subexpressions = []
+        if not start_from_monomials:
+            mono_eqs = self.poly_to_mono_matrix @ sp.Matrix(self.post_collision_symbols)
+            subexpressions = [Assignment(c, v) for c, v in zip(self.post_collision_monomial_symbols, mono_eqs)]
+
+        main_assignments = []
+        for exp in self.central_moment_exponents:
+            eq = self.central_moment_from_cumulants(exp, self.mono_base_post)
+            k_symbol = statistical_quantity_symbol(central_moment_base, exp)
+            main_assignments.append(Assignment(k_symbol, eq))
+
+        symbol_gen = SymbolGen(subexpression_base)
+        ac = AssignmentCollection(main_assignments, subexpressions=subexpressions,
+                                  subexpression_symbol_generator=symbol_gen)
+
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    def cumulant_from_central_moments(self, cumulant_exponents, moment_symbol_base):
+        dim = self.dim
+        wave_numbers = WAVE_NUMBER_SYMBOLS[:dim]
+        K = sp.Function('K')
+
+        u_symbols = self.equilibrium_velocity
+        rho = self.equilibrium_density
+
+        C = sum(w * u for w, u in zip(wave_numbers, u_symbols)) + sp.log(K(*wave_numbers))
+
+        diff_args = chain.from_iterable([var, i] for var, i in zip(wave_numbers, cumulant_exponents))
+        cumulant = C.diff(*diff_args)
+
+        derivatives = cumulant.atoms(sp.Derivative)
+        derivative_subs = [(d, derivative_as_statistical_quantity(d, wave_numbers, moment_symbol_base))
+                           for d in derivatives]
+        derivative_subs = sorted(derivative_subs, key=lambda x: count_derivatives(x[0]), reverse=True)
+        derivative_subs.append((K(*wave_numbers), statistical_quantity_symbol(moment_symbol_base, (0,) * dim)))
+
+        cumulant = cumulant.subs(derivative_subs)
+
+        value_subs = {x: 0 for x in wave_numbers}
+        cumulant = cumulant.subs(value_subs)
+
+        return (rho * cumulant).collect(rho)
+
+    def central_moment_from_cumulants(self, moment_exponents, cumulant_symbol_base):
+        dim = len(moment_exponents)
+        wave_numbers = WAVE_NUMBER_SYMBOLS[:dim]
+        C = sp.Function('C')
+
+        u_symbols = self.equilibrium_velocity
+        rho = self.equilibrium_density
+
+        K = sp.exp(C(*wave_numbers) - sum(w * u for w,
+                                          u in zip(wave_numbers, u_symbols)))
+
+        diff_args = chain.from_iterable([var, i] for var, i in zip(wave_numbers, moment_exponents))
+        moment = K.diff(*diff_args)
+
+        derivatives = moment.atoms(sp.Derivative)
+
+        derivative_subs = [(d, derivative_as_statistical_quantity(d, wave_numbers, 'c')) for d in derivatives]
+        derivative_subs = sorted(derivative_subs, key=lambda x: count_derivatives(x[0]), reverse=True)
+        derivative_subs.append((C(*wave_numbers), statistical_quantity_symbol('c', (0,) * dim)))
+
+        moment = moment.subs(derivative_subs)
+
+        value_subs = [(x, 0) for x in wave_numbers]
+
+        moment = moment.subs(value_subs)
+
+        c_indices = [(0,) * dim] + [moment_index_from_derivative(d, wave_numbers) for d in derivatives]
+        moment = moment.subs([(statistical_quantity_symbol('c', idx),
+                               statistical_quantity_symbol(cumulant_symbol_base, idx) / rho)
+                              for idx in c_indices])
+
+        return moment.expand().collect(rho)
+
+    @property
+    def _default_simplification(self):
+        simplification = SimplificationStrategy()
+        return simplification
+# end class CentralMomentsToCumulantsByGeneratingFunc

src/lbmpy/moment_transforms/rawmomenttransforms.py

+from functools import partial
+import sympy as sp
+
+from pystencils import Assignment, AssignmentCollection
+from pystencils.simp import (
+    SimplificationStrategy, add_subexpressions_for_divisions, add_subexpressions_for_constants,
+    insert_aliases, insert_constants)
+from pystencils.simp.assignment_collection import SymbolGen
+
+from lbmpy.moments import (
+    moment_matrix, monomial_to_polynomial_transformation_matrix, non_aliased_polynomial_raw_moments)
+from lbmpy.moments import statistical_quantity_symbol as sq_sym
+
+from .abstractmomenttransform import (
+    AbstractMomentTransform,
+    PRE_COLLISION_RAW_MOMENT, POST_COLLISION_RAW_MOMENT,
+    PRE_COLLISION_MONOMIAL_RAW_MOMENT, POST_COLLISION_MONOMIAL_RAW_MOMENT
+)
+
+
+class AbstractRawMomentTransform(AbstractMomentTransform):
+    """Abstract base class for all transformations between population space
+    and raw-moment space."""
+
+    def __init__(self, stencil, moment_polynomials,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 pre_collision_symbol_base=PRE_COLLISION_RAW_MOMENT,
+                 post_collision_symbol_base=POST_COLLISION_RAW_MOMENT,
+                 pre_collision_monomial_symbol_base=PRE_COLLISION_MONOMIAL_RAW_MOMENT,
+                 post_collision_monomial_symbol_base=POST_COLLISION_MONOMIAL_RAW_MOMENT,
+                 **kwargs):
+        super(AbstractRawMomentTransform, self).__init__(
+            stencil, equilibrium_density, equilibrium_velocity,
+            moment_polynomials=moment_polynomials,
+            pre_collision_symbol_base=pre_collision_symbol_base,
+            post_collision_symbol_base=post_collision_symbol_base,
+            pre_collision_monomial_symbol_base=pre_collision_monomial_symbol_base,
+            post_collision_monomial_symbol_base=post_collision_monomial_symbol_base,
+            **kwargs
+        )
+
+        assert len(self.moment_polynomials) == self.q, 'Number of moments must match stencil'
+
+    def _rm_background_shift(self, raw_moments):
+        if self.background_distribution is not None:
+            shift = self.background_distribution.moments(raw_moments)
+        else:
+            shift = (0,) * self.q
+        return sp.Matrix(shift)
+# end class AbstractRawMomentTransform
+
+
+class PdfsToMomentsByMatrixTransform(AbstractRawMomentTransform):
+    """Transform between populations and moment space spanned by a polynomial
+    basis, using matrix-vector multiplication."""
+
+    def __init__(self, stencil, moment_polynomials,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 conserved_quantity_equations=None,
+                 **kwargs):
+        super(PdfsToMomentsByMatrixTransform, self).__init__(
+            stencil, moment_polynomials, equilibrium_density, equilibrium_velocity,
+            conserved_quantity_equations=conserved_quantity_equations, **kwargs)
+
+        self.moment_matrix = moment_matrix(self.moment_polynomials, stencil)
+        self.inv_moment_matrix = self.moment_matrix.inv()
+
+    @property
+    def absorbs_conserved_quantity_equations(self):
+        return False
+
+    def forward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_f_to_M',
+                          return_monomials=False):
+        r"""Returns an assignment collection containing equations for pre-collision polynomial
+        moments, expressed in terms of the pre-collision populations by matrix-multiplication.
+
+        The moment transformation matrix :math:`M` provided by :func:`lbmpy.moments.moment_matrix` is
+        used to compute the pre-collision moments as :math:`\mathbf{M} = M \cdot \mathbf{f}`,
+        which is returned element-wise.
+
+        Args:
+            pdf_symbols: List of symbols that represent the pre-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            return_monomials: Return equations for monomial moments. Use only when specifying
+                              ``moment_exponents`` in constructor!
+        """
+        simplification = self._get_simp_strategy(simplification, 'forward')
+
+        if return_monomials:
+            assert len(self.moment_exponents) == self.q, "Could not derive invertible monomial transform." \
+                f"Expected {self.q} monomials, but got {len(self.moment_exponents)}."
+            mm = moment_matrix(self.moment_exponents, self.stencil)
+            background_shift = self._rm_background_shift(self.moment_exponents)
+            pre_collision_moments = self.pre_collision_monomial_symbols
+        else:
+            mm = self.moment_matrix
+            background_shift = self._rm_background_shift(self.moment_polynomials)
+            pre_collision_moments = self.pre_collision_symbols
+
+        f_to_m_vec = mm * sp.Matrix(pdf_symbols) + background_shift
+        main_assignments = [Assignment(m, eq) for m, eq in zip(pre_collision_moments, f_to_m_vec)]
+
+        symbol_gen = SymbolGen(symbol=subexpression_base)
+        ac = AssignmentCollection(main_assignments, subexpression_symbol_generator=symbol_gen)
+
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    def backward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_k_to_f',
+                           start_from_monomials=False):
+        r"""Returns an assignment collection containing equations for post-collision populations,
+        expressed in terms of the post-collision polynomial moments by matrix-multiplication.
+
+        The moment transformation matrix :math:`M` provided by :func:`lbmpy.moments.moment_matrix` is
+        inverted and used to compute the pre-collision moments as
+        :math:`\mathbf{f}^{\ast} = M^{-1} \cdot \mathbf{M}_{\mathrm{post}}`, which is returned element-wise.
+
+        **Simplifications**
+
+        If simplification is enabled, the equations for populations :math:`f_i` and :math:`f_{\bar{i}}`
+        of opposite stencil directions :math:`\mathbf{c}_i` and :math:`\mathbf{c}_{\bar{i}} = - \mathbf{c}_i`
+        are split into their symmetric and antisymmetric parts :math:`f_i^{\mathrm{sym}}, f_i^{\mathrm{anti}}`, such
+        that
+
+        .. math::
+
+            f_i = f_i^{\mathrm{sym}} + f_i^{\mathrm{anti}}
+
+            f_{\bar{i}} = f_i^{\mathrm{sym}} - f_i^{\mathrm{anti}}
+
+
+        Args:
+            pdf_symbols: List of symbols that represent the post-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            start_from_monomials: Return equations for monomial moments. Use only when specifying
+                                  ``moment_exponents`` in constructor!
+        """
+        simplification = self._get_simp_strategy(simplification, 'backward')
+
+        if start_from_monomials:
+            assert len(self.moment_exponents) == self.q, "Could not derive invertible monomial transform." \
+                f"Expected {self.q} monomials, but got {len(self.moment_exponents)}."
+            mm_inv = moment_matrix(self.moment_exponents, self.stencil).inv()
+            background_shift = self._rm_background_shift(self.moment_exponents)
+            post_collision_moments = self.post_collision_monomial_symbols
+        else:
+            mm_inv = self.inv_moment_matrix
+            background_shift = self._rm_background_shift(self.moment_polynomials)
+            post_collision_moments = self.post_collision_symbols
+
+        symbol_gen = SymbolGen(subexpression_base)
+
+        subexpressions = [Assignment(xi, m - s)
+                          for xi, m, s in zip(symbol_gen, post_collision_moments, background_shift)]
+
+        m_to_f_vec = mm_inv * sp.Matrix([s.lhs for s in subexpressions])
+        main_assignments = [Assignment(f, eq) for f, eq in zip(pdf_symbols, m_to_f_vec)]
+
+        ac = AssignmentCollection(main_assignments, subexpressions=subexpressions,
+                                  subexpression_symbol_generator=symbol_gen)
+
+        ac.add_simplification_hint('stencil', self.stencil)
+        ac.add_simplification_hint('post_collision_pdf_symbols', pdf_symbols)
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    #   ----------------------------- Private Members -----------------------------
+
+    @property
+    def _default_simplification(self):
+        forward_simp = SimplificationStrategy()
+        # forward_simp.add(substitute_moments_in_conserved_quantity_equations)
+        forward_simp.add(insert_aliases)
+        forward_simp.add(add_subexpressions_for_divisions)
+
+        from lbmpy.methods.momentbased.momentbasedsimplifications import split_pdf_main_assignments_by_symmetry
+
+        backward_simp = SimplificationStrategy()
+        backward_simp.add(insert_aliases)
+        backward_simp.add(split_pdf_main_assignments_by_symmetry)
+        backward_simp.add(add_subexpressions_for_constants)
+
+        return {
+            'forward': forward_simp,
+            'backward': backward_simp
+        }
+
+
+# end class PdfsToMomentsByMatrixTransform
+
+class PdfsToMomentsByChimeraTransform(AbstractRawMomentTransform):
+    """Transform between populations and moment space spanned by a polynomial
+    basis, using the raw-moment chimera transform in the forward direction and
+    matrix-vector multiplication in the backward direction."""
+
+    def __init__(self, stencil, moment_polynomials,
+                 equilibrium_density,
+                 equilibrium_velocity,
+                 conserved_quantity_equations=None,
+                 **kwargs):
+
+        if moment_polynomials:
+            #   Remove aliases
+            moment_polynomials = non_aliased_polynomial_raw_moments(moment_polynomials, stencil)
+
+        super(PdfsToMomentsByChimeraTransform, self).__init__(
+            stencil, moment_polynomials, equilibrium_density, equilibrium_velocity,
+            conserved_quantity_equations=conserved_quantity_equations, **kwargs)
+
+        self.inv_moment_matrix = moment_matrix(self.moment_exponents, self.stencil).inv()
+        self.mono_to_poly_matrix = monomial_to_polynomial_transformation_matrix(self.moment_exponents,
+                                                                                self.moment_polynomials)
+        self.poly_to_mono_matrix = self.mono_to_poly_matrix.inv()
+
+    @property
+    def absorbs_conserved_quantity_equations(self):
+        return True
+
+    def get_cq_to_moment_symbols_dict(self, moment_symbol_base):
+        """Returns a dictionary mapping the density and velocity symbols to the correspondig
+        zeroth- and first-order raw moment symbols"""
+        if self.cqe is None:
+            return dict()
+
+        rho = self.equilibrium_density
+        u = self.equilibrium_velocity
+        cq_symbols_to_moments = dict()
+        if isinstance(rho, sp.Symbol) and rho in self.cqe.defined_symbols:
+            cq_symbols_to_moments[rho] = sq_sym(moment_symbol_base, (0,) * self.dim)
+        for d, u_sym in enumerate(u):
+            if isinstance(u_sym, sp.Symbol) and u_sym in self.cqe.defined_symbols:
+                cq_symbols_to_moments[u_sym] = sq_sym(moment_symbol_base, tuple(
+                    (1 if i == d else 0) for i in range(self.dim)))
+        return cq_symbols_to_moments
+
+    def forward_transform(self, pdf_symbols, simplification=True,
+                          subexpression_base='sub_f_to_m',
+                          return_monomials=False):
+        r"""Returns an assignment collection containing equations for pre-collision polynomial
+        moments, expressed in terms of the pre-collision populations, using the raw moment chimera transform.
+
+        The chimera transform for raw moments is given by :cite:`geier2015` :
+
+        .. math::
+
+            f_{xyz} &:= f_i \text{ such that } c_i = (x,y,z)^T \\
+            m_{xy|\gamma} &:= \sum_{z \in \{-1, 0, 1\} } f_{xyz} \cdot z^{\gamma} \\
+            m_{x|\beta \gamma} &:= \sum_{y \in \{-1, 0, 1\}} m_{xy|\gamma} \cdot y^{\beta} \\
+            m_{\alpha \beta \gamma} &:= \sum_{x \in \{-1, 0, 1\}} m_{x|\beta \gamma} \cdot x^{\alpha}
+
+
+        The obtained raw moments are afterward combined to the desired polynomial moments.
+
+        **Conserved Quantity Equations**
+
+        If given, this transform absorbs the conserved quantity equations and simplifies them
+        using the monomial raw moment equations, if simplification is enabled.
+
+        **De-Aliasing**
+
+        If more than :math:`q` monomial moments are extracted from the polynomial set, the polynomials
+        are de-aliased by eliminating aliases according to the stencil
+        using `lbmpy.moments.non_aliased_polynomial_raw_moments`.
+
+        **Simplification**
+
+        If simplification is enabled, the absorbed conserved quantity equations are - if possible -
+        rewritten using the monomial symbols. If the conserved quantities originate somewhere else
+        than in the lower-order moments (like from an external field), they are not affected by this
+        simplification. Furthermore, aliases and constants are propagated in the chimera equations.
+
+        Args:
+            pdf_symbols: List of symbols that represent the pre-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            return_monomials: Return equations for monomial moments. Use only when specifying
+                              ``moment_exponents`` in constructor!
+        """
+
+        simplification = self._get_simp_strategy(simplification, 'forward')
+        monomial_symbol_base = self.mono_base_pre
+
+        def _partial_kappa_symbol(fixed_directions, remaining_exponents):
+            fixed_str = '_'.join(str(direction) for direction in fixed_directions).replace('-', 'm')
+            exp_str = '_'.join(str(exp) for exp in remaining_exponents).replace('-', 'm')
+            return sp.Symbol(f"partial_{monomial_symbol_base}_{fixed_str}_e_{exp_str}")
+
+        partial_sums_dict = dict()
+        monomial_eqs = []
+
+        def collect_partial_sums(exponents, dimension=0, fixed_directions=tuple()):
+            if dimension == self.dim:
+                #   Base Case
+                if fixed_directions in self.stencil:
+                    return pdf_symbols[self.stencil.index(fixed_directions)]
+                else:
+                    return 0
+            else:
+                #   Recursive Case
+                summation = sp.sympify(0)
+                for d in [-1, 0, 1]:
+                    next_partial = collect_partial_sums(
+                        exponents, dimension=dimension + 1, fixed_directions=fixed_directions + (d,))
+                    summation += next_partial * d ** exponents[dimension]
+
+                if dimension == 0:
+                    lhs_symbol = sq_sym(monomial_symbol_base, exponents)
+                    monomial_eqs.append(Assignment(lhs_symbol, summation))
+                else:
+                    lhs_symbol = _partial_kappa_symbol(fixed_directions, exponents[dimension:])
+                    partial_sums_dict[lhs_symbol] = summation
+                return lhs_symbol
+
+        for e in self.moment_exponents:
+            collect_partial_sums(e)
+
+        main_assignments = self.cqe.main_assignments.copy() if self.cqe is not None else []
+        subexpressions = self.cqe.subexpressions.copy() if self.cqe is not None else []
+        subexpressions += [Assignment(lhs, rhs) for lhs, rhs in partial_sums_dict.items()]
+
+        if return_monomials:
+            shift = self._rm_background_shift(self.moment_exponents)
+            main_assignments += [Assignment(a.lhs, a.rhs + s) for a, s in zip(monomial_eqs, shift)]
+        else:
+            subexpressions += monomial_eqs
+            moment_eqs = self.mono_to_poly_matrix * sp.Matrix(self.pre_collision_monomial_symbols)
+            moment_eqs += self._rm_background_shift(self.moment_polynomials)
+            main_assignments += [Assignment(m, v) for m, v in zip(self.pre_collision_symbols, moment_eqs)]
+
+        symbol_gen = SymbolGen(subexpression_base)
+        ac = AssignmentCollection(main_assignments, subexpressions=subexpressions,
+                                  subexpression_symbol_generator=symbol_gen)
+        ac.add_simplification_hint('cq_symbols_to_moments', self.get_cq_to_moment_symbols_dict(monomial_symbol_base))
+
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    def backward_transform(self, pdf_symbols, simplification=True,
+                           subexpression_base='sub_k_to_f',
+                           start_from_monomials=False):
+        r"""Returns an assignment collection containing equations for post-collision populations,
+        expressed in terms of the post-collision polynomial moments by matrix-multiplication.
+
+        The post-collision monomial moments :math:`\mathbf{m}_{\mathrm{post}}` are first obtained from the polynomials.
+        Then, the monomial transformation matrix :math:`M_r` provided by :func:`lbmpy.moments.moment_matrix`
+        is inverted and used to compute the post-collision populations as
+        :math:`\mathbf{f}_{\mathrm{post}} = M_r^{-1} \cdot \mathbf{m}_{\mathrm{post}}`.
+
+        **De-Aliasing**
+
+        See `PdfsToMomentsByChimeraTransform.forward_transform`.
+
+        **Simplifications**
+
+        If simplification is enabled, the equations for populations :math:`f_i` and :math:`f_{\bar{i}}`
+        of opposite stencil directions :math:`\mathbf{c}_i` and :math:`\mathbf{c}_{\bar{i}} = - \mathbf{c}_i`
+        are split into their symmetric and antisymmetric parts :math:`f_i^{\mathrm{sym}}, f_i^{\mathrm{anti}}`, such
+        that
+
+        .. math::
+
+            f_i = f_i^{\mathrm{sym}} + f_i^{\mathrm{anti}}
+
+            f_{\bar{i}} = f_i^{\mathrm{sym}} - f_i^{\mathrm{anti}}
+
+
+        Args:
+            pdf_symbols: List of symbols that represent the post-collision populations
+            simplification: Simplification specification. See :class:`AbstractMomentTransform`
+            subexpression_base: The base name used for any subexpressions of the transformation.
+            start_from_monomials: Return equations for monomial moments. Use only when specifying
+                                  ``moment_exponents`` in constructor!
+        """
+
+        simplification = self._get_simp_strategy(simplification, 'backward')
+
+        post_collision_moments = self.post_collision_symbols
+        post_collision_monomial_moments = self.post_collision_monomial_symbols
+
+        symbol_gen = SymbolGen(subexpression_base)
+
+        subexpressions = []
+        if not start_from_monomials:
+            background_shift = self._rm_background_shift(self.moment_polynomials)
+            shift_equations = [Assignment(xi, m - s)
+                               for xi, m, s in zip(symbol_gen, post_collision_moments, background_shift)]
+            raw_moment_eqs = self.poly_to_mono_matrix * sp.Matrix([s.lhs for s in shift_equations])
+            subexpressions += shift_equations
+            subexpressions += [Assignment(rm, v) for rm, v in zip(post_collision_monomial_moments, raw_moment_eqs)]
+        else:
+            background_shift = self._rm_background_shift(self.moment_exponents)
+            shift_equations = [Assignment(xi, m - s)
+                               for xi, m, s in zip(symbol_gen, post_collision_monomial_moments, background_shift)]
+            subexpressions += shift_equations
+            post_collision_monomial_moments = [s.lhs for s in shift_equations]
+
+        rm_to_f_vec = self.inv_moment_matrix * sp.Matrix(post_collision_monomial_moments)
+        main_assignments = [Assignment(f, eq) for f, eq in zip(pdf_symbols, rm_to_f_vec)]
+
+        ac = AssignmentCollection(main_assignments, subexpressions=subexpressions,
+                                  subexpression_symbol_generator=symbol_gen)
+        ac.add_simplification_hint('stencil', self.stencil)
+        ac.add_simplification_hint('post_collision_pdf_symbols', pdf_symbols)
+        if simplification:
+            ac = simplification.apply(ac)
+        return ac
+
+    #   ----------------------------- Private Members -----------------------------
+
+    @property
+    def _default_simplification(self):
+        from lbmpy.methods.momentbased.momentbasedsimplifications import (
+            substitute_moments_in_conserved_quantity_equations,
+            split_pdf_main_assignments_by_symmetry
+        )
+
+        cq = (self.equilibrium_density,) + self.equilibrium_velocity
+        fw_skip = cq + self.pre_collision_monomial_symbols
+
+        forward_simp = SimplificationStrategy()
+        forward_simp.add(substitute_moments_in_conserved_quantity_equations)
+        forward_simp.add(partial(insert_aliases, skip=fw_skip))
+        forward_simp.add(partial(insert_constants, skip=fw_skip))
+
+        bw_skip = self.post_collision_monomial_symbols
+
+        backward_simp = SimplificationStrategy()
+        backward_simp.add(partial(insert_aliases, skip=bw_skip))
+        backward_simp.add(split_pdf_main_assignments_by_symmetry)
+        backward_simp.add(add_subexpressions_for_constants)
+
+        return {
+            'forward': forward_simp,
+            'backward': backward_simp
+        }
+
+# end class PdfsToMomentsByChimeraTransform

lbmpy/moments.py → src/lbmpy/moments.py

@@ -86,8 +86,8 @@ def moment_permutations(exponent_tuple):
 def moments_of_order(order, dim=3, include_permutations=True):
     """All tuples of length 'dim' which sum equals 'order'"""
     for item in __fixed_sum_tuples(dim, order, ordered=not include_permutations):
-        assert(len(item) == dim)
-        assert(sum(item) == order)
+        assert len(item) == dim
+        assert sum(item) == order
         yield item
 
 
@@ -110,6 +110,14 @@ def extend_moments_with_permutations(exponent_tuples):
     return __unique(all_moments)
 
 
+def contained_moments(exponent_tuple, min_order=0, exclude_original=True):
+    """Returns all moments contained in exponent_tuple, in the sense that their exponents are less than or
+    equal to the corresponding exponents in exponent_tuple."""
+    return [t for t
+            in itertools.product(*(range(i + 1) for i in exponent_tuple))
+            if sum(t) >= min_order and (not exclude_original or t != exponent_tuple)]
+
+
 # ------------------------------ Representation Conversions ------------------------------------------------------------
 
 
@@ -151,7 +159,7 @@ def polynomial_to_exponent_representation(polynomial, dim=3):
     summands = [polynomial] if polynomial.func != sp.Add else polynomial.args
     for expr in summands:
         if len(expr.atoms(sp.Symbol) - set(MOMENT_SYMBOLS)) > 0:
-            raise ValueError("Invalid moment polynomial: " + str(expr))
+            raise ValueError(f"Invalid moment polynomial: {str(expr)}")
         c, x_exp, y_exp, z_exp = sp.Wild('c'), sp.Wild('xexp'), sp.Wild('yexp'), sp.Wild('zc')
         match_res = expr.match(c * x**x_exp * y**y_exp * z**z_exp)
         assert match_res[x_exp].is_integer and match_res[y_exp].is_integer and match_res[z_exp].is_integer
@@ -171,6 +179,10 @@ def moment_sort_key(moment):
     return get_order(moment), len(moment.atoms(sp.Symbol)), len(mom_str), mom_str
 
 
+def exponent_tuple_sort_key(x):
+    return moment_sort_key(exponent_to_polynomial_representation(x))
+
+
 def sort_moments_into_groups_of_same_order(moments):
     """Returns a dictionary mapping the order (int) to a list of moments with that order."""
     result = defaultdict(list)
@@ -182,6 +194,10 @@ def sort_moments_into_groups_of_same_order(moments):
 # -------------------- Common Function working with exponent tuples and polynomial moments -----------------------------
 
 
+def statistical_quantity_symbol(name, exponents):
+    return sp.Symbol(f'{name}_{"".join(str(i) for i in exponents)}')
+
+
 def is_even(moment):
     """
     A moment is considered even when under sign reversal nothing changes i.e. :math:`m(-x,-y,-z) = m(x,y,z)`
@@ -209,7 +225,7 @@ def is_even(moment):
 
 def get_moment_indices(moment_exponent_tuple):
     """Returns indices for a given exponent tuple:
-    
+
     Example:
         >>> get_moment_indices((2,1,0))
         [0, 0, 1]
@@ -273,6 +289,62 @@ def non_aliased_moment(moment_tuple: Sequence[int]) -> Tuple[int, ...]:
     return tuple(result)
 
 
+def aliases_from_moment_list(moment_exponents, stencil):
+    """Takes a list of moment exponent tuples and finds aliases within it
+    according the given stencil. Two moments are aliases of each other on a stencil
+    if they produce the same coefficients in the moment sum over all populations.
+
+    Apart from the obvious aliases (e.g. ``(4,0,0)`` to ``(2,0,0)``, etc), there are aliasing
+    effects for example on the D3Q15 stencil, where some third and fourth order moments
+    are the same.
+    """
+    mm = moment_matrix(moment_exponents, stencil)
+    rows_dict = dict()
+    aliases = dict()
+    for r, moment in enumerate(moment_exponents):
+        row = tuple(mm[r, :])
+        if row in rows_dict:
+            aliases[moment] = rows_dict[row]
+        else:
+            rows_dict[row] = moment
+    return aliases
+
+
+def non_aliased_polynomial_raw_moments(polys, stencil, nested=False):
+    """Takes a (potentially nested) list of raw moment polynomials and rewrites them by eliminating
+    any aliased monomials. 
+
+    All polynomials are expanded and occuring monomials are collected. Using `aliases_from_moment_list`,
+    aliases are eliminated and substituted in the polynomials. 
+
+    Attention: Only use this method for monomials in raw moment space. It will produce wrong results
+    if used for central moments, since there is no direct aliasing in central moment space!
+    """
+    dim = len(stencil[0])
+    if nested:
+        polys_unnested = list(itertools.chain.from_iterable(polys))
+    else:
+        polys_unnested = polys
+    monos = sorted(extract_monomials(polys_unnested, dim=dim), key=exponent_tuple_sort_key)
+    aliases = aliases_from_moment_list(monos, stencil)
+
+    if not aliases:  # Stop early if there are no aliases
+        return polys
+
+    def nonalias_polynomial(poly):
+        exponents = polynomial_to_exponent_representation(poly, dim)
+        exponents_unaliased = [(coeff, aliases.get(m, m)) for coeff, m in exponents]
+        return sum(coeff * exponent_to_polynomial_representation(m) for coeff, m in exponents_unaliased)
+
+    if nested:
+        output_polys = []
+        for group in polys:
+            output_polys.append(list(map(nonalias_polynomial, group)))
+        return output_polys
+    else:
+        return list(map(nonalias_polynomial, polys))
+
+
 def is_bulk_moment(moment, dim):
     """The bulk moment is x**2+y**2+z**2"""
     if type(moment) is not tuple:
@@ -321,10 +393,10 @@ def discrete_moment(func, moment, stencil, shift_velocity=None):
         func: list of distribution functions for each direction
         moment: can either be a exponent tuple, or a sympy polynomial expression
         stencil: sequence of directions
-        shift_velocity: velocity vector u to compute central moments, the lattice velocity is replaced by
-                        (lattice_velocity - shift_velocity)
+        shift_velocity: velocity vector :math:`\mathbf{u}` to compute central moments, the lattice
+                        velocity is replaced by (lattice_velocity - shift_velocity)
     """
-    assert len(stencil) == len(func)
+    assert stencil.Q == len(func)
     if shift_velocity is None:
         shift_velocity = (0,) * len(stencil[0])
     res = 0
@@ -365,7 +437,33 @@ def moment_matrix(moments, stencil, shift_velocity=None):
                 evaluated = evaluated.subs(var, stencil_entry - shift)
             return evaluated
 
-    return sp.Matrix(len(moments), len(stencil), generator)
+    return sp.Matrix(len(moments), stencil.Q, generator)
+
+
+def set_up_shift_matrix(moments, stencil, velocity_symbols=sp.symbols("u_:3")):
+    """
+    Sets up a shift matrix to shift raw moments to central moment space.
+
+    Args:
+        - moments: Sequence of polynomials or sequence of exponent tuples, sorted
+                   ascendingly by moment order.
+        - stencil: Nested tuple of lattice velocities
+        - velocity_symbols: Sequence of symbols corresponding to the shift velocity
+    """
+    dim = len(stencil[0])
+    if len(velocity_symbols) > dim:
+        velocity_symbols = velocity_symbols[:dim]
+
+    M = moment_matrix(moments, stencil, shift_velocity=None)
+    MN = moment_matrix(moments, stencil, shift_velocity=velocity_symbols)
+
+    N = sp.simplify(MN * M.inv())
+
+    assert N.is_lower, "Calculating the shift matrix gave not a lower diagonal matrix. Thus it failed"
+    assert sum(N[i, i] for i in range(stencil.Q)) == stencil.Q, "Calculating the shift matrix failed. " \
+        "There are entries on the diagonal which are not equal to one"
+
+    return N
 
 
 def gram_schmidt(moments, stencil, weights=None):
@@ -376,15 +474,15 @@ def gram_schmidt(moments, stencil, weights=None):
         moments: sequence of moments, either in tuple or polynomial form
         stencil: stencil as sequence of directions
         weights: optional weights, that define the scalar product which is used for normalization.
-                 Scalar product :math:`< a,b > = \sum a_i b_i w_i` with weights :math:`w_i`.
+                 Scalar product :math:`< \mathbf{a},\mathbf{b} > = \sum a_i b_i w_i` with weights :math:`w_i`.
                  Passing no weights sets all weights to 1.
     Returns:
         set of orthogonal moments in polynomial form
     """
     if weights is None:
-        weights = sp.eye(len(stencil))
+        weights = sp.eye(stencil.Q)
     if type(weights) is list:
-        assert len(weights) == len(stencil)
+        assert len(weights) == stencil.Q
         weights = sp.diag(*weights)
 
     if type(moments[0]) is tuple:
@@ -401,8 +499,8 @@ def gram_schmidt(moments, stencil, weights=None):
             prev_element = orthogonalized_vectors[j]
             denominator = prev_element.dot(weights * prev_element)
             if denominator == 0:
-                raise ValueError("Not an independent set of vectors given: "
-                                 "vector %d is dependent on previous vectors" % (i,))
+                raise ValueError(f"Not an independent set of vectors given: "
+                                 f"vector {i} is dependent on previous vectors")
             overlap = current_element.dot(weights * prev_element) / denominator
             current_element -= overlap * prev_element
             moments[i] -= overlap * moments[j]
@@ -415,22 +513,19 @@ def get_default_moment_set_for_stencil(stencil):
     """
     Returns a sequence of moments that are commonly used to construct a LBM equilibrium for the given stencil
     """
-    from lbmpy.stencils import get_stencil
-    from pystencils.stencil import have_same_entries
-
     to_poly = exponents_to_polynomial_representations
 
-    if have_same_entries(stencil, get_stencil("D2Q9")):
+    if stencil.D == 2 and stencil.Q == 9:
         return sorted(to_poly(moments_up_to_component_order(2, dim=2)), key=moment_sort_key)
 
     all27_moments = moments_up_to_component_order(2, dim=3)
-    if have_same_entries(stencil, get_stencil("D3Q27")):
+    if stencil.D == 3 and stencil.Q == 27:
         return sorted(to_poly(all27_moments), key=moment_sort_key)
-    if have_same_entries(stencil, get_stencil("D3Q19")):
+    if stencil.D == 3 and stencil.Q == 19:
         non_matched_moments = [(1, 2, 2), (1, 1, 2), (2, 2, 2), (1, 1, 1)]
         moments19 = set(all27_moments) - set(extend_moments_with_permutations(non_matched_moments))
         return sorted(to_poly(moments19), key=moment_sort_key)
-    if have_same_entries(stencil, get_stencil("D3Q15")):
+    if stencil.D == 3 and stencil.Q == 15:
         x, y, z = MOMENT_SYMBOLS
         non_matched_moments = [(1, 2, 0), (2, 2, 0), (1, 1, 2), (1, 2, 2), (2, 2, 2)]
         additional_moments = (6 * (x ** 2 * y ** 2 + x ** 2 * z ** 2 + y ** 2 * z ** 2),
@@ -440,6 +535,15 @@ def get_default_moment_set_for_stencil(stencil):
                               )
         to_remove = set(extend_moments_with_permutations(non_matched_moments))
         return sorted(to_poly(set(all27_moments) - to_remove) + additional_moments, key=moment_sort_key)
+    if stencil.D == 3 and stencil.Q == 7:
+        x, y, z = MOMENT_SYMBOLS
+        non_matched_moments = [(1, 1, 0), (1, 1, 1), (1, 1, 2), (1, 2, 0),
+                               (1, 2, 2), (2, 0, 0), (2, 2, 0), (2, 2, 2)]
+        additional_moments = (x ** 2 - y ** 2,
+                              x ** 2 - z ** 2,
+                              x ** 2 + y ** 2 + z ** 2)
+        to_remove = set(extend_moments_with_permutations(non_matched_moments))
+        return sorted(to_poly(set(all27_moments) - to_remove) + additional_moments, key=moment_sort_key)
 
     raise NotImplementedError("No default moment system available for this stencil - define matched moments yourself")
 
@@ -453,10 +557,10 @@ def extract_monomials(sequence_of_polynomials, dim=3):
         dim: length of returned exponent tuples
 
     >>> x, y, z = MOMENT_SYMBOLS
-    >>> extract_monomials([x**2 + y**2 + y, y + y**2])
-    {(0, 2, 0), (0, 1, 0), (2, 0, 0)}
-    >>> extract_monomials([x**2 + y**2 + y, y + y**2], dim=2)
-    {(0, 1), (2, 0), (0, 2)}
+    >>> extract_monomials([x**2 + y**2 + y, y + y**2]) == {(0, 1, 0),(0, 2, 0),(2, 0, 0)}
+    True
+    >>> extract_monomials([x**2 + y**2 + y, y + y**2], dim=2) == {(0, 1), (0, 2), (2, 0)}
+    True
     """
     monomials = set()
     for polynomial in sequence_of_polynomials:
@@ -474,13 +578,13 @@ def monomial_to_polynomial_transformation_matrix(monomials, polynomials):
         polynomials: sequence of polynomials in the MOMENT_SYMBOLS
 
     >>> x, y, z = MOMENT_SYMBOLS
-    >>> polys = [7 * x**2 + 3 * x + 2 * y **2, \
-                 9 * x**2 - 5 * x]
+    >>> polys = [7 * x**2 + 3 * x + 2 * y **2, 9 * x**2 - 5 * x]
     >>> mons = list(extract_monomials(polys, dim=2))
+    >>> mons.sort()
     >>> monomial_to_polynomial_transformation_matrix(mons, polys)
     Matrix([
-    [7,  3, 2],
-    [9, -5, 0]])
+    [2,  3, 7],
+    [0, -5, 9]])
     """
     dim = len(monomials[0])
 
@@ -492,6 +596,33 @@ def monomial_to_polynomial_transformation_matrix(monomials, polynomials):
     return result
 
 
+def central_moment_reduced_monomial_to_polynomial_matrix(polynomials, stencil, velocity_symbols=None):
+    r"""
+    Returns a transformation matrix from a reduced set of monomial central moments to a set of polynomial
+    central moments.
+
+    Use for a set of :math:`q` central moment polynomials that reduces to a too-large set of :math:`r > q` 
+    monomials (as given by `extract_monomials`). Reduces the monomials by eliminating aliases in raw moment
+    space, and computes a reduced polynomialization matrix :math:`\mathbf{P}^{r}` that computes the given
+    polynomials from the reduced set of monomials.
+    """
+    dim = len(stencil[0])
+    if velocity_symbols is None:
+        velocity_symbols = sp.symbols(f"u_:{dim}")
+
+    reduced_polynomials = non_aliased_polynomial_raw_moments(polynomials, stencil)
+    reduced_monomials = sorted(extract_monomials(reduced_polynomials), key=exponent_tuple_sort_key)
+
+    assert len(reduced_monomials) == stencil.Q, "Could not extract a base set of correct size from the polynomials."
+
+    N = set_up_shift_matrix(reduced_monomials, stencil, velocity_symbols=velocity_symbols)
+    N_P = set_up_shift_matrix(polynomials, stencil, velocity_symbols=velocity_symbols)
+    P_mono = monomial_to_polynomial_transformation_matrix(reduced_monomials, reduced_polynomials)
+
+    P_reduced = (N_P * P_mono * N.inv()).expand()
+    return P_reduced, reduced_monomials
+
+
 # ---------------------------------- Visualization ---------------------------------------------------------------------
 
 
@@ -515,21 +646,20 @@ def moment_equality_table(stencil, discrete_equilibrium=None, continuous_equilib
     import ipy_table
     from lbmpy.continuous_distribution_measures import continuous_moment
 
-    dim = len(stencil[0])
-    u = sp.symbols("u_:{dim}".format(dim=dim))
+    u = sp.symbols(f"u_:{stencil.D}")
     if discrete_equilibrium is None:
         from lbmpy.maxwellian_equilibrium import discrete_maxwellian_equilibrium
         discrete_equilibrium = discrete_maxwellian_equilibrium(stencil, c_s_sq=sp.Rational(1, 3), compressible=True,
                                                                u=u, order=truncate_order)
     if continuous_equilibrium is None:
         from lbmpy.maxwellian_equilibrium import continuous_maxwellian_equilibrium
-        continuous_equilibrium = continuous_maxwellian_equilibrium(dim=dim, u=u, c_s_sq=sp.Rational(1, 3))
+        continuous_equilibrium = continuous_maxwellian_equilibrium(dim=stencil.D, u=u, c_s_sq=sp.Rational(1, 3))
 
     table = []
     matched_moments = 0
     non_matched_moments = 0
 
-    moments_list = [list(moments_of_order(o, dim, include_permutations=False)) for o in range(max_order + 1)]
+    moments_list = [list(moments_of_order(o, stencil.D, include_permutations=False)) for o in range(max_order + 1)]
 
     colors = dict()
     nr_of_columns = max([len(v) for v in moments_list]) + 1
@@ -540,11 +670,11 @@ def moment_equality_table(stencil, discrete_equilibrium=None, continuous_equilib
 
     for order, moments in enumerate(moments_list):
         row = [' '] * nr_of_columns
-        row[0] = '%d' % (order,)
+        row[0] = f'{order}'
         for moment, col_idx in zip(moments, range(1, len(row))):
             multiplicity = moment_multiplicity(moment)
             dm = discrete_moment(discrete_equilibrium, moment, stencil)
-            cm = continuous_moment(continuous_equilibrium, moment, symbols=sp.symbols("v_0 v_1 v_2")[:dim])
+            cm = continuous_moment(continuous_equilibrium, moment, symbols=sp.symbols("v_0 v_1 v_2")[:stencil.D])
             difference = sp.simplify(dm - cm)
             if truncate_order:
                 difference = sp.simplify(remove_higher_order_terms(difference, symbols=u, order=truncate_order))
@@ -555,7 +685,7 @@ def moment_equality_table(stencil, discrete_equilibrium=None, continuous_equilib
                 colors[(order + 1, col_idx)] = 'lightGreen'
                 matched_moments += multiplicity
 
-            row[col_idx] = '%s  x %d' % (moment, moment_multiplicity(moment))
+            row[col_idx] = f'{moment}  x {moment_multiplicity(moment)}'
 
         table.append(row)
 
@@ -564,8 +694,8 @@ def moment_equality_table(stencil, discrete_equilibrium=None, continuous_equilib
     for cell_idx, color in colors.items():
         ipy_table.set_cell_style(cell_idx[0], cell_idx[1], color=color)
 
-    print("Matched moments %d - non matched moments %d - total %d" %
-          (matched_moments, non_matched_moments, matched_moments + non_matched_moments))
+    print(f"Matched moments {matched_moments} - non matched moments {non_matched_moments} "
+          f"- total {matched_moments + non_matched_moments}")
 
     return table_display
 
@@ -596,8 +726,8 @@ def moment_equality_table_by_stencil(name_to_stencil_dict, moments, truncate_ord
 
     colors = {}
     for stencil_idx, stencil in enumerate(stencils):
-        dim = len(stencil[0])
-        u = sp.symbols("u_:{dim}".format(dim=dim))
+        dim = stencil.D
+        u = sp.symbols(f"u_:{dim}")
         discrete_equilibrium = discrete_maxwellian_equilibrium(stencil, c_s_sq=sp.Rational(1, 3), compressible=True,
                                                                u=u, order=truncate_order)
         continuous_equilibrium = continuous_maxwellian_equilibrium(dim=dim, u=u, c_s_sq=sp.Rational(1, 3))
@@ -654,7 +784,7 @@ def __unique_permutations(elements: Sequence[T]) -> Iterable[T]:
 
     """
     if len(elements) == 1:
-        yield (elements[0],)
+        yield elements[0],
     else:
         unique_elements = set(elements)
         for first_element in unique_elements:

src/lbmpy/non_newtonian_models.py

+from dataclasses import dataclass
+import sympy as sp
+
+from lbmpy.relaxationrates import get_shear_relaxation_rate, lattice_viscosity_from_relaxation_rate, \
+    relaxation_rate_from_lattice_viscosity
+from lbmpy.utils import extract_shear_relaxation_rate, frobenius_norm, second_order_moment_tensor
+from pystencils import Assignment
+
+
+@dataclass
+class CassonsParameters:
+    yield_stress: float
+    """
+    The yield stress controls the intensity of non-Newtonian behavior.
+    """
+    omega_min: float = 0.2
+    """
+    The minimal shear relaxation rate that is used as a lower bound
+    """
+    omega_max: float = 1.98
+    """
+    The maximal shear relaxation rate that is used as an upper bound
+    """
+
+
+def add_cassons_model(collision_rule, parameter: CassonsParameters, omega_output_field=None):
+    r""" Adds the Cassons model to a lattice Boltzmann collision rule that can be found here :cite:`Casson`
+
+    The only parameter of the model is the so-called yield_stress. The main idea is that no strain rate is
+    observed below some stress. However, this leads to the problem that the modified relaxation rate might no longer
+    lead to stable LBM simulations. Thus, an upper and lower limit for the shear relaxation rate must be given.
+    All the parameters are combined in the `CassonsParameters` dataclass
+
+    """
+    yield_stress = parameter.yield_stress
+    omega_min = parameter.omega_min
+    omega_max = parameter.omega_max
+
+    method = collision_rule.method
+    equilibrium = method.equilibrium_distribution
+
+    omega_s = get_shear_relaxation_rate(method)
+    omega_s, found_symbolic_shear_relaxation = extract_shear_relaxation_rate(collision_rule, omega_s)
+
+    if not found_symbolic_shear_relaxation:
+        raise ValueError("For the Cassons model the shear relaxation rate has to be a symbol or it has to be "
+                         "assigned to a single equation in the assignment list")
+
+    sigma, theta, rhs, mu, tau, adapted_omega = sp.symbols("sigma theta rhs mu tau omega_new")
+
+    rho = equilibrium.density if equilibrium.compressible else equilibrium.background_density
+    f_neq = sp.Matrix(method.pre_collision_pdf_symbols) - method.get_equilibrium_terms()
+
+    second_invariant_strain_rate_tensor = frobenius_norm(second_order_moment_tensor(f_neq, method.stencil))
+    eta = lattice_viscosity_from_relaxation_rate(omega_s)
+
+    one = sp.Rational(1, 1)
+
+    # rhs of equation 14 in https://doi.org/10.1007/s10955-005-8415-x
+    # Note that C_2 / C_4 = 3 for all configurations thus we directly insert it here
+    eq14 = one / (one - theta) * (one + sp.sqrt(theta * (one + rho / eta * sp.Rational(1, 6) * (one - theta))))
+
+    new_omega = one / tau
+    omega_cond = sp.Piecewise((omega_min, new_omega < omega_min), (omega_max, new_omega > omega_max), (new_omega, True))
+
+    eqs = [Assignment(sigma, second_invariant_strain_rate_tensor),
+           Assignment(theta, yield_stress / sigma),
+           Assignment(rhs, eq14),
+           Assignment(mu, eta * rhs ** 2),
+           Assignment(tau, one / relaxation_rate_from_lattice_viscosity(mu)),
+           Assignment(adapted_omega, omega_cond)]
+
+    collision_rule = collision_rule.new_with_substitutions({omega_s: adapted_omega}, substitute_on_lhs=False)
+    collision_rule.subexpressions += eqs
+    collision_rule.topological_sort(sort_subexpressions=True, sort_main_assignments=False)
+
+    if omega_output_field:
+        collision_rule.main_assignments.append(Assignment(omega_output_field.center, adapted_omega))
+
+    return collision_rule

src/lbmpy/oldroydb.py

+import pystencils as ps
+import sympy as sp
+import numpy as np
+
+from pystencils.boundaries.boundaryconditions import Boundary
+from pystencils.stencil import inverse_direction_string, direction_string_to_offset
+
+
+class OldroydB:
+    def __init__(self, dim, u, tau, F, tauflux, tauface, lambda_p, eta_p, vof=True):
+        assert not ps.FieldType.is_staggered(u)
+        assert not ps.FieldType.is_staggered(tau)
+        assert not ps.FieldType.is_staggered(F)
+        assert ps.FieldType.is_staggered(tauflux)
+        assert ps.FieldType.is_staggered(tauface)
+
+        self.dim = dim
+        self.u = u
+        self.tau = tau
+        self.F = F
+        self.tauflux = tauflux
+        self.tauface_field = tauface
+        self.lambda_p = lambda_p
+        self.eta_p = eta_p
+
+        full_stencil = ["C"] + self.tauflux.staggered_stencil + \
+            list(map(inverse_direction_string, self.tauflux.staggered_stencil))
+        self.stencil = tuple(map(lambda d: tuple(ps.stencil.direction_string_to_offset(d, self.dim)), full_stencil))
+        full_stencil = ["C"] + self.tauface_field.staggered_stencil + \
+            list(map(inverse_direction_string, self.tauface_field.staggered_stencil))
+        self.force_stencil = tuple(map(lambda d: tuple(ps.stencil.direction_string_to_offset(d, self.dim)),
+                                   full_stencil))
+
+        self.disc = ps.fd.FVM1stOrder(self.tau, self._flux(), self._source())
+        if vof:
+            self.vof = ps.fd.VOF(self.tauflux, self.u, self.tau)
+        else:
+            self.vof = None
+
+    def _flux(self):
+        return [self.tau.center_vector.applyfunc(lambda t: t * self.u.center_vector[i]) for i in range(self.dim)]
+
+    def _source(self):
+        gradu = sp.Matrix([[ps.fd.diff(self.u.center_vector[j], i) for j in range(self.dim)] for i in range(self.dim)])
+        gamma = gradu + gradu.transpose()
+        return self.tau.center_vector * gradu + gradu.transpose() * self.tau.center_vector + \
+            (self.eta_p * gamma - self.tau.center_vector) / self.lambda_p
+
+    def tauface(self):
+        return ps.AssignmentCollection([ps.Assignment(self.tauface_field.staggered_vector_access(d),
+                                        (self.tau.center_vector + self.tau.neighbor_vector(d)) / 2)
+                                        for d in self.tauface_field.staggered_stencil])
+
+    def force(self):
+        full_stencil = self.tauface_field.staggered_stencil + \
+            list(map(inverse_direction_string, self.tauface_field.staggered_stencil))
+        dtau = sp.Matrix([sum([sum([
+                          self.tauface_field.staggered_access(d, (i, j)) * direction_string_to_offset(d)[i]
+                          for i in range(self.dim)]) / sp.Matrix(direction_string_to_offset(d)).norm()
+                         for d in full_stencil]) for j in range(self.dim)])
+        A0 = sum([sp.Matrix(direction_string_to_offset(d)).norm() for d in full_stencil])
+        return ps.AssignmentCollection(ps.Assignment(self.F.center_vector, dtau / A0 * 2 * self.dim))
+
+    def flux(self):
+        if self.vof is not None:
+            return self.vof
+        else:
+            return self.disc.discrete_flux(self.tauflux)
+
+    def continuity(self):
+        cont = self.disc.discrete_continuity(self.tauflux)
+        tau_copy = sp.Matrix(self.dim, self.dim, lambda i, j: sp.Symbol("tau_old_%d_%d" % (i, j)))
+        tau_subs = {self.tau.center_vector[i, j]: tau_copy[i, j] for i in range(self.dim) for j in range(self.dim)}
+        return [ps.Assignment(tau_copy[i, j], self.tau.center_vector[i, j])
+                for i in range(self.dim) for j in range(self.dim)] + \
+               [ps.Assignment(a.lhs, a.rhs.subs(tau_subs)) for a in cont]
+
+
+class Flux(Boundary):
+    inner_or_boundary = True  # call the boundary condition with the fluid cell
+    single_link = False  # needs to be called for all directional fluxes
+
+    def __init__(self, stencil, value=None):
+        self.stencil = stencil
+        self.value = value
+
+    def __call__(self, field, direction_symbol, **kwargs):
+        assert ps.FieldType.is_staggered(field)
+
+        assert all([s == 0 for s in self.stencil[0]])
+        accesses = [field.staggered_vector_access(ps.stencil.offset_to_direction_string(d))
+                    for d in self.stencil[1:]]
+        conds = [sp.Equality(direction_symbol, d + 1) for d in range(len(accesses))]
+
+        if self.value is None:
+            val = sp.Matrix(np.zeros(accesses[0].shape, dtype=int))
+        else:
+            val = self.value
+
+        # use conditional
+        conditional = None
+        for a, c, d in zip(accesses, conds, self.stencil[1:]):
+            d = ps.stencil.offset_to_direction_string(d)
+            assignments = []
+            for i in range(len(a)):
+                fac = 1
+                if ps.FieldType.is_staggered_flux(field) and type(a[i]) is sp.Mul and a[i].args[0] == -1:
+                    fac = a[i].args[0]
+                    a[i] *= a[i].args[0]
+                assignments.append(ps.Assignment(a[i], fac * val[i]))
+            if len(assignments) > 0:
+                conditional = ps.astnodes.Conditional(c,
+                                                      ps.astnodes.Block(assignments),
+                                                      conditional)
+        return [conditional]
+
+    def __hash__(self):
+        return hash((Flux, self.stencil, self.value))
+
+    def __eq__(self, other):
+        return type(other) is Flux and other.stencil == self.stencil and self.value == other.value
+
+
+class Extrapolation(Boundary):
+    inner_or_boundary = True  # call the boundary condition with the fluid cell
+    single_link = False  # needs to be called for all directional fluxes
+
+    def __init__(self, stencil, src_field, order):
+        self.stencil = stencil
+        self.src = src_field
+        if order == 0:
+            self.weights = (1,)
+        elif order == 1:
+            self.weights = (sp.Rational(3, 2), - sp.Rational(1, 2))
+        elif order == 2:
+            self.weights = (sp.Rational(15, 8), - sp.Rational(10, 8), sp.Rational(3, 8))
+        else:
+            raise NotImplementedError("weights are not known for extrapolation orders > 2")
+
+    def __call__(self, field, direction_symbol, **kwargs):
+        assert ps.FieldType.is_staggered(field)
+
+        assert all([s == 0 for s in self.stencil[0]])
+        accesses = [field.staggered_vector_access(ps.stencil.offset_to_direction_string(d))
+                    for d in self.stencil[1:]]
+        conds = [sp.Equality(direction_symbol, d + 1) for d in range(len(accesses))]
+
+        # use conditional
+        conditional = None
+        for a, c, o in zip(accesses, conds, self.stencil[1:]):
+            assignments = []
+            src = [self.src.neighbor_vector(tuple([-1 * n * i for i in o])) for n in range(len(self.weights))]
+            interp = self.weights[0] * src[0]
+            for i in range(1, len(self.weights)):
+                interp += self.weights[i] * src[i]
+            for i in range(len(a)):
+                fac = 1
+                if ps.FieldType.is_staggered_flux(field) and type(a[i]) is sp.Mul and a[i].args[0] == -1:
+                    fac = a[i].args[0]
+                    a[i] *= a[i].args[0]
+                assignments.append(ps.Assignment(a[i], fac * interp[i]))
+            if len(assignments) > 0:
+                conditional = ps.astnodes.Conditional(ps.data_types.type_all_numbers(c, "int"),
+                                                      ps.astnodes.Block(assignments),
+                                                      conditional)
+        return [conditional]
+
+    def __hash__(self):
+        return hash((Extrapolation, self.stencil, self.src, self.weights))
+
+    def __eq__(self, other):
+        return type(other) is Extrapolation and other.stencil == self.stencil and \
+            other.src == self.src and other.weights == self.weights
+
+
+class ForceOnBoundary(Boundary):
+    inner_or_boundary = False  # call the boundary condition with the boundary cell
+    single_link = False  # needs to be called for all directional fluxes
+
+    def __init__(self, stencil, force_field, name=None):
+        self.stencil = stencil
+        self.force_field = force_field
+        super(ForceOnBoundary).__init__(name)
+
+        assert not ps.FieldType.is_staggered(force_field)
+
+    def __call__(self, face_stress_field, direction_symbol, **kwargs):
+        assert ps.FieldType.is_staggered(face_stress_field)
+
+        assert all([s == 0 for s in self.stencil[0]])
+        accesses = [face_stress_field.staggered_vector_access(ps.stencil.offset_to_direction_string(d))
+                    for d in self.stencil[1:]]
+        conds = [sp.Equality(direction_symbol, d + 1) for d in range(len(accesses))]
+
+        # use conditional
+        conditional = None
+        for a, c, o in zip(accesses, conds, self.stencil[1:]):
+            assignments = ps.Assignment(self.force_field.center_vector,
+                                        self.force_field.center_vector + 1 * a.transpose() * sp.Matrix(o))
+            conditional = ps.astnodes.Conditional(ps.data_types.type_all_numbers(c, "int"),
+                                                  ps.astnodes.Block(assignments),
+                                                  conditional)
+        return [conditional]
+
+    def __hash__(self):
+        return hash((ForceOnBoundary, self.stencil, self.force_field))
+
+    def __eq__(self, other):
+        return type(other) is ForceOnBoundary and other.stencil == self.stencil and \
+            other.force_field == self.force_field