From 869614f6ced1242e9f84fdb4256e82de5a93cc2e Mon Sep 17 00:00:00 2001
From: Martin Bauer <martin.bauer@fau.de>
Date: Thu, 19 Oct 2017 17:27:21 +0200
Subject: [PATCH] Advanced boundary handling for lbmpy

---
 boundaries/boundary_kernel.py    |  4 +-
 boundaries/boundaryconditions.py | 63 ++++++++++++++++++++++++++------
 boundaries/boundaryhandling.py   |  2 +
 boundaries/handlinginterface.py  | 49 ++++++++++++++++++-------
 chapman_enskog/steady_state.py   | 12 ++++--
 forcemodels.py                   |  2 -
 parallel/boundaryhandling.py     |  4 ++
 plot2d.py                        | 55 +++++++++++++++++++++++++++-
 8 files changed, 157 insertions(+), 34 deletions(-)

diff --git a/boundaries/boundary_kernel.py b/boundaries/boundary_kernel.py
index 2c17bf0f..a7e86cc0 100644
--- a/boundaries/boundary_kernel.py
+++ b/boundaries/boundary_kernel.py
@@ -57,8 +57,8 @@ def generateIndexBoundaryKernel(pdfField, indexArr, lbMethod, boundaryFunctor, t
     dirSymbol = TypedSymbol("dir", indexArr.dtype.fields['dir'][0])
     boundaryEqList = [sp.Eq(dirSymbol, indexField[0]('dir'))]
     if createInitializationKernel:
-        boundaryEqList += boundaryFunctor.additionalDataInit(pdfField=pdfField, directionSymbol=dirSymbol,
-                                                             lbMethod=lbMethod, indexField=indexField)
+        boundaryEqList += boundaryFunctor.additionalDataInitKernelEquations(pdfField=pdfField, directionSymbol=dirSymbol,
+                                                                            lbMethod=lbMethod, indexField=indexField)
     else:
         boundaryEqList += boundaryFunctor(pdfField=pdfField, directionSymbol=dirSymbol, lbMethod=lbMethod,
                                           indexField=indexField)
diff --git a/boundaries/boundaryconditions.py b/boundaries/boundaryconditions.py
index 94525cf0..19d2d690 100644
--- a/boundaries/boundaryconditions.py
+++ b/boundaries/boundaryconditions.py
@@ -29,15 +29,26 @@ class Boundary(object):
 
     @property
     def additionalData(self):
+        """Return a list of (name, type) tuples for additional data items required in this boundary
+        These data items can either be initialized in separate kernel see additionalDataKernelInit or by 
+        Python callbacks - see additionalDataCallback """
         return []
 
+    @property
+    def additionalDataInitCallback(self):
+        """Return a callback function called with (x, y, [z]), and returning a dict of data-name to data for each
+         element that should be initialized"""
+        return None
+
+    @property
+    def additionalDataInitKernelEquations(self):
+        """Should return callback that returns kernel equations for the boundary data initialization kernel"""
+        return None
+
     @property
     def name(self):
         return type(self).__name__
 
-    def additionalDataInit(self, idxArray):
-        return []
-
 
 class NoSlip(Boundary):
     """No-Slip, (half-way) simple bounce back boundary condition, enforcing zero velocity at obstacle"""
@@ -61,8 +72,11 @@ class NoSlipFullWay(Boundary):
     def additionalData(self):
         return [('lastValue', createTypeFromString("double"))]
 
-    def additionalDataInit(self, pdfField, directionSymbol, indexField, **kwargs):
-        return [sp.Eq(indexField('lastValue'), pdfField(directionSymbol))]
+    @property
+    def additionalDataInitKernelEquations(self):
+        def kernelEqGetter(pdfField, directionSymbol, indexField, **kwargs):
+            return [sp.Eq(indexField('lastValue'), pdfField(directionSymbol))]
+        return kernelEqGetter
 
     def __call__(self, pdfField, directionSymbol, lbMethod, indexField, **kwargs):
         neighbor = offsetFromDir(directionSymbol, lbMethod.dim)
@@ -82,20 +96,47 @@ class UBB(Boundary):
 
     """Velocity bounce back boundary condition, enforcing specified velocity at obstacle"""
 
-    def __init__(self, velocity, adaptVelocityToForce=False):
+    def __init__(self, velocity, adaptVelocityToForce=False, dim=None):
+        """
+        
+        :param velocity: can either be a constant, an access into a field, or a callback function.
+                         The callback functions gets a numpy record array with members, 'x','y','z', 'dir' (direction) 
+                         and 'velocity' which has to be set to the desired velocity of the corresponding link
+        :param adaptVelocityToForce:
+        """
         self._velocity = velocity
         self._adaptVelocityToForce = adaptVelocityToForce
+        if callable(self._velocity) and not dim:
+            raise ValueError("When using a velocity callback the dimension has to be specified with the dim parameter")
+        elif not callable(self._velocity):
+            dim = len(velocity)
+        self.dim = dim
 
-    def __call__(self, pdfField, directionSymbol, lbMethod, **kwargs):
-        vel = self._velocity
+    @property
+    def additionalData(self):
+        if callable(self._velocity):
+            return [('vel_%d' % (i,), createTypeFromString("double")) for i in range(self.dim)]
+        else:
+            return []
+
+    @property
+    def additionalDataInitCallback(self):
+        if callable(self._velocity):
+            return self._velocity
+
+    def __call__(self, pdfField, directionSymbol, lbMethod, indexField, **kwargs):
+        velFromIdxField = callable(self._velocity)
+        vel = [indexField('vel_%d' % (i,)) for i in range(self.dim)] if velFromIdxField else self._velocity
         direction = directionSymbol
 
-        assert len(vel) == lbMethod.dim, \
-            "Dimension of velocity (%d) does not match dimension of LB method (%d)" % (len(vel), lbMethod.dim)
+        assert self.dim == lbMethod.dim, "Dimension of UBB (%d) does not match dimension of method (%d)" \
+                                         % (self.dim, lbMethod.dim)
+
         neighbor = offsetFromDir(direction, lbMethod.dim)
         inverseDir = invDir(direction)
 
-        velocity = tuple(v_i.getShifted(*neighbor) if isinstance(v_i, Field.Access) else v_i for v_i in vel)
+        velocity = tuple(v_i.getShifted(*neighbor) if isinstance(v_i, Field.Access) and not velFromIdxField else v_i
+                         for v_i in vel)
 
         if self._adaptVelocityToForce:
             cqc = lbMethod.conservedQuantityComputation
diff --git a/boundaries/boundaryhandling.py b/boundaries/boundaryhandling.py
index bf74fe73..f055f2d7 100644
--- a/boundaries/boundaryhandling.py
+++ b/boundaries/boundaryhandling.py
@@ -7,6 +7,8 @@ class BoundaryHandling(PeriodicityHandling, GenericBoundaryHandling):  # importa
 
     def __init__(self, pdfField, domainShape, lbMethod, ghostLayers=1, target='cpu', openMP=True, flagDtype=np.uint32):
         shapeWithGl = [a + 2 * ghostLayers for a in domainShape]
+        self.domainShape = domainShape
+        self.domainShapeWithGhostLayers = shapeWithGl
         flagFieldInterface = NumpyFlagFieldInterface(shapeWithGl, flagDtype)
 
         GenericBoundaryHandling.__init__(self, flagFieldInterface, pdfField, lbMethod, ghostLayers, target, openMP)
diff --git a/boundaries/handlinginterface.py b/boundaries/handlinginterface.py
index 31d7279f..8f85e006 100644
--- a/boundaries/handlinginterface.py
+++ b/boundaries/handlinginterface.py
@@ -40,10 +40,11 @@ class FlagFieldInterface(object):
 class GenericBoundaryHandling(object):
 
     class BoundaryInfo(object):
-        def __init__(self, kernel, ast, indexArray):
+        def __init__(self, kernel, ast, indexArray, idxArrayForExecution):
             self.kernel = kernel
             self.ast = ast
             self.indexArray = indexArray
+            self.idxArrayForExecution = idxArrayForExecution  # is different for GPU kernels
 
     def __init__(self, flagFieldInterface, pdfField, lbMethod, ghostLayers=1, target='cpu', openMP=True):
         """
@@ -86,7 +87,7 @@ class GenericBoundaryHandling(object):
             self.prepare()
         for boundary in self._boundaryInfos.values():
             if boundary.kernel:
-                boundary.kernel(**kwargs)
+                boundary.kernel(indexField=boundary.idxArrayForExecution, **kwargs)
 
     def getBoundaryIndexArray(self, boundaryObject):
         return self._boundaryInfos[boundaryObject].indexArray
@@ -100,6 +101,30 @@ class GenericBoundaryHandling(object):
     def _invalidateCache(self):
         self._dirty = True
 
+    def triggerReinitializationOfBoundaryData(self, **kwargs):
+        if self._dirty:
+            self.prepare()
+            return
+        else:
+            for boundaryObject, boundaryInfo in self._boundaryInfos.items():
+                self.__boundaryDataInitialization(boundaryInfo.indexArray, boundaryObject, **kwargs)
+                if self._target == 'gpu':
+                    import pycuda.gpuarray as gpuarray
+                    boundaryInfo.idxArrayForExecution = gpuarray.to_gpu(boundaryInfo.indexArray)
+
+    def __boundaryDataInitialization(self, idxArray, boundaryObject, **kwargs):
+        if boundaryObject.additionalDataInitCallback:
+            #TODO x,y,z coordinates should be transformed here
+            boundaryObject.additionalDataInitCallback(idxArray, **kwargs)
+
+        if boundaryObject.additionalDataInitKernelEquations:
+            initKernelAst = generateIndexBoundaryKernel(self._symbolicPdfField, idxArray, self._lbMethod,
+                                                        boundaryObject, target='cpu',
+                                                        createInitializationKernel=True)
+            from pystencils.cpu import makePythonFunction as makePythonCpuFunction
+            initKernel = makePythonCpuFunction(initKernelAst, {'indexField': idxArray, 'pdfs': self._pdfField})
+            initKernel()
+
     def prepare(self):
         """Compiles boundary kernels according to flag field. When setting boundaries the cache is automatically
         invalidated, so there is no need to call this function manually, as long as the flag field is not manually
@@ -127,29 +152,25 @@ class GenericBoundaryHandling(object):
                     extendedIdxField[prop] = idxArray[prop]
 
                 idxArray = extendedIdxField
-                if boundaryObject.additionalDataInit:
-                    initKernelAst = generateIndexBoundaryKernel(self._symbolicPdfField, idxArray, self._lbMethod,
-                                                                boundaryObject, target='cpu',
-                                                                createInitializationKernel=True)
-                    from pystencils.cpu import makePythonFunction as makePythonCpuFunction
-                    initKernel = makePythonCpuFunction(initKernelAst, {'indexField': idxArray, 'pdfs': self._pdfField})
-                    initKernel()
+                self.__boundaryDataInitialization(idxArray, boundaryObject)
 
             ast = generateIndexBoundaryKernel(self._symbolicPdfField, idxArray, self._lbMethod, boundaryObject,
                                               target=self._target)
+
             if self._target == 'cpu':
                 from pystencils.cpu import makePythonFunction as makePythonCpuFunction, addOpenMP
                 addOpenMP(ast, numThreads=self._openMP)
-                kernel = makePythonCpuFunction(ast, {'indexField': idxArray})
+                idxArrayForExecution = idxArray
+                kernel = makePythonCpuFunction(ast)
             elif self._target == 'gpu':
                 from pystencils.gpucuda import makePythonFunction as makePythonGpuFunction
                 import pycuda.gpuarray as gpuarray
-                idxGpuField = gpuarray.to_gpu(idxArray)
-                kernel = makePythonGpuFunction(ast, {'indexField': idxGpuField})
+                idxArrayForExecution = gpuarray.to_gpu(idxArray)
+                kernel = makePythonGpuFunction(ast)
             else:
                 assert False
 
-            boundaryInfo = GenericBoundaryHandling.BoundaryInfo(kernel, ast, idxArray)
+            boundaryInfo = GenericBoundaryHandling.BoundaryInfo(kernel, ast, idxArray, idxArrayForExecution)
             self._boundaryInfos[boundaryObject] = boundaryInfo
 
         self._dirty = False
@@ -181,7 +202,7 @@ class GenericBoundaryHandling(object):
         self._dirty = True
 
     def getMask(self, boundaryObject):
-        return np.logical_and(self._flagFieldInterface.array, self._flagFieldInterface.getFlag(boundaryObject))
+        return np.bitwise_and(self._flagFieldInterface.array, self._flagFieldInterface.getFlag(boundaryObject))
 
     @property
     def flagField(self):
diff --git a/chapman_enskog/steady_state.py b/chapman_enskog/steady_state.py
index 6f63ed82..8547846b 100644
--- a/chapman_enskog/steady_state.py
+++ b/chapman_enskog/steady_state.py
@@ -29,7 +29,9 @@ class SteadyStateChapmanEnskogAnalysis(object):
 
         # Perform the analysis
         self.tayloredEquation = self._createTaylorExpandedEquation()
-        self.pdfHierarchy = self._createPdfHierarchy(self.tayloredEquation)
+        insertedHierarchy, rawHierarchy = self._createPdfHierarchy(self.tayloredEquation)
+        self.pdfHierarchy = insertedHierarchy
+        self.pdfHierarchyRaw = rawHierarchy
         self.recombinedEq = self._recombinePdfs(self.pdfHierarchy)
 
         symbolsToValues = self._getSymbolsToValuesDict()
@@ -61,16 +63,18 @@ class SteadyStateChapmanEnskogAnalysis(object):
                                                         pdfs=(['f', 0, self.order + 1],), commutative=False)
         chapmanEnskogHierarchy = [chapmanEnskogHierarchy[i] for i in range(self.order + 1)]
 
