Commit ed36b869 authored by Martin Bauer's avatar Martin Bauer
Browse files

Data handling - flag array support

parent c35df1c3
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:
......
......@@ -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)
......
......@@ -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)
......
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)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment