Refactoring force models

lbmpy/boundaries/boundaryconditions.py

+import abc
 from warnings import warn
 
 from lbmpy.advanced_streaming.utility import AccessPdfValues, Timestep
@@ -13,7 +14,7 @@ from pystencils.stencil import offset_to_direction_string, direction_string_to_o
 import sympy as sp
 
 
-class LbBoundary:
+class LbBoundary(abc.ABC):
     """Base class that all boundaries should derive from.
 
     Args:

lbmpy/creationfunctions.py

@@ -537,6 +537,7 @@ def force_model_from_string(force_model_name, force_values):
         'edm': forcemodels.EDM,
         'schiller': forcemodels.Schiller,
         'cumulant': cumulant_force_model.CenteredCumulantForceModel,
+        'He': forcemodels.He
     }
     if force_model_name.lower() not in force_model_dict:
         raise ValueError("Unknown force model %s" % (force_model_name,))

lbmpy/forcemodels.py

@@ -7,7 +7,7 @@ Get started:
 ------------
 
 This module offers different models to introduce a body force in the lattice Boltzmann scheme.
-If you don't know which model to choose, use :class:`lbmpy.forcemodels.Schiller`.
+If you don't know which model to choose, use :class:`lbmpy.forcemodels.Guo`.
 For incompressible collision models the :class:`lbmpy.forcemodels.Buick` model can be better.
 
 
@@ -18,15 +18,16 @@ Force models add a term :math:`C_F` to the collision equation:
 
 .. math ::
 
-    f(\pmb{x} + c_q \Delta t, t + \Delta t) - f(\pmb{x},t) = \Omega(f, f^{(eq)})
+    f(\pmb{x} + c_q \Delta t, t + \Delta t) - f(\pmb{x},t) = \Omega(f, f^{(\mathrm{eq})})
                                                             + \underbrace{F_q}_{\mbox{forcing term}}
 
 The form of this term depends on the concrete force model: the first moment of this forcing term is equal
-to the acceleration :math:`\pmb{a}` for all force models.
+to the acceleration :math:`\pmb{a}` for all force models. Here :math:`\mathbf{F}` is the D dimensional force vector,
+which defines the force for each spatial dircetion.
 
 .. math ::
 
-    \sum_q \pmb{c}_q F_q = \pmb{a}
+    \sum_q \pmb{c}_q \mathbf{F} = \pmb{a}
 
 
 The second order moment is different for the forcing models - if it is zero the model is suited for
@@ -35,9 +36,9 @@ should be chosen.
 
 .. math ::
 
-    \sum_q c_{qi} c_{qj} f_q = F_i u_j + F_j u_i  \hspace{1cm} \mbox{for Guo, Luo models}
+    \sum_q c_{qi} c_{qj} f_q &= F_i u_j + F_j u_i  &\qquad \mbox{for Guo, Luo models}
     
-    \sum_q c_{qi} c_{qj} f_q = 0  \hspace{1cm} \mbox{for Simple, Buick}
+    \sum_q c_{qi} c_{qj} f_q &= 0  &\qquad \mbox{for Simple, Buick}
     
 Models with zero second order moment have:
 
@@ -57,9 +58,9 @@ For all force models the computation of the macroscopic velocity has to be adapt
     
     .. math ::
     
-        \pmb{u} = \sum_q \pmb{c}_q f_q + S_{macro}
+        \pmb{u} &= \sum_q \pmb{c}_q f_q + S_{\mathrm{macro}}
         
-        S_{macro} = \frac{\Delta t}{2} \sum_q F_q
+        S_{\mathrm{macro}} &= \frac{\Delta t}{2 \cdot \rho} \sum_q F_q
 
 
 Some models also shift the velocity entering the equilibrium distribution.
@@ -85,86 +86,151 @@ second moment nonzero   :class:`Luo`         :class:`Guo`
 =====================  ====================  =================
   
 """
-
+import abc
 import sympy as sp
 
+from pystencils.sympyextensions import scalar_product, remove_higher_order_terms
+
 from lbmpy.relaxationrates import get_bulk_relaxation_rate, get_shear_relaxation_rate
 
 
