lbmpy: Refactored boundary handling

- boundary handling more flexible now: boundaries return LIST of equations now
- new pressure boundary

boundaries.py

@@ -3,12 +3,14 @@ import itertools
 import numpy as np
 import sympy as sp
 
+from lbmpy.densityVelocityExpressions import getDensityVelocityExpressions
 from lbmpy.stencils import getStencil
 from pystencils.backends.cbackend import CustomCppCode
 from pystencils.types import TypedSymbol
 from pystencils.field import Field
 from pystencils.ast import Node, Block, SympyAssignment, LoopOverCoordinate, KernelFunction
-from pystencils.transformations import moveConstantsBeforeLoop, resolveFieldAccesses
+from pystencils.transformations import moveConstantsBeforeLoop, resolveFieldAccesses, typingFromSympyInspection, \
+    typeAllEquations
 
 INV_DIR_SYMBOL = TypedSymbol("invDir", "int")
 WEIGHTS_SYMBOL = TypedSymbol("weights", "double")
@@ -66,7 +68,12 @@ def generateBoundaryHandling(pdfField, indexArr, latticeModel, boundaryFunctor):
     dirSymbol = TypedSymbol("dir", "int")
     cellLoopBody.append(SympyAssignment(dirSymbol, indexField[0](dim)))
 
-    cellLoopBody.append(boundaryFunctor(pdfField, dirSymbol, latticeModel))
+    boundaryEqList = boundaryFunctor(pdfField, dirSymbol, latticeModel)
+    typeInfos = typingFromSympyInspection(boundaryEqList, pdfField.dtype)
+    fieldsRead, fieldsWritten, assignments = typeAllEquations(boundaryEqList, typeInfos)
+
+    for be in assignments:
+        cellLoopBody.append(be)
 
     functionBody = Block([cellLoop])
     ast = KernelFunction(functionBody, [pdfField, indexField])
