Improved lbmpy boundary setup

- easy and efficient access to fluid, boundary and link positions
  through BoundaryDataSetter object
- more code reuse for serial and parallel setups

boundaries/boundaryhandling.py

@@ -9,20 +9,11 @@ class BoundaryHandling(PeriodicityHandling, GenericBoundaryHandling):  # importa
         shapeWithGl = [a + 2 * ghostLayers for a in domainShape]
         self.domainShape = domainShape
         self.domainShapeWithGhostLayers = shapeWithGl
-        flagFieldInterface = NumpyFlagFieldInterface(shapeWithGl, flagDtype)
+        flagInterface = NumpyFlagFieldInterface(shapeWithGl, flagDtype)
 
-        GenericBoundaryHandling.__init__(self, flagFieldInterface, pdfField, lbMethod, ghostLayers, target, openMP)
+        GenericBoundaryHandling.__init__(self, flagInterface, pdfField, lbMethod, None, ghostLayers, target, openMP)
         PeriodicityHandling.__init__(self, list(domainShape) + [len(lbMethod.stencil)])
 
-    def setBoundary(self, boundaryObject, indexExpr=None, maskCallback=None, sliceNormalizationGhostLayers=1):
-        mask = None
-        if maskCallback is not None:
-            gridParams = [np.arange(start=-sliceNormalizationGhostLayers, stop=s - sliceNormalizationGhostLayers) + 0.5
-                          for s in self.flagField.shape]
-            indices = np.meshgrid(*gridParams, indexing='ij')
-            mask = maskCallback(*indices)
-        return GenericBoundaryHandling.setBoundary(self, boundaryObject, indexExpr, mask)
-
     def __call__(self, *args, **kwargs):
         for cls in BoundaryHandling.__bases__:
             cls.__call__(self, *args, **kwargs)

boundaries/handlinginterface.py

@@ -4,6 +4,7 @@ from pystencils.field import Field
 from pystencils.slicing import normalizeSlice
 from lbmpy.boundaries.boundary_kernel import generateIndexBoundaryKernel
 from lbmpy.boundaries.createindexlist import createBoundaryIndexList
+from pystencils.cache import memorycache
 
 
 class FlagFieldInterface(object):
@@ -37,20 +38,63 @@ class FlagFieldInterface(object):
         return name
 
 
+class BoundaryDataSetter:
+
+    def __init__(self, indexArray, offset, stencil, ghostLayers):
+        self.indexArray = indexArray
+        self.offset = offset
+        self.stencil = np.array(stencil)
+        arrFieldNames = indexArray.dtype.names
+        self.dim = 3 if 'z' in arrFieldNames else 2
+        assert 'x' in arrFieldNames and 'y' in arrFieldNames and 'dir' in arrFieldNames, str(arrFieldNames)
+        self.boundaryDataNames = set(self.indexArray.dtype.names) - set(['x', 'y', 'z', 'dir'])
+        self.coordMap = {0: 'x', 1: 'y', 2: 'z'}
+        self.ghostLayers = ghostLayers
+
+    def fluidCellPositions(self, coord):
+        assert coord < self.dim
+        return self.indexArray[self.coordMap[coord]] + self.offset[coord] - self.ghostLayers
+
+    @memorycache()
+    def linkOffsets(self):
+        return self.stencil[self.indexArray['dir']]
+
+    @memorycache()
+    def linkPositions(self, coord):
+        return self.fluidCellPositions(coord) + 0.5 * self.linkOffsets()[:, coord]
+
+    @memorycache()
+    def boundaryCellPositions(self, coord):
+        return self.fluidCellPositions(coord) + self.linkOffsets()[:, coord]
+
+    def __setitem__(self, key, value):
+        if key not in self.boundaryDataNames:
+            raise KeyError("Invalid boundary data name %s. Allowed are %s" % (key, self.boundaryDataNames))
+        self.indexArray[key] = value
+
+    def __getitem__(self, item):
+        if item not in self.boundaryDataNames:
+            raise KeyError("Invalid boundary data name %s. Allowed are %s" % (item, self.boundaryDataNames))
+        return self.indexArray[item]
+
+
 class GenericBoundaryHandling(object):
 
     class BoundaryInfo(object):