-class Simple:
+class AbstractForceModel(abc.ABC):
+    r"""
+    Abstract base class for all force models. All force models have to implement the __call__, which should return a
+    q dimensional vector added to the PDFs in the population space. If an MRT method is used, it is also possible
+    to apply the force directly in the moment space. This is done by additionally providing the function
+    moment_space_forcing. The MRT method will check if it is available and apply the force directly in the moment
+    space. For MRT methods in the central moment space the central_moment_space_forcing function can be provided,
+    which shifts the force vector to the central moment space. Applying the force in the collision space has the
+    advantage of saving FLOPs. Furthermore, it is sometimes easier to apply the correct force vector in the
+    collision space because often, the relaxation matrix has to be taken into account.
+
+    Args:
+        force: force vector of size dim which contains the force for each spatial dimension.
+    """
+
+    def __init__(self, force):
+        self._force = force
+
+        # The following booleans should indicate if a force model is has the function moment_space_forcing which
+        # transfers the forcing terms to the moment space or central_moment_space_forcing which transfers them to the
+        # central moment space
+        self.has_moment_space_forcing = hasattr(self, "moment_space_forcing")
+        self.has_central_moment_space_forcing = hasattr(self, "central_moment_space_forcing")
+
+    def __call__(self, lb_method):
+        r"""
+        This function returns a vector of size q which is added to the PDFs in the PDF space. It has to be implemented
+        by all forcing models.
+
+        Args:
+            lb_method: LB method, see lbmpy.creationfunctions.create_lb_method
+        """
+        raise NotImplementedError("Force model class has to overwrite __call__")
+
+    def macroscopic_velocity_shift(self, density):
+        r"""
+        macroscopic velocity shift by :math:`\frac{\Delta t}{2 \cdot \rho}`
+
+        Args:
+            density: Density symbol which is needed for the shift
+        """
+        return default_velocity_shift(density, self._force)
+
+    def macroscopic_momentum_density_shift(self, *_):
+        r"""
+        macroscopic momentum density shift by :math:`\frac{\Delta t}{2}`
+        """
+        return default_momentum_density_shift(self._force)
+
+    def equilibrium_velocity_shift(self, density):
+        r"""
+        Some models also shift the velocity entering the equilibrium distribution. By default the shift is zero
+
+        Args:
+            density: Density symbol which is needed for the shift
+        """
+        return [0] * len(self._force)
+
+
+class Simple(AbstractForceModel):
     r"""
     A simple force model which introduces the following additional force term in the
-    collision process :math:`\frac{w_q}{c_s^2} c_{qi} \; a_i` (often: force = rho * acceleration)
+    collision process :math:`\frac{w_q}{c_s^2} \mathbf{c_{q}} \cdot \mathbf{F}` (often: force = rho * acceleration
+    where rho is the zeroth moment to be consistent with the above definition)
     Should only be used with constant forces!
-    Shifts the macroscopic velocity by F/2, but does not change the equilibrium velocity.
+    Shifts the macroscopic velocity by :math:`\frac{\mathbf{F}}{2}`, but does not change the equilibrium velocity.
     """
 
     def __init__(self, force):
-        self._force = force
+        super(Simple, self).__init__(force)
 
-    def __call__(self, lb_method, **kwargs):
+    def __call__(self, lb_method):
         assert len(self._force) == lb_method.dim
+        cs_sq = sp.Rational(1, 3)  # squared speed of sound
 
-        def scalar_product(a, b):
-            return sum(a_i * b_i for a_i, b_i in zip(a, b))
-
-        return [3 * w_i * scalar_product(self._force, direction)
+        return [(w_i / cs_sq) * 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)
-
-    def macroscopic_momentum_density_shift(self, density):
-        return default_momentum_density_shift(self._force)
+    def central_moment_space_forcing(self, lb_method):
+        moments = (lb_method.moment_matrix * sp.Matrix(self(lb_method))).expand()
+        return lb_method.shift_matrix * moments
 
 
-class Luo:
+class Luo(AbstractForceModel):
     r"""Force model by Luo :cite:`luo1993lattice`.
 
-    Shifts the macroscopic velocity by F/2, but does not change the equilibrium velocity.
+    Shifts the macroscopic velocity by :math:`\frac{\mathbf{F}}{2}`, but does not change the equilibrium velocity.
     """
 
     def __init__(self, force):
-        self._force = force
+        super(Luo, self).__init__(force)
 
     def __call__(self, lb_method):
         u = sp.Matrix(lb_method.first_order_equilibrium_moment_symbols)
         force = sp.Matrix(self._force)
 
+        cs_sq = sp.Rational(1, 3)  # squared speed of sound
+
         result = []
         for direction, w_i in zip(lb_method.stencil, lb_method.weights):
             direction = sp.Matrix(direction)
-            result.append(3 * w_i * force.dot(direction - u + 3 * direction * direction.dot(u)))
+            first_summand = (direction - u) / cs_sq
+            second_summand = (direction * direction.dot(u)) / cs_sq ** 2
+
+            fq = w_i * force.dot(first_summand + second_summand)
+
+            result.append(fq.simplify())
         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)
