Bugfix in equationcollection

boundaries/boundaryconditions.py

+import sympy as sp
+from lbmpy.boundaries.boundaryhandling import offsetFromDir, weightOfDirection, invDir
+
+
+def noSlip(pdfField, direction, latticeModel):
+    neighbor = offsetFromDir(direction, latticeModel.dim)
+    inverseDir = invDir(direction)
+    return [sp.Eq(pdfField[neighbor](inverseDir), pdfField(direction))]
+
+
+def ubb(pdfField, direction, latticeModel, velocity):
+    neighbor = offsetFromDir(direction, latticeModel.dim)
+    inverseDir = invDir(direction)
+
+    velTerm = 6 * sum([d_i * v_i for d_i, v_i in zip(neighbor, velocity)]) * weightOfDirection(direction)
+    return [sp.Eq(pdfField[neighbor](inverseDir),
+                  pdfField(direction) - velTerm)]
+
+
+def fixedDensity(pdfField, direction, latticeModel, density):
+    from lbmpy_old.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))]
+

boundaries/boundaryhandling.py

+import sympy as sp
+import numpy as np
+from pystencils import TypedSymbol, Field
+from pystencils.backends.cbackend import CustomCppCode
+from pystencils.ast import Block, SympyAssignment, LoopOverCoordinate, KernelFunction
+from pystencils.transformations import moveConstantsBeforeLoop, resolveFieldAccesses, typingFromSympyInspection, \
+    typeAllEquations
+from pystencils.cpu import makePythonFunction
+
+
+INV_DIR_SYMBOL = TypedSymbol("invDir", "int")
+WEIGHTS_SYMBOL = TypedSymbol("weights", "double")
+
+
+class BoundaryHandling:
+    def __init__(self, symbolicPdfField, domainShape, latticeModel, ghostLayers=1):
+        self._symbolicPdfField = symbolicPdfField
+        self._shapeWithGhostLayers = [d + 2 * ghostLayers for d in domainShape]
+        self._fluidFlag = 2 ** 31
+        self.flagField = np.full(self._shapeWithGhostLayers, self._fluidFlag, dtype=np.int32)
+        self._ghostLayers = ghostLayers
+        self._latticeModel = latticeModel
+        self._boundaryFunctions = []
+        self._nameToIndex = {}
+        self._boundarySweeps = []
+
+    def addBoundary(self, boundaryFunction, name=None):
+        if name is None:
+            name = boundaryFunction.__name__
+
+        self._nameToIndex[name] = len(self._boundaryFunctions)
+        self._boundaryFunctions.append(boundaryFunction)
+
+    def invalidateIndexCache(self):
+        self._boundarySweeps = []
+
+    def clear(self):
+        np.fill(self._fluidFlag)
+        self.invalidateIndexCache()
+
+    def getFlag(self, name):
+        return 2 ** self._nameToIndex[name]
+
+    def setBoundary(self, name, indexExpr, clearOtherBoundaries=True):
+        if not isinstance(name, str):
+            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)
+        if clearOtherBoundaries:
+            self.flagField[indexExpr] = flag
+        else:
+            # clear fluid flag
+            np.bitwise_and(self.flagField[indexExpr], np.invert(self._fluidFlag), self.flagField[indexExpr])
+            # add new boundary flag
+            np.bitwise_or(self.flagField[indexExpr], flag, self.flagField[indexExpr])
+
+        self.invalidateIndexCache()
+
+    def prepare(self):
+        self.invalidateIndexCache()
+        for boundaryIdx, boundaryFunc in enumerate(self._boundaryFunctions):
+            idxField = createBoundaryIndexList(self.flagField, self._ghostLayers, self._latticeModel.stencil,
+                                               2 ** boundaryIdx, self._fluidFlag)
+            ast = generateBoundaryHandling(self._symbolicPdfField, idxField, self._latticeModel, boundaryFunc)
+            self._boundarySweeps.append(makePythonFunction(ast, {'indexField': idxField}))
+
+    def __call__(self, **kwargs):
+        if len(self._boundarySweeps) == 0:
+            self.prepare()
+        for boundarySweep in self._boundarySweeps:
+            boundarySweep(**kwargs)
+
+
+# -------------------------------------- Helper Functions --------------------------------------------------------------
+
+
+def offsetSymbols(dim):
+    return [TypedSymbol("c_%d" % (d,), "int") for d in range(dim)]
+
+
+def offsetFromDir(dirIdx, dim):
+    return tuple([sp.IndexedBase(symbol, shape=(1,))[dirIdx] for symbol in offsetSymbols(dim)])
+
+
+def invDir(dirIdx):
+    return sp.IndexedBase(INV_DIR_SYMBOL, shape=(1,))[dirIdx]
+
+
+def weightOfDirection(dirIdx):
+    return sp.IndexedBase(WEIGHTS_SYMBOL, shape=(1,))[dirIdx]
+
+
+# ------------------------------------- Kernel Generation --------------------------------------------------------------
+
+class LatticeModelInfo(CustomCppCode):
+    def __init__(self, latticeModel):
+        stencil = latticeModel.stencil
+        symbolsDefined = set(offsetSymbols(latticeModel.dim) + [INV_DIR_SYMBOL, WEIGHTS_SYMBOL])
+
+        offsetSym = offsetSymbols(latticeModel.dim)
+        code = "\n"
+        for i in range(latticeModel.dim):
+            offsetStr = ", ".join([str(d[i]) for d in stencil])
+            code += "const int %s [] = { %s };\n" % (offsetSym[i].name, offsetStr)
+
+        invDirs = []
+        for direction in stencil:
+            inverseDir = tuple([-i for i in direction])
+            invDirs.append(str(stencil.index(inverseDir)))
+
+        code += "static const int %s [] = { %s };\n" % (INV_DIR_SYMBOL.name, ", ".join(invDirs))
+        weights = [str(w.evalf()) for w in latticeModel.weights]
+        code += "static const double %s [] = { %s };\n" % (WEIGHTS_SYMBOL.name, ",".join(weights))
+        super(LatticeModelInfo, self).__init__(code, symbolsRead=set(), symbolsDefined=symbolsDefined)
+
+
+def generateBoundaryHandling(pdfField, indexArr, latticeModel, boundaryFunctor):
+    dim = latticeModel.dim
+
+    cellLoopBody = Block([])
+    cellLoop = LoopOverCoordinate(cellLoopBody, coordinateToLoopOver=0, start=0, stop=indexArr.shape[0])
+
+    indexField = Field.createFromNumpyArray("indexField", indexArr, indexDimensions=1)
+
+    coordinateSymbols = [TypedSymbol(name, "int") for name in ['x', 'y', 'z']]
+    for d in range(dim):
+        cellLoopBody.append(SympyAssignment(coordinateSymbols[d], indexField[0](d)))
+    dirSymbol = TypedSymbol("dir", "int")
+    cellLoopBody.append(SympyAssignment(dirSymbol, indexField[0](dim)))
+
+    boundaryEqList = boundaryFunctor(pdfField, dirSymbol, latticeModel)
+    if type(boundaryEqList) is tuple:
+        boundaryEqList, additionalNodes = boundaryEqList
+    else:
+        additionalNodes = []
+
+    typeInfos = typingFromSympyInspection(boundaryEqList, pdfField.dtype)
+    fieldsRead, fieldsWritten, assignments = typeAllEquations(boundaryEqList, typeInfos)
+    fieldsAccessed = fieldsRead.union(fieldsWritten) - set([indexField])
+
+    for be in assignments:
+        cellLoopBody.append(be)
+
+    functionBody = Block([cellLoop])
+    ast = KernelFunction(functionBody, [pdfField, indexField])
+
+    if len(additionalNodes) > 0:
+        loops = ast.atoms(LoopOverCoordinate)
+        assert len(loops) == 1
+        loop = list(loops)[0]
+        for node in additionalNodes:
+            loop.body.append(node)
+
+    functionBody.insertFront(LatticeModelInfo(latticeModel))
+
+    fixedCoordinateMapping = {f.name: coordinateSymbols[:dim] for f in fieldsAccessed}
+    resolveFieldAccesses(ast, set(['indexField']), fieldToFixedCoordinates=fixedCoordinateMapping)
+    moveConstantsBeforeLoop(ast)
+    return ast

