new lbm module: momentbased methods & new conserved quantity computation

forcemodels.py

@@ -16,11 +16,10 @@ class Simple:
     def __init__(self, force):
         self._force = force
 
-    def __call__(self, abstractLbmMethod, **kwargs):
-        dim = len(stencil[0])
-        assert len(self._force) == dim
+    def __call__(self, lbmMethod, **kwargs):
+        assert len(self._force) == lbmMethod.dim
         return [3 * w_i * sum([d_i * f_i for d_i, f_i in zip(direction, self._force)])
-                for direction, w_i in zip(stencil, weights)]
+                for direction, w_i in zip(lbmMethod.stencil, lbmMethod.weights)]
 
 
 class Luo:
@@ -36,18 +35,18 @@ class Luo:
     def __init__(self, force):
         self._force = force
 
-    def __call__(self, abstractLbmMethod, firstOrderMoments):
-        u = firstOrderMoments
+    def __call__(self, lbmMethod):
+        u = sp.Matrix(lbmMethod.firstOrderEquilibriumMomentSymbols)
         force = sp.Matrix(self._force)
 
         result = []
-        for direction, w_i in zip(stencil, weights):
+        for direction, w_i in zip(lbmMethod.stencil, lbmMethod.weights):
             direction = sp.Matrix(direction)
             result.append(3 * w_i * force.dot(direction - u + 3 * direction * direction.dot(u)))
         return result
 
-    def macroscopicVelocity(self, vel, density):
-        return defaultVelocityShift(vel, density, self._force)
+    def macroscopicVelocityShift(self, density):
+        return defaultVelocityShift(density, self._force)
 
 
 class Guo:
@@ -60,25 +59,26 @@ class Guo:
 
     Adapts the calculation of the macroscopic velocity as well as the equilibrium velocity (both shifted by F/2)!
     """
-    def __init__(self, force, viscosityRelaxationRate):
+    def __init__(self, force):
         self._force = force
-        self._viscosityRelaxationRate = viscosityRelaxationRate
 
-    def __call__(self, abstractLbmMethod):
+    def __call__(self, lbmMethod):
         luo = Luo(self._force)
-        correctionFactor = (1 - sp.Rational(1, 2) * self._viscosityRelaxationRate)
-        return [correctionFactor * t for t in luo(latticeModel)]
+        u = lbmMethod.firstOrderEquilibriumMomentSymbols
+        shearRelaxationRate = lbmMethod.getShearRelaxationRate()
+        correctionFactor = (1 - sp.Rational(1, 2) * shearRelaxationRate)
+        return [correctionFactor * t for t in luo(lbmMethod)]
 
-    def macroscopicVelocity(self, vel, density):
-        return defaultVelocityShift(vel, density, self._force)
+    def macroscopicVelocityShift(self, density):
+        return defaultVelocityShift(density, self._force)
 
-    def equilibriumVelocity(self, vel, density):
-        return defaultVelocityShift(vel, density, self._force)
+    def equilibriumVelocityShift(self, density):
+        return defaultVelocityShift(density, self._force)
 
 
 # --------------------------------  Helper functions  ------------------------------------------------------------------
 
 
-def defaultVelocityShift(velocity, density, force):
-    return [v_i + f_i / (2 * density) for v_i, f_i in zip(velocity, force)]
+def defaultVelocityShift(density, force):
+    return [f_i / (2 * density) for f_i in force]
 

methods/__init__.py

+from lbmpy.methods.abstractlbmmethod import AbstractLbmMethod
+from lbmpy.methods.momentbased import MomentBasedLbmMethod, RelaxationInfo
+from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation, DensityVelocityComputation

methods/abstractlbmmethod.py

@@ -3,6 +3,9 @@ import sympy as sp
 
 
 class AbstractLbmMethod(metaclass=abc.ABCMeta):
+    """
+    Abstract base class for all LBM methods
+    """
 
     def __init__(self, stencil):
         self._stencil = stencil
@@ -29,29 +32,20 @@ class AbstractLbmMethod(metaclass=abc.ABCMeta):
     # ------------------------- Abstract Methods & Properties ----------------------------------------------------------
 
     @abc.abstractproperty
-    def availableMacroscopicQuantities(self):
-        """Returns a dict of string to symbol(s), where each entry is a computable macroscopic quantity"""
-        pass
+    def conservedQuantityComputation(self):
+        """Returns an instance of class :class:`lbmpy.methods.AbstractConservedQuantityComputation`"""
 
     @abc.abstractproperty
-    def conservedQuantitiesSymbols(self):
-        """Returns symbols representing conserved quantities (subset of macroscopic quantities)"""
-        return
-
-    @abc.abstractmethod
-    def getMacroscopicQuantitiesEquations(self, macroscopicQuantities):
-        """Returns equation collection defining conserved quantities. The passed symbols have to be keys in
-        of the dict returned by :func:`availableMacroscopicQuantities`"""
-        pass
+    def weights(self):
+        """Returns a sequence of weights, one for each lattice direction"""
 
     @abc.abstractmethod
     def getEquilibrium(self):
         """Returns equation collection, to compute equilibrium values.
         The equations have the post collision symbols as left hand sides and are
         functions of the conserved quantities"""
-        return
 
     @abc.abstractmethod
     def getCollisionRule(self):
         """Returns an equation collection defining the collision operator."""
-        return
+

methods/defaultMacroscopicValueEquations.py → methods/conservedquantitycomputation.py

+import abc
 import sympy as sp
 from pystencils.equationcollection import EquationCollection
 
 
+class AbstractConservedQuantityComputation(metaclass=abc.ABCMeta):
+    """
+    This class defines how conserved quantities are computed as functions of the pdfs.
+    Conserved quantities are used:
+        - for output
+        - as input to the equilibrium in the collision step
+
+    Depending on the method they might also be computed slightly different, e.g. due to a force model.
+
+    An additional method describes how to get the conserved quantities for the equilibrium for initialization.
+    In most cases the inputs can be used directly, but for some methods they have to be altered slightly.
+    For example in zero centered hydrodynamic schemes with force model, the density has
+    to be decreased by one, and the given velocity has to be shifted dependent on the force.
+    """
+
+    @abc.abstractproperty
+    def conservedQuantities(self):
+        """
+        Dict, mapping names (symbol) to dimensionality (int)
+        For example: {'density' : 1, 'velocity' : 3}
+        The naming strings can be used in :func:`outputEquationsFromPdfs`
+        and :func:`equilibriumInputEquationsFromInitValues`
+        """
+
+    def definedSymbols(self, order='all'):
+        """
+        Returns a dict, mapping names of conserved quantities to their symbols
+        """
+
+    @abc.abstractproperty
+    def defaultValues(self):
+        """
+        Returns a dict of symbol to default value, where "default" means that
+        the equilibrium simplifies to the weights if these values are inserted.
+        Hydrodynamic example: rho=1, u_i = 0
+        """
+
+    @abc.abstractmethod
+    def equilibriumInputEquationsFromPdfs(self, pdfs):
+        """
+        Returns an equation collection that defines all necessary quantities to compute the equilibrium as functions
+        of the pdfs.
+        For hydrodynamic LBM schemes this is usually the density and velocity.
+
+        :param pdfs: values or symbols for the pdf values
+        """
+
+    @abc.abstractmethod
+    def outputEquationsFromPdfs(self, pdfs, outputQuantityNames):
+        """
+        Returns an equation collection that defines conserved quantities for output. These conserved quantities might
+        be slightly different that the ones used as input for the equilibrium e.g. due to a force model.
+        :param pdfs: values for the pdf entries
+        :param outputQuantityNames: list of conserved quantity names, defining which parameters should be written out.
+                                    See :func:`conservedQuantities`
+        """
+
+    @abc.abstractmethod
+    def equilibriumInputEquationsFromInitValues(self, **kwargs):
+        """
+        Returns an equation collection that defines all necessary quantities to compute the equilibrium as function of
+        given conserved quantities. Parameters can be names that are given by
+        symbol names of :func:`conservedQuantities`.
+        For all parameters not specified each implementation should use sensible defaults. For example hydrodynamic
+        schemes use density=1 and velocity=0.
+        """
+
+
+class DensityVelocityComputation(AbstractConservedQuantityComputation):
+    def __init__(self, stencil, compressible, forceModel=None,
+                 zerothOrderMomentSymbol=sp.Symbol("rho"),
+                 firstOrderMomentSymbols=sp.symbols("u_:3")):
+        dim = len(stencil[0])
+        self._stencil = stencil
+        self._compressible = compressible
+        self._forceModel = forceModel
+        self._symbolOrder0 = zerothOrderMomentSymbol
+        self._symbolsOrder1 = firstOrderMomentSymbols[:dim]
+
+    @property
+    def conservedQuantities(self):
+        return {'density': 1,
+                'velocity': 3}
+
+    def definedSymbols(self, order='all'):
+        if order == 'all':
+            return {'density': self._symbolOrder0,
+                    'velocity': self._symbolsOrder1}
+        elif order == 0:
+            return 'density', self._symbolOrder0
+        elif order == 1:
+            return 'velocity', self._symbolsOrder1
+        else:
+            return None
+
+    @property
+    def defaultValues(self):
+        result = {self._symbolOrder0: 1}
+        for s in self._symbolsOrder1:
+            result[s] = 0
+        return result
+
+    def equilibriumInputEquationsFromPdfs(self, pdfs):
+        dim = len(self._stencil[0])
+        eqColl = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs, self._symbolOrder0,
+                                                          self._symbolsOrder1[:dim])
+        if self._compressible:
+            eqColl = divideFirstOrderMomentsByRho(eqColl, dim)
+
+        eqColl = applyForceModelShift('equilibriumVelocityShift', dim, eqColl, self._forceModel, self._compressible)
+        return eqColl
+
+    def equilibriumInputEquationsFromInitValues(self, density=1, velocity=[0, 0, 0]):
+        dim = len(self._stencil[0])
+        zerothOrderMoment = density
+        firstOrderMoments = velocity[:dim]
+        velOffset = [0] * dim
+        if self._compressible:
+            if self._forceModel and hasattr(self._forceModel, 'macroscopicVelocityShift'):
+                velOffset = self._forceModel.macroscopicVelocityShift(zerothOrderMoment)
+        else:
+            if self._forceModel and hasattr(self._forceModel, 'macroscopicVelocityShift'):
+                velOffset = self._forceModel.macroscopicVelocityShift(sp.Rational(1, 1))
+            zerothOrderMoment -= sp.Rational(1, 1)
+        firstOrderMoments = [a - b for a, b in zip(firstOrderMoments, velOffset)]
+
+        eqs = [sp.Eq(self._symbolOrder0, zerothOrderMoment)]
+        eqs += [sp.Eq(l, r) for l, r in zip(self._symbolsOrder1, firstOrderMoments)]
+        return EquationCollection(eqs, [])
+
+    def outputEquationsFromPdfs(self, pdfs, outputQuantityNames):
+        outputQuantityNames = set(outputQuantityNames)
+
+        dim = len(self._stencil[0])
+        eqColl = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs, self._symbolOrder0,
+                                                          self._symbolsOrder1[:dim])
+        if self._compressible:
+            eqColl = divideFirstOrderMomentsByRho(eqColl, dim)
+        else:
+            eqColl = addDensityOffset(eqColl)
+
+        eqColl = applyForceModelShift('macroscopicVelocityShift', dim, eqColl, self._forceModel, self._compressible)
+
+        nameToSymbol = {'density': self._symbolOrder0,
+                        'velocity': self._symbolsOrder1}
+        return eqColl.extract({nameToSymbol[e] for e in outputQuantityNames})
+
+    def __repr__(self):
+        return "ConservedValueComputation for %s" % (", " .join(self.conservedQuantities.keys()),)
+
+# -----------------------------------------  Helper functions ----------------------------------------------------------
+
+
 def getEquationsForZerothAndFirstOrderMoment(stencil, symbolicPdfs, symbolicZerothMoment, symbolicFirstMoments):
     """
     Returns an equation system that computes the zeroth and first order moments with the least amount of operations
