boundaryhandling.py

import numpy as np
import sympy as sp
from pystencils.assignment import Assignment
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 create_type


class FlagInterface:
    FLAG_DTYPE = np.uint32

    def __init__(self, dataHandling, flagFieldName):
        self.flagFieldName = flagFieldName
        self.domainFlag = self.FLAG_DTYPE(1 << 0)
        self._nextFreeFlag = 1
        self.dataHandling = dataHandling

        # Add flag field to data handling if it does not yet exist
        if dataHandling.hasData(self.flagFieldName):
            raise ValueError("There is already a boundary handling registered at the data handling."
                             "If you want to add multiple handlings, choose a different name.")

        dataHandling.addArray(self.flagFieldName, dtype=self.FLAG_DTYPE, cpu=True, gpu=False)
        ffGhostLayers = dataHandling.ghostLayersOfField(self.flagFieldName)
        for b in dataHandling.iterate(ghostLayers=ffGhostLayers):
            b[self.flagFieldName].fill(self.domainFlag)

    def allocateNextFlag(self):
        result = self.FLAG_DTYPE(1 << self._nextFreeFlag)
        self._nextFreeFlag += 1
        return result


class BoundaryHandling:

    def __init__(self, dataHandling, fieldName, stencil, name="boundaryHandling", flagInterface=None,
                 target='cpu', openMP=True):
        assert dataHandling.hasData(fieldName)

        self._dataHandling = dataHandling
        self._fieldName = fieldName
        self._indexArrayName = name + "IndexArrays"
        self._target = target
        self._openMP = openMP
        self._boundaryObjectToBoundaryInfo = {}
        self.stencil = stencil
        self._dirty = True
        self.flagInterface = flagInterface if flagInterface is not None else FlagInterface(dataHandling, name + "Flags")

        gpu = self._target == 'gpu'
        dataHandling.addCustomClass(self._indexArrayName, self.IndexFieldBlockData, cpu=True, gpu=gpu)

    @property
    def dataHandling(self):
        return self._dataHandling

    def getFlag(self, boundaryObj):
        return self._boundaryObjectToBoundaryInfo[boundaryObj].flag

    @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.flagInterface.flagFieldName

    def getBoundaryNameToFlagDict(self):
        result = {bObj.name: bInfo.flag for bObj, bInfo in self._boundaryObjectToBoundaryInfo.items()}
        result['domain'] = self.flagInterface.domainFlag
        return result

    def getMask(self, sliceObj, boundaryObj, inverse=False):
        if isinstance(boundaryObj, str) and boundaryObj.lower() == 'domain':
            flag = self.flagInterface.domainFlag
        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,
                    replace=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() == 'domain':
            flag = self.flagInterface.domainFlag
        else:
            flag = self._addBoundary(boundaryObject)

        for b in self._dataHandling.iterate(sliceObj, ghostLayers=ghostLayers, innerGhostLayers=innerGhostLayers):
            flagArr = b[self.flagInterface.flagFieldName]
            if maskCallback is not None:
                mask = maskCallback(*b.midpointArrays)
                if replace:
                    flagArr[mask] = flag
                else:
                    np.bitwise_or(flagArr, flag, where=mask, out=flagArr)
                    np.bitwise_and(flagArr, ~self.flagInterface.domainFlag, where=mask, out=flagArr)
            else:
                if replace:
                    flagArr.fill(flag)
                else:
                    np.bitwise_or(flagArr, flag, out=flagArr)
                    np.bitwise_and(flagArr, ~self.flagInterface.domainFlag, out=flagArr)

        self._dirty = True

        return flag

    def setBoundaryWhereFlagIsSet(self, boundaryObject, flag):
        self._addBoundary(boundaryObject, flag)
        self._dirty = True
        return flag

    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.flagInterface.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]
                kwargs['indexField'] = idxArr
                dataUsedInKernel = (p.fieldName
                                    for p in self._boundaryObjectToBoundaryInfo[bObj].kernel.parameters
                                    if p.isFieldPtrArgument and p.fieldName not in kwargs)
                kwargs.update({name: b[name] for name in dataUsedInKernel})

                self._boundaryObjectToBoundaryInfo[bObj].kernel(**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 'domain' for domain 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()) + ['domain']
        elif not hasattr(boundaries, "__len__"):
            boundaries = [boundaries]

        masksToName = {}
        for b in boundaries:
            if b == 'domain':
                masksToName[self.flagInterface.domainFlag] = 'domain'
            else:
                masksToName[self._boundaryObjectToBoundaryInfo[b].flag] = b.name

        writer = self.dataHandling.vtkWriterFlags(fileName, self.flagInterface.flagFieldName,
                                                  masksToName, ghostLayers=ghostLayers)
        writer(1)

    # ------------------------------ Implementation Details ------------------------------------------------------------

    def _addBoundary(self, boundaryObject, flag=None):
        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)
            if flag is None:
                flag = self.flagInterface.allocateNextFlag()
            boundaryInfo = self.BoundaryInfo(boundaryObject, flag=flag, kernel=ast.compile())
            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.flagInterface.flagFieldName)
        for b in dh.iterate(ghostLayers=ffGhostLayers):
            flagArr = b[self.flagInterface.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.flagInterface.domainFlag,
                                                  bInfo.boundaryObject, ffGhostLayers)
                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 nonBoundaryCellPositions(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.nonBoundaryCellPositions(coord) + 0.5 * self.linkOffsets()[:, coord]

    @memorycache()
    def boundaryCellPositions(self, coord):
        return self.nonBoundaryCellPositions(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, symbols_read=set(),
                                                 symbols_defined=set(offsetSymbols + [self.INV_DIR_SYMBOL]))

    @staticmethod
    def _offsetSymbols(dim):
        return [TypedSymbol("c_%d" % (d,), create_type(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.numpy_dtype
    dirSymbol = TypedSymbol("dir", indexArrDtype.fields['dir'][0])
    elements += [Assignment(dirSymbol, indexField[0]('dir'))]
    elements += boundaryFunctor(field, directionSymbol=dirSymbol, indexField=indexField)
    return createIndexedKernel(elements, [indexField], target=target, cpuOpenMP=openMP)