methods/conservedquantitycomputation.py

@@ -136,12 +136,13 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
         eqs += [sp.Eq(l, r) for l, r in zip(self._symbolsOrder1, firstOrderMoments)]
         return EquationCollection(eqs, [])
 
-    def outputEquationsFromPdfs(self, pdfs, outputQuantityNames):
-        outputQuantityNames = set(outputQuantityNames)
-
+    def outputEquationsFromPdfs(self, pdfs, outputQuantityNamesToSymbols):
         dim = len(self._stencil[0])
-        eqColl = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs, self._symbolOrder0,
-                                                          self._symbolsOrder1[:dim])
+
+        symbolsToExtract = set()
+
+        eqColl = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs, self._symbolOrder0, self._symbolsOrder1)
+
         if self._compressible:
             eqColl = divideFirstOrderMomentsByRho(eqColl, dim)
         else:
@@ -149,18 +150,18 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
 
         eqColl = applyForceModelShift('macroscopicVelocityShift', dim, eqColl, self._forceModel, self._compressible)
 
-        nameToSymbol = {'density': self._symbolOrder0,
-                        'velocity': self._symbolsOrder1}
-
-        symbolsToExtract = set()
-        for e in outputQuantityNames:
-            symbol = nameToSymbol[e]
-            if hasattr(symbol, "__len__"):
-                symbolsToExtract.update(symbol)
-            else:
-                symbolsToExtract.add(symbol)
-
-        return eqColl.extract(symbolsToExtract)
+        mainEquations = []
+        if 'density' in outputQuantityNamesToSymbols:
+            densityOutputSymbol = outputQuantityNamesToSymbols['density']
+            mainEquations.append(sp.Eq(densityOutputSymbol, self._symbolOrder0))
+            symbolsToExtract.add(densityOutputSymbol)
+        if 'velocity' in outputQuantityNamesToSymbols:
+            velOutputSymbols = outputQuantityNamesToSymbols['velocity']
+            mainEquations += [sp.Eq(a, b) for a, b in zip(velOutputSymbols, self._symbolsOrder1)]
+            symbolsToExtract.update(velOutputSymbols)
+
+        eqColl = eqColl.copy(mainEquations, eqColl.allEquations)
+        return eqColl.newWithoutUnusedSubexpressions()
 
     def __repr__(self):
         return "ConservedValueComputation for %s" % (", " .join(self.conservedQuantities.keys()),)

