Merge branch 'improvedMomentMethod' into 'master'

Improved Moment-Based Method

See merge request pycodegen/lbmpy!91

doc/sphinx/api.rst

@@ -9,6 +9,7 @@ API Reference
    methodcreation.rst
    stencils.rst
    methods.rst
+   moment_transforms.rst
    maxwellian_equilibrium.rst
    continuous_distribution_measures.rst
    moments.rst

doc/sphinx/moment_transforms.rst

+*******************************************
+Moment Transforms (lbmpy.moment_transforms)
+*******************************************
+
+Abstract Base Class
+===================
+
+.. autoclass:: lbmpy.moment_transforms.AbstractMomentTransform
+    :members:
+
+
+Moment Space Transforms
+===========================
+
+By Matrix-Vector Multiplication
+-------------------------------
+
+.. autoclass:: lbmpy.moment_transforms.PdfsToMomentsByMatrixTransform
+    :members:
+
+By Chimera-Transform
+--------------------
+
+.. autoclass:: lbmpy.moment_transforms.PdfsToMomentsByChimeraTransform
+    :members:
+
+
+Central Moment Space Transforms
+===============================
+
+.. autoclass:: lbmpy.moment_transforms.PdfsToCentralMomentsByMatrix
+    :members:
+
+.. autoclass:: lbmpy.moment_transforms.FastCentralMomentTransform
+    :members:
+
+.. autoclass:: lbmpy.moment_transforms.PdfsToCentralMomentsByShiftMatrix
+    :members:
+
+Cumulant Space Transforms
+=========================
+
+.. autoclass:: lbmpy.methods.centeredcumulant.cumulant_transform.CentralMomentsToCumulantsByGeneratingFunc
+    :members:
+

lbmpy/creationfunctions.py

@@ -111,7 +111,7 @@ Simplifications / Transformations:
   is built into the kernel. This parameters specifies if the domain is periodic in (x,y,z) direction. Even if the
   periodicity is built into the kernel, the fields have one ghost layer to be consistent with other functions.
 - ``pre_simplification=True``: Simplifications applied during the derivaton of the collision rule for cumulant LBMs
-  For details see :mod:`lbmpy.methods.momentbased.moment_transforms`.
+  For details see :mod:`lbmpy.moment_transforms`.
     
 
 Field size information:
@@ -202,7 +202,7 @@ from lbmpy.methods import (create_mrt_orthogonal, create_mrt_raw, create_central
                            create_srt, create_trt, create_trt_kbc)
 from lbmpy.methods.abstractlbmethod import RelaxationInfo
 from lbmpy.methods.centeredcumulant import CenteredCumulantBasedLbMethod
-from lbmpy.methods.momentbased.moment_transforms import PdfsToCentralMomentsByShiftMatrix
+from lbmpy.moment_transforms import PdfsToMomentsByChimeraTransform, PdfsToCentralMomentsByShiftMatrix
 from lbmpy.methods.centeredcumulant.cumulant_transform import CentralMomentsToCumulantsByGeneratingFunc
 from lbmpy.methods.creationfunctions import (
     create_with_monomial_cumulants, create_with_polynomial_cumulants, create_with_default_polynomial_cumulants)
@@ -220,7 +220,7 @@ from pystencils import Assignment, AssignmentCollection, create_kernel
 from pystencils.cache import disk_cache_no_fallback
 from pystencils.data_types import collate_types
 from pystencils.field import Field, get_layout_of_array
-from pystencils.simp import sympy_cse
+from pystencils.simp import sympy_cse, SimplificationStrategy
 from pystencils.stencil import have_same_entries
 
 
@@ -364,8 +364,10 @@ def create_lb_collision_rule(lb_method=None, optimization=None, **kwargs):
 
     if opt_params['simplification'] == 'auto':
         simplification = create_simplification_strategy(lb_method, split_inner_loop=split_inner_loop)
-    else:
+    elif callable(opt_params['simplification']):
         simplification = opt_params['simplification']
+    else:
+        simplification = SimplificationStrategy()
     collision_rule = simplification(collision_rule)
 
     if params['fluctuating']:
@@ -418,7 +420,9 @@ def create_lb_method(**params):
         'equilibrium_order': params['equilibrium_order'],
         'force_model': force_model,
         'maxwellian_moments': params['maxwellian_moments'],
-        'c_s_sq': params['c_s_sq']
+        'c_s_sq': params['c_s_sq'],
+        'moment_transform_class': params['moment_transform_class'],
+        'central_moment_transform_class': params['central_moment_transform_class'],
     }
 
     cumulant_params = {
@@ -582,6 +586,8 @@ def update_with_default_parameters(params, opt_params=None, fail_on_unknown_para
         'initial_velocity': None,
 
         'galilean_correction': False,  # only available for D3Q27 cumulant methods
+
+        'moment_transform_class': PdfsToMomentsByChimeraTransform,
         'central_moment_transform_class': PdfsToCentralMomentsByShiftMatrix,
         'cumulant_transform_class': CentralMomentsToCumulantsByGeneratingFunc,
 
@@ -644,6 +650,18 @@ def update_with_default_parameters(params, opt_params=None, fail_on_unknown_para
 
             del params['relaxation_rate']
 
+    #   By default, do not derive moment equations for SRT and TRT methods
+    if 'method' not in params or params['method'][:3].lower() in ('srt', 'trt'):
+        if 'moment_transform_class' not in params:
+            params['moment_transform_class'] = None
+
+    #   Workaround until entropic method supports relaxation in subexpressions
+    #   and the problem with RNGs in the assignment collection has been solved
+    if (('entropic' in params and params['entropic'])
+            or ('fluctuating' in params and params['fluctuating'])):
+        if 'moment_transform_class' not in params:
+            params['moment_transform_class'] = None
+
     #   for backwards compatibility
     if 'kernel_type' in params:
         kernel_type_to_streaming_pattern = {

lbmpy/forcemodels.py

@@ -98,6 +98,7 @@ class Simple:
     Should only be used with constant forces!
     Shifts the macroscopic velocity by F/2, but does not change the equilibrium velocity.
     """
+
     def __init__(self, force):
         self._force = force
 
@@ -110,6 +111,9 @@ class Simple:
         return [3 * w_i * scalar_product(self._force, direction)
                 for direction, w_i in zip(lb_method.stencil, lb_method.weights)]
 
+    def moment_space_forcing(self, lb_method):
+        return (lb_method.moment_matrix * sp.Matrix(self(lb_method))).expand()
+
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
@@ -122,6 +126,7 @@ class Luo:
 
     Shifts the macroscopic velocity by F/2, but does not change the equilibrium velocity.
     """
+
     def __init__(self, force):
         self._force = force
 
@@ -135,6 +140,9 @@ class Luo:
             result.append(3 * w_i * force.dot(direction - u + 3 * direction * direction.dot(u)))
         return result
 
+    def moment_space_forcing(self, lb_method):
+        return (lb_method.moment_matrix * sp.Matrix(self(lb_method))).expand()
+
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
@@ -147,6 +155,7 @@ class Guo:
     Force model by Guo  :cite:`guo2002discrete`
     Adapts the calculation of the macroscopic velocity as well as the equilibrium velocity (both shifted by F/2)!
     """
+
     def __init__(self, force):
         self._force = force
 
@@ -154,10 +163,18 @@ class Guo:
         luo = Luo(self._force)
 
         shear_relaxation_rate = get_shear_relaxation_rate(lb_method)
-        assert len(set(lb_method.relaxation_rates)) == 1, "Guo only works for SRT, use Schiller instead"
+        assert len(set(lb_method.relaxation_rates)) == 1, (
+            "In population space, guo only works for SRT, use Schiller instead")
         correction_factor = (1 - sp.Rational(1, 2) * shear_relaxation_rate)
         return [correction_factor * t for t in luo(lb_method)]
 
+    def moment_space_forcing(self, lb_method):
+        luo = Luo(self._force)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.relaxation_matrix
+        moments = correction_factor * (lb_method.moment_matrix * sp.Matrix(luo(lb_method)))
+        return moments.expand()
+
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
@@ -173,13 +190,14 @@ class Schiller:
     Force model by Schiller  :cite:`schiller2008thermal`, equation 4.67
     Equivalent to Guo but not restricted to SRT.
     """
+
     def __init__(self, force):
         self._force = force
 
     def __call__(self, lb_method):
         u = sp.Matrix(lb_method.first_order_equilibrium_moment_symbols)
         force = sp.Matrix(self._force)
-        
+
         uf = u.dot(force) * sp.eye(len(force))
         omega = get_shear_relaxation_rate(lb_method)
         omega_bulk = get_bulk_relaxation_rate(lb_method)
@@ -192,7 +210,7 @@ class Schiller:
             tr = sp.trace(G * (direction * direction.transpose() - sp.Rational(1, 3) * sp.eye(len(force))))
             result.append(3 * w_i * (force.dot(direction) + sp.Rational(3, 2) * tr))
         return result
-    
+
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 

lbmpy/methods/centeredcumulant/centered_cumulants.py

@@ -3,11 +3,7 @@ import sympy as sp
 
 from pystencils.stencil import have_same_entries
 
-from lbmpy.moments import MOMENT_SYMBOLS, moment_sort_key, exponent_to_polynomial_representation
-
-
-def exponent_tuple_sort_key(x):
-    return moment_sort_key(exponent_to_polynomial_representation(x))
+from lbmpy.moments import MOMENT_SYMBOLS
 
 
 def get_default_polynomial_cumulants_for_stencil(stencil):

lbmpy/methods/centeredcumulant/centeredcumulantmethod.py

@@ -14,18 +14,17 @@ from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantity
 from lbmpy.moments import (
     moments_up_to_order, get_order,
     monomial_to_polynomial_transformation_matrix,
-    moment_sort_key, exponent_to_polynomial_representation, extract_monomials, MOMENT_SYMBOLS,
+    moment_sort_key, exponent_tuple_sort_key,
+    exponent_to_polynomial_representation, extract_monomials, MOMENT_SYMBOLS,
     statistical_quantity_symbol)
 
 #   Local Imports
 
-from lbmpy.methods.centeredcumulant.centered_cumulants import exponent_tuple_sort_key
-
-from lbmpy.methods.centeredcumulant.cumulant_transform import (
+from .cumulant_transform import (
     PRE_COLLISION_CUMULANT, POST_COLLISION_CUMULANT,
     CentralMomentsToCumulantsByGeneratingFunc)
 
-from lbmpy.methods.momentbased.moment_transforms import (
+from lbmpy.moment_transforms import (
     PRE_COLLISION_CENTRAL_MOMENT, POST_COLLISION_CENTRAL_MOMENT,
     PdfsToCentralMomentsByShiftMatrix)
 
@@ -119,7 +118,7 @@ class CenteredCumulantBasedLbMethod(AbstractLbMethod):
     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
-    is governed by subclasses of :class:`lbmpy.methods.momentbased.moment_transforms.AbstractMomentTransform`
+    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.
 

lbmpy/methods/centeredcumulant/cumulant_transform.py

@@ -5,12 +5,13 @@ from pystencils import Assignment, AssignmentCollection
 from pystencils.simp import SimplificationStrategy, add_subexpressions_for_divisions
 from pystencils.simp.assignment_collection import SymbolGen
 
-from lbmpy.moments import moments_up_to_order, get_order, statistical_quantity_symbol
+from lbmpy.moments import (
+    moments_up_to_order, get_order, statistical_quantity_symbol, exponent_tuple_sort_key
+)
 
 from itertools import product, chain
 
-from lbmpy.methods.centeredcumulant.centered_cumulants import exponent_tuple_sort_key
-from lbmpy.methods.momentbased.moment_transforms import (
+from lbmpy.moment_transforms import (
     AbstractMomentTransform, PRE_COLLISION_CENTRAL_MOMENT, POST_COLLISION_CENTRAL_MOMENT
 )
 
@@ -43,9 +44,10 @@ def count_derivatives(derivative):
 
 class CentralMomentsToCumulantsByGeneratingFunc(AbstractMomentTransform):
 
-    def __init__(self, stencil, cumulants, equilibrium_density, equilibrium_velocity, **kwargs):
+    def __init__(self, stencil, cumulant_exponents, equilibrium_density, equilibrium_velocity, **kwargs):
         super(CentralMomentsToCumulantsByGeneratingFunc, self).__init__(
-            stencil, cumulants, equilibrium_density, equilibrium_velocity, **kwargs)
+            stencil, equilibrium_density, equilibrium_velocity,  
+            moment_exponents=cumulant_exponents, **kwargs)
 
         self.cumulant_exponents = self.moment_exponents
         self.central_moment_exponents = self.compute_required_central_moments()

lbmpy/methods/centeredcumulant/simplification.py

-import sympy as sp
-from pystencils.sympyextensions import is_constant
-
-#   Subexpression Insertion
-
-
-def insert_subexpressions(ac, selection_callback, skip=set()):
-    i = 0
-    while i < len(ac.subexpressions):
-        exp = ac.subexpressions[i]
-        if exp.lhs not in skip and selection_callback(exp):
-            ac = ac.new_with_inserted_subexpression(exp.lhs)
-        else:
-            i += 1
-
-    return ac
-
-
-def insert_aliases(ac, **kwargs):
-    return insert_subexpressions(ac, lambda x: isinstance(x.rhs, sp.Symbol), **kwargs)
-
-
-def insert_zeros(ac, **kwargs):
-    zero = sp.Integer(0)
-    return insert_subexpressions(ac, lambda x: x.rhs == zero, **kwargs)
-
-
-def insert_constants(ac, **kwargs):
-    return insert_subexpressions(ac, lambda x: is_constant(x.rhs), **kwargs)
-
-
-def insert_symbol_times_minus_one(ac, **kwargs):
-    def callback(exp):
-        rhs = exp.rhs
-        minus_one = sp.Integer(-1)
-        return isinstance(rhs, sp.Mul) and \
-            len(rhs.args) == 2 and \
-            (rhs.args[0] == minus_one or rhs.args[1] == minus_one)
-    return insert_subexpressions(ac, callback, **kwargs)
-
-
-def insert_constant_multiples(ac, **kwargs):
-    def callback(exp):
-        rhs = exp.rhs
-        return isinstance(rhs, sp.Mul) and \
-            len(rhs.args) == 2 and \
-            (is_constant(rhs.args[0]) or is_constant(rhs.args[1]))
-    return insert_subexpressions(ac, callback, **kwargs)
-
-
-def insert_constant_additions(ac, **kwargs):
-    def callback(exp):
-        rhs = exp.rhs
-        return isinstance(rhs, sp.Add) and \
-            len(rhs.args) == 2 and \
-            (is_constant(rhs.args[0]) or is_constant(rhs.args[1]))
-    return insert_subexpressions(ac, callback, **kwargs)
-
-
-def insert_squares(ac, **kwargs):
-    two = sp.Integer(2)
-
-    def callback(exp):
-        rhs = exp.rhs
-        return isinstance(rhs, sp.Pow) and rhs.args[1] == two
-    return insert_subexpressions(ac, callback, **kwargs)
-
-
-def bind_symbols_to_skip(insertion_function, skip):
-    return lambda ac: insertion_function(ac, skip=skip)

lbmpy/methods/creationfunctions.py

@@ -19,7 +19,7 @@ from lbmpy.methods.conservedquantitycomputation import DensityVelocityComputatio
 
 from lbmpy.methods.momentbased.momentbasedmethod import MomentBasedLbMethod
 from lbmpy.methods.momentbased.centralmomentbasedmethod import CentralMomentBasedLbMethod
-from lbmpy.methods.momentbased.moment_transforms import FastCentralMomentTransform
+from lbmpy.moment_transforms import PdfsToCentralMomentsByShiftMatrix, PdfsToMomentsByChimeraTransform
 
 from lbmpy.moments import (
     MOMENT_SYMBOLS, discrete_moment, exponents_to_polynomial_representations,
@@ -36,7 +36,8 @@ from pystencils.sympyextensions import common_denominator
 def create_with_discrete_maxwellian_eq_moments(stencil, moment_to_relaxation_rate_dict, compressible=False,
                                                force_model=None, equilibrium_order=2, c_s_sq=sp.Rational(1, 3),
                                                central_moment_space=False,
-                                               central_moment_transform_class=FastCentralMomentTransform):
+                                               moment_transform_class=None,
+                                               central_moment_transform_class=PdfsToCentralMomentsByShiftMatrix):
     r"""Creates a moment-based LBM by taking a list of moments with corresponding relaxation rate.
 
     These moments are relaxed against the moments of the discrete Maxwellian distribution.
@@ -53,13 +54,15 @@ def create_with_discrete_maxwellian_eq_moments(stencil, moment_to_relaxation_rat
         force_model: force model instance, or None if no external forces
         equilibrium_order: approximation order of macroscopic velocity :math:`\mathbf{u}` in the equilibrium
         c_s_sq: Speed of sound squared
-        central_moment_space: If set to True and instance of
-        `lbmpy.methods.momentbased.centralmomentbasedmethod.CentralMomentBasedLbMethod` is returned.
-        Thus the collision will be performed in the central moment space.
-        central_moment_transform_class: class to transform PDFs to the central moment space.
+        central_moment_space: If set to True, an instance of 
+                              :class:`lbmpy.methods.momentbased.CentralMomentBasedLbMethod` is returned, 
+                              and the the collision will be performed in the central moment space.
+        moment_transform_class: Class implementing the transform from populations to moment space.
+        central_moment_transform_class: Class implementing the transform from populations to central moment space.
 
     Returns:
-        `lbmpy.methods.momentbased.MomentBasedLbMethod` instance
+        Instance of either :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` or 
+        :class:`lbmpy.methods.momentbased.CentralMomentBasedLbMethod` 
     """
     if isinstance(stencil, str):
         stencil = get_stencil(stencil)
@@ -84,13 +87,14 @@ def create_with_discrete_maxwellian_eq_moments(stencil, moment_to_relaxation_rat
         return CentralMomentBasedLbMethod(stencil, rr_dict, density_velocity_computation,
                                           force_model, central_moment_transform_class)
     else:
-        return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model)
+        return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model, moment_transform_class)
 
 
 def create_with_continuous_maxwellian_eq_moments(stencil, moment_to_relaxation_rate_dict, compressible=False,
                                                  force_model=None, equilibrium_order=2, c_s_sq=sp.Rational(1, 3),
                                                  central_moment_space=False,
-                                                 central_moment_transform_class=FastCentralMomentTransform):
+                                                 moment_transform_class=None,
+                                                 central_moment_transform_class=PdfsToCentralMomentsByShiftMatrix):
     r"""
     Creates a moment-based LBM by taking a list of moments with corresponding relaxation rate. These moments are
     relaxed against the moments of the continuous Maxwellian distribution.
@@ -109,13 +113,15 @@ def create_with_continuous_maxwellian_eq_moments(stencil, moment_to_relaxation_r
         force_model: force model instance, or None if no external forces
         equilibrium_order: approximation order of macroscopic velocity :math:`\mathbf{u}` in the equilibrium
         c_s_sq: Speed of sound squared
-        central_moment_space: If set to True and instance of
-        `lbmpy.methods.momentbased.centralmomentbasedmethod.CentralMomentBasedLbMethod` is returned.
-        Thus the collision will be performend in the central moment space.
-        central_moment_transform_class: class to transform PDFs to the central moment space.
+        central_moment_space: If set to True, an instance of 
+                              :class:`lbmpy.methods.momentbased.CentralMomentBasedLbMethod` is returned, 
+                              and the the collision will be performed in the central moment space.
+        moment_transform_class: Class implementing the transform from populations to moment space.
+        central_moment_transform_class: Class implementing the transform from populations to central moment space.
 
     Returns:
-        `lbmpy.methods.momentbased.MomentBasedLbMethod` instance
+        Instance of either :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` or 
+        :class:`lbmpy.methods.momentbased.CentralMomentBasedLbMethod` 
     """
     if isinstance(stencil, str):
         stencil = get_stencil(stencil)
@@ -148,11 +154,11 @@ def create_with_continuous_maxwellian_eq_moments(stencil, moment_to_relaxation_r
         return CentralMomentBasedLbMethod(stencil, rr_dict, density_velocity_computation,
                                           force_model, central_moment_transform_class)
     else:
-        return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model)
+        return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model, moment_transform_class)
 
 
 def create_generic_mrt(stencil, moment_eq_value_relaxation_rate_tuples, compressible=False,
-                       force_model=None):
+                       force_model=None, moment_transform_class=PdfsToMomentsByChimeraTransform):
     r"""
     Creates a generic moment-based LB method.
 
@@ -168,11 +174,12 @@ def create_generic_mrt(stencil, moment_eq_value_relaxation_rate_tuples, compress
     for moment, eq_value, rr in moment_eq_value_relaxation_rate_tuples:
         moment = sp.sympify(moment)
         rr_dict[moment] = RelaxationInfo(eq_value, rr)
-    return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model)
+    return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model, moment_transform_class)
 
 
 def create_from_equilibrium(stencil, equilibrium, moment_to_relaxation_rate_dict,
-                            compressible=False, force_model=None):
+                            compressible=False, force_model=None,
+                            moment_transform_class=PdfsToMomentsByChimeraTransform):
     r"""
     Creates a moment-based LB method using a given equilibrium distribution function
 
@@ -196,7 +203,7 @@ def create_from_equilibrium(stencil, equilibrium, moment_to_relaxation_rate_dict
 
     rr_dict = OrderedDict([(mom, RelaxationInfo(discrete_moment(equilibrium, mom, stencil).expand(), rr))
                            for mom, rr in zip(mom_to_rr_dict.keys(), mom_to_rr_dict.values())])
-    return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model)
+    return MomentBasedLbMethod(stencil, rr_dict, density_velocity_computation, force_model, moment_transform_class)
 
 
 # ------------------------------------ SRT / TRT/ MRT Creators ---------------------------------------------------------
@@ -554,7 +561,7 @@ def mrt_orthogonal_modes_literature(stencil, is_weighted):
 def create_centered_cumulant_model(stencil, cumulant_to_rr_dict, force_model=None,
                                    equilibrium_order=None, c_s_sq=sp.Rational(1, 3),
                                    galilean_correction=False,
-                                   central_moment_transform_class=FastCentralMomentTransform,
+                                   central_moment_transform_class=PdfsToCentralMomentsByShiftMatrix,
                                    cumulant_transform_class=CentralMomentsToCumulantsByGeneratingFunc):
     r"""Creates a cumulant lattice Boltzmann model.
 
@@ -570,7 +577,7 @@ def create_centered_cumulant_model(stencil, cumulant_to_rr_dict, force_model=Non
         galilean_correction: special correction for D3Q27 cumulant collisions. See Appendix H in
                              :cite:`geier2015`. Implemented in :mod:`lbmpy.methods.centeredcumulant.galilean_correction`
         central_moment_transform_class: Class which defines the transformation to the central moment space
-                                        (see :mod:`lbmpy.methods.momentbased.moment_transforms`)
+                                        (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`)
 

lbmpy/methods/momentbased/__init__.py

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

lbmpy/methods/momentbased/centralmomentbasedmethod.py

@@ -5,8 +5,8 @@ from pystencils import Assignment, AssignmentCollection
 
 from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
 from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
-from lbmpy.methods.momentbased.moment_transforms import (FastCentralMomentTransform,
-                                                         PRE_COLLISION_CENTRAL_MOMENT, POST_COLLISION_CENTRAL_MOMENT)
+from lbmpy.moment_transforms import (FastCentralMomentTransform,
+                                     PRE_COLLISION_CENTRAL_MOMENT, POST_COLLISION_CENTRAL_MOMENT)
 
 from lbmpy.moments import (polynomial_to_exponent_representation, MOMENT_SYMBOLS, moment_matrix, set_up_shift_matrix,
                            statistical_quantity_symbol)

lbmpy/methods/momentbased/momentbasedmethod.py

@@ -6,12 +6,15 @@ 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
-from pystencils import Assignment
 from pystencils.sympyextensions import subs_additive
+from pystencils import Assignment, AssignmentCollection
+
+from lbmpy.moment_transforms import PdfsToMomentsByChimeraTransform
 
 
 class MomentBasedLbMethod(AbstractLbMethod):
-    def __init__(self, stencil, moment_to_relaxation_info_dict, conserved_quantity_computation=None, force_model=None):
+    def __init__(self, stencil, moment_to_relaxation_info_dict, conserved_quantity_computation=None, force_model=None,
+                 moment_transform_class=PdfsToMomentsByChimeraTransform):
         """
         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
@@ -35,11 +38,17 @@ class MomentBasedLbMethod(AbstractLbMethod):
         self._momentToRelaxationInfoDict = OrderedDict(moment_to_relaxation_info_dict.items())
         self._conservedQuantityComputation = conserved_quantity_computation
         self._weights = None
+        self._moment_transform_class = moment_transform_class
 
     @property
     def force_model(self):
         return self._forceModel
 
+    @property
+    def moment_space_collision(self):
+        """Returns whether collision is derived in terms of moments or in terms of populations only."""
+        return (self._moment_transform_class is not None)
+
     @property
     def relaxation_info_dict(self):
         return self._momentToRelaxationInfoDict
@@ -78,11 +87,21 @@ class MomentBasedLbMethod(AbstractLbMethod):
         assert len(weights) == len(self.stencil)
         self._weights = weights
 
-    def get_equilibrium(self, conserved_quantity_equations=None, include_force_terms=False):
+    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))
-        return self._collision_rule_with_relaxation_matrix(relaxation_matrix,
-                                                           conserved_quantity_equations=conserved_quantity_equations,
-                                                           include_force_terms=include_force_terms)
+        ac = self._collision_rule_with_relaxation_matrix(relaxation_matrix,
+                                                         conserved_quantity_equations=conserved_quantity_equations,
+                                                         include_force_terms=include_force_terms,
+                                                         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)
+            else:
+                return ac.new_without_subexpressions()
+        else:
+            return ac
 
     def get_equilibrium_terms(self):
         equilibrium = self.get_equilibrium()
@@ -90,9 +109,10 @@ class MomentBasedLbMethod(AbstractLbMethod):
 
     def get_collision_rule(self, conserved_quantity_equations=None, pre_simplification=True):
         d = self.relaxation_matrix
-        relaxation_rate_sub_expressions, d = self._generate_relaxation_matrix(d, pre_simplification)
+        relaxation_rate_sub_expressions, d = self._generate_relaxation_matrix(d, True)
         ac = self._collision_rule_with_relaxation_matrix(d, relaxation_rate_sub_expressions,
-                                                         True, conserved_quantity_equations)
+                                                         True, conserved_quantity_equations,
+                                                         pre_simplification=pre_simplification)
         return ac
 
     def set_zeroth_moment_relaxation_rate(self, relaxation_rate):
@@ -198,36 +218,89 @@ class MomentBasedLbMethod(AbstractLbMethod):
             weights.append(value)
         return weights
 
+    def _bound_symbols_cqc(self, conserved_quantity_equations=None):
+        f = self.pre_collision_pdf_symbols
+        cqe = conserved_quantity_equations
+
+        if cqe is None:
+            cqe = self._conservedQuantityComputation.equilibrium_input_equations_from_pdfs(f)
+
+        return cqe.bound_symbols
+