-
-    def macroscopic_momentum_density_shift(self, density):
-        return default_momentum_density_shift(self._force)
+    def central_moment_space_forcing(self, lb_method):
+        moments = (lb_method.moment_matrix * sp.Matrix(self(lb_method))).expand()
+        return lb_method.shift_matrix * moments
 
 
-class Guo:
+class Guo(AbstractForceModel):
     r"""
     Force model by Guo  :cite:`guo2002discrete`
-    Adapts the calculation of the macroscopic velocity as well as the equilibrium velocity (both shifted by F/2)!
+    Adapts the calculation of the macroscopic velocity as well as the equilibrium velocity
+    (both shifted by :math:`\frac{\mathbf{F}}{2}`)!
     """
 
     def __init__(self, force):
-        self._force = force
+        super(Guo, self).__init__(force)
 
     def __call__(self, lb_method):
         luo = Luo(self._force)
 
         shear_relaxation_rate = get_shear_relaxation_rate(lb_method)
         assert len(set(lb_method.relaxation_rates)) == 1, (
-            "In population space, guo only works for SRT, use Schiller instead")
+            "In population space, guo only works for SRT. Use the function moment_space_forcing as call operator "
+            "instead")
         correction_factor = (1 - sp.Rational(1, 2) * shear_relaxation_rate)
         return [correction_factor * t for t in luo(lb_method)]
 
@@ -175,24 +241,84 @@ class Guo:
         moments = correction_factor * (lb_method.moment_matrix * sp.Matrix(luo(lb_method)))
         return moments.expand()
 
-    def macroscopic_velocity_shift(self, density):
+    def central_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 = (lb_method.moment_matrix * sp.Matrix(luo(lb_method)))
+        central_moments = correction_factor * (lb_method.shift_matrix * moments)
+        return central_moments.expand()
+
+    def equilibrium_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
-    def macroscopic_momentum_density_shift(self, density):
-        return default_momentum_density_shift(self._force)
+
+class He(AbstractForceModel):
+    r"""
+    Force model by He  :cite:`HeForce`
+    Adapts the calculation of the macroscopic velocity as well as the equilibrium velocity
+    (both shifted by :math:`\frac{\mathbf{F}}{2}`)!
+    """
+
+    def __init__(self, force):
+        super(He, self).__init__(force)
+
+    def forcing_terms(self, lb_method):
+        rho = lb_method.zeroth_order_equilibrium_moment_symbol
+        u = sp.Matrix(lb_method.first_order_equilibrium_moment_symbols)
+        force = sp.Matrix(self._force) / rho
+
+        cs_sq = sp.Rational(1, 3)  # squared speed of sound
+        eq_terms = lb_method.get_equilibrium_terms()
+
+        result = []
+        for direction, eq in zip(lb_method.stencil, eq_terms):
+            direction = sp.Matrix(direction)
+            eu_dot_f = (direction - u).dot(force)
+            result.append(eq * eu_dot_f / (rho * cs_sq))
+
+        return result
+
+    def __call__(self, lb_method):
+        fq = self.forcing_terms(lb_method)
+
+        shear_relaxation_rate = get_shear_relaxation_rate(lb_method)
+        assert len(set(lb_method.relaxation_rates)) == 1, (
+            "In population space, He only works for SRT. Use the function moment_space_forcing as call operator "
+            "instead")
+        correction_factor = (1 - sp.Rational(1, 2) * shear_relaxation_rate)
+        return [correction_factor * t for t in fq]
+
+    def moment_space_forcing(self, lb_method):
+        fq = self.forcing_terms(lb_method)
+        u = lb_method.first_order_equilibrium_moment_symbols
+        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(fq))
+        moments = moments.expand()
+        return [remove_higher_order_terms(moment, order=2, symbols=u) for moment in moments]
+
+    def central_moment_space_forcing(self, lb_method):
+        fq = self.forcing_terms(lb_method)
+        u = lb_method.first_order_equilibrium_moment_symbols
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.relaxation_matrix
+        moments = (lb_method.moment_matrix * sp.Matrix(fq))
+        central_moments = correction_factor * (lb_method.shift_matrix * moments)
+        return [remove_higher_order_terms(central_moment, order=2, symbols=u) for central_moment in central_moments]
 
     def equilibrium_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
 
