DataHandling for pystencils

parallel/blockiteration.py

-import numpy as np
-import waLBerla as wlb
-from pystencils.slicing import normalizeSlice
-
-
-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]
-    sliceObj = normalizeSlice(sliceObj, domainExtent)
-    targetCellBB = wlb.CellInterval.fromSlice(sliceObj)
-    targetCellBB.shift(*[a - sliceNormalizationGhostLayers for a in domainCellBB.min])
-
-    for block in blocks:
-        intersection = blocks.getBlockCellBB(block).getExpanded(ghostLayers)
-        intersection.intersect(targetCellBB)
-        if intersection.empty():
-            continue
-
-        localTargetBB = blocks.transformGlobalToLocal(block, intersection)
-        localTargetBB.shift(ghostLayers, ghostLayers, ghostLayers)
-        localSlice = localTargetBB.toSlice()
-
-        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

 import numpy as np
 import waLBerla as wlb
 from lbmpy.boundaries.handlinginterface import GenericBoundaryHandling, FlagFieldInterface
-from lbmpy.parallel.blockiteration import slicedBlockIteration
+from pystencils.parallel.blockiteration import slicedBlockIteration
 
 
 class BoundaryHandling(object):
@@ -61,8 +61,8 @@ class BoundaryHandling(object):
         for block in self._blocks:
             block[self._boundaryId].reserveFlag(boundaryObject)
 
-        for block, localSlice in slicedBlockIteration(self._blocks, indexExpr, gl, sliceNormalizationGhostLayers):
-            block[self._boundaryId].setBoundary(boundaryObject, indexExpr=localSlice, maskCallback=maskCallback)
+        for iter in slicedBlockIteration(self._blocks, indexExpr, gl, sliceNormalizationGhostLayers):
+            block[self._boundaryId].setBoundary(boundaryObject, indexExpr=iter.localSlice, maskCallback=maskCallback)
 
 
 # ----------------------------------------------------------------------------------------------------------------------

parallel/scenarios.py

@@ -5,7 +5,7 @@ from lbmpy.creationfunctions import createLatticeBoltzmannFunction, updateWithDe
     switchToSymbolicRelaxationRatesForEntropicMethods
 from lbmpy.macroscopic_value_kernels import compileMacroscopicValuesGetter, compileMacroscopicValuesSetter
 from lbmpy.parallel.boundaryhandling import BoundaryHandling
-from lbmpy.parallel.blockiteration import slicedBlockIteration
+from pystencils.parallel.blockiteration import slicedBlockIteration
 from lbmpy.boundaries import NoSlip, UBB
 from pystencils.slicing import sliceFromDirection
 
@@ -170,13 +170,13 @@ class Scenario(object):
         else:
             raise ValueError("density has to be a number or a callback function")
 
-        for block, (x, y, z), localSlice in slicedBlockIteration(self.blocks, indexExpr, 1, 1):
+        for it in slicedBlockIteration(self.blocks, indexExpr, 1, 1):
             if density:
-                densityArr = wlb.field.toArray(block[self.densityFieldId], withGhostLayers=True)
-                densityArr[localSlice] = densityCallback(x, y, z) if densityCallback else densityValue
+                densityArr = wlb.field.toArray(it.block[self.densityFieldId], withGhostLayers=True)
+                densityArr[it.localSlice] = densityCallback(*it.midpointArrays) if densityCallback else densityValue
             if velocity:
-                velArr = wlb.field.toArray(block[self.velocityFieldId], withGhostLayers=True)
-                velArr[localSlice, :] = velocityCallback(x, y, z) if velocityCallback else velocityValue
+                velArr = wlb.field.toArray(it.block[self.velocityFieldId], withGhostLayers=True)
+                velArr[it.localSlice, :] = velocityCallback(*it.midpointArrays) if velocityCallback else velocityValue
 
         self._setMacroscpicValues()