diff --git a/datahandling.py b/datahandling.py
index b64300ac87cd5bde97f8916711c3b88f52cba67e..1c7c6e20f88edb54e911e75a8f42531b0e8d3c91 100644
--- a/datahandling.py
+++ b/datahandling.py
@@ -1,5 +1,10 @@
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:
diff --git a/field.py b/field.py
index a3ffc22bf9709a8560e9514af2bc129d4de46c66..a8e0e7274726f100fb74a7d63aa8b19989a106c8 100644
--- a/field.py
+++ b/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)
diff --git a/finitedifferences.py b/finitedifferences.py
index aa3e9d344bf189092596f4b369ab20864b06be53..a60bb8b9cb2b81a61023613f9ac2d51c191888e1 100644
--- a/finitedifferences.py
+++ b/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)
diff --git a/parallel/datahandling.py b/parallel/datahandling.py
index 7aa029f16404e9494f6a3d7a5f5eccf9c9968e2a..55cb5fc8a2b8c7e4ed01f8853b09559cf0df8a18 100644
--- a/parallel/datahandling.py
+++ b/parallel/datahandling.py
@@ -1,6 +1,6 @@
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)