Implemented CM in momentbased.py

lbmpy/creationfunctions.py

@@ -179,6 +179,7 @@ For example, to modify the AST one can run::
 
 """
 from copy import copy
+from warnings import warn
 
 import sympy as sp
 
@@ -522,6 +523,7 @@ def update_with_default_parameters(params, opt_params=None, fail_on_unknown_para
         'force_model': 'none',
         'force': (0, 0, 0),
         'maxwellian_moments': True,
+        'central_moments': False,
         'cumulant': False,
         'initial_velocity': None,
 
@@ -613,6 +615,16 @@ def update_with_default_parameters(params, opt_params=None, fail_on_unknown_para
             raise ValueError("Unknown optimization parameter(s): " + ",".join(unknown_opt_params))
 
     params_result = copy(default_method_description)
+
+    # TODO: the equlibrium order should be set automatically for all methods
+    if 'central_moments' in params and 'equilibrium_order' in params:
+        if params['central_moments'] and params['equilibrium_order'] < 4:
+            warn("The central moments should be used with an equilibrium order depending on the highest order moment")
+
+    if 'central_moments' in params and 'equilibrium_order' not in params:
+        if params['central_moments']:
+            params['equilibrium_order'] = 4
+
     params_result.update(params)
     opt_params_result = copy(default_optimization_description)
     opt_params_result.update(opt_params)

lbmpy/methods/momentbased.py

@@ -5,7 +5,7 @@ import sympy as sp
 from lbmpy.maxwellian_equilibrium import get_weights
 from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
 from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
-from lbmpy.moments import MOMENT_SYMBOLS, moment_matrix, shift_matrix
+from lbmpy.moments import MOMENT_SYMBOLS, moment_matrix, set_up_shift_matrix
 from pystencils import Assignment
 from pystencils.sympyextensions import subs_additive
 
@@ -17,6 +17,10 @@ class MomentBasedLbMethod(AbstractLbMethod):
         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.
+        The collision can also be applied in the central moment space. For this Moment based LBM needs to set up a
+        shift matrix which is used to shift from moment space to central moment space. The shift matrix can be set up by
+        calling set_shift_matrix which either sets up the matrix based on the moments and the stencil of
+        Moment based LBM or takes an individual shift matrix as argument which is used for further calculations.
 
         Args:
             stencil: see :func:`lbmpy.stencils.get_stencil`
@@ -35,6 +39,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
         self._momentToRelaxationInfoDict = OrderedDict(moment_to_relaxation_info_dict.items())
         self._conservedQuantityComputation = conserved_quantity_computation
         self._weights = None
+        self._shift_matrix = None
 
     @property
     def force_model(self):
@@ -74,13 +79,23 @@ class MomentBasedLbMethod(AbstractLbMethod):
             self._weights = self._compute_weights()
         return self._weights
 
+    def set_shift_matrix(self, shift_matrix=None):
+        if shift_matrix:
+            self._shift_matrix = shift_matrix
+        else:
+            self._shift_matrix = set_up_shift_matrix(self.moments, self.stencil)
+
     def override_weights(self, weights):
         assert len(weights) == len(self.stencil)
         self._weights = weights
 
     def get_equilibrium(self, conserved_quantity_equations=None, include_force_terms=False):
         relaxation_matrix = sp.eye(len(self.relaxation_rates))
-        return self._collision_rule_with_relaxation_matrix(relaxation_matrix,
+        if self.shift_matrix:
+            sm = sp.eye(len(self.relaxation_rates))
+        else:
+            sm = None
+        return self._collision_rule_with_relaxation_matrix(relaxation_matrix, sm,
                                                            conserved_quantity_equations=conserved_quantity_equations,
                                                            include_force_terms=include_force_terms)
 
@@ -90,8 +105,9 @@ class MomentBasedLbMethod(AbstractLbMethod):
 
     def get_collision_rule(self, conserved_quantity_equations=None, keep_rrs_symbolic=True):
         d = self.relaxation_matrix
+        sm = self.shift_matrix
         relaxation_rate_sub_expressions, d = self._generate_relaxation_matrix(d, keep_rrs_symbolic)
-        ac = self._collision_rule_with_relaxation_matrix(d, relaxation_rate_sub_expressions,
+        ac = self._collision_rule_with_relaxation_matrix(d, sm, relaxation_rate_sub_expressions,
                                                          True, conserved_quantity_equations)
         return ac
 
@@ -120,13 +136,21 @@ class MomentBasedLbMethod(AbstractLbMethod):
     def collision_matrix(self):
         pdfs_to_moments = self.moment_matrix
         d = self.relaxation_matrix
-        return pdfs_to_moments.inv() * d * pdfs_to_moments
+        sm = self._shift_matrix
+        if sm:
+            return pdfs_to_moments.inv() * sm.inv() * d * sm * pdfs_to_moments
+        else:
+            return pdfs_to_moments.inv() * d * pdfs_to_moments
 
     @property
     def inverse_collision_matrix(self):
         pdfs_to_moments = self.moment_matrix
         inverse_relaxation_matrix = self.relaxation_matrix.inv()
-        return pdfs_to_moments.inv() * inverse_relaxation_matrix * pdfs_to_moments
+        sm = self._shift_matrix
+        if sm:
+            return pdfs_to_moments.inv() * sm.inv() * inverse_relaxation_matrix * sm * pdfs_to_moments
+        else:
+            return pdfs_to_moments.inv() * inverse_relaxation_matrix * pdfs_to_moments
 
     @property
     def moment_matrix(self):
@@ -134,7 +158,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
 
     @property
     def shift_matrix(self):
-        return shift_matrix(self.moments, self.stencil)
+        return self._shift_matrix
 
     @property
     def relaxation_matrix(self):
@@ -202,13 +226,18 @@ class MomentBasedLbMethod(AbstractLbMethod):
             weights.append(value)
         return weights
 
-    def _collision_rule_with_relaxation_matrix(self, d, additional_subexpressions=(), include_force_terms=True,
+    def _collision_rule_with_relaxation_matrix(self, d, sm=None, additional_subexpressions=(), include_force_terms=True,
                                                conserved_quantity_equations=None):
         f = sp.Matrix(self.pre_collision_pdf_symbols)
         pdf_to_moment = self.moment_matrix
         m_eq = sp.Matrix(self.moment_equilibrium_values)
 
-        collision_rule = f + pdf_to_moment.inv() * d * (m_eq - pdf_to_moment * f)
+        if sm:
+            print("Using collision with shift matrix")
+            collision_rule = f + pdf_to_moment.inv() * sm.inv() * d * sm * (m_eq - pdf_to_moment * f)
+        else:
+            print("Using collision without shift matrix")
+            collision_rule = f + pdf_to_moment.inv() * d * (m_eq - pdf_to_moment * f)
         collision_eqs = [Assignment(lhs, rhs) for lhs, rhs in zip(self.post_collision_pdf_symbols, collision_rule)]
 
         if conserved_quantity_equations is None:

lbmpy/moments.py

@@ -369,9 +369,9 @@ def moment_matrix(moments, stencil, shift_velocity=None):
     return sp.Matrix(len(moments), len(stencil), generator)
 
 
-def shift_matrix(moments, stencil):
+def set_up_shift_matrix(moments, stencil):
     """