-class Schiller:
+class Schiller(AbstractForceModel):
     r"""
     Force model by Schiller  :cite:`schiller2008thermal`, equation 4.67
     Equivalent to Guo but not restricted to SRT.
     """
 
     def __init__(self, force):
-        self._force = force
+        super(Schiller, self).__init__(force)
 
     def __call__(self, lb_method):
         u = sp.Matrix(lb_method.first_order_equilibrium_moment_symbols)
@@ -211,61 +337,70 @@ class Schiller:
             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)
 
-    def macroscopic_momentum_density_shift(self, density):
-        return default_momentum_density_shift(self._force)
-
-
-class Buick:
+class Buick(AbstractForceModel):
     r"""
     This force model :cite:`buick2000gravity` has a force term with zero second moment. 
     It is suited for incompressible lattice models.
     """
 
     def __init__(self, force):
-        self._force = force
+        super(Buick, self).__init__(force)
 
     def __call__(self, lb_method, **kwargs):
         simple = Simple(self._force)
 
         shear_relaxation_rate = get_shear_relaxation_rate(lb_method)
-        assert len(set(lb_method.relaxation_rates)) == 1, "Buick only works for SRT"
+        assert len(set(lb_method.relaxation_rates)) == 1, (
+            "In population space, buick only works for SRT. Use the function moment_space_forcing as call operator "
+            "instead")
         correction_factor = (1 - sp.Rational(1, 2) * shear_relaxation_rate)
         return [correction_factor * t for t in simple(lb_method)]
 
-    def macroscopic_velocity_shift(self, density):
-        return default_velocity_shift(density, self._force)
+    def moment_space_forcing(self, lb_method):
+        simple = Simple(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(simple(lb_method)))
+        return moments.expand()
 
-    def macroscopic_momentum_density_shift(self, density):
-        return default_momentum_density_shift(self._force)
+    def central_moment_space_forcing(self, lb_method):
+        simple = Simple(self._force)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.relaxation_matrix
+        moments = (lb_method.moment_matrix * sp.Matrix(simple(lb_method)))
+        central_moments = correction_factor * (lb_method.shift_matrix * moments)
+        return central_moments.expand()
 
     def equilibrium_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
 
-class EDM:
+class EDM(AbstractForceModel):
     r"""Exact differencing force model"""
 
     def __init__(self, force):
-        self._force = force
+        super(EDM, self).__init__(force)
 
-    def __call__(self, lb_method, **kwargs):
+    def __call__(self, lb_method):
         cqc = lb_method.conserved_quantity_computation
         rho = cqc.zeroth_order_moment_symbol if cqc.compressible else 1
         u = cqc.first_order_moment_symbols
 
+        equilibrium_terms = lb_method.get_equilibrium_terms()
+
         shifted_u = (u_i + f_i / rho for u_i, f_i in zip(u, self._force))
-        eq_terms = lb_method.get_equilibrium_terms()
-        shifted_eq = eq_terms.subs({u_i: su_i for u_i, su_i in zip(u, shifted_u)})
-        return shifted_eq - eq_terms
+        shifted_eq = equilibrium_terms.subs({u_i: su_i for u_i, su_i in zip(u, shifted_u)})
+        return shifted_eq - equilibrium_terms
 
-    def macroscopic_velocity_shift(self, density):
-        return default_velocity_shift(density, self._force)
+    def moment_space_forcing(self, lb_method):
+        moments = lb_method.moment_matrix * self(lb_method)
+        return moments.expand()
 
-    def macroscopic_momentum_density_shift(self, density):
-        return default_momentum_density_shift(self._force)
+    def central_moment_space_forcing(self, lb_method):
+        moments = lb_method.moment_matrix * self(lb_method)
+        central_moments = lb_method.shift_matrix * moments.expand()
+        return central_moments.expand()
 
 
 # --------------------------------  Helper functions  ------------------------------------------------------------------

lbmpy/methods/centeredcumulant/force_model.py

-from lbmpy.forcemodels import default_velocity_shift
+from lbmpy.forcemodels import AbstractForceModel
 
 
 #   =========================== Centered Cumulant Force Model ==========================================================
 
 
-class CenteredCumulantForceModel:
+class CenteredCumulantForceModel(AbstractForceModel):
     """
     A force model to be used with the centered cumulant-based LB Method.
     It only shifts the macroscopic and equilibrium velocities and does not introduce a forcing term to the
@@ -16,14 +16,9 @@ class CenteredCumulantForceModel:
     """
 
     def __init__(self, force):
-        self._force = force
         self.override_momentum_relaxation_rate = 2
 
-    def __call__(self, lb_method, **kwargs):
-        raise Exception('This force model does not provide a forcing term.')
-
-    def macroscopic_velocity_shift(self, density):
-        return default_velocity_shift(density, self._force)
+        super(CenteredCumulantForceModel, self).__init__(force)
 
-    def equilibrium_velocity_shift(self, density):
-        return default_velocity_shift(density, self._force)
+    def __call__(self, lb_method):
+        raise Exception('This force model does not provide a forcing term.')

lbmpy/methods/conservedquantitycomputation.py

@@ -140,23 +140,23 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
         if self._compressible:
             eq_coll = divide_first_order_moments_by_rho(eq_coll, dim)
 
-        eq_coll = apply_force_model_shift('equilibrium_velocity_shift', dim, eq_coll,
-                                          self._forceModel, self._compressible)
+        if self._forceModel is not None:
+            eq_coll = apply_force_model_shift(self._forceModel.equilibrium_velocity_shift,
+                                              dim, eq_coll, self._compressible)
         return eq_coll
 
     def equilibrium_input_equations_from_init_values(self, density=1, velocity=(0, 0, 0)):
         dim = len(self._stencil[0])
-        zeroth_order_moment = density
+        zeroth_order_moment = density if self._compressible else density - sp.Rational(1, 1)
         first_order_moments = velocity[:dim]
         vel_offset = [0] * dim
 
         if self._compressible:
-            if self._forceModel and hasattr(self._forceModel, 'macroscopic_velocity_shift'):
+            if self._forceModel is not None:
                 vel_offset = self._forceModel.macroscopic_velocity_shift(zeroth_order_moment)
         else:
-            if self._forceModel and hasattr(self._forceModel, 'macroscopic_velocity_shift'):
+            if self._forceModel is not None:
                 vel_offset = self._forceModel.macroscopic_velocity_shift(sp.Rational(1, 1))
-            zeroth_order_moment -= sp.Rational(1, 1)
         eqs = [Assignment(self._symbolOrder0, zeroth_order_moment)]
 
         first_order_moments = [a - b for a, b in zip(first_order_moments, vel_offset)]
@@ -176,7 +176,8 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
         else:
             ac = add_density_offset(ac)
 
-        ac = apply_force_model_shift('macroscopic_velocity_shift', dim, ac, self._forceModel, self._compressible)
+        if self._forceModel is not None:
+            ac = apply_force_model_shift(self._forceModel.macroscopic_velocity_shift, dim, ac, self._compressible)
 
         main_assignments = []
         eqs = OrderedDict([(eq.lhs, eq.rhs) for eq in ac.all_assignments])
@@ -210,8 +211,9 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
             mom_density_eq_coll = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
                                                                                   self._symbolOrder0,
                                                                                   self._symbolsOrder1)
-            mom_density_eq_coll = apply_force_model_shift('macroscopic_momentum_density_shift', dim, 
-                                                          mom_density_eq_coll, self._forceModel, self._compressible)
+            if self._forceModel is not None:
+                mom_density_eq_coll = apply_force_model_shift(self._forceModel.macroscopic_momentum_density_shift, dim,
+                                                              mom_density_eq_coll, self._compressible)
 
             for sym, val in zip(momentum_density_output_symbols, mom_density_eq_coll.main_assignments[1:]):
                 main_assignments.append(Assignment(sym, val.rhs))
@@ -287,11 +289,13 @@ def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs, symb
                 p[dim * i + j] += f * int(offset[i]) * int(offset[j])
 
     plus_terms = [set(filter_out_plus_terms(u_i).args) for u_i in u]