-        def __init__(self, kernel, ast, indexArray, idxArrayForExecution):
+        def __init__(self, kernel, ast, indexArray, idxArrayForExecution, boundaryDataSetter):
             self.kernel = kernel
             self.ast = ast
             self.indexArray = indexArray
             self.idxArrayForExecution = idxArrayForExecution  # is different for GPU kernels
+            self.boundaryDataSetter = boundaryDataSetter
 
-    def __init__(self, flagFieldInterface, pdfField, lbMethod, ghostLayers=1, target='cpu', openMP=True):
+    def __init__(self, flagFieldInterface, pdfField, lbMethod, offset=None, ghostLayers=1, target='cpu', openMP=True):
         """
         :param flagFieldInterface: implementation of FlagFieldInterface
         :param pdfField: numpy array
         :param lbMethod:
+        :param offset: offset that is added to all coordinates when calling callback functions to set up geometry
+                       or boundary data. This is used for waLBerla simulations to pass the block offset  
         :param target: 'cpu' or 'gpu'
         :param openMP:
         """
@@ -61,7 +105,7 @@ class GenericBoundaryHandling(object):
         self._openMP = openMP
         self.ghostLayers = ghostLayers
         self._dirty = False
-
+        self.offset = offset if offset else (0,) * lbMethod.dim
         self._boundaryInfos = {}  # mapping of boundary object to boundary info
         self._symbolicPdfField = Field.createFromNumpyArray('pdfs', pdfField, indexDimensions=1)
 
@@ -98,33 +142,17 @@ class GenericBoundaryHandling(object):
         self._dirty = False
         self._boundaryInfos = {}
 
-    def _invalidateCache(self):
-        self._dirty = True
-
     def triggerReinitializationOfBoundaryData(self, **kwargs):
         if self._dirty:
             self.prepare()
             return
         else:
             for boundaryObject, boundaryInfo in self._boundaryInfos.items():
-                self.__boundaryDataInitialization(boundaryInfo.indexArray, boundaryObject, **kwargs)
+                self.__boundaryDataInitialization(boundaryInfo, boundaryObject, **kwargs)
                 if self._target == 'gpu':
                     import pycuda.gpuarray as gpuarray
                     boundaryInfo.idxArrayForExecution = gpuarray.to_gpu(boundaryInfo.indexArray)
 
-    def __boundaryDataInitialization(self, idxArray, boundaryObject, **kwargs):
-        if boundaryObject.additionalDataInitCallback:
-            #TODO x,y,z coordinates should be transformed here
-            boundaryObject.additionalDataInitCallback(idxArray, **kwargs)
-
-        if boundaryObject.additionalDataInitKernelEquations:
-            initKernelAst = generateIndexBoundaryKernel(self._symbolicPdfField, idxArray, self._lbMethod,
-                                                        boundaryObject, target='cpu',
-                                                        createInitializationKernel=True)
-            from pystencils.cpu import makePythonFunction as makePythonCpuFunction
-            initKernel = makePythonCpuFunction(initKernelAst, {'indexField': idxArray, 'pdfs': self._pdfField})
-            initKernel()
-
     def prepare(self):
         """Compiles boundary kernels according to flag field. When setting boundaries the cache is automatically
         invalidated, so there is no need to call this function manually, as long as the flag field is not manually
@@ -152,7 +180,6 @@ class GenericBoundaryHandling(object):
                     extendedIdxField[prop] = idxArray[prop]
 
                 idxArray = extendedIdxField
-                self.__boundaryDataInitialization(idxArray, boundaryObject)
 
             ast = generateIndexBoundaryKernel(self._symbolicPdfField, idxArray, self._lbMethod, boundaryObject,
                                               target=self._target)
@@ -170,15 +197,30 @@ class GenericBoundaryHandling(object):
             else:
                 assert False
 
-            boundaryInfo = GenericBoundaryHandling.BoundaryInfo(kernel, ast, idxArray, idxArrayForExecution)
+            boundaryDataSetter = BoundaryDataSetter(idxArray, self.offset, self._lbMethod.stencil, self.ghostLayers)
+            boundaryInfo = GenericBoundaryHandling.BoundaryInfo(kernel, ast, idxArray,
+                                                                idxArrayForExecution, boundaryDataSetter)
             self._boundaryInfos[boundaryObject] = boundaryInfo
 
+            if boundaryObject.additionalData:
+                self.__boundaryDataInitialization(boundaryInfo, boundaryObject)
+
         self._dirty = False
 
     def reserveFlag(self, boundaryObject):
         self._flagFieldInterface.getFlag(boundaryObject)
 
-    def setBoundary(self, boundaryObject, indexExpr=None, maskArr=None):
+    def setBoundary(self, boundaryObject, indexExpr=None, maskCallback=None):
+        mask = None
+        flagField = self._flagFieldInterface.array
+        if maskCallback is not None:
+            gridParams = [offset + np.arange(-self.ghostLayers, s - self.ghostLayers) + 0.5
+                          for s, offset in zip(flagField.shape, self.offset)]
+            indices = np.meshgrid(*gridParams, indexing='ij')
+            mask = maskCallback(*indices)
+        return self.setBoundaryWithMaskArray(boundaryObject, indexExpr, mask)
+
+    def setBoundaryWithMaskArray(self, boundaryObject, indexExpr=None, maskArr=None):
         """
         Sets boundary in a rectangular region (slice)
 
@@ -207,3 +249,19 @@ class GenericBoundaryHandling(object):
     @property
     def flagField(self):
         return self._flagFieldInterface.array
+
+    def _invalidateCache(self):
+        self._dirty = True
+
+    def __boundaryDataInitialization(self, boundaryInfo, boundaryObject, **kwargs):
+        if boundaryObject.additionalDataInitCallback:
+            boundaryObject.additionalDataInitCallback(boundaryInfo.boundaryDataSetter, **kwargs)
+
+        if boundaryObject.additionalDataInitKernelEquations:
+            initKernelAst = generateIndexBoundaryKernel(self._symbolicPdfField, boundaryInfo.indexArray, self._lbMethod,
+                                                        boundaryObject, target='cpu',
+                                                        createInitializationKernel=True)
+            from pystencils.cpu import makePythonFunction as makePythonCpuFunction
+            initKernel = makePythonCpuFunction(initKernelAst, {'indexField': boundaryInfo.indexArray,
+                                                               'pdfs': self._pdfField})
+            initKernel()
\ No newline at end of file

chapman_enskog/steady_state.py

@@ -64,7 +64,7 @@ class SteadyStateChapmanEnskogAnalysis(object):
         chapmanEnskogHierarchy = [chapmanEnskogHierarchy[i] for i in range(self.order + 1)]
 
         insertedHierarchy = []
-        rawHierarchy = []j
+        rawHierarchy = []
         substitutionDict = {}
         for ceEq, f_i in zip(chapmanEnskogHierarchy, self.fSyms):
             newEq = -1 / self.collisionOpSym * (ceEq - self.collisionOpSym * f_i)

parallel/blockiteration.py

@@ -3,14 +3,28 @@ import waLBerla as wlb
 from pystencils.slicing import normalizeSlice
 
 
