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

@@ -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. 
@@ -201,11 +201,11 @@ 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,
+                                   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,
+                                   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,
@@ -304,7 +304,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.
 
@@ -318,6 +318,7 @@ 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
     """
@@ -325,7 +326,7 @@ def create_mrt_raw(stencil, relaxation_rates, continuous_equilibrium=True, **kwa
     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:
@@ -333,7 +334,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, fraction_field=None, **kwargs):
     r"""
     Creates moment based LB method where the collision takes place in the central moment space.
 
@@ -346,6 +347,8 @@ 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.
+        fraction_field: fraction field for the PSM method
     Returns:
         :class:`lbmpy.methods.momentbased.CentralMomentBasedLbMethod` instance
     """
@@ -367,7 +370,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 is not None:
+        relaxation_rates_modifier = (1.0 - fraction_field.center)
+        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)
@@ -455,7 +462,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
@@ -476,6 +483,7 @@ 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)
     check_and_set_mrt_space(CollisionSpace.RAW_MOMENTS)
@@ -507,7 +515,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,
@@ -518,8 +527,7 @@ def create_mrt_orthogonal(stencil, relaxation_rates, continuous_equilibrium=True
 
 
 # ----------------------------------------- Cumulant method creators ---------------------------------------------------
-
-def create_cumulant(stencil, relaxation_rates, cumulant_groups, **kwargs):
+def create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments=True, fraction_field=None, **kwargs):
     r"""Creates a cumulant-based lattice Boltzmann method.
 
     Args:
@@ -531,12 +539,19 @@ def create_cumulant(stencil, relaxation_rates, cumulant_groups, **kwargs):
                           that the force is applied correctly to the momentum groups
         cumulant_groups: Nested sequence of polynomial expressions defining the cumulants to be relaxed. All cumulants 
                          within one group are relaxed with the same relaxation rate.
+        conserved_moments: If lower order moments are conserved or not.
         kwargs: See :func:`create_with_continuous_maxwellian_equilibrium`
 
     Returns:
         :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod` instance
     """
-    cumulant_to_rr_dict = _get_relaxation_info_dict(relaxation_rates, cumulant_groups, stencil.D)
+    cumulant_to_rr_dict = _get_relaxation_info_dict(relaxation_rates, cumulant_groups, stencil.D, conserved_moments)
+
+    if fraction_field is not None:
+        relaxation_rates_modifier = (1.0 - fraction_field.center)
+        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:
@@ -547,7 +562,7 @@ def create_cumulant(stencil, relaxation_rates, cumulant_groups, **kwargs):
                                                          **kwargs)
 
 
-def create_with_monomial_cumulants(stencil, relaxation_rates, **kwargs):
+def create_with_monomial_cumulants(stencil, relaxation_rates, conserved_moments=True, **kwargs):
     r"""Creates a cumulant lattice Boltzmann model using the given stencil's canonical monomial cumulants.
 
     Args:
@@ -557,6 +572,7 @@ 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
+        conserved_moments: If lower order moments are conserved or not.
         kwargs: See :func:`create_cumulant`
 
     Returns:
@@ -565,10 +581,10 @@ def create_with_monomial_cumulants(stencil, relaxation_rates, **kwargs):
     # 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, **kwargs)
+    return create_cumulant(stencil, relaxation_rates, cumulant_groups, conserved_moments, **kwargs)
 
 
-def create_with_default_polynomial_cumulants(stencil, relaxation_rates, **kwargs):
+def create_with_default_polynomial_cumulants(stencil, relaxation_rates, conserved_moments=True, **kwargs):
     r"""Creates a cumulant lattice Boltzmann model based on the default polynomial set of :cite:`geier2015`.
 
     Args: See :func:`create_cumulant`.
@@ -578,10 +594,11 @@ def create_with_default_polynomial_cumulants(stencil, relaxation_rates, **kwargs
     """
     # Get polynomial groups
     cumulant_groups = cascaded_moment_sets_literature(stencil)
-    return create_cumulant(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:
@@ -591,6 +608,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()
 
@@ -599,7 +617,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
@@ -618,12 +638,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:
@@ -647,15 +673,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 "
@@ -664,6 +700,9 @@ 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
 
 

lbmpy/methods/cumulantbased/__init__.py → 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']
+__all__ = ['CumulantBasedLbMethod', 'add_galilean_correction', 'add_fourth_order_correction']

lbmpy/methods/cumulantbased/cumulant_simplifications.py → src/lbmpy/methods/cumulantbased/cumulant_simplifications.py

@@ -16,11 +16,8 @@ def insert_logs(ac, **kwargs):
 def insert_log_products(ac, **kwargs):
     def callback(asm):
         rhs = asm.rhs
-        if isinstance(rhs, sp.log):
+        if rhs.find(sp.log):
             return True
-        if isinstance(rhs, sp.Mul):
-            if any(isinstance(arg, sp.log) for arg in rhs.args):
-                return True
         return False
 
     return insert_subexpressions(ac, callback, **kwargs)

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

@@ -23,7 +23,7 @@ class CumulantBasedLbMethod(AbstractLbMethod):
     This class implements cumulant-based lattice boltzmann methods which relax all the non-conserved quantities
     as either monomial or polynomial cumulants. It is mostly inspired by the work presented in :cite:`geier2015`.
 
-    This method is implemented 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.
@@ -124,7 +124,7 @@ class CumulantBasedLbMethod(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
 

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/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

@@ -16,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:
@@ -39,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)
 
@@ -48,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
 
@@ -174,7 +175,14 @@ class MomentBasedLbMethod(AbstractLbMethod):
 
     def get_collision_rule(self, conserved_quantity_equations: AssignmentCollection = None,
                            pre_simplification: bool = True) -> LbmCollisionRule:
-        rr_sub_expressions, d = self._generate_symbolic_relaxation_matrix()
+
+        if self.fraction_field is not None:
+            relaxation_rates_modifier = (1.0 - self.fraction_field.center)
+            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,
@@ -378,7 +386,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/moment_transforms/cumulanttransforms.py → src/lbmpy/moment_transforms/cumulanttransforms.py

@@ -67,10 +67,10 @@ class CentralMomentsToCumulantsByGeneratingFunc(AbstractMomentTransform):
         self.cumulant_exponents = self.moment_exponents
         self.cumulant_polynomials = self.moment_polynomials
 
-        if(len(self.cumulant_exponents) != stencil.Q):
+        if len(self.cumulant_exponents) != stencil.Q:
             raise ValueError("Number of cumulant exponent tuples must match stencil size.")
 
-        if(len(self.cumulant_polynomials) != stencil.Q):
+        if len(self.cumulant_polynomials) != stencil.Q:
             raise ValueError("Number of cumulant polynomials must match stencil size.")
 
         self.central_moment_exponents = self.compute_required_central_moments()

lbmpy/moment_transforms/rawmomenttransforms.py → src/lbmpy/moment_transforms/rawmomenttransforms.py

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

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

@@ -86,8 +86,8 @@ def moment_permutations(exponent_tuple):
 def moments_of_order(order, dim=3, include_permutations=True):
     """All tuples of length 'dim' which sum equals 'order'"""
     for item in __fixed_sum_tuples(dim, order, ordered=not include_permutations):
-        assert(len(item) == dim)
-        assert(sum(item) == order)
+        assert len(item) == dim
+        assert sum(item) == order
         yield item