Data handling - flag array support

datahandling.py

 import numpy as np
 from abc import ABC, abstractmethod, abstractproperty
+
+from collections import defaultdict
+from contextlib import contextmanager
+
+from lbmpy.boundaries.periodicityhandling import PeriodicityHandling
 from pystencils import Field, makeSlice
 from pystencils.parallel.blockiteration import BlockIterationInfo
 from pystencils.slicing import normalizeSlice, removeGhostLayers
@@ -11,6 +16,61 @@ except ImportError:
     gpuarray = None
 
 
+class WalberlaFlagInterface:
+    def __init__(self, flagField):
+        self.flagField = flagField
+
+    def registerFlag(self, flagName):
+        return self.flagField.registerFlag(flagName)
+
+    def flag(self, flagName):
+        return self.flagField.flag(flagName)
+
+    def flagName(self, flag):
+        return self.flagField.flagName(flag)
+
+    @property
+    def flags(self):
+        return self.flagField.flags
+
+
+class PythonFlagInterface:
+    def __init__(self):
+        self.nameToFlag = {}
+        self.flagToName = {}
+        self.nextFreeBit = 0
+
+    def registerFlag(self, flagName):
+        assert flagName not in self.nameToFlag
+        flag = 1 << self.nextFreeBit
+        self.nextFreeBit += 1
+        self.flagToName[flag] = flagName
+        self.nameToFlag[flagName] = flag
+        return flag
+
+    def flag(self, flagName):
+        return self.nameToFlag[flagName]
+
+    def flagName(self, flag):
+        return self.flagToName[flag]
+
+    @property
+    def flags(self):
+        return tuple(self.nameToFlag.keys())
+
+
+class FlagArray(np.ndarray):
+    def __new__(cls, inputArray, flagInterface):
+        obj = np.asarray(inputArray).view(cls)
+        obj.flagInterface = flagInterface
+        assert inputArray.dtype.kind in ('u', 'i'), "FlagArrays can only be created from integer arrays"
+        return obj
+
+    def __array_finalize__(self, obj):
+        if obj is None: return
+        self.flagInterface = getattr(obj, 'flagInterface', None)
+
+
 class DataHandling(ABC):
     """
     Manages the storage of arrays and maps them to a symbolic field.
@@ -21,13 +81,16 @@ class DataHandling(ABC):
     'gather' function that has collects (parts of the) distributed data on a single process.
     """
 
+    def __init__(self):
+        self._preAccessFunctions = defaultdict(list)
+        self._postAccessFunctions = defaultdict(list)
+
     # ---------------------------- Adding and accessing data -----------------------------------------------------------
 
     @property
     @abstractmethod
     def dim(self):
         """Dimension of the domain, either 2 or 3"""
-        pass
 
     @abstractmethod
     def add(self, name, fSize=1, dtype=np.float64, latexName=None, ghostLayers=None, layout=None, cpu=True, gpu=False):
@@ -48,7 +111,6 @@ class DataHandling(ABC):
         :param cpu: allocate field on the CPU
         :param gpu: allocate field on the GPU
         """
-        pass
 
     @abstractmethod
     def addLike(self, name, nameOfTemplateField, latexName=None, cpu=True, gpu=False):
@@ -60,13 +122,18 @@ class DataHandling(ABC):
         :param cpu: see 'add' method
         :param gpu: see 'add' method
         """
-        pass
+
+    def addFlagArray(self, name, dtype=np.int32, latexName=None, ghostLayers=None):
+        """
+        Adds a flag array (of integer type) where each bit is interpreted as a boolean
+        Flag arrays additionally store a mapping of name to bit nr, which is accessible as arr.flagInterface.
+        For parameter documentation see 'add()' function.
+        """
 
     @property
     @abstractmethod
     def fields(self):
         """Dictionary mapping data name to symbolic pystencils field - use this to create pystencils kernels"""
-        pass
 
     @abstractmethod
     def access(self, name, sliceObj=None, innerGhostLayers=None, outerGhostLayers=0):