-def slicedBlockIteration(blocks, indexExpr=None, ghostLayers=1, sliceNormalizationGhostLayers=1):
-    if indexExpr is None:
-        indexExpr = [slice(None, None, None)] * 3
+def slicedBlockIteration(blocks, sliceObj=None, ghostLayers=1, sliceNormalizationGhostLayers=1,
+                         withIndexArrays=False):
+    """
+    Iterates of all blocks that have an intersection with the given slice object.
+    For these blocks a slice object in local block coordinates is given
+    
+    :param blocks: waLBerla block data structure
+    :param sliceObj: a slice (i.e. rectangular subregion), can be created with makeSlice[]
+    :param ghostLayers: how many ghost layers are included in the local slice and the optional index arrays
+    :param sliceNormalizationGhostLayers: slices can have relative coordinates e.g. makeSlice[0.2, :, :]
+                                          when computing absolute values, the domain size is needed. This parameter 
+                                          specifies how many ghost layers are taken into account for this operation.
+    :param withIndexArrays: if true index arrays [x,y,z] are yielded as last parameters. These arrays contain the
+                            cell midpoints in global coordinates
+    """
+    if sliceObj is None:
+        sliceObj = [sliceObj(None, None, None)] * 3
 
     domainCellBB = blocks.getDomainCellBB()
     domainExtent = [s + 2 * sliceNormalizationGhostLayers for s in domainCellBB.size]
-    indexExpr = normalizeSlice(indexExpr, domainExtent)
-    targetCellBB = wlb.CellInterval.fromSlice(indexExpr)
+    sliceObj = normalizeSlice(sliceObj, domainExtent)
+    targetCellBB = wlb.CellInterval.fromSlice(sliceObj)
     targetCellBB.shift(*[a - sliceNormalizationGhostLayers for a in domainCellBB.min])
 
     for block in blocks:
@@ -19,12 +33,15 @@ def slicedBlockIteration(blocks, indexExpr=None, ghostLayers=1, sliceNormalizati
         if intersection.empty():
             continue
 
-        meshGridParams = [offset + 0.5 + np.arange(width)
-                          for offset, width in zip(intersection.min, intersection.size)]
-
-        indexArrays = np.meshgrid(*meshGridParams, indexing='ij')
-
         localTargetBB = blocks.transformGlobalToLocal(block, intersection)
         localTargetBB.shift(ghostLayers, ghostLayers, ghostLayers)
         localSlice = localTargetBB.toSlice()
-        yield block, indexArrays, localSlice
+
+        if withIndexArrays:
+            meshGridParams = [offset + 0.5 + np.arange(width)
+                              for offset, width in zip(intersection.min, intersection.size)]
+            indexArrays = np.meshgrid(*meshGridParams, indexing='ij')
+            yield block, localSlice, indexArrays
+        else:
+            yield block, localSlice
+

parallel/boundaryhandling.py

@@ -10,8 +10,9 @@ class BoundaryHandling(object):
             flagFieldInterface = WalberlaFlagFieldInterface(block, flagFieldId)
             pdfField = wlb.field.toArray(block[pdfFieldId], withGhostLayers=True)
             ghostLayers = block[pdfFieldId].nrOfGhostLayers
-            return GenericBoundaryHandling(flagFieldInterface, pdfField, lbMethod, ghostLayers=ghostLayers,
-                                           target=target, openMP=openMP)
+            offset = blocks.getBlockCellBB(block).min
+            return GenericBoundaryHandling(flagFieldInterface, pdfField, lbMethod, offset=offset,
+                                           ghostLayers=ghostLayers, target=target, openMP=openMP)
 
         self._boundaryId = boundaryId
         self._pdfFieldId = pdfFieldId
@@ -41,13 +42,11 @@ class BoundaryHandling(object):
             return
 
         gl = self._blocks[0][self._boundaryId].ghostLayers
-        ngl = sliceNormalizationGhostLayers
         for block in self._blocks:
             block[self._boundaryId].reserveFlag(boundaryObject)
 
-        for block, indexArrs, localSlice in slicedBlockIteration(self._blocks, indexExpr, gl, ngl):
-            mask = maskCallback(*indexArrs) if maskCallback else None
-            block[self._boundaryId].setBoundary(boundaryObject, indexExpr=localSlice, maskArr=mask)
+        for block, localSlice in slicedBlockIteration(self._blocks, indexExpr, gl, sliceNormalizationGhostLayers):
+            block[self._boundaryId].setBoundary(boundaryObject, indexExpr=localSlice, maskCallback=maskCallback)
 
 
 # ----------------------------------------------------------------------------------------------------------------------