From a1bb764f6268313a7ed9aabda9be4689ce0c328b Mon Sep 17 00:00:00 2001
From: Martin Bauer <martin.bauer@fau.de>
Date: Tue, 17 Oct 2017 16:55:42 +0200
Subject: [PATCH] started with generalized steady state Chapman Enskog
 expansion

---
 chapman_enskog/chapman_enskog.py        |  12 +-
 chapman_enskog/derivative.py            |  11 +-
 chapman_enskog/steady_state.py          | 142 ++++++++++++++++++++++++
 methods/conservedquantitycomputation.py |   4 +
 methods/momentbased.py                  |  12 ++
 parallel/scenarios.py                   |   4 +-
 6 files changed, 174 insertions(+), 11 deletions(-)
 create mode 100644 chapman_enskog/steady_state.py

diff --git a/chapman_enskog/chapman_enskog.py b/chapman_enskog/chapman_enskog.py
index 5b90433b..2b7a8ba9 100644
--- a/chapman_enskog/chapman_enskog.py
+++ b/chapman_enskog/chapman_enskog.py
@@ -25,12 +25,12 @@ def getExpandedName(originalObject, number):
     return originalObject.func(newName)
 
 
-def expandedSymbol(name, subscript=None, superscript=None):
+def expandedSymbol(name, subscript=None, superscript=None, **kwargs):
     if subscript is not None:
         name += "_{%s}" % (subscript,)
     if superscript is not None:
         name += "^{(%s)}" % (superscript,)
-    return sp.Symbol(name)
+    return sp.Symbol(name, **kwargs)
 
 # --------------------------------   Summation Convention  -------------------------------------------------------------
 
@@ -383,7 +383,7 @@ def getTaylorExpandedLbEquation(pdfSymbolName="f", pdfsAfterCollisionOperator="\
 
 
 def useChapmanEnskogAnsatz(equation, timeDerivativeOrders=(1, 3), spatialDerivativeOrders=(1, 2),
-                           pdfs=(['f', 0, 3], ['\Omega f', 1, 3])):
+                           pdfs=(['f', 0, 3], ['\Omega f', 1, 3]), **kwargs):
 
     t, eps = sp.symbols("t epsilon")
 
