diff --git a/boundaries/boundaryconditions.py b/boundaries/boundaryconditions.py
new file mode 100644
index 0000000000000000000000000000000000000000..48ab5c894bb10e96c38d5f47d48dd83d79b42359
--- /dev/null
+++ b/boundaries/boundaryconditions.py
@@ -0,0 +1,48 @@
+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))]
+
diff --git a/boundaries/boundaryhandling.py b/boundaries/boundaryhandling.py
new file mode 100644
index 0000000000000000000000000000000000000000..dba1bd2db4eed5af19674fd04ac7b7bb93d1818e
--- /dev/null
+++ b/boundaries/boundaryhandling.py
@@ -0,0 +1,166 @@
+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
diff --git a/methods/conservedquantitycomputation.py b/methods/conservedquantitycomputation.py
index 749832c10df2999d76b109270c26599f1a935d80..a65d2a46804c73f6c5e444be135db7935f94172a 100644
--- a/methods/conservedquantitycomputation.py
+++ b/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()),)
diff --git a/updatekernels.py b/updatekernels.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b8c8b6a413ccb2d3b571e94cde038eb18091542
--- /dev/null
+++ b/updatekernels.py
@@ -0,0 +1,71 @@
+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)
+