updatekernels.py

+import numpy as np
+from pystencils import Field
+from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor
+
+
+# -------------------------------------------- LBM Kernel Creation -----------------------------------------------------
+
+
+def createLBMKernel(collisionRule, inputField, outputField, accessor):
+    """
+    Replaces the pre- and post collision symbols in the collision rule by field accesses.
+
+    :param collisionRule:  instance of LbmCollisionRule, defining the collision step
+    :param inputField: field used for reading pdf values
+    :param outputField: field used for writing pdf values (may be the same as input field for certain accessors)
+    :param accessor: instance of PdfFieldAccessor, defining where to read and write values
+                     to create e.g. a fused stream-collide kernel
+    :return: LbmCollisionRule where pre- and post collision symbols have been replaced
+    """
+    method = collisionRule.method
+    preCollisionSymbols = method.preCollisionSymbols
+    postCollisionSymbols = method.postCollisionPdfSymbols
+    substitutions = {}
+
+    inputAccesses = accessor.read(inputField, method.stencil)
+    outputAccesses = accessor.write(outputField, method.stencil)
+
+    for (idx, offset), inputAccess, outputAccess in zip(enumerate(method.stencil), inputAccesses, outputAccesses):
+        substitutions[preCollisionSymbols[idx]] = inputAccess
+        substitutions[postCollisionSymbols[idx]] = outputAccess
+
+    return collisionRule.copyWithSubstitutionsApplied(substitutions)
+
+
+def createStreamPullKernel(collisionRule, numpyField=None, srcFieldName="src", dstFieldName="dst",
+                           genericLayout='numpy', genericFieldType=np.float64):
+    """
+    Implements a stream-pull scheme, where values are read from source and written to destination field
+    :param collisionRule: a collision rule created by lbm method
+    :param numpyField: optional numpy field for PDFs. Used to create a kernel of fixed loop bounds and strides
+                       if None, a generic kernel is created
+    :param srcFieldName: name of the pdf source field
+    :param dstFieldName: name of the pdf destination field
+    :param genericLayout: if no numpyField is given to determine the layout, a variable sized field with the given
+                          genericLayout is used
+    :param genericFieldType: if no numpyField is given, this data type is used for the fields
+    :return: lbm update rule, where generic pdf references are replaced by field accesses
+    """
+    dim = collisionRule.method.dim
+    if numpyField is not None:
+        assert len(numpyField.shape) == dim + 1
+
+    if numpyField is None:
+        src = Field.createGeneric(srcFieldName, dim, indexDimensions=1, layout=genericLayout, dtype=genericFieldType)
+        dst = Field.createGeneric(dstFieldName, dim, indexDimensions=1, layout=genericLayout, dtype=genericFieldType)
+    else:
+        src = Field.createFromNumpyArray(srcFieldName, numpyField, indexDimensions=1)
+        dst = Field.createFromNumpyArray(dstFieldName, numpyField, indexDimensions=1)
+
+    return createLBMKernel(collisionRule, src, dst, StreamPullTwoFieldsAccessor)
+
+
+# ------------------------------------------- Add output fields to kernel ----------------------------------------------
+
+
+def addOutputFieldForConservedQuantities(collisionRule, conservedQuantitiesToOutputFieldDict):
+    method = collisionRule.method
+    cqc = method.conservedQuantityComputation.outputEquationsFromPdfs(method.preCollisionPdfSymbols,
+                                                                      conservedQuantitiesToOutputFieldDict)
+    return collisionRule.merge(cqc)
+