@@ -79,7 +146,6 @@ class DataHandling(ABC):
         :param outerGhostLayers: how many ghost layers at the domain border to include
         Yields a numpy array with local part of data and a BlockIterationInfo object containing geometric information
         """
-        pass
 
     @abstractmethod
     def gather(self, name, sliceObj=None, allGather=False):
@@ -92,8 +158,20 @@ class DataHandling(ABC):
         :param allGather: if False only the root process receives the result, if True all processes
         :return: generator expression yielding the gathered field, the gathered field does not include any ghost layers
         """
-        pass
 
+    def registerPreAccessFunction(self, name, function):
+        self._preAccessFunctions[name].append(function)
+
+    def registerPostAccessFunction(self, name, function):
+        self._postAccessFunctions[name].append(function)
+
+    @contextmanager
+    def accessWrapper(self, name):
+        for func in self._preAccessFunctions[name]:
+            func()
+        yield
+        for func in self._postAccessFunctions[name]:
+            func()
     # ------------------------------- CPU/GPU transfer -----------------------------------------------------------------
 
     @abstractmethod
@@ -136,12 +214,19 @@ class SerialDataHandling(DataHandling):
         :param defaultGhostLayers: nr of ghost layers used if not specified in add() method
         :param defaultLayout: layout used if no layout is given to add() method
         """
+        super(SerialDataHandling, self).__init__()
         self._domainSize = tuple(domainSize)
         self.defaultGhostLayers = defaultGhostLayers
         self.defaultLayout = defaultLayout
         self._fields = DotDict()
         self.cpuArrays = DotDict()
         self.gpuArrays = DotDict()
+        #if periodicity is None or periodicity is False:
+        #    periodicity = [False] * self.dim
+        #if periodicity is True:
+        #    periodicity = [True] * self.dim
+        #
+        #self._periodicity = periodicity
         self._fieldInformation = {}
 
     @property
@@ -193,27 +278,36 @@ class SerialDataHandling(DataHandling):
         self.fields[name] = Field.createFixedSize(latexName, shape=kwargs['shape'], indexDimensions=indexDimensions,
                                                   dtype=kwargs['dtype'], layout=kwargs['order'])
 
+    def addFlagArray(self, name, dtype=np.int32, latexName=None, ghostLayers=None):
+        self.add(name, 1, dtype, latexName, ghostLayers, layout='AoS', cpu=True, gpu=False)
+        self.cpuArrays[name] = FlagArray(self.cpuArrays[name], PythonFlagInterface())
+
     def addLike(self, name, nameOfTemplateField, latexName=None, cpu=True, gpu=False):
+        if hasattr(self.fields[nameOfTemplateField], 'flagInterface'):
+            raise ValueError("addLike() does not work for flag arrays")
         self.add(name,latexName=latexName, cpu=cpu, gpu=gpu, **self._fieldInformation[nameOfTemplateField])
 
     def access(self, name, sliceObj=None, outerGhostLayers=0, **kwargs):
         if sliceObj is None:
             sliceObj = [slice(None, None)] * self.dim
-        arr = self.cpuArrays[name]
-        glToRemove = self._fieldInformation[name]['ghostLayers'] - outerGhostLayers
-        assert glToRemove >= 0
-        arr = removeGhostLayers(arr, indexDimensions=self.fields[name].indexDimensions, ghostLayers=glToRemove)
-        sliceObj = normalizeSlice(sliceObj, arr.shape[:self.dim])
-        yield arr[sliceObj], BlockIterationInfo(None, tuple(s.start for s in sliceObj), sliceObj)
 
-    def gather(self, name, sliceObj=None, **kwargs):
-        gls = self._fieldInformation[name]['ghostLayers']
-        arr = self.cpuArrays[name]
-        arr = removeGhostLayers(arr, indexDimensions=self.fields[name].indexDimensions, ghostLayers=gls)
+        with self.accessWrapper(name):
+            arr = self.cpuArrays[name]
+            glToRemove = self._fieldInformation[name]['ghostLayers'] - outerGhostLayers
+            assert glToRemove >= 0
+            arr = removeGhostLayers(arr, indexDimensions=self.fields[name].indexDimensions, ghostLayers=glToRemove)
+            sliceObj = normalizeSlice(sliceObj, arr.shape[:self.dim])
+            yield arr[sliceObj], BlockIterationInfo(None, tuple(s.start for s in sliceObj), sliceObj)
 
-        if sliceObj is not None:
-            arr = arr[sliceObj]
-        yield arr
+    def gather(self, name, sliceObj=None, **kwargs):
+        with self.accessWrapper(name):
+            gls = self._fieldInformation[name]['ghostLayers']
+            arr = self.cpuArrays[name]
+            arr = removeGhostLayers(arr, indexDimensions=self.fields[name].indexDimensions, ghostLayers=gls)
+
+            if sliceObj is not None:
+                arr = arr[sliceObj]
+            yield arr
 
     def swap(self, name1, name2, gpu=False):
         if not gpu:

field.py

@@ -218,6 +218,10 @@ class Field(object):
         offsetList[coordId] = offset
         return Field.Access(self, tuple(offsetList))
 
+    def neighbors(self, stencil):
+        return [ self.__getitem__(dir) for s in stencil]
+
+    @property
     def center(self):
         center = tuple([0] * self.spatialDimensions)
         return Field.Access(self, center)

finitedifferences.py

@@ -174,7 +174,7 @@ class Advection(sp.Function):
 def advection(scalar, vector, idx=None):
     assert isinstance(scalar, Field), "Advected scalar has to be a pystencils.Field"
 
-    args = [scalar.center(), vector if not isinstance(vector, Field) else vector.center()]
+    args = [scalar.center, vector if not isinstance(vector, Field) else vector.center]
     if idx is not None:
         args.append(idx)
     return Advection(*args)
@@ -210,7 +210,7 @@ class Diffusion(sp.Function):
 
 def diffusion(scalar, diffusionCoeff, idx=None):
     assert isinstance(scalar, Field), "Advected scalar has to be a pystencils.Field"
-    args = [scalar.center(), diffusionCoeff if not isinstance(diffusionCoeff, Field) else diffusionCoeff.center()]
+    args = [scalar.center, diffusionCoeff if not isinstance(diffusionCoeff, Field) else diffusionCoeff.center]
     if idx is not None:
         args.append(idx)
     return Diffusion(*args)
@@ -231,7 +231,7 @@ class Transient(sp.Function):
 
 
 def transient(scalar, idx=None):
-    args = [scalar.center()]
+    args = [scalar.center]
     if idx is not None:
         args.append(idx)
     return Transient(*args)
@@ -249,12 +249,12 @@ class Discretization2ndOrder:
             if isinstance(expr.diffusionCoeff, Field):
                 firstDiffs = [offset *
                               (scalar.neighbor(c, offset) * expr.diffusionCoeff.neighbor(c, offset) -
-                               scalar.center() * expr.diffusionCoeff.center())
+                               scalar.center * expr.diffusionCoeff.center())
                               for offset in [-1, 1]]
             else:
                 firstDiffs = [offset *
                               (scalar.neighbor(c, offset) * expr.diffusionCoeff -
-                               scalar.center() * expr.diffusionCoeff)
+                               scalar.center * expr.diffusionCoeff)
                               for offset in [-1, 1]]
             result += firstDiffs[1] - firstDiffs[0]
         return result / (self.dx**2)

parallel/datahandling.py

 import numpy as np
 from pystencils import Field, makeSlice
-from pystencils.datahandling import DataHandling
+from pystencils.datahandling import DataHandling, FlagArray, WalberlaFlagInterface
 from pystencils.parallel.blockiteration import slicedBlockIteration
 from pystencils.utils import DotDict
 import waLBerla as wlb
@@ -20,6 +20,7 @@ class ParallelDataHandling(DataHandling):
                     waLBerla always uses three dimensions, so if dim=2 the extend of the
                     z coordinate of blocks has to be 1
         """
+        super(ParallelDataHandling, self).__init__()
         assert dim in (2, 3)
         self.blocks = blocks
         self.defaultGhostLayers = defaultGhostLayers
@@ -41,6 +42,68 @@ class ParallelDataHandling(DataHandling):
         return self._fields
 
     def add(self, name, fSize=1, dtype=np.float64, latexName=None, ghostLayers=None, layout=None, cpu=True, gpu=False):
+        return self._add(name, fSize, dtype, latexName, ghostLayers, layout, cpu, gpu, flagField=False)
+
+    def addLike(self, name, nameOfTemplateField, latexName=None, cpu=True, gpu=False):
+        assert not self._fieldInformation[nameOfTemplateField]['flagField']
+        self._add(name,latexName=latexName, cpu=cpu, gpu=gpu, **self._fieldInformation[nameOfTemplateField])
+
+    def addFlagArray(self, name, dtype=np.int32, latexName=None, ghostLayers=None):
+        return self._add(name, dtype=dtype, latexName=latexName, ghostLayers=ghostLayers, flagField=True)
+
+    def swap(self, name1, name2, gpu=False):
+        if gpu:
+            name1 = self.GPU_DATA_PREFIX + name1
+            name2 = self.GPU_DATA_PREFIX + name2
+        for block in self.blocks:
+            block[name1].swapDataPointers(block[name2])
+
+    def access(self, name, sliceObj=None, innerGhostLayers=None, outerGhostLayers=0):
+        fieldInfo = self._fieldInformation[name]
+        with self.accessWrapper(name):
+            if innerGhostLayers is None:
+                innerGhostLayers = fieldInfo['ghostLayers']
+
+            if outerGhostLayers is None:
+                outerGhostLayers = fieldInfo['ghostLayers']
+
+            for iterInfo in slicedBlockIteration(self.blocks, sliceObj, innerGhostLayers, outerGhostLayers):
+                arr = wlb.field.toArray(iterInfo.block[name], withGhostLayers=innerGhostLayers)[iterInfo.localSlice]
+                if fieldInfo['flagField']:
+                    arr = FlagArray(arr, WalberlaFlagInterface(iterInfo.block[name]))
+                if self.fields[name].indexDimensions == 0:
+                    arr = arr[..., 0]
+                if self.dim == 2:
+                    arr = arr[:, :, 0]
+                yield arr, iterInfo
+
+    def gather(self, name, sliceObj=None, allGather=False):
+        with self.accessWrapper(name):
+            if sliceObj is None:
+                sliceObj = makeSlice[:, :, :]
+            for array in wlb.field.gatherGenerator(self.blocks, name, sliceObj, allGather):
+                if self.fields[name].indexDimensions == 0:
+                    array = array[..., 0]
+                if self.dim == 2:
+                    array = array[:, :, 0]
+                yield array
+
+    def toCpu(self, name):
+        wlb.cuda.copyFieldToCpu(self.blocks, self.GPU_DATA_PREFIX + name, name)
+
+    def toGpu(self, name):
+        wlb.cuda.copyFieldToGpu(self.blocks, self.GPU_DATA_PREFIX + name, name)
+
+    def allToCpu(self):
+        for cpuName, gpuName in self._cpuGpuPairs:
+            wlb.cuda.copyFieldToCpu(self.blocks, gpuName, cpuName)
+
+    def allToGpu(self):
+        for cpuName, gpuName in self._cpuGpuPairs:
+            wlb.cuda.copyFieldToGpu(self.blocks, gpuName, cpuName)
+
+    def _add(self, name, fSize=1, dtype=np.float64, latexName=None, ghostLayers=None, layout=None,
+            cpu=True, gpu=False, flagField=False):
         if ghostLayers is None:
             ghostLayers = self.defaultGhostLayers
         if layout is None:
@@ -57,19 +120,27 @@ class ParallelDataHandling(DataHandling):
         self._fieldInformation[name] = {'ghostLayers': ghostLayers,
                                         'fSize': fSize,
                                         'layout': layout,
-                                        'dtype': dtype}
+                                        'dtype': dtype,
+                                        'flagField': flagField}
 
         layoutMap = {'fzyx': wlb.field.Layout.fzyx, 'zyxf': wlb.field.Layout.zyxf,
                      'SoA': wlb.field.Layout.fzyx,  'AoS': wlb.field.Layout.zyxf}
-        if cpu:
-            wlb.field.addToStorage(self.blocks, name, dtype, fSize=fSize, layout=layoutMap[layout],
-                                   ghostLayers=ghostLayers)
-        if gpu:
-            wlb.cuda.addGpuFieldToStorage(self.blocks, self.GPU_DATA_PREFIX+name, dtype, fSize=fSize,
-                                          usePitchedMem=False, ghostLayers=ghostLayers, layout=layoutMap[layout])
 
-        if cpu and gpu:
-            self._cpuGpuPairs.append((name, self.GPU_DATA_PREFIX + name))
+        if flagField:
+            assert not gpu
+            assert np.dtype(dtype).kind in ('u', 'i'), "FlagArrays can only be created from integer arrays"
+            nrOfBits = np.dtype(dtype).itemsize * 8
+            wlb.field.addFlagFieldToStorage(self.blocks, name, nrOfBits, ghostLayers)
+        else:
+            if cpu:
+                wlb.field.addToStorage(self.blocks, name, dtype, fSize=fSize, layout=layoutMap[layout],
+                                       ghostLayers=ghostLayers)
+            if gpu:
+                wlb.cuda.addGpuFieldToStorage(self.blocks, self.GPU_DATA_PREFIX+name, dtype, fSize=fSize,
+                                              usePitchedMem=False, ghostLayers=ghostLayers, layout=layoutMap[layout])
+
+            if cpu and gpu:
+                self._cpuGpuPairs.append((name, self.GPU_DATA_PREFIX + name))
 
         blockBB = self.blocks.getBlockCellBB(self.blocks[0])
         shape = tuple(s + 2 * ghostLayers for s in blockBB.size)
@@ -79,52 +150,3 @@ class ParallelDataHandling(DataHandling):
 
         assert all(f.name != latexName for f in self.fields.values()), "Symbolic field with this name already exists"
         self.fields[name] = Field.createFixedSize(latexName, shape, indexDimensions, dtype, layout)
-
-    def addLike(self, name, nameOfTemplateField, latexName=None, cpu=True, gpu=False):
-        self.add(name,latexName=latexName, cpu=cpu, gpu=gpu, **self._fieldInformation[nameOfTemplateField])
-
-    def swap(self, name1, name2, gpu=False):
-        if gpu:
-            name1 = self.GPU_DATA_PREFIX + name1
-            name2 = self.GPU_DATA_PREFIX + name2
-        for block in self.blocks:
-            block[name1].swapDataPointers(block[name2])
-
-    def access(self, name, sliceObj=None, innerGhostLayers=None, outerGhostLayers=0):
-        if innerGhostLayers is None:
-            innerGhostLayers = self._fieldInformation[name]['ghostLayers']
-
-        if outerGhostLayers is None:
-            outerGhostLayers = self._fieldInformation[name]['ghostLayers']
-
-        for iterInfo in slicedBlockIteration(self.blocks, sliceObj, innerGhostLayers, outerGhostLayers):
-            arr = wlb.field.toArray(iterInfo.block[name], withGhostLayers=innerGhostLayers)[iterInfo.localSlice]
-            if self.fields[name].indexDimensions == 0:
-                arr = arr[..., 0]
-            if self.dim == 2:
-                arr = arr[:, :, 0]
-            yield arr, iterInfo
-
-    def gather(self, name, sliceObj=None, allGather=False):
-        if sliceObj is None:
-            sliceObj = makeSlice[:, :, :]
-        for array in wlb.field.gatherGenerator(self.blocks, name, sliceObj, allGather):
-            if self.fields[name].indexDimensions == 0:
-                array = array[..., 0]
-            if self.dim == 2:
-                array = array[:, :, 0]
-            yield array
-
-    def toCpu(self, name):
-        wlb.cuda.copyFieldToCpu(self.blocks, self.GPU_DATA_PREFIX + name, name)
-
-    def toGpu(self, name):
-        wlb.cuda.copyFieldToGpu(self.blocks, self.GPU_DATA_PREFIX + name, name)
-
-    def allToCpu(self):
-        for cpuName, gpuName in self._cpuGpuPairs:
-            wlb.cuda.copyFieldToCpu(self.blocks, gpuName, cpuName)
-
-    def allToGpu(self):
-        for cpuName, gpuName in self._cpuGpuPairs:
-            wlb.cuda.copyFieldToGpu(self.blocks, gpuName, cpuName)