@@ -11,7 +165,7 @@ def getEquationsForZerothAndFirstOrderMoment(stencil, symbolicPdfs, symbolicZero
     .. math :
 
         \rho = \sum_{d \in S} f_d
-        \u_j = \sum_{d \in S} f_d u_jd
+        u_j = \sum_{d \in S} f_d u_jd
 
     :param stencil: called :math:`S` above
     :param symbolicPdfs: called :math:`f` above
@@ -69,8 +223,8 @@ def divideFirstOrderMomentsByRho(equationCollection, dim):
     rhoInv = sp.Symbol("rhoInv")
     newSubExpression = sp.Eq(rhoInv, 1 / rho)
     newFirstOrderMomentEq = [sp.Eq(eq.lhs, eq.rhs * rhoInv) for eq in oldEqs[1:dim+1]]
-    newEqs = oldEqs[0] + newFirstOrderMomentEq + oldEqs[dim+1:]
-    return equationCollection.createNewWithAdditionalSubexpressions(newEqs, newSubExpression)
+    newEqs = [oldEqs[0]] + newFirstOrderMomentEq + oldEqs[dim+1:]
+    return equationCollection.newWithAdditionalSubexpressions(newEqs, [newSubExpression])
 
 
 def addDensityOffset(equationCollection, offset=sp.Rational(1, 1)):
@@ -78,11 +232,11 @@ def addDensityOffset(equationCollection, offset=sp.Rational(1, 1)):
     Assumes that first equation is the density (zeroth moment). Changes the density equations by adding offset to it.
     """
     oldEqs = equationCollection.mainEquations
-    newDensity = sp.Eq(oldEqs[0].lhs, oldEqs[1].rhs + offset)
-    return equationCollection.createNewWithAdditionalSubexpressions([newDensity] + oldEqs[1:], [])
+    newDensity = sp.Eq(oldEqs[0].lhs, oldEqs[0].rhs + offset)
+    return equationCollection.newWithAdditionalSubexpressions([newDensity] + oldEqs[1:], [])
 
 
-def applyForceModelShift(shiftMemberName, dim, equationCollection, forceModel, compressible):
+def applyForceModelShift(shiftMemberName, dim, equationCollection, forceModel, compressible, reverse=False):
     """
     Modifies the first order moment equations in equationCollection according to the force model shift.
     It is applied if force model has a method named shiftMemberName. The equations 1: dim+1 of the passed
@@ -93,41 +247,15 @@ def applyForceModelShift(shiftMemberName, dim, equationCollection, forceModel, c
         density = oldEqs[0].lhs if compressible else sp.Rational(1, 1)
         oldVelEqs = oldEqs[1:dim + 1]
         shiftFunc = getattr(forceModel, shiftMemberName)
-        shiftedVels = shiftFunc([eq.rhs for eq in oldVelEqs], density)
-        shiftedVelocityEqs = [sp.Eq(oldEq.lhs, shiftedVel) for oldEq, shiftedVel in zip(oldVelEqs, shiftedVels)]
+        velOffsets = shiftFunc(density)
+        if reverse:
+            velOffsets = [-v for v in velOffsets]
+        shiftedVelocityEqs = [sp.Eq(oldEq.lhs, oldEq.rhs + offset) for oldEq, offset in zip(oldVelEqs, velOffsets)]
         newEqs = [oldEqs[0]] + shiftedVelocityEqs + oldEqs[dim + 1:]
-        return equationCollection.createNewWithAdditionalSubexpressions(newEqs, [])
+        return equationCollection.newWithAdditionalSubexpressions(newEqs, [])
     else:
         return equationCollection
 
 
-# --------------------------- Density / Velocity definitions with force shift ------------------------------------------
-
-
-def densityVelocityExpressionsForEquilibrium(stencil, symbolicPdfs, compressible, symbolicDensity,
-                                             symbolicVelocities, forceModel=None):
-    """
-    Returns an equation collection, containing equations to compute density, velocity and pressure from pdf values
-    """
 
-    eqColl = getEquationsForZerothAndFirstOrderMoment(stencil, symbolicPdfs, symbolicDensity, symbolicVelocities)
-    if compressible:
-        eqColl = divideFirstOrderMomentsByRho(eqColl)
-    if forceModel is not None:
-        eqColl = applyForceModelShift('equilibriumVelocity', len(stencil[0]), eqColl, forceModel, compressible)
-    return eqColl
-
-
-def densityVelocityExpressionsForOutput(stencil, symbolicPdfs, compressible, symbolicDensity,
-                                        symbolicVelocities, forceModel=None):
-    """
-    Returns an equation collection, containing equations to compute density, velocity and pressure from pdf values
-    """
 
-    eqColl = getEquationsForZerothAndFirstOrderMoment(stencil, symbolicPdfs, symbolicDensity, symbolicVelocities)
-    if compressible:
-        eqColl = divideFirstOrderMomentsByRho(eqColl)
-        eqColl = addDensityOffset(eqColl)
-    if forceModel is not None:
-        eqColl = applyForceModelShift('macroscopicVelocity', len(stencil[0]), eqColl, forceModel, compressible)
-    return eqColl

methods/hydrodynamiclbmmethod.py

-import sympy as sp
-from lbmpy.methods.abstractlbmmethod import AbstractLbmMethod
-from pystencils.equationcollection.equationcollection import EquationCollection
-
-
-class HydrodynamicLbmMethod(AbstractLbmMethod):
-
-    def __init__(self, stencil, forceModel, compressible, viscosityRelaxationRate, c_s_sq):
-        super(HydrodynamicLbmMethod, self).__init__(stencil)
-        self._compressible = compressible
-        self._forceModel = forceModel
-        self._viscosityRelaxationRate = viscosityRelaxationRate
-        self._c_s_sq = c_s_sq
-
-    @property
-    def densitySymbol(self):
-        return sp.Symbol('rho')
-
-    @property
-    def pressureSymbol(self):
-        return sp.Symbol("p")
-
-    @property
-    def velocitySymbols(self):
-        return sp.symbols("u_:%d" % (self.dim,))
-
-    @property
-    def availableMacroscopicQuantities(self):
-        return {'density': self.densitySymbol,
-                'pressure': self.pressureSymbol,
-                'velocity': self.velocitySymbols}
-
-    @property
-    def viscosityRelaxationRate(self):
-        return self._viscosityRelaxationRate
-
-    @property
-    def conservedQuantitiesSymbols(self):
-        return self.availableMacroscopicQuantities.values()
-
-    def getMacroscopicQuantitiesEquations(self, sequenceOfSymbols):
-        # TODO shift of velocity by force model
-        eqColl = _densityVelocityExpressions(self.stencil, self.preCollisionPdfSymbols, self._compressible,
-                                             self._c_s_sq, self.densitySymbol, self.pressureSymbol,
-                                             self.velocitySymbols)
-
-        return eqColl.extract(sequenceOfSymbols)
-
-
-def _densityVelocityExpressions(stencil, symbolicPdfs, compressible, c_s_sq,
-                                symbolicDensity, symbolicPressure, symbolicVelocities):
-    """
-    Returns an equation collection, containing equations to compute density, velocity and pressure from pdf values
-    """
-    def filterOutPlusTerms(expr):
-        result = 0
-        for term in expr.args:
-            if not type(term) is sp.Mul:
-                result += term
-        return result
-
-    dim = len(stencil[0])
-
-    subexpressions = []
-    pdfSum = sum(symbolicPdfs)
-    u = [0] * dim
-    for f, offset in zip(symbolicPdfs, stencil):
-        for i in range(dim):
-            u[i] += f * int(offset[i])
-
-    plusTerms = [set(filterOutPlusTerms(u_i).args) for u_i in u]
-    for i in range(dim):
-        rhs = plusTerms[i]
-        for j in range(i):
-            rhs -= plusTerms[j]
-        eq = sp.Eq(sp.Symbol("vel%dTerm" % (i,)), sum(rhs))
-        subexpressions.append(eq)
-
-    for subexpression in subexpressions:
-        pdfSum = pdfSum.subs(subexpression.rhs, subexpression.lhs)
-
-    for i in range(dim):
-        u[i] = u[i].subs(subexpressions[i].rhs, subexpressions[i].lhs)
-
-    equations = []
-    if compressible:
-        rho_inv = sp.Symbol("rho_inv")
-        subexpressions.append(sp.Eq(rho_inv, 1 / symbolicDensity))
-        equations += [sp.Eq(symbolicDensity, pdfSum)]
-        equations += [sp.Eq(symbolicPressure, pdfSum * c_s_sq)]
-        equations += [sp.Eq(u_i_sym, u_i * rho_inv) for u_i_sym, u_i in zip(symbolicVelocities, u)]
-    else:
-        subexpressions.append(sp.Eq(symbolicDensity, sp.Rational(1, 1)))
-        subexpressions.append(sp.Eq(symbolicPressure, pdfSum * c_s_sq))
-        equations += [sp.Eq(u_i_sym, u_i) for u_i_sym, u_i in zip(symbolicVelocities, u)]
-
-    return EquationCollection(equations, subexpressions)

methods/momentbased.py

+import sympy as sp
+import itertools
+import collections
+from collections import namedtuple
+from lbmpy.methods.abstractlbmmethod import AbstractLbmMethod
+from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
+from lbmpy.moments import MOMENT_SYMBOLS, momentMatrix, exponentToPolynomialRepresentation, isShearMoment
+from pystencils.equationcollection import EquationCollection
+
+"""
+Ways to create method:
+    - moment (polynomial or tuple) mapped to relaxation rate
+    - moment matrix & relaxation vector
+    - createSRT, createTRT, createMRT
+"""
+
+RelaxationInfo = namedtuple('Relaxationinfo', ['equilibriumValue', 'relaxationRate'])
+
+
+class MomentBasedLbmMethod(AbstractLbmMethod):
+
+    def __init__(self, stencil, momentToRelaxationInfoDict, conservedQuantityComputation, forceModel=None):
+        """
+
+        :param stencil:
+        :param momentToRelaxationInfoDict:
+        :param conservedQuantityComputation: instance of :class:`lbmpy.methods.AbstractConservedQuantityComputation`
+        :param forceModel:
+        """
+        super(MomentBasedLbmMethod, self).__init__(stencil)
+
+        assert isinstance(conservedQuantityComputation, AbstractConservedQuantityComputation)
+
+        moments = []
+        relaxationRates = []
+        equilibriumMoments = []
+        for moment, relaxInfo in momentToRelaxationInfoDict.items():
+            moments.append(moment)
+            relaxationRates.append(relaxInfo.relaxationRate)
+            equilibriumMoments.append(relaxInfo.equilibriumValue)
+
+        self._forceModel = forceModel
+        self._momentToRelaxationInfoDict = momentToRelaxationInfoDict
+        self._momentMatrix = momentMatrix(moments, self.stencil)
+        self._relaxationRates = sp.Matrix(relaxationRates)
+        self._equilibriumMoments = sp.Matrix(equilibriumMoments)
+        self._conservedQuantityComputation = conservedQuantityComputation
+
+        symbolsInEquilibriumMoments = self._equilibriumMoments.atoms(sp.Symbol)
+        conservedQuantities = set()
+        for v in self._conservedQuantityComputation.definedSymbols().values():
+            if isinstance(v, collections.Sequence):
+                conservedQuantities.update(v)
+            else:
+                conservedQuantities.add(v)
+        undefinedEquilibriumSymbols = symbolsInEquilibriumMoments - conservedQuantities
+        assert len(undefinedEquilibriumSymbols) == 0, "Undefined symbol(s) in equilibrium moment: %s" % \
+                                                      (undefinedEquilibriumSymbols, )
+
+        self._weights = self._computeWeights()
+
+    @property
+    def zerothOrderEquilibriumMomentSymbol(self,):
+        return self._conservedQuantityComputation.definedSymbols(order=0)[1]
+
+    @property
+    def firstOrderEquilibriumMomentSymbols(self,):
+        return self._conservedQuantityComputation.definedSymbols(order=1)[1]
+
+    @property
+    def weights(self):
+        return self._weights
+
+    def _computeWeights(self):
+        replacements = self._conservedQuantityComputation.defaultValues
+        eqColl = self.getEquilibrium().newWithSubstitutionsApplied(replacements).insertSubexpressions()
+        weights = []
+        for eq in eqColl.mainEquations:
+            value = eq.rhs.expand()
+            assert len(value.atoms(sp.Symbol)) == 0, "Failed to compute weights"
+            weights.append(value)
+        return weights
+
+    def getShearRelaxationRate(self):
+        """
+        Assumes that all shear moments are relaxed with same rate - returns this rate
+        Shear moments in 3D are: x*y, x*z and y*z - in 2D its only x*y
+        The shear relaxation rate determines the viscosity in hydrodynamic LBM schemes
+        """
+        relaxationRates = set()
+        for moment, relaxInfo in self._momentToRelaxationInfoDict.items():
+            if isShearMoment(moment):
+                relaxationRates.add(relaxInfo.relaxationRate)
+        if len(relaxationRates) == 1:
+            return relaxationRates.pop()
+        else:
+            if len(relaxationRates) > 1:
+                raise ValueError("Shear moments are relaxed with different relaxation times: %s" % (relaxationRates,))
+            else:
+                raise NotImplementedError("Shear moments seem to be not relaxed separately - "
+                                          "Can not determine their relaxation rate automatically")
+
+    def getEquilibrium(self):
+        D = sp.eye(len(self._relaxationRates))
+        return self._getCollisionRuleWithRelaxationMatrix(D)
+
+    def getCollisionRule(self):
+        D = sp.diag(*self._relaxationRates)
+        eqColl = self._getCollisionRuleWithRelaxationMatrix(D)
+        if self._forceModel is not None:
+            forceModelTerms = self._forceModel(self)
+            newEqs = [sp.Eq(eq.lhs, eq.rhs + fmt) for eq, fmt in zip(eqColl.mainEquations, forceModelTerms)]
+            eqColl = eqColl.newWithAdditionalSubexpressions(newEqs, [])
+        return eqColl
+
+    @property
+    def conservedQuantityComputation(self):
+        return self._conservedQuantityComputation
+
+    def _getCollisionRuleWithRelaxationMatrix(self, D):
+        f = sp.Matrix(self.preCollisionPdfSymbols)
+        M = self._momentMatrix
+        m_eq = self._equilibriumMoments
+
+        collisionRule = f + M.inv() * D * (m_eq - M * f)
+        collisionEqs = [sp.Eq(lhs, rhs) for lhs, rhs in zip(self.postCollisionPdfSymbols, collisionRule)]
+
+        eqValueEqs = self._conservedQuantityComputation.equilibriumInputEquationsFromPdfs(f)
+        simplificationHints = eqValueEqs.simplificationHints