+
+    velo_terms = sp.symbols(f"vel:{dim}Term")
     for i in range(dim):
         rhs = plus_terms[i]
         for j in range(i):
             rhs -= plus_terms[j]
-        eq = Assignment(sp.Symbol("vel%dTerm" % (i,)), sum(rhs))
+        eq = Assignment(velo_terms[i], sum(rhs))
         subexpressions.append(eq)
 
     for subexpression in subexpressions:
@@ -334,23 +338,27 @@ def add_density_offset(assignment_collection, offset=sp.Rational(1, 1)):
     return assignment_collection.copy([new_density] + old_eqs[1:])
 
 
-def apply_force_model_shift(shift_member_name, dim, assignment_collection, force_model, compressible, reverse=False):
+def apply_force_model_shift(shift_func, dim, assignment_collection, compressible, reverse=False):
     """
     Modifies the first order moment equations in assignment collection according to the force model shift.
     It is applied if force model has a method named shift_member_name. The equations 1: dim+1 of the passed
     equation collection are assumed to be the velocity equations.
+
+    Args:
+        shift_func: shift function which is applied. See lbmpy.forcemodels.AbstractForceModel for details
+        dim: number of spatial dimensions
+        assignment_collection: assignment collection containing the conserved quantity computation
+        compressible: True if a compressible LB method is used. Otherwise the Helmholtz decomposition was applied
+                      for rho
+        reverse: If True the sign of the shift is flipped
     """
-    if force_model is not None and hasattr(force_model, shift_member_name):
-        old_eqs = assignment_collection.main_assignments
-        density = old_eqs[0].lhs if compressible else sp.Rational(1, 1)
-        old_vel_eqs = old_eqs[1:dim + 1]
-        shift_func = getattr(force_model, shift_member_name)
-        vel_offsets = shift_func(density)
-        if reverse:
-            vel_offsets = [-v for v in vel_offsets]
-        shifted_velocity_eqs = [Assignment(old_eq.lhs, old_eq.rhs + offset)
-                                for old_eq, offset in zip(old_vel_eqs, vel_offsets)]
-        new_eqs = [old_eqs[0]] + shifted_velocity_eqs + old_eqs[dim + 1:]
-        return assignment_collection.copy(new_eqs)
-    else:
-        return assignment_collection
+    old_eqs = assignment_collection.main_assignments
+    density = old_eqs[0].lhs if compressible else sp.Rational(1, 1)
+    old_vel_eqs = old_eqs[1:dim + 1]
+    vel_offsets = shift_func(density)
+    if reverse:
+        vel_offsets = [-v for v in vel_offsets]
+    shifted_velocity_eqs = [Assignment(old_eq.lhs, old_eq.rhs + offset)
+                            for old_eq, offset in zip(old_vel_eqs, vel_offsets)]
+    new_eqs = [old_eqs[0]] + shifted_velocity_eqs + old_eqs[dim + 1:]
+    return assignment_collection.copy(new_eqs)

lbmpy/methods/creationfunctions.py

@@ -310,7 +310,16 @@ def create_central_moment(stencil, relaxation_rates, maxwellian_moments=False, *
     """
     if isinstance(stencil, str):
         stencil = get_stencil(stencil)
+
+    dim = len(stencil[0])
+
     moments = get_default_moment_set_for_stencil(stencil)
+    # x, y, z = MOMENT_SYMBOLS
+    # if dim == 2:
+    #     diagonal_viscous_moments = [x ** 2 + y ** 2, x ** 2 - y ** 2]
+    # for i, d in enumerate(MOMENT_SYMBOLS[:dim]):
+    #     moments[moments.index(d ** 2)] = diagonal_viscous_moments[i]
+
     sorted_moments = sort_moments_into_groups_of_same_order(moments)
     if len(relaxation_rates) == len(sorted_moments) - 2:
         relaxation_rates = [0, 0, *relaxation_rates]

lbmpy/methods/momentbased/momentbasedmethod.py

@@ -30,6 +30,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
                                             This determines how conserved quantities are computed, and defines
                                             the symbols used in the equilibrium moments like e.g. density and velocity
             force_model: force model instance, or None if no forcing terms are required
+            moment_transform_class: transformation class to transform PDFs to the moment space.