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

@@ -49,6 +49,9 @@ class LatticeBoltzmannStep:
                                                  default_target=target,
                                                  parallel=False)
 
+        if lbm_config:
+            method_parameters['stencil'] = lbm_config.stencil
+
         if 'stencil' not in method_parameters:
             method_parameters['stencil'] = LBStencil(Stencil.D2Q9) \
                 if data_handling.dim == 2 else LBStencil(Stencil.D3Q27)

lbmpy/macroscopic_value_kernels.py → src/lbmpy/macroscopic_value_kernels.py

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

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

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

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

@@ -76,38 +76,49 @@ def discrete_maxwellian_equilibrium(stencil, rho=sp.Symbol("rho"), u=sp.symbols(
     """
     weights = get_weights(stencil, c_s_sq)
     assert stencil.Q == len(weights)
-
     u = u[:stencil.D]
 
+    res = [discrete_equilibrium(e_q, u, rho, w_q, order, c_s_sq, compressible) for w_q, e_q in zip(weights, stencil)]
+    return tuple(res)
+
+
+def discrete_equilibrium(v=sp.symbols("v_:3"), u=sp.symbols("u_:3"), rho=sp.Symbol("rho"), weight=sp.Symbol("w"),
+                         order=2, c_s_sq=sp.Symbol("c_s") ** 2, compressible=True):
+    """
+    Returns the common discrete LBM equilibrium depending on the mesoscopic velocity and the directional lattice weight
+
+    Args:
+        v: symbols for mesoscopic velocity
+        u: symbols for macroscopic velocity
+        rho: sympy symbol for the density
+        weight: symbol for stencil weights
+        order: highest order of velocity terms (for hydrodynamics order 2 is sufficient)
+        c_s_sq: square of speed of sound
+        compressible: compressibility
+    """
     rho_outside = rho if compressible else sp.Rational(1, 1)
     rho_inside = rho if not compressible else sp.Rational(1, 1)
 
-    res = []
-    for w_q, e_q in zip(weights, stencil):
-        e_times_u = 0
-        for c_q_alpha, u_alpha in zip(e_q, u):
-            e_times_u += c_q_alpha * u_alpha
-
-        fq = rho_inside + e_times_u / c_s_sq
+    e_times_u = 0
+    for c_q_alpha, u_alpha in zip(v, u):
+        e_times_u += c_q_alpha * u_alpha
 
-        if order <= 1:
-            res.append(fq * rho_outside * w_q)
-            continue
+    fq = rho_inside + e_times_u / c_s_sq
 
-        u_times_u = 0
-        for u_alpha in u:
-            u_times_u += u_alpha * u_alpha
-        fq += sp.Rational(1, 2) / c_s_sq**2 * e_times_u ** 2 - sp.Rational(1, 2) / c_s_sq * u_times_u
+    if order <= 1:
+        return fq * rho_outside * weight
 
-        if order <= 2:
-            res.append(fq * rho_outside * w_q)
-            continue
+    u_times_u = 0
+    for u_alpha in u:
+        u_times_u += u_alpha * u_alpha
+    fq += sp.Rational(1, 2) / c_s_sq**2 * e_times_u ** 2 - sp.Rational(1, 2) / c_s_sq * u_times_u
 
-        fq += sp.Rational(1, 6) / c_s_sq**3 * e_times_u**3 - sp.Rational(1, 2) / c_s_sq**2 * u_times_u * e_times_u
+    if order <= 2:
+        return fq * rho_outside * weight
 
-        res.append(sp.expand(fq * rho_outside * w_q))
+    fq += sp.Rational(1, 6) / c_s_sq**3 * e_times_u**3 - sp.Rational(1, 2) / c_s_sq**2 * u_times_u * e_times_u
 
-    return tuple(res)
+    return sp.expand(fq * rho_outside * weight)
 
 
 @disk_cache

lbmpy/methods/__init__.py → src/lbmpy/methods/__init__.py

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

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

@@ -5,6 +5,7 @@ import sympy as sp
 from sympy.core.numbers import Zero
 
 from pystencils import Assignment, AssignmentCollection
+from lbmpy.stencils import LBStencil
 
 RelaxationInfo = namedtuple('RelaxationInfo', ['equilibrium_value', 'relaxation_rate'])
 
@@ -21,7 +22,7 @@ class LbmCollisionRule(AssignmentCollection):
 class AbstractLbMethod(abc.ABC):
     """Abstract base class for all LBM methods."""
 
-    def __init__(self, stencil):
+    def __init__(self, stencil: LBStencil):
         self._stencil = stencil
 
     @property
@@ -32,17 +33,17 @@ class AbstractLbMethod(abc.ABC):
     @property
     def dim(self):
         """The method's spatial dimensionality"""
-        return len(self.stencil[0])
+        return self._stencil.D
 
     @property
     def pre_collision_pdf_symbols(self):
         """Tuple of symbols representing the pdf values before collision"""
-        return sp.symbols(f"f_:{len(self.stencil)}")
+        return sp.symbols(f"f_:{self._stencil.Q}")
 
     @property
     def post_collision_pdf_symbols(self):
         """Tuple of symbols representing the pdf values after collision"""
-        return sp.symbols(f"d_:{len(self.stencil)}")
+        return sp.symbols(f"d_:{self._stencil.Q}")
 
     @property
     @abc.abstractmethod
@@ -66,10 +67,10 @@ class AbstractLbMethod(abc.ABC):
         return d
 
     @property
-    def subs_dict_relxation_rate(self):
-        """returns a dictonary which maps the replaced numerical relaxation rates to its original numerical value"""
+    def subs_dict_relaxation_rate(self):
+        """returns a dictionary which maps the replaced numerical relaxation rates to its original numerical value"""
         result = dict()
-        for i in range(len(self.stencil)):
+        for i in range(self._stencil.Q):
             result[self.symbolic_relaxation_matrix[i, i]] = self.relaxation_matrix[i, i]
         return result
 
@@ -98,26 +99,33 @@ class AbstractLbMethod(abc.ABC):
 
     # -------------------------------- Helper Functions ----------------------------------------------------------------
 
-    def _generate_symbolic_relaxation_matrix(self):
+    def _generate_symbolic_relaxation_matrix(self, relaxation_rates=None, relaxation_rates_modifier=None):
         """
         This function replaces the numbers in the relaxation matrix with symbols in this case, and returns also
         the subexpressions, that assign the number to the newly introduced symbol
         """
-        rr = [self.relaxation_matrix[i, i] for i in range(self.relaxation_matrix.rows)]
-        unique_relaxation_rates = set()
+        rr = relaxation_rates if relaxation_rates is not None else self.relaxation_rates
         subexpressions = {}
+        symbolic_relaxation_rates = list()
         for relaxation_rate in rr:
-            if relaxation_rate not in unique_relaxation_rates:
-                relaxation_rate = sp.sympify(relaxation_rate)
-                # special treatment for zero, sp.Zero would be an integer ..
+            relaxation_rate = sp.sympify(relaxation_rate)
+            if isinstance(relaxation_rate, sp.Symbol):
+                symbolic_relaxation_rate = relaxation_rate
+            else:
                 if isinstance(relaxation_rate, Zero):
                     relaxation_rate = 0.0
-                if not isinstance(relaxation_rate, sp.Symbol):
-                    rt_symbol = sp.Symbol(f"rr_{len(subexpressions)}")
-                    subexpressions[relaxation_rate] = rt_symbol
-            unique_relaxation_rates.add(relaxation_rate)
+                if relaxation_rate in subexpressions:
+                    symbolic_relaxation_rate = subexpressions[relaxation_rate]
+                else:
+                    symbolic_relaxation_rate = sp.Symbol(f"rr_{len(subexpressions)}")
+                    subexpressions[relaxation_rate] = symbolic_relaxation_rate
+            symbolic_relaxation_rates.append(symbolic_relaxation_rate)
 
-        new_rr = [subexpressions[sp.sympify(e)] if sp.sympify(e) in subexpressions else e for e in rr]
         substitutions = [Assignment(e[1], e[0]) for e in subexpressions.items()]
+        if relaxation_rates_modifier is not None:
+            symbolic_relaxation_rates = [r * relaxation_rates_modifier for r in symbolic_relaxation_rates]
+        else:
+            for srr in symbolic_relaxation_rates:
+                assert isinstance(srr, sp.Symbol)
 
-        return substitutions, sp.diag(*new_rr)
+        return substitutions, sp.diag(*symbolic_relaxation_rates)

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

@@ -14,15 +14,15 @@ from lbmpy.equilibrium import ContinuousHydrodynamicMaxwellian, DiscreteHydrodyn
 
 from lbmpy.methods.default_moment_sets import cascaded_moment_sets_literature
 
-from lbmpy.methods.centeredcumulant import CenteredCumulantBasedLbMethod
-from lbmpy.methods.centeredcumulant.cumulant_transform import CentralMomentsToCumulantsByGeneratingFunc
+from lbmpy.moment_transforms import CentralMomentsToCumulantsByGeneratingFunc
 
-from lbmpy.methods.conservedquantitycomputation import DensityVelocityComputation
+from .conservedquantitycomputation import DensityVelocityComputation
 
-from lbmpy.methods.momentbased.momentbasedmethod import MomentBasedLbMethod
-from lbmpy.methods.momentbased.centralmomentbasedmethod import CentralMomentBasedLbMethod
+from .momentbased.momentbasedmethod import MomentBasedLbMethod
+from .momentbased.centralmomentbasedmethod import CentralMomentBasedLbMethod
+from .cumulantbased import CumulantBasedLbMethod
 from lbmpy.moment_transforms import (
-    AbstractMomentTransform, PdfsToCentralMomentsByShiftMatrix, PdfsToMomentsByChimeraTransform)
+    AbstractMomentTransform, BinomialChimeraTransform, PdfsToMomentsByChimeraTransform)
 from lbmpy.moment_transforms.rawmomenttransforms import AbstractRawMomentTransform
 from lbmpy.moment_transforms.centralmomenttransforms import AbstractCentralMomentTransform
 
@@ -58,14 +58,14 @@ class CollisionSpaceInfo:
     """
     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.PdfsToCentralMomentsByShiftMatrix`
+    :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.methods.centeredcumulant.cumulant_transform.CentralMomentsToCumulantsByGeneratingFunc`
+    :class:`lbmpy.moment_transforms.CentralMomentsToCumulantsByGeneratingFunc`
     if `CollisionSpace.CUMULANTS` is given, or `None` otherwise.
     """
 
@@ -74,7 +74,7 @@ class CollisionSpaceInfo:
             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 = PdfsToCentralMomentsByShiftMatrix
+            self.central_moment_transform_class = BinomialChimeraTransform
         if self.collision_space == CollisionSpace.CUMULANTS and self.cumulant_transform_class is None:
             self.cumulant_transform_class = CentralMomentsToCumulantsByGeneratingFunc
 
@@ -162,7 +162,7 @@ def create_with_continuous_maxwellian_equilibrium(stencil, moment_to_relaxation_
 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):
+                            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. 
@@ -202,19 +202,24 @@ def create_from_equilibrium(stencil, equilibrium, conserved_quantity_computation
     if cspace.collision_space == CollisionSpace.POPULATIONS:
         return MomentBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
                                    force_model=force_model, zero_centered=zero_centered,
+                                   fraction_field=fraction_field,
                                    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:
-        raise NotImplementedError("Creating a cumulant method from general equilibria is not supported yet.")
-        # return CenteredCumulantBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
-        #                                      force_model=force_model, zero_centered=zero_centered)
+        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 ---------------------------------------------------------
@@ -239,7 +244,7 @@ def create_srt(stencil, relaxation_rate, continuous_equilibrium=True, **kwargs):
         :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
     """
     continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
-    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.POPULATIONS))
+    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:
@@ -268,7 +273,7 @@ def create_trt(stencil, relaxation_rate_even_moments, relaxation_rate_odd_moment
         :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
     """
     continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
-    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.POPULATIONS))
+    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])
@@ -303,7 +308,7 @@ def create_trt_with_magic_number(stencil, relaxation_rate, magic_number=sp.Ratio
                       relaxation_rate_odd_moments=rr_odd, **kwargs)
 
 