-        result = []
+        insertedHierarchy = []
+        rawHierarchy = []j
         substitutionDict = {}
         for ceEq, f_i in zip(chapmanEnskogHierarchy, self.fSyms):
             newEq = -1 / self.collisionOpSym * (ceEq - self.collisionOpSym * f_i)
+            rawHierarchy.append(newEq)
             newEq = expandUsingLinearity(newEq.subs(substitutionDict), functions=self.fSyms + [self.forceSym])
             if newEq:
                 substitutionDict[f_i] = newEq
-            result.append(newEq)
+            insertedHierarchy.append(newEq)
 
-        return result
+        return insertedHierarchy, rawHierarchy
 
     def _recombinePdfs(self, pdfHierarchy):
         return sum(pdfHierarchy[i] * self.eps**(i-1) for i in range(1, self.order+1))
diff --git a/forcemodels.py b/forcemodels.py
index 39cc4e04..dbc93edf 100644
--- a/forcemodels.py
+++ b/forcemodels.py
@@ -161,8 +161,6 @@ class Buick(object):
     r"""
     This force model :cite:`buick2000gravity` has a force term with zero second moment. 
     It is suited for incompressible lattice models.
-
-    
     """
 
     def __init__(self, force):
diff --git a/parallel/boundaryhandling.py b/parallel/boundaryhandling.py
index 92ec8465..3e59ec0a 100644
--- a/parallel/boundaryhandling.py
+++ b/parallel/boundaryhandling.py
@@ -30,6 +30,10 @@ class BoundaryHandling(object):
         for block in self._blocks:
             block[self._boundaryId].prepare()
 
