import numpy as np
import sympy as sp
from pystencils import Field, TypedSymbol, createIndexedKernel
from pystencils.backends.cbackend import CustomCppCode
from pystencils.boundaries.createindexlist import numpyDataTypeForBoundaryObject, createBoundaryIndexArray
from pystencils.cache import memorycache
from pystencils.data_types import createType
class BoundaryHandling:
def __init__(self, dataHandling, fieldName, stencil, name="boundaryHandling", target='cpu', openMP=True):
assert dataHandling.hasData(fieldName)
self._dataHandling = dataHandling
self._fieldName = fieldName
self._flagFieldName = name + "Flags"
self._indexArrayName = name + "IndexArrays"
self._target = target
self._openMP = openMP
self._boundaryObjectToBoundaryInfo = {}
self._fluidFlag = 1 << 0
self._nextFreeFlag = 1
self.stencil = stencil
self._dirty = True
# Add flag field to data handling if it does not yet exist
if dataHandling.hasData(self._flagFieldName) or dataHandling.hasData(self._indexArrayName):
raise ValueError("There is already a boundary handling registered at the data handling."
"If you want to add multiple handlings, choose a different name.")
gpu = self._target == 'gpu'
dataHandling.addArray(self._flagFieldName, dtype=np.uint32, cpu=True, gpu=False)
dataHandling.addCustomClass(self._indexArrayName, self.IndexFieldBlockData, cpu=True, gpu=gpu)
ffGhostLayers = self._dataHandling.ghostLayersOfField(self._flagFieldName)
for b in self._dataHandling.iterate(ghostLayers=ffGhostLayers):
b[self._flagFieldName].fill(self._fluidFlag)
@property
def dataHandling(self):
return self._dataHandling
@property
def shape(self):
return self._dataHandling.shape
@property
def dim(self):
return self._dataHandling.dim
@property
def boundaryObjects(self):
return tuple(self._boundaryObjectToName.keys())
@property
def flagArrayName(self):
return self._flagFieldName
def getBoundaryNameToFlagDict(self):
result = {bObj.name: bInfo.flag for bObj, bInfo in self._boundaryObjectToBoundaryInfo.items()}
result['fluid'] = self._fluidFlag
return result
def getMask(self, sliceObj, boundaryObj, inverse=False):
if isinstance(boundaryObj, str) and boundaryObj.lower() == 'fluid':
flag = self._fluidFlag
else:
flag = self._boundaryObjectToBoundaryInfo[boundaryObj].flag
arr = self.dataHandling.gatherArray(self.flagArrayName, sliceObj)
if arr is None:
return None
else:
result = np.bitwise_and(arr, flag)
if inverse:
result = np.logical_not(result)
return result
def setBoundary(self, boundaryObject, sliceObj=None, maskCallback=None, ghostLayers=True, innerGhostLayers=True):
"""
Sets boundary using either a rectangular slice, a boolean mask or a combination of both
:param boundaryObject: instance of a boundary object that should be set
:param sliceObj: a slice object (can be created with makeSlice[]) that selects a part of the domain where
the boundary should be set. If none, the complete domain is selected which makes only sense
if a maskCallback is passed. The slice can have ':' placeholders, which are interpreted
depending on the 'includeGhostLayers' parameter i.e. if it is True, the slice extends
into the ghost layers
:param maskCallback: callback function getting x,y (z) parameters of the cell midpoints and returning a
boolean mask with True entries where boundary cells should be set.
The x, y, z arrays have 2D/3D shape such that they can be used directly
to create the boolean return array. i.e return x < 10 sets boundaries in cells with
midpoint x coordinate smaller than 10.
:param ghostLayers see DataHandling.iterate()
"""
if isinstance(boundaryObject, str) and boundaryObject.lower() == 'fluid':
flag = self._fluidFlag
else:
flag = self._getFlagForBoundary(boundaryObject)
for b in self._dataHandling.iterate(sliceObj, ghostLayers=ghostLayers, innerGhostLayers=innerGhostLayers):
flagArr = b[self._flagFieldName]
if maskCallback is not None:
mask = maskCallback(*b.midpointArrays)
flagArr[mask] = flag
else:
flagArr.fill(flag)
self._dirty = True
def prepare(self):
if not self._dirty:
return
self._createIndexFields()
self._dirty = False
def triggerReinitializationOfBoundaryData(self, **kwargs):
if self._dirty:
self.prepare()
else:
ffGhostLayers = self._dataHandling.ghostLayersOfField(self._flagFieldName)
for b in self._dataHandling.iterate(ghostLayers=ffGhostLayers):
for bObj, setter in b[self._indexArrayName].boundaryObjectToDataSetter.items():
self._boundaryDataInitialization(bObj, setter, **kwargs)
def __call__(self, **kwargs):
if self._dirty:
self.prepare()
for b in self._dataHandling.iterate(gpu=self._target == 'gpu'):
for bObj, idxArr in b[self._indexArrayName].boundaryObjectToIndexList.items():
kwargs[self._fieldName] = b[self._fieldName]
self._boundaryObjectToBoundaryInfo[bObj].kernel(indexField=idxArr, **kwargs)
def geometryToVTK(self, fileName='geometry', boundaries='all', ghostLayers=False):
"""
Writes a VTK field where each cell with the given boundary is marked with 1, other cells are 0
This can be used to display the simulation geometry in Paraview
:param fileName: vtk filename
:param boundaries: boundary object, or special string 'fluid' for fluid cells or special string 'all' for all
boundary conditions.
can also be a sequence, to write multiple boundaries to VTK file
:param ghostLayers: number of ghost layers to write, or True for all, False for none
"""
if boundaries == 'all':
boundaries = list(self._boundaryObjectToBoundaryInfo.keys()) + ['fluid']
elif not hasattr(boundaries, "__len__"):
boundaries = [boundaries]
masksToName = {}
for b in boundaries:
if b == 'fluid':
masksToName[self._fluidFlag] = 'fluid'
else:
masksToName[self._boundaryObjectToBoundaryInfo[b].flag] = b.name
writer = self.dataHandling.vtkWriterFlags(fileName, self._flagFieldName, masksToName, ghostLayers=ghostLayers)
writer(1)
# ------------------------------ Implementation Details ------------------------------------------------------------
def _getFlagForBoundary(self, boundaryObject):
if boundaryObject not in self._boundaryObjectToBoundaryInfo:
symbolicIndexField = Field.createGeneric('indexField', spatialDimensions=1,
dtype=numpyDataTypeForBoundaryObject(boundaryObject, self.dim))
ast = self._createBoundaryKernel(self._dataHandling.fields[self._fieldName],
symbolicIndexField, boundaryObject)
boundaryInfo = self.BoundaryInfo(boundaryObject, flag=1 << self._nextFreeFlag, kernel=ast.compile())
self._nextFreeFlag += 1
self._boundaryObjectToBoundaryInfo[boundaryObject] = boundaryInfo
return self._boundaryObjectToBoundaryInfo[boundaryObject].flag
def _createBoundaryKernel(self, symbolicField, symbolicIndexField, boundaryObject):
return createBoundaryKernel(symbolicField, symbolicIndexField, self.stencil, boundaryObject,
target=self._target, openMP=self._openMP)
def _createIndexFields(self):
dh = self._dataHandling
ffGhostLayers = dh.ghostLayersOfField(self._flagFieldName)
for b in dh.iterate(ghostLayers=ffGhostLayers):
flagArr = b[self._flagFieldName]
pdfArr = b[self._fieldName]
indexArrayBD = b[self._indexArrayName]
indexArrayBD.clear()
for bInfo in self._boundaryObjectToBoundaryInfo.values():
idxArr = createBoundaryIndexArray(flagArr, self.stencil, bInfo.flag, self._fluidFlag,
bInfo.boundaryObject, dh.ghostLayersOfField(self._flagFieldName))
if idxArr.size == 0:
continue
boundaryDataSetter = BoundaryDataSetter(idxArr, b.offset, self.stencil, ffGhostLayers, pdfArr)
indexArrayBD.boundaryObjectToIndexList[bInfo.boundaryObject] = idxArr
indexArrayBD.boundaryObjectToDataSetter[bInfo.boundaryObject] = boundaryDataSetter
self._boundaryDataInitialization(bInfo.boundaryObject, boundaryDataSetter)
def _boundaryDataInitialization(self, boundaryObject, boundaryDataSetter, **kwargs):
if boundaryObject.additionalDataInitCallback:
boundaryObject.additionalDataInitCallback(boundaryDataSetter, **kwargs)
if self._target == 'gpu':
self._dataHandling.toGpu(self._indexArrayName)
class BoundaryInfo(object):
def __init__(self, boundaryObject, flag, kernel):
self.boundaryObject = boundaryObject
self.flag = flag
self.kernel = kernel
class IndexFieldBlockData:
def __init__(self, *args, **kwargs):
self.boundaryObjectToIndexList = {}
self.boundaryObjectToDataSetter = {}
def clear(self):
self.boundaryObjectToIndexList.clear()
self.boundaryObjectToDataSetter.clear()
@staticmethod
def toCpu(gpuVersion, cpuVersion):
gpuVersion = gpuVersion.boundaryObjectToIndexList
cpuVersion = cpuVersion.boundaryObjectToIndexList
for obj, cpuArr in cpuVersion.values():
gpuVersion[obj].get(cpuArr)
@staticmethod
def toGpu(gpuVersion, cpuVersion):
from pycuda import gpuarray
gpuVersion = gpuVersion.boundaryObjectToIndexList
cpuVersion = cpuVersion.boundaryObjectToIndexList
for obj, cpuArr in cpuVersion.items():
if obj not in gpuVersion:
gpuVersion[obj] = gpuarray.to_gpu(cpuArr)
else:
gpuVersion[obj].set(cpuArr)
class BoundaryDataSetter:
def __init__(self, indexArray, offset, stencil, ghostLayers, pdfArray):
self.indexArray = indexArray
self.offset = offset
self.stencil = np.array(stencil)
self.pdfArray = pdfArray.view()
self.pdfArray.flags.writeable = False
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 + 0.5
@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 BoundaryOffsetInfo(CustomCppCode):
# --------------------------- Functions to be used by boundaries --------------------------
@staticmethod
def offsetFromDir(dirIdx, dim):
return tuple([sp.IndexedBase(symbol, shape=(1,))[dirIdx]
for symbol in BoundaryOffsetInfo._offsetSymbols(dim)])
@staticmethod
def invDir(dirIdx):
return sp.IndexedBase(BoundaryOffsetInfo.INV_DIR_SYMBOL, shape=(1,))[dirIdx]
# ---------------------------------- Internal ---------------------------------------------
def __init__(self, stencil):
dim = len(stencil[0])
offsetSym = BoundaryOffsetInfo._offsetSymbols(dim)
code = "\n"
for i in range(dim):
offsetStr = ", ".join([str(d[i]) for d in stencil])
code += "const int64_t %s [] = { %s };\n" % (offsetSym[i].name, offsetStr)
invDirs = []
for direction in stencil:
inverseDir = tuple([-i for i in direction])
invDirs.append(str(stencil.index(inverseDir)))
code += "const int %s [] = { %s };\n" % (self.INV_DIR_SYMBOL.name, ", ".join(invDirs))
offsetSymbols = BoundaryOffsetInfo._offsetSymbols(dim)
super(BoundaryOffsetInfo, self).__init__(code, symbolsRead=set(),
symbolsDefined=set(offsetSymbols + [self.INV_DIR_SYMBOL]))
@staticmethod
def _offsetSymbols(dim):
return [TypedSymbol("c_%d" % (d,), createType(np.int64)) for d in range(dim)]
INV_DIR_SYMBOL = TypedSymbol("invDir", "int")
def createBoundaryKernel(field, indexField, stencil, boundaryFunctor, target='cpu', openMP=True):
elements = [BoundaryOffsetInfo(stencil)]
indexArrDtype = indexField.dtype.numpyDtype
dirSymbol = TypedSymbol("dir", indexArrDtype.fields['dir'][0])
elements += [sp.Eq(dirSymbol, indexField[0]('dir'))]
elements += boundaryFunctor(field, directionSymbol=dirSymbol, indexField=indexField)
return createIndexedKernel(elements, [indexField], target=target, cpuOpenMP=openMP)