-def create_mrt_raw(stencil, relaxation_rates, continuous_equilibrium=True, **kwargs):
+def create_mrt_raw(stencil, relaxation_rates, continuous_equilibrium=True, conserved_moments=True, **kwargs):
     r"""
     Creates a MRT method using non-orthogonalized moments.
 
@@ -317,14 +322,15 @@ def create_mrt_raw(stencil, relaxation_rates, continuous_equilibrium=True, **kwa
         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)
-    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.RAW_MOMENTS))
+    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)
+    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:
@@ -332,7 +338,7 @@ def create_mrt_raw(stencil, relaxation_rates, continuous_equilibrium=True, **kwa
 
 
 def create_central_moment(stencil, relaxation_rates, nested_moments=None,
-                          continuous_equilibrium=True, **kwargs):
+                          continuous_equilibrium=True, conserved_moments=True, **kwargs):
     r"""
     Creates moment based LB method where the collision takes place in the central moment space.
 
@@ -345,11 +351,16 @@ def create_central_moment(stencil, relaxation_rates, nested_moments=None,
         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]
@@ -362,7 +373,11 @@ def create_central_moment(stencil, relaxation_rates, nested_moments=None,
     if not nested_moments:
         nested_moments = cascaded_moment_sets_literature(stencil)
 
-    rr_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments, stencil.D)
+    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)
@@ -392,7 +407,7 @@ def create_trt_kbc(dim, shear_relaxation_rate, higher_order_relaxation_rate, met
                         used to compute the equilibrium moments.
     """
     continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
-    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.POPULATIONS))
+    check_and_set_mrt_space(CollisionSpace.POPULATIONS)
 
     def product(iterable):
         return reduce(operator.mul, iterable, 1)
@@ -450,7 +465,7 @@ def create_trt_kbc(dim, shear_relaxation_rate, higher_order_relaxation_rate, met
 
 
 def create_mrt_orthogonal(stencil, relaxation_rates, continuous_equilibrium=True, weighted=None,
-                          nested_moments=None, **kwargs):
+                          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
@@ -471,9 +486,10 @@ def create_mrt_orthogonal(stencil, relaxation_rates, continuous_equilibrium=True
         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)
-    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.RAW_MOMENTS))
+    check_and_set_mrt_space(CollisionSpace.RAW_MOMENTS)
 
     if weighted:
         weights = get_weights(stencil, sp.Rational(1, 3))
@@ -502,7 +518,8 @@ def create_mrt_orthogonal(stencil, relaxation_rates, continuous_equilibrium=True
         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)
+    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,
@@ -513,59 +530,8 @@ def create_mrt_orthogonal(stencil, relaxation_rates, continuous_equilibrium=True
 
 
 # ----------------------------------------- Cumulant method creators ---------------------------------------------------
-
-
-def create_centered_cumulant_model(stencil, cumulant_to_rr_dict, force_model=None,
-                                   zero_centered=True,
-                                   c_s_sq=sp.Rational(1, 3),
-                                   galilean_correction=False,
-                                   collision_space_info=CollisionSpaceInfo(CollisionSpace.CUMULANTS)):
-    r"""Creates a cumulant lattice Boltzmann model.
-
-    Args:
-        stencil: instance of :class:`lbmpy.stencils.LBStencil`
-        cumulant_to_rr_dict: dict that has as many entries as the stencil. Each cumulant, which can be
-                             represented by an exponent tuple or in polynomial form is mapped to a relaxation rate.
-                             See :func:`lbmpy.methods.default_moment_sets.cascaded_moment_sets_literature`
-        force_model: force model used for the collision. For cumulant LB method a good choice is
-                     `lbmpy.methods.centeredcumulant.CenteredCumulantForceModel`
-        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).
-        c_s_sq: Speed of sound squared
-        galilean_correction: special correction for D3Q27 cumulant collisions. See Appendix H in
-                             :cite:`geier2015`. Implemented in :mod:`lbmpy.methods.centeredcumulant.galilean_correction`
-        central_moment_transform_class: Class which defines the transformation to the central moment space
-                                        (see :mod:`lbmpy.moment_transforms`)
-        cumulant_transform_class: Class which defines the transformation from the central moment space to the
-                                  cumulant space (see :mod:`lbmpy.methods.centeredcumulant.cumulant_transform`)
-
-    Returns:
-        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
-    """
-
-    assert len(cumulant_to_rr_dict) == stencil.Q, \
-        "The number of moments has to be equal to the number of stencil entries"
-    assert collision_space_info.collision_space == CollisionSpace.CUMULANTS
-
-    # CQC
-    cqc = DensityVelocityComputation(stencil, True, zero_centered, force_model=force_model, c_s_sq=c_s_sq)
-
-    equilibrium = ContinuousHydrodynamicMaxwellian(dim=stencil.D, compressible=True,
-                                                   deviation_only=False,
-                                                   order=None,
-                                                   c_s_sq=c_s_sq)
-
-    cspace = collision_space_info
-    return CenteredCumulantBasedLbMethod(stencil, equilibrium, cumulant_to_rr_dict,
-                                         conserved_quantity_computation=cqc, force_model=force_model,
-                                         zero_centered=zero_centered,
-                                         galilean_correction=galilean_correction,
-                                         central_moment_transform_class=cspace.central_moment_transform_class,
-                                         cumulant_transform_class=cspace.cumulant_transform_class)
-
-
-def create_with_polynomial_cumulants(stencil, relaxation_rates, cumulant_groups, **kwargs):
-    r"""Creates a cumulant lattice Boltzmann model based on a default polynomial set.
+def create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments=True, **kwargs):
+    r"""Creates a cumulant-based lattice Boltzmann method.
 
     Args:
         stencil: instance of :class:`lbmpy.stencils.LBStencil`
@@ -576,17 +542,31 @@ def create_with_polynomial_cumulants(stencil, relaxation_rates, cumulant_groups,
                           that the force is applied correctly to the momentum groups
         cumulant_groups: Nested sequence of polynomial expressions defining the cumulants to be relaxed. All cumulants 
                          within one group are relaxed with the same relaxation rate.
-        kwargs: See :func:`create_centered_cumulant_model`
+        conserved_moments: If lower order moments are conserved or not.
+        kwargs: See :func:`create_with_continuous_maxwellian_equilibrium`
 
     Returns:
-        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
+        :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
     """
-    cumulant_to_rr_dict = _get_relaxation_info_dict(relaxation_rates, cumulant_groups, stencil.D)
-    return create_centered_cumulant_model(stencil, cumulant_to_rr_dict, **kwargs)
+    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, **kwargs):
-    r"""Creates a cumulant lattice Boltzmann model based on a default polinomial set.
+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`
@@ -595,32 +575,33 @@ def create_with_monomial_cumulants(stencil, relaxation_rates, **kwargs):
                           used for determine the viscosity of the simulation. All other cumulants are relaxed with one.
                           If a cumulant force model is provided the first order cumulants are relaxed with two to ensure
                           that the force is applied correctly to the momentum groups
-        kwargs: See :func:`create_centered_cumulant_model`
+        conserved_moments: If lower order moments are conserved or not.
+        kwargs: See :func:`create_cumulant`
 
     Returns:
-        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
+        :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
     """
     # 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)
 
-    return create_with_polynomial_cumulants(stencil, relaxation_rates, cumulant_groups, **kwargs)
 
+def create_with_default_polynomial_cumulants(stencil, relaxation_rates, conserved_moments=True, **kwargs):
+    r"""Creates a cumulant lattice Boltzmann model based on the default polynomial set of :cite:`geier2015`.
 
-def create_with_default_polynomial_cumulants(stencil, relaxation_rates, **kwargs):
-    r"""Creates a cumulant lattice Boltzmann model based on a default polynomial set.
-
-    Args: See :func:`create_with_polynomial_cumulants`.
+    Args: See :func:`create_cumulant`.
 
     Returns:
-        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
+        :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
     """
     # Get polynomial groups
     cumulant_groups = cascaded_moment_sets_literature(stencil)
-    return create_with_polynomial_cumulants(stencil, relaxation_rates, cumulant_groups, **kwargs)
+    return create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments, **kwargs)
 
 
-def _get_relaxation_info_dict(relaxation_rates, nested_moments, dim):
+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:
@@ -630,6 +611,7 @@ def _get_relaxation_info_dict(relaxation_rates, nested_moments, dim):
                           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()
 
@@ -638,7 +620,9 @@ def _get_relaxation_info_dict(relaxation_rates, nested_moments, dim):
             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
@@ -657,12 +641,18 @@ def _get_relaxation_info_dict(relaxation_rates, nested_moments, dim):
     if len(relaxation_rates) == 1:
         for group in nested_moments:
             for moment in group:
-                if get_order(moment) <= 1:
-                    result[moment] = 0.0
-                elif is_shear_moment(moment, dim):
-                    result[moment] = relaxation_rates[0]
+                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:
-                    result[moment] = 1.0
+                    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:
@@ -686,15 +676,25 @@ def _get_relaxation_info_dict(relaxation_rates, nested_moments, dim):
                     rr = next(rr_iter)
                 next_rr = False
                 for moment in group:
-                    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
+                    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:
-                        next_rr = True
-                        result[moment] = rr
+                        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 "
@@ -703,7 +703,18 @@ def _get_relaxation_info_dict(relaxation_rates, nested_moments, dim):
                              "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 ----------------------------------------------
 
 

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.""")

lbmpy/methods/centeredcumulant/centeredcumulantmethod.py → src/lbmpy/methods/cumulantbased/cumulantbasedmethod.py

-from pystencils.simp.simplifications import sympy_cse
 import sympy as sp
 from collections import OrderedDict
-from warnings import warn, filterwarnings
+from typing import Set
+from warnings import filterwarnings
 
 from pystencils import Assignment, AssignmentCollection
-from pystencils.simp.assignment_collection import SymbolGen
-from pystencils.stencil import have_same_entries
-from pystencils.cache import disk_cache
 from pystencils.sympyextensions import is_constant
+from pystencils.simp import apply_to_all_assignments
 
-from lbmpy.enums import Stencil
-from lbmpy.stencils import LBStencil
 from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
 from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
 
-from lbmpy.moments import (moments_up_to_order, get_order, moment_matrix,
-                           monomial_to_polynomial_transformation_matrix,
-                           moment_sort_key, exponent_tuple_sort_key,
-                           exponent_to_polynomial_representation, extract_monomials, MOMENT_SYMBOLS,
-                           statistical_quantity_symbol)
-
-from lbmpy.forcemodels import Luo, Simple
-
-#   Local Imports
-
-from .cumulant_transform import (
-    PRE_COLLISION_CUMULANT, POST_COLLISION_CUMULANT,
-    CentralMomentsToCumulantsByGeneratingFunc)
+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,
-    PdfsToCentralMomentsByShiftMatrix)
-
-from lbmpy.methods.centeredcumulant.force_model import CenteredCumulantForceModel
-from lbmpy.methods.centeredcumulant.galilean_correction import (
-    contains_corrected_polynomials,
-    add_galilean_correction,
-    get_galilean_correction_terms)
-
-
-#   ============================ Cached Transformations ================================================================
-
-@disk_cache
-def cached_forward_transform(transform_obj, *args, **kwargs):
-    return transform_obj.forward_transform(*args, **kwargs)
-
-
-@disk_cache
-def cached_backward_transform(transform_obj, *args, **kwargs):
-    return transform_obj.backward_transform(*args, **kwargs)
-
-
-#   ============================ Lower Order Central Moment Collision ==================================================
-
-
-@disk_cache
-def relax_lower_order_central_moments(moment_indices, pre_collision_values,
-                                      relaxation_rates, equilibrium_values,
-                                      post_collision_base=POST_COLLISION_MONOMIAL_CENTRAL_MOMENT):
-    post_collision_symbols = [statistical_quantity_symbol(post_collision_base, i) for i in moment_indices]
-    equilibrium_vec = sp.Matrix(equilibrium_values)
-    moment_vec = sp.Matrix(pre_collision_values)
-    relaxation_matrix = sp.diag(*relaxation_rates)
-    moment_vec = moment_vec + relaxation_matrix * (equilibrium_vec - moment_vec)
-    main_assignments = [Assignment(s, eq) for s, eq in zip(post_collision_symbols, moment_vec)]
-
-    return AssignmentCollection(main_assignments)
-
-
-#   ============================ Polynomial Cumulant Collision =========================================================
-
-@disk_cache
-def relax_polynomial_cumulants(monomial_exponents, polynomials, relaxation_rates, equilibrium_values,
-                               pre_simplification,
-                               galilean_correction_terms=None,
-                               pre_collision_base=PRE_COLLISION_CUMULANT,
-                               post_collision_base=POST_COLLISION_CUMULANT,
-                               subexpression_base='sub_col'):
-    mon_to_poly_matrix = monomial_to_polynomial_transformation_matrix(monomial_exponents, polynomials)
-    mon_vec = sp.Matrix([statistical_quantity_symbol(pre_collision_base, exp) for exp in monomial_exponents])
-    equilibrium_vec = sp.Matrix(equilibrium_values)
-    relaxation_matrix = sp.diag(*relaxation_rates)
-
-    subexpressions = []
-
-    poly_vec = mon_to_poly_matrix * mon_vec
-    relaxed_polys = poly_vec + relaxation_matrix * (equilibrium_vec - poly_vec)
-
-    if galilean_correction_terms is not None:
-        relaxed_polys = add_galilean_correction(relaxed_polys, polynomials, galilean_correction_terms)
-        subexpressions = galilean_correction_terms.all_assignments
-
-    relaxed_monos = mon_to_poly_matrix.inv() * relaxed_polys
-
-    main_assignments = [Assignment(statistical_quantity_symbol(post_collision_base, exp), v)
-                        for exp, v in zip(monomial_exponents, relaxed_monos)]
-
-    symbol_gen = SymbolGen(subexpression_base)
-    ac = AssignmentCollection(
-        main_assignments, subexpressions=subexpressions, subexpression_symbol_generator=symbol_gen)
-    if pre_simplification == 'default_with_cse':
-        ac = sympy_cse(ac)
-    return ac
+    CentralMomentsToCumulantsByGeneratingFunc,
+    BinomialChimeraTransform)
 
 
-#   =============================== LB Method Implementation ===========================================================
-
-class CenteredCumulantBasedLbMethod(AbstractLbMethod):
+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`.
 
-    Conserved quantities are relaxed in central moment space. This method supports an implicit forcing scheme
-    through :class:`lbmpy.methods.centeredcumulant.CenteredCumulantForceModel` where forces are applied by
-    shifting the central-moment frame of reference by :math:`F/2` and then relaxing the first-order central
-    moments with a relaxation rate of two. This corresponds to the change-of-sign described in the paper.
-    Classical forcing schemes can still be applied.
-
-    The galilean correction described in :cite:`geier2015` is also available for the D3Q27 lattice.
-
-    This method is implemented modularily as the transformation from populations to central moments to cumulants
+    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.
@@ -136,7 +40,6 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
         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.
-        galilean_correction: if set to True the galilean_correction is applied to a D3Q27 cumulant method
         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
@@ -144,53 +47,36 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
 
     def __init__(self, stencil, equilibrium, relaxation_dict,
                  conserved_quantity_computation=None,
-                 force_model=None, zero_centered=False,
-                 galilean_correction=False,
-                 central_moment_transform_class=PdfsToCentralMomentsByShiftMatrix,
+                 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(CenteredCumulantBasedLbMethod, self).__init__(stencil)
+        super(CumulantBasedLbMethod, self).__init__(stencil)
 
         if force_model is not None:
-            assert (isinstance(force_model, CenteredCumulantForceModel)
-                    or isinstance(force_model, Simple)
-                    or isinstance(force_model, Luo)), "Given force model currently not supported."
-
-        for m in moments_up_to_order(1, dim=self.dim):
-            if exponent_to_polynomial_representation(m) not in relaxation_dict.keys():
-                raise ValueError(f'No relaxation info given for conserved cumulant {m}!')
+            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._galilean_correction = galilean_correction
-
-        if galilean_correction:
-            if not have_same_entries(stencil, LBStencil(Stencil.D3Q27)):
-                raise ValueError("Galilean Correction only available for D3Q27 stencil")
-
-            if not contains_corrected_polynomials(relaxation_dict):
-                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!")
-
         self._cumulant_transform_class = cumulant_transform_class
         self._central_moment_transform_class = central_moment_transform_class
 
-        self.force_model_rr_override = False
-        if isinstance(self._force_model, CenteredCumulantForceModel) and \
-                self._force_model.override_momentum_relaxation_rate is not None:
-            self.set_first_moment_relaxation_rate(self._force_model.override_momentum_relaxation_rate)
-            self.force_model_rr_override = True
-
     @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.
@@ -228,15 +114,7 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
     @property
     def cumulant_equilibrium_values(self):
         """Equilibrium values of this method's :attr:`cumulants`."""
-        cumulants = self.cumulants
-        equilibria = list(self._equilibrium.cumulants(cumulants, rescale=True))
-
-        for i, c in enumerate(cumulants):
-            if get_order(c) == 0:
-                equilibria[i] = self.zeroth_order_equilibrium_moment_symbol
-            elif get_order(c) == 1:
-                equilibria[i] = sp.Integer(0)
-        return tuple(equilibria)
+        return self._equilibrium.cumulants(self.cumulants, rescale=True)
 
     @property
     def relaxation_rates(self):
@@ -251,7 +129,7 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
     @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
+        which is the area under its curve
         (see :attr:`lbmpy.equilibrium.AbstractEquilibrium.zeroth_order_moment_symbol`)."""
         return self._equilibrium.zeroth_order_moment_symbol
 
@@ -261,19 +139,11 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
         which is its mean value. (see :attr:`lbmpy.equilibrium.AbstractEquilibrium.first_order_moment_symbols`)."""
         return self._equilibrium.first_order_moment_symbols
 
-    @property
-    def galilean_correction(self):
-        """Whether or not the gallilean correction is included in the collision equations."""
-        return self._galilean_correction
-
     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):
-        if self.force_model_rr_override:
-            warn("Overwriting first-order relaxation rates governed by CenteredCumulantForceModel "
-                 "might break your forcing scheme.")
         for e in MOMENT_SYMBOLS[:self.dim]:
             assert e in self._relaxation_dict, \
                 "First cumulants are not relaxed separately by this method"
@@ -285,9 +155,8 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
         self.set_first_moment_relaxation_rate(relaxation_rate)
 
     def set_force_model(self, force_model):
-        assert (isinstance(force_model, CenteredCumulantForceModel)
-                or isinstance(force_model, Simple)
-                or isinstance(force_model, Luo)), "Given force model currently not supported."
+        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):
@@ -326,11 +195,6 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
                 'eq_value': sp.latex(eq_value),
                 'nb': 'style="border:none"',
             }
-            order = get_order(cumulant)
-            if order <= 1:
-                vals['cumulant'] += ' (central moment)'
-                if order == 1 and self.force_model_rr_override:
-                    vals['rr'] += ' (overridden by force model)'
             html += """<tr {nb}>
                             <td {nb}>${cumulant}$</td>
                             <td {nb}>${eq_value}$</td>
@@ -353,10 +217,11 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
         assert len(weights) == len(self.stencil)
         self._weights = weights
 
-    def get_equilibrium(self, conserved_quantity_equations=None, subexpressions=False, pre_simplification=False,
-                        keep_cqc_subexpressions=True):
+    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
+        The equations have the post collision symbols as left-hand sides and are
         functions of the conserved quantities
 
         Args:
@@ -367,34 +232,46 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
             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 = {c: RelaxationInfo(info.equilibrium_value, sp.Integer(1))
-                       for c, info in self.relaxation_info_dict.items()}
-        ac = self._centered_cumulant_collision_rule(
-            r_info_dict, conserved_quantity_equations, pre_simplification, include_galilean_correction=False)
+        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)
-                return ac.new_without_subexpressions(subexpressions_to_keep=bs)
+                ac = expand_all_assignments(ac.new_without_subexpressions(subexpressions_to_keep=bs))
+                return ac.new_without_unused_subexpressions()
             else:
-                return ac.new_without_subexpressions()
+                ac = expand_all_assignments(ac.new_without_subexpressions())
+                return ac.new_without_unused_subexpressions()
         else:
-            return ac
+            return ac.new_without_unused_subexpressions()
 
-    def get_equilibrium_terms(self):
+    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=None, pre_simplification=False):
+    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(
-            self.relaxation_info_dict, conserved_quantity_equations, pre_simplification, True,
-            symbolic_relaxation_rates=True)
+        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=None):
+    def _bound_symbols_cqc(self, conserved_quantity_equations: AssignmentCollection = None) -> Set[sp.Symbol]:
         f = self.pre_collision_pdf_symbols
         cqe = conserved_quantity_equations
 
@@ -421,12 +298,11 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
 
         return [w for w in weights]
 
-    def _centered_cumulant_collision_rule(self, cumulant_to_relaxation_info_dict,
-                                          conserved_quantity_equations=None,
-                                          pre_simplification=False,
-                                          include_force_terms=False,
-                                          include_galilean_correction=True,
-                                          symbolic_relaxation_rates=False):
+    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
@@ -438,15 +314,17 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
         velocity = self.first_order_equilibrium_moment_symbols
         cqe = conserved_quantity_equations
 
-        relaxation_info_dict = dict()
-        subexpressions_relaxation_rates = []
+        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, sd = self._generate_symbolic_relaxation_matrix()
-            for i, cumulant in enumerate(cumulant_to_relaxation_info_dict):
-                relaxation_info_dict[cumulant] = RelaxationInfo(cumulant_to_relaxation_info_dict[cumulant][0],
-                                                                sd[i, i])
+            subexpressions_relaxation_rates, d = self._generate_symbolic_relaxation_matrix(relaxation_rates=rrs)
         else:
-            relaxation_info_dict = cumulant_to_relaxation_info_dict
+            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)
@@ -454,6 +332,8 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
         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:
@@ -463,70 +343,56 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
         else:
             background_distribution = None
 
-        #   1) Extract Monomial Cumulants for the higher-order polynomials
-        polynomial_cumulants = relaxation_info_dict.keys()
-        polynomial_cumulants = sorted(list(polynomial_cumulants), key=moment_sort_key)
-        higher_order_polynomials = [p for p in polynomial_cumulants if get_order(p) > 1]
-        monomial_cumulants = sorted(list(extract_monomials(
-            higher_order_polynomials, dim=self.dim)), key=exponent_tuple_sort_key)
-
-        #   2) Get Forward and Backward Transformations between central moment and cumulant space,
+        #   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, monomial_cumulants, density, velocity)
-        k_to_c_eqs = cached_forward_transform(k_to_c_transform, simplification=pre_simplification)
-        c_post_to_k_post_eqs = cached_backward_transform(
-            k_to_c_transform, simplification=pre_simplification, omit_conserved_moments=True)
+        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
-        assert len(central_moments) == stencil.Q, 'Number of required central moments must match stencil size.'
 
-        #   3) Get Forward Transformation from PDFs to 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 = cached_forward_transform(
-            pdfs_to_k_transform, f, simplification=pre_simplification, return_monomials=True)
-
-        #   4) Add relaxation rules for lower order moments
-        lower_order_moments = moments_up_to_order(1, dim=self.dim)
-        lower_order_moment_symbols = [statistical_quantity_symbol(PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT, exp)
-                                      for exp in lower_order_moments]
-
-        lower_order_relaxation_infos = [relaxation_info_dict[exponent_to_polynomial_representation(e)]
-                                        for e in lower_order_moments]
-        lower_order_relaxation_rates = [info.relaxation_rate for info in lower_order_relaxation_infos]
-        lower_order_equilibrium = [info.equilibrium_value for info in lower_order_relaxation_infos]
-
-        lower_order_moment_collision_eqs = relax_lower_order_central_moments(
-            lower_order_moments, tuple(lower_order_moment_symbols),
-            tuple(lower_order_relaxation_rates), tuple(lower_order_equilibrium))
-
-        #   5) Add relaxation rules for higher-order, polynomial cumulants
-        poly_relaxation_infos = [relaxation_info_dict[c] for c in higher_order_polynomials]
-        poly_relaxation_rates = [info.relaxation_rate for info in poly_relaxation_infos]
-        poly_equilibrium = [info.equilibrium_value for info in poly_relaxation_infos]
-
-        if self._galilean_correction and include_galilean_correction:
-            galilean_correction_terms = get_galilean_correction_terms(
-                relaxation_info_dict, density, velocity)
-        else:
-            galilean_correction_terms = None
-
-        cumulant_collision_eqs = relax_polynomial_cumulants(
-            tuple(monomial_cumulants), tuple(higher_order_polynomials),
-            tuple(poly_relaxation_rates), tuple(poly_equilibrium),
-            pre_simplification,
-            galilean_correction_terms=galilean_correction_terms)
-
-        #   6) Get backward transformation from central moments to PDFs
+        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 = cached_backward_transform(
-            pdfs_to_k_transform, d, simplification=pre_simplification, start_from_monomials=True)
+        k_post_to_pdfs_eqs = pdfs_to_k_transform.backward_transform(
+            d, simplification=pre_simplification, start_from_monomials=True)
 
-        #   7) That's all. Now, put it all together.
+        #   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,
-                    lower_order_moment_collision_eqs, cumulant_collision_eqs, c_post_to_k_post_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
@@ -534,18 +400,8 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
         simplification_hints = cqe.simplification_hints.copy()
         simplification_hints.update(self._cqc.defined_symbols())
         simplification_hints['relaxation_rates'] = [rr for rr in self.relaxation_rates]
-
-        #   8) Maybe add forcing terms if CenteredCumulantForceModel was not used
-        if self._force_model is not None and \
-                not isinstance(self._force_model, CenteredCumulantForceModel) and include_force_terms:
-            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
+        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)

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

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

lbmpy/methods/momentbased/centralmomentbasedmethod.py → 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 FastCentralMomentTransform
+from lbmpy.moment_transforms import BinomialChimeraTransform
 
 from lbmpy.moments import MOMENT_SYMBOLS, moment_matrix, set_up_shift_matrix
 
@@ -53,8 +55,8 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
 
     def __init__(self, stencil, equilibrium, relaxation_dict,
                  conserved_quantity_computation=None,
-                 force_model=None, zero_centered=False,
-                 central_moment_transform_class=FastCentralMomentTransform):
+                 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)
 
@@ -63,6 +65,7 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
         self._cqc = conserved_quantity_computation
         self._force_model = force_model
         self._zero_centered = zero_centered
+        self._fraction_field = fraction_field
         self._weights = None
         self._central_moment_transform_class = central_moment_transform_class
 
@@ -71,6 +74,10 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
         """Force model employed by this method."""
         return self._force_model
 
+    @property
+    def fraction_field(self):
+        return self._fraction_field
+
     @property
     def relaxation_info_dict(self):
         """Dictionary mapping this method's moments to their relaxation rates and equilibrium values.
@@ -102,7 +109,7 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
     @property
     def moment_equilibrium_values(self):
         """Equilibrium values of this method's :attr:`moments`."""
-        return self._equilibrium.central_moments(self.moments)
+        return self._equilibrium.central_moments(self.moments, self.first_order_equilibrium_moment_symbols)
 
     @property
     def relaxation_rates(self):
@@ -152,15 +159,15 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
         self._force_model = force_model
 
     @property
-    def moment_matrix(self):
+    def moment_matrix(self) -> sp.Matrix:
         return moment_matrix(self.moments, self.stencil)
 
     @property
-    def shift_matrix(self):
+    def shift_matrix(self) -> sp.Matrix:
         return set_up_shift_matrix(self.moments, self.stencil)
 
     @property
-    def relaxation_matrix(self):
+    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]
@@ -218,10 +225,11 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
 
     #   ----------------------- Overridden Abstract Members --------------------------
 
-    def get_equilibrium(self, conserved_quantity_equations=None, subexpressions=False, pre_simplification=False,
-                        keep_cqc_subexpressions=True, include_force_terms=False):
+    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
+        The equations have the post collision symbols as left-hand sides and are
         functions of the conserved quantities
 
         Args:
@@ -232,34 +240,48 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
             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 = {c: RelaxationInfo(info.equilibrium_value, sp.Integer(1))
-                       for c, info in self.relaxation_info_dict.items()}
-        ac = self._central_moment_collision_rule(r_info_dict, conserved_quantity_equations, pre_simplification,
+        _, 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)
-                return ac.new_without_subexpressions(subexpressions_to_keep=bs)
+                ac = expand_all_assignments(ac.new_without_subexpressions(subexpressions_to_keep=bs))
+                return ac.new_without_unused_subexpressions()
             else:
-                return ac.new_without_subexpressions()
+                ac = expand_all_assignments(ac.new_without_subexpressions())
+                return ac.new_without_unused_subexpressions()
         else:
-            return ac
+            return ac.new_without_unused_subexpressions()
 
-    def get_equilibrium_terms(self):
+    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=None, pre_simplification=False):
+    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(self.relaxation_info_dict, conserved_quantity_equations,
-                                                   pre_simplification, True, symbolic_relaxation_rates=True)
+        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=None):
+    def _bound_symbols_cqc(self, conserved_quantity_equations: AssignmentCollection = None) -> Set[sp.Symbol]:
         f = self.pre_collision_pdf_symbols
         cqe = conserved_quantity_equations
 
@@ -285,11 +307,11 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
 
         return [w for w in weights]
 
-    def _central_moment_collision_rule(self, moment_to_relaxation_info_dict,
-                                       conserved_quantity_equations=None,
-                                       pre_simplification=False,
-                                       include_force_terms=False,
-                                       symbolic_relaxation_rates=False):
+    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

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

@@ -286,7 +286,7 @@ def _get_relaxation_rates(collision_rule):
     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_relxation_rate.items():
+    for symbolic_rr, rr in method.subs_dict_relaxation_rate.items():
         if omega_s == rr:
             omega_s = symbolic_rr
 

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

 from collections import OrderedDict
+from typing import Iterable, Set
 
 import sympy as sp
 import numpy as np
@@ -15,8 +16,8 @@ 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
+    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:
@@ -38,7 +39,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
 
     def __init__(self, stencil, equilibrium, relaxation_dict,
                  conserved_quantity_computation=None, force_model=None, zero_centered=False,
-                 moment_transform_class=PdfsToMomentsByChimeraTransform):
+                 fraction_field=None, moment_transform_class=PdfsToMomentsByChimeraTransform):
         assert isinstance(conserved_quantity_computation, AbstractConservedQuantityComputation)
         super(MomentBasedLbMethod, self).__init__(stencil)
 
@@ -47,6 +48,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
         self._cqc = conserved_quantity_computation
         self._force_model = force_model
         self._zero_centered = zero_centered
+        self._fraction_field = fraction_field
         self._weights = None
         self._moment_transform_class = moment_transform_class
 
@@ -55,6 +57,10 @@ class MomentBasedLbMethod(AbstractLbMethod):
         """Force model employed by this method."""
         return self._force_model
 
+    @property
+    def fraction_field(self):
+        return self._fraction_field
+
     @property
     def relaxation_info_dict(self):
         """Dictionary mapping this method's moments to their relaxation rates and equilibrium values.
@@ -127,31 +133,64 @@ class MomentBasedLbMethod(AbstractLbMethod):
         assert len(weights) == len(self.stencil)
         self._weights = weights
 
-    def get_equilibrium(self, conserved_quantity_equations=None, include_force_terms=False,
-                        pre_simplification=False, subexpressions=False, keep_cqc_subexpressions=True):
-        relaxation_matrix = sp.eye(len(self.relaxation_rates))
-        ac = self._collision_rule_with_relaxation_matrix(relaxation_matrix,
-                                                         conserved_quantity_equations=conserved_quantity_equations,
+    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)
-                return ac.new_without_subexpressions(subexpressions_to_keep=bs)
+                ac = ac.new_without_subexpressions(subexpressions_to_keep=bs)
+                return ac.new_without_unused_subexpressions()
             else:
-                return ac.new_without_subexpressions()
+                ac = ac.new_without_subexpressions()
+                return ac.new_without_unused_subexpressions()
         else:
-            return ac
+            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=None, pre_simplification=True):
-        relaxation_rate_sub_expressions, d = self._generate_symbolic_relaxation_matrix()
-        ac = self._collision_rule_with_relaxation_matrix(d, relaxation_rate_sub_expressions,
-                                                         True, conserved_quantity_equations,
+    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
 
@@ -179,27 +218,27 @@ class MomentBasedLbMethod(AbstractLbMethod):
             self._cqc.set_force_model(force_model)
 
     @property
-    def collision_matrix(self):
+    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):
+    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):
+    def moment_matrix(self) -> sp.Matrix:
         return moment_matrix(self.moments, self.stencil)
 
     @property
-    def is_orthogonal(self):
+    def is_orthogonal(self) -> bool:
         return (self.moment_matrix * self.moment_matrix.T).is_diagonal()
 
     @property
-    def is_weighted_orthogonal(self):
+    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))
 
@@ -273,7 +312,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
 
         return [w for w in weights]
 
-    def _bound_symbols_cqc(self, conserved_quantity_equations=None):
+    def _bound_symbols_cqc(self, conserved_quantity_equations: AssignmentCollection = None) -> Set[sp.Symbol]:
         f = self.pre_collision_pdf_symbols
         cqe = conserved_quantity_equations
 
@@ -282,8 +321,13 @@ class MomentBasedLbMethod(AbstractLbMethod):
 
         return cqe.bound_symbols
 
-    def _collision_rule_with_relaxation_matrix(self, d, additional_subexpressions=(), include_force_terms=True,
-                                               conserved_quantity_equations=None, pre_simplification=False):
+    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)
 
@@ -346,7 +390,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
             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 customly required, derive equations entirely in
+            #   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"

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

@@ -4,10 +4,12 @@ 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, 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
@@ -41,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)
@@ -318,6 +320,45 @@ def split_pdf_main_assignments_by_symmetry(ac: AssignmentCollection):
     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 --------------------------------------------------------------
 
 
@@ -344,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