@@ -406,9 +406,9 @@ def useChapmanEnskogAnsatz(equation, timeDerivativeOrders=(1, 3), spatialDerivat
     for pdfName, startOrder, stopOrder in pdfs:
         if isinstance(pdfName, sp.Symbol):
             pdfName = pdfName.name
-        expandedPdfSymbols += [expandedSymbol(pdfName, superscript=i) for i in range(startOrder, stopOrder)]
-        subsDict[sp.Symbol(pdfName)] = sum(eps ** i * expandedSymbol(pdfName, superscript=i)
-                                           for i in range(startOrder, stopOrder))
+        expandedPdfSymbols += [expandedSymbol(pdfName, superscript=i, **kwargs) for i in range(startOrder, stopOrder)]
+        subsDict[sp.Symbol(pdfName, **kwargs)] = sum(eps ** i * expandedSymbol(pdfName, superscript=i, **kwargs)
+                                                     for i in range(startOrder, stopOrder))
         maxExpansionOrder = max(maxExpansionOrder, stopOrder)
     equation = equation.subs(subsDict)
     equation = expandUsingLinearity(equation, functions=expandedPdfSymbols).expand().collect(eps)
diff --git a/chapman_enskog/derivative.py b/chapman_enskog/derivative.py
index e54537dc..ddd0bee6 100644
--- a/chapman_enskog/derivative.py
+++ b/chapman_enskog/derivative.py
@@ -78,7 +78,6 @@ class Diff(sp.Expr):
         - all simplifications have to be done manually
         - each Diff has a Chapman Enskog expansion order index: 'ceIdx'
     """
-    is_commutative = True
     is_number = False
     is_Rational = False
 
@@ -87,6 +86,14 @@ class Diff(sp.Expr):
             return sp.Rational(0, 1)
         return sp.Expr.__new__(cls, argument.expand(), sp.sympify(label), sp.sympify(ceIdx), **kwargs)
 
+    @property
+    def is_commutative(self):
+        anyNonCommutative = any(not s.is_commutative for s in self.atoms(sp.Symbol))
+        if anyNonCommutative:
+            return False
+        else:
+            return True
+
     def getArgRecursive(self):
         """Returns the argument the derivative acts on, for nested derivatives the inner argument is returned"""
         if not isinstance(self.arg, Diff):
@@ -213,7 +220,7 @@ def expandUsingLinearity(expr, functions=None, constants=None):
 
     if isinstance(expr, Diff):
         arg = expandUsingLinearity(expr.arg, functions)
-        if arg.func == sp.Add:
+        if hasattr(arg, 'func') and arg.func == sp.Add:
             result = 0
             for a in arg.args:
                 result += Diff(a, label=expr.label, ceIdx=expr.ceIdx).splitLinear(functions)
diff --git a/chapman_enskog/steady_state.py b/chapman_enskog/steady_state.py
new file mode 100644
index 00000000..6f63ed82
--- /dev/null
+++ b/chapman_enskog/steady_state.py
@@ -0,0 +1,142 @@
+import sympy as sp
+from pystencils.sympyextensions import normalizeProduct
+from lbmpy.chapman_enskog.derivative import Diff, DiffOperator, expandUsingLinearity, normalizeDiffOrder
+from lbmpy.chapman_enskog.chapman_enskog import expandedSymbol, useChapmanEnskogAnsatz
+
+
+class SteadyStateChapmanEnskogAnalysis(object):
+
+    def __init__(self, method, forceModelClass=None, order=4):
+        self.method = method
+        self.dim = method.dim
+        self.order = order
+        self.physicalVariables = list(sp.Matrix(self.method.momentEquilibriumValues).atoms(sp.Symbol))  # rho, u..
+        self.eps = sp.Symbol("epsilon")
+
+        self.fSym = sp.Symbol("f", commutative=False)
+        self.fSyms = [expandedSymbol("f", superscript=i, commutative=False) for i in range(order + 1)]
+        self.collisionOpSym = sp.Symbol("A", commutative=False)
+        self.forceSym = sp.Symbol("F_q", commutative=False)
+        self.velocitySyms = sp.Matrix([expandedSymbol("c", subscript=i, commutative=False) for i in range(self.dim)])
+
+        self.F_q = [0] * len(self.method.stencil)
+        self.forceModel = None
+        if forceModelClass:
+            accelerationSymbols = sp.symbols("a_:%d" % (self.dim,), commutative=False)
+            self.physicalVariables += accelerationSymbols
+            self.forceModel = forceModelClass(accelerationSymbols)
+            self.F_q = self.forceModel(self.method)
+
+        # Perform the analysis
+        self.tayloredEquation = self._createTaylorExpandedEquation()
+        self.pdfHierarchy = self._createPdfHierarchy(self.tayloredEquation)
+        self.recombinedEq = self._recombinePdfs(self.pdfHierarchy)
+
+        symbolsToValues = self._getSymbolsToValuesDict()
+        self.continuityEquation = self._computeContinuityEquation(self.recombinedEq, symbolsToValues)
+        self.momentumEquations = [self._computeMomentumEquation(self.recombinedEq, symbolsToValues, h)
+                                  for h in range(self.dim)]
+
+    def _createTaylorExpandedEquation(self):
+        """
+        Creates a generic, Taylor expanded lattice Boltzmann update equation with collision and force term.
+        Collision operator and force terms are represented symbolically.
+        """
+        c = self.velocitySyms
+        Dx = sp.Matrix([DiffOperator(label=l) for l in range(self.dim)])
+
+        differentialOperator = sum((self.eps * c.dot(Dx)) ** n / sp.factorial(n)
+                                   for n in range(1, self.order + 1))
+        taylorExpansion = DiffOperator.apply(differentialOperator.expand(), self.fSym)
+
+        fNonEq = self.fSym - self.fSyms[0]
+        return taylorExpansion + self.collisionOpSym * fNonEq - self.eps * self.forceSym
+
+    def _createPdfHierarchy(self, tayloredEquation):
+        """
+        Expresses the expanded pdfs f^1, f^2, ..  as functions of the equilibrium f^0.
+        Returns a list where element [1] is the equation for f^1 etc.
+        """
+        chapmanEnskogHierarchy = useChapmanEnskogAnsatz(tayloredEquation, spatialDerivativeOrders=None,
+                                                        pdfs=(['f', 0, self.order + 1],), commutative=False)
+        chapmanEnskogHierarchy = [chapmanEnskogHierarchy[i] for i in range(self.order + 1)]
+
+        result = []
+        substitutionDict = {}
+        for ceEq, f_i in zip(chapmanEnskogHierarchy, self.fSyms):
+            newEq = -1 / self.collisionOpSym * (ceEq - self.collisionOpSym * f_i)
+            newEq = expandUsingLinearity(newEq.subs(substitutionDict), functions=self.fSyms + [self.forceSym])
+            if newEq:
+                substitutionDict[f_i] = newEq
+            result.append(newEq)
+
+        return result
+
+    def _recombinePdfs(self, pdfHierarchy):
+        return sum(pdfHierarchy[i] * self.eps**(i-1) for i in range(1, self.order+1))
+
+    def _computeContinuityEquation(self, recombinedEq, symbolsToValues):
+        return self._computeMoments(recombinedEq, symbolsToValues)
+
+    def _computeMomentumEquation(self, recombinedEq, symbolsToValues, coordinate):
+        eq = sp.expand(self.velocitySyms[coordinate] * recombinedEq)
+
+        result = self._computeMoments(eq, symbolsToValues)
+        if self.forceModel and hasattr(self.forceModel, 'equilibriumVelocityShift'):
+            compressible = self.method.conservedQuantityComputation.compressible
+            shift = self.forceModel.equilibriumVelocityShift(sp.Symbol("rho") if compressible else 1)
+            result += shift[coordinate]
+        return result
+
+    def _getSymbolsToValuesDict(self):
+        result = {1 / self.collisionOpSym: self.method.inverseCollisionMatrix,
+                  self.forceSym: sp.Matrix(self.forceModel(self.method)) if self.forceModel else 0,
+                  self.fSyms[0]: self.method.getEquilibriumTerms()}
+        for i, c_i in enumerate(self.velocitySyms):
+            result[c_i] = sp.Matrix([d[i] for d in self.method.stencil])
+
+        return result
+
+    def _computeMoments(self, recombinedEq, symbolsToValues):
+        eq = recombinedEq.expand()
+        assert eq.func is sp.Add
+
+        newProducts = []
+        for product in eq.args:
+            assert product.func is sp.Mul
+
+            derivative = None
+
+            newProd = 1
+            for arg in reversed(normalizeProduct(product)):
+                if isinstance(arg, Diff):
+                    assert derivative is None, "More than one derivative term in the product"
+                    derivative = arg
+                    arg = arg.getArgRecursive()  # new argument is inner part of derivative
+
+                if arg in symbolsToValues:
+                    arg = symbolsToValues[arg]
+
+                haveShape = hasattr(arg, 'shape') and hasattr(newProd, 'shape')
+                if haveShape and arg.shape == newProd.shape and arg.shape[1] == 1:
+                    newProd = sp.matrix_multiply_elementwise(newProd, arg)
+                else:
+                    newProd = arg * newProd
+
+            newProd = sp.expand(sum(newProd))
+
+            if derivative is not None:
+                newProd = derivative.changeArgRecursive(newProd)
+
+            newProducts.append(newProd)
+
+        return normalizeDiffOrder(expandUsingLinearity(sp.Add(*newProducts), functions=self.physicalVariables))
+
+
+if __name__ == '__main__':
+    from lbmpy.creationfunctions import createLatticeBoltzmannMethod
+    from lbmpy.forcemodels import Buick
+
+    method = createLatticeBoltzmannMethod(stencil='D2Q9', method='srt')
+    result = SteadyStateChapmanEnskogAnalysis(method, Buick, 3)
+
diff --git a/methods/conservedquantitycomputation.py b/methods/conservedquantitycomputation.py
index b6ea416d..10aa1c49 100644
--- a/methods/conservedquantitycomputation.py
+++ b/methods/conservedquantitycomputation.py
@@ -93,6 +93,10 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
         return {'density': 1,
                 'velocity': len(self._stencil[0])}
 
+    @property
+    def compressible(self):
+        return self._compressible
+
     def definedSymbols(self, order='all'):
         if order == 'all':
             return {'density': self._symbolOrder0,
diff --git a/methods/momentbased.py b/methods/momentbased.py
index e1e8da99..6a30ccab 100644
--- a/methods/momentbased.py
+++ b/methods/momentbased.py
@@ -110,6 +110,18 @@ class MomentBasedLbMethod(AbstractLbMethod):
             newEntry = RelaxationInfo(prevEntry[0], relaxationRate)
             self._momentToRelaxationInfoDict[e] = newEntry
 
+    @property
+    def collisionMatrix(self):
+        M = self.momentMatrix
+        D = sp.diag(*self.relaxationRates)
+        return M.inv() * D * M
+
+    @property
+    def inverseCollisionMatrix(self):
+        M = self.momentMatrix
+        Dinv = sp.diag(*[1/e for e in self.relaxationRates])
+        return M.inv() * Dinv * M
+
     @property
     def momentMatrix(self):
         return momentMatrix(self.moments, self.stencil)
diff --git a/parallel/scenarios.py b/parallel/scenarios.py
index acc7bc12..5c88699c 100644
--- a/parallel/scenarios.py
+++ b/parallel/scenarios.py
@@ -1,4 +1,3 @@
-import sympy as sp
 import numpy as np
 import numbers
 import waLBerla as wlb
@@ -7,7 +6,6 @@ from lbmpy.creationfunctions import createLatticeBoltzmannFunction, updateWithDe
 from lbmpy.macroscopic_value_kernels import compileMacroscopicValuesGetter, compileMacroscopicValuesSetter
 from lbmpy.parallel.boundaryhandling import BoundaryHandling
 from lbmpy.parallel.blockiteration import slicedBlockIteration
-from pystencils.field import createNumpyArrayWithLayout, getLayoutOfArray
 
 
 class Scenario(object):
@@ -38,7 +36,7 @@ class Scenario(object):
             switchToSymbolicRelaxationRatesForEntropicMethods(methodParameters, kernelParams)
             methodParameters['optimizationParams'] = optimizationParams
             self._lbmKernel = createLatticeBoltzmannFunction(**methodParameters)
-        else:
+        else:for
             self._lbmKernel = lbmKernel
 
         # ----- Add fields
-- 
GitLab