-    Returns shift matrix to shift raw moments to central moment space
+    Sets up a shift matrix to shift raw moments to central moment space
     """
     x, y, z = MOMENT_SYMBOLS
     dim = len(stencil[0])

lbmpy/simplificationfactory.py

@@ -16,22 +16,27 @@ def create_simplification_strategy(lb_method, split_inner_loop=False):
     expand = apply_to_all_assignments(sp.expand)
 
     if isinstance(lb_method, MomentBasedLbMethod):
-        if len(set(lb_method.relaxation_rates)) <= 2:
-            s.add(expand)
-            s.add(replace_second_order_velocity_products)
-            s.add(expand)
-            s.add(factor_relaxation_rates)
-            s.add(replace_density_and_velocity)
-            s.add(replace_common_quadratic_and_constant_term)
-            s.add(factor_density_after_factoring_relaxation_times)
-            s.add(subexpression_substitution_in_main_assignments)
-            if split_inner_loop:
-                s.add(create_lbm_split_groups)
-            s.add(add_subexpressions_for_divisions)
-        else:
-            s.add(subexpression_substitution_in_main_assignments)
-            if split_inner_loop:
-                s.add(create_lbm_split_groups)
+        if not lb_method.shift_matrix:
+            if len(set(lb_method.relaxation_rates)) <= 2:
+                s.add(expand)
+                s.add(replace_second_order_velocity_products)
+                s.add(expand)
+                s.add(factor_relaxation_rates)
+                s.add(replace_density_and_velocity)
+                s.add(replace_common_quadratic_and_constant_term)
+                s.add(factor_density_after_factoring_relaxation_times)
+                s.add(subexpression_substitution_in_main_assignments)
+                if split_inner_loop:
+                    s.add(create_lbm_split_groups)
+                s.add(add_subexpressions_for_divisions)
+            else:
+                s.add(subexpression_substitution_in_main_assignments)
+                if split_inner_loop:
+                    s.add(create_lbm_split_groups)
+        # else:
+        #     s.add(expand)
+        #     s.add(factor_relaxation_rates)
+        #     s.add(add_subexpressions_for_divisions)
     elif isinstance(lb_method, CumulantBasedLbMethod):
         s.add(expand)
         s.add(factor_relaxation_rates)

lbmpy_tests/test_moments.py

@@ -111,7 +111,7 @@ def test_shift_matrix():
         method = create_lb_method(stencil=stencil, method='mrt_raw', compressible=True, equilibrium_order=6)
         moments = method.moments
 
-        sm = shift_matrix(moments, stencil)
+        sm = set_up_shift_matrix(moments, stencil)
 
         m_eq = sp.Matrix(method.moment_equilibrium_values)
         m_eq_rel = sp.simplify(sm * m_eq)
@@ -119,6 +119,6 @@ def test_shift_matrix():
         m_eq_sol = sp.Matrix(method.moment_equilibrium_values)
 
         for i in range(0, len(stencil)):
-            m_eq_sol[i] = m_eq_sol[i].subs({sp.Symbol("u_0"):0, sp.Symbol("u_1"):0, sp.Symbol("u_2"):0})
+            m_eq_sol[i] = m_eq_sol[i].subs({sp.Symbol("u_0"): 0, sp.Symbol("u_1"): 0, sp.Symbol("u_2"): 0})
 
         assert m_eq_rel == m_eq_sol