@@ -94,29 +101,59 @@ def createBoundaryIndexList(flagFieldArr, nrOfGhostLayers, stencil, boundaryMask
 
 def noSlip(pdfField, direction, latticeModel):
     neighbor = offsetFromDir(direction, latticeModel.dim)
-    current = tuple([0] * latticeModel.dim)
     inverseDir = invDir(direction)
-    return SympyAssignment(pdfField[neighbor](inverseDir), pdfField[current](direction))
+    return [sp.Eq(pdfField[neighbor](inverseDir), pdfField(direction))]
 
 
 def ubb(pdfField, direction, latticeModel, velocity):
     neighbor = offsetFromDir(direction, latticeModel.dim)
-    current = tuple([0] * latticeModel.dim)
     inverseDir = invDir(direction)
 
     velTerm = 6 * sum([d_i * v_i for d_i, v_i in zip(neighbor, velocity)]) * weightOfDirection(direction)
-    return SympyAssignment(pdfField[neighbor](inverseDir),
-                           pdfField[current](direction) - velTerm)
+    return [sp.Eq(pdfField[neighbor](inverseDir),
+                  pdfField(direction) - velTerm)]
+
+
+def fixedDensity(pdfField, direction, latticeModel, density):
+    from lbmpy.equilibria import standardDiscreteEquilibrium
+    neighbor = offsetFromDir(direction, latticeModel.dim)
+    inverseDir = invDir(direction)
+    stencil = latticeModel.stencil
+
+    if not latticeModel.compressible:
+        density -= 1
+
+    eqParams = {'stencil': stencil,
+                'order': 2,
+                'c_s_sq': sp.Rational(1, 3),
+                'compressible': latticeModel.compressible,
+                'rho': density}
+
+    u = sp.Matrix(latticeModel.symbolicVelocity)
+    symmetricEq = (standardDiscreteEquilibrium(u=u, **eqParams) + standardDiscreteEquilibrium(u=-u, **eqParams)) / 2
+
+    subExpr1, rhoExpr, subExpr2, uExprs = getDensityVelocityExpressions(stencil,
+                                                                        [pdfField(i) for i in range(len(stencil))],
+                                                                        latticeModel.compressible)
+    subExprs = subExpr1 + [rhoExpr] + subExpr2 + uExprs
+
+    conditions = [(eq_i, sp.Equality(direction, i)) for i, eq_i in enumerate(symmetricEq)] + [(0, True)]
+    eq_component = sp.Piecewise(*conditions)
 
+    return subExprs + [sp.Eq(pdfField[neighbor](inverseDir),
+                             2 * eq_component - pdfField(direction))]
 
 if __name__ == "__main__":
-    import lbmpy.collisionoperator as coll
-    lm = coll.makeSRT(getStencil("D3Q19"))
+    from lbmpy.latticemodel import makeSRT
+    from pystencils.cpu import generateC
+    import functools
+    lm = makeSRT(getStencil("D3Q19"))
     pdfField = Field.createGeneric("pdfField", lm.dim, indexDimensions=1)
 
     indexArr = np.array([[1, 1, 1, 1], [1, 2, 1,  1], [2, 1, 1, 1]], dtype=np.int32)
 
-    ast = generateBoundaryHandling(pdfField, indexArr, lm, noSlip)
+    pressureBoundary = functools.partial(fixedDensity, density=1.0)
+    ast = generateBoundaryHandling(pdfField, indexArr, lm, pressureBoundary)
 
-    print(ast.generateC())
+    print(generateC(ast))
 

boundaryhandling.py

@@ -34,7 +34,14 @@ class BoundaryHandling:
 
     def setBoundary(self, name, indexExpr):
         if not isinstance(name, str):
-            name = name.__name__
+            function = name
+            if hasattr(function, '__name__'):
+                name = function.__name__
+            else:
+                name = function.name
+
+            if function not in self._boundaryFunctions:
+                self.addBoundary(function, name)
 
         flag = self.getFlag(name)
         self.flagField[indexExpr] = flag

equilibria.py

@@ -39,7 +39,7 @@ getWeights.weights = {
 }
 
 
-@diskcache
+#@diskcache
 def standardDiscreteEquilibrium(stencil, rho=None, u=None, order=2, c_s_sq=None, compressible=True):
     """
     Returns the common quadratic LBM equilibrium as a list of sympy expressions
@@ -56,9 +56,9 @@ def standardDiscreteEquilibrium(stencil, rho=None, u=None, order=2, c_s_sq=None,
 
     D = len(e[0])
 
-    if not rho:
+    if rho is None:
         rho = getSymbolicDensity()
-    if not u:
+    if u is None:
         u = getSymbolicVelocityVector(D, "u")
 
     if not c_s_sq:

helperkernels.py

+import numpy as np
 
 
 def initializePdfField(latticeModel, pdfArray):
@@ -14,4 +15,11 @@ def initializePdfField(latticeModel, pdfArray):
             raise NotImplementedError()
 
 
-
+def computeVelocity(latticeModel, pdfArray):
+    vel = np.zeros(pdfArray.shape[:-1] + (2,))
+    for i, dir in enumerate(latticeModel.stencil):
+        if dir[0] != 0:
+            vel[:, :, 0] += dir[0] * pdfArray[:, :, i]
+        if dir[1] != 0:
+            vel[:, :, 1] += dir[1] * pdfArray[:, :, i]
+    return vel

plot2d.py

+from matplotlib.pyplot import *
+import numpy as np
+from numpy.linalg import norm
+
+
+def removeGhostLayers(field):
+    return field[1:-1, 1:-1]
+
+
+def vectorField(field, step=2, **kwargs):
+    field = removeGhostLayers(field)
+    veln = field.swapaxes(0, 1)
+    quiver(veln[::step, ::step, 0], veln[::step, ::step, 1], **kwargs)
+
+
+def vectorFieldMagnitude(field, **kwargs):
+    field = norm(field, axis=2, ord=2)
+    scalarField(field, **kwargs)
+
+
+def scalarField(field, **kwargs):
+    field = removeGhostLayers(field)
+    field = np.swapaxes(field, 0, 1)
+    imshow(field, origin='lower', **kwargs)
+