+    def triggerReinitializationOfBoundaryData(self, **kwargs):
+        for block in self._blocks:
+            block[self._boundaryId].triggerReinitializationOfBoundaryData()
+
     def setBoundary(self, boundaryObject, indexExpr=None, maskCallback=None, sliceNormalizationGhostLayers=1):
         if indexExpr is None:
             indexExpr = [slice(None, None, None)] * self.dim
diff --git a/plot2d.py b/plot2d.py
index e8e02274..ad2f1a72 100644
--- a/plot2d.py
+++ b/plot2d.py
@@ -1 +1,54 @@
-from pystencils.plot2d import *
\ No newline at end of file
+from pystencils.plot2d import *
+from pystencils.slicing import normalizeSlice
+
+
+def plotBoundaryHandling(boundaryHandling, slice=None, boundaryNameToColor=None, legend=True):
+    """
+    Shows boundary cells
+
+    :param boundaryHandling: instance of :class:`lbmpy.boundaries.BoundaryHandling`
+    :param boundaryNameToColor: optional dictionary mapping boundary names to colors
+    """
+    import matplotlib
+    import matplotlib.pyplot as plt
+
+    boundaryHandling.prepare()
+
+    if len(boundaryHandling.flagField.shape) != 2 and slice is None:
+        raise ValueError("To plot 3D boundary handlings a slice has to be passed")
+
+    if boundaryNameToColor:
+        fixedColors = boundaryNameToColor
+    else:
+        fixedColors = {
+            'fluid': '#56b4e9',
+            'NoSlip': '#999999',
+            'UBB': '#d55e00',
+            'FixedDensity': '#009e73',
+        }
+
+    boundaryNames = []
+    flagValues = []
+    for name, flag in sorted(boundaryHandling.getBoundaryNameToFlagDict().items(), key=lambda l: l[1]):
+        boundaryNames.append(name)
+        flagValues.append(flag)
+    defaultCycler = matplotlib.rcParams['axes.prop_cycle']
+    colorValues = [fixedColors[name] if name in fixedColors else cycle['color']
+                   for cycle, name in zip(defaultCycler, boundaryNames)]
+
+    cmap = matplotlib.colors.ListedColormap(colorValues)
+    bounds = np.array(flagValues, dtype=float) - 0.5
+    bounds = list(bounds) + [bounds[-1] + 1]
+    norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)
+
+    flagField = boundaryHandling.flagField
+    if slice:
+        flagField = flagField[normalizeSlice(slice, flagField.shape)]
+    flagField = flagField.swapaxes(0, 1)
+    plt.imshow(flagField, interpolation='none', origin='lower',
+               cmap=cmap, norm=norm)
+
+    patches = [matplotlib.patches.Patch(color=color, label=name) for color, name in zip(colorValues, boundaryNames)]
+    plt.axis('equal')
+    if legend:
+        plt.legend(handles=patches, bbox_to_anchor=(1.02, 0.5), loc=2, borderaxespad=0.)
\ No newline at end of file
-- 
GitLab