lbmpy: boundary index list creation with Cython

boundaries/boundaryconditions.py

@@ -2,14 +2,14 @@ import sympy as sp
 from lbmpy.boundaries.boundaryhandling import offsetFromDir, weightOfDirection, invDir
 
 
-def noSlip(pdfField, direction, latticeModel):
-    neighbor = offsetFromDir(direction, latticeModel.dim)
+def noSlip(pdfField, direction, lbmMethod):
+    neighbor = offsetFromDir(direction, lbmMethod.dim)
     inverseDir = invDir(direction)
     return [sp.Eq(pdfField[neighbor](inverseDir), pdfField(direction))]
 
 
-def ubb(pdfField, direction, latticeModel, velocity):
-    neighbor = offsetFromDir(direction, latticeModel.dim)
+def ubb(pdfField, direction, lbmMethod, velocity):
+    neighbor = offsetFromDir(direction, lbmMethod.dim)
     inverseDir = invDir(direction)
 
     velTerm = 6 * sum([d_i * v_i for d_i, v_i in zip(neighbor, velocity)]) * weightOfDirection(direction)

boundaries/createindexlist.py

+import numpy as np
+import itertools
+
+#try:
+if True:
+    import pyximport;
+    pyximport.install()
+    from lbmpy.boundaries.createindexlistcython import createBoundaryIndexList2D, createBoundaryIndexList3D
+    cythonFuncsAvailable = True
+#except Exception:
+#    cythonFuncsAvailable = False
+#    createBoundaryIndexList2D = None
+#    createBoundaryIndexList3D = None
+
+
+def _createBoundaryIndexListPython(flagFieldArr, nrOfGhostLayers, boundaryMask, fluidMask, stencil):
+    result = []
+    gl = nrOfGhostLayers
+    for cell in itertools.product(*[range(gl, i-gl) for i in flagFieldArr.shape]):
+        if not flagFieldArr[cell] & fluidMask:
+            continue
+        for dirIdx, direction in enumerate(stencil):
+            neighborCell = tuple([cell_i + dir_i for cell_i, dir_i in zip(cell, direction)])
+            if flagFieldArr[neighborCell] & boundaryMask:
+                result.append(list(cell) + [dirIdx])
+
+    return np.array(result, dtype=np.int32)
+
+
+def createBoundaryIndexList(flagField, stencil, boundaryMask, fluidMask, nrOfGhostLayers=1):
+    if cythonFuncsAvailable:
+        stencil = np.array(stencil, dtype=np.int32)
+        if len(flagField.shape) == 2:
+            return np.array(createBoundaryIndexList2D(flagField, nrOfGhostLayers, boundaryMask, fluidMask, stencil))
+        elif len(flagField.shape) == 3:
+            return np.array(createBoundaryIndexList3D(flagField, nrOfGhostLayers, boundaryMask, fluidMask, stencil))
+        else:
+            raise ValueError("Flag field has to be a 2 or 3 dimensional numpy array")
+    else:
+        return _createBoundaryIndexListPython(flagField, nrOfGhostLayers, boundaryMask, fluidMask, stencil)

boundaries/createindexlistcython.pyx

+import cython
+
+ctypedef fused IntegerType:
+    short
+    int
+    long
+    unsigned short
+    unsigned int
+    unsigned long
+
+@cython.boundscheck(False) # turn off bounds-checking for entire function
+@cython.wraparound(False)  # turn off negative index wrapping for entire function
+def createBoundaryIndexList2D(object[IntegerType, ndim=2] flagField,
+                              int nrOfGhostLayers, IntegerType boundaryMask, IntegerType fluidMask,
+                              object[int, ndim=2] stencil):
+    cdef int xs, ys, x, y
+    cdef int dirIdx, numDirections, dx, dy
+
+    xs, ys = flagField.shape
+    boundaryIndexList = []
+    numDirections = stencil.shape[0]
+
+    for x in range(nrOfGhostLayers,xs-nrOfGhostLayers):
+        for y in range(nrOfGhostLayers,ys-nrOfGhostLayers):
+            if flagField[x,y] & fluidMask:
+                for dirIdx in range(1, numDirections):
+                    dx = stencil[dirIdx,0]
+                    dy = stencil[dirIdx,1]
+                    if flagField[x+dx, y+dy] & boundaryMask:
+                        boundaryIndexList.append((x,y, dirIdx))
+    return boundaryIndexList
+
+
+@cython.boundscheck(False) # turn off bounds-checking for entire function
+@cython.wraparound(False)  # turn off negative index wrapping for entire function
+def createBoundaryIndexList3D(object[IntegerType, ndim=3] flagField,
+                              int nrOfGhostLayers, IntegerType boundaryMask, IntegerType fluidMask,
+                              object[int, ndim=2] stencil):
+    cdef int xs, ys, zs, x, y, z
+    cdef int dirIdx, numDirections, dx, dy, dz
+
+    xs, ys, zs = flagField.shape
+    boundaryIndexList = []
+    numDirections = stencil.shape[0]
+
+    for x in range(nrOfGhostLayers,xs-nrOfGhostLayers):
+        for y in range(nrOfGhostLayers,ys-nrOfGhostLayers):
+            for z in range(nrOfGhostLayers, zs-nrOfGhostLayers):
+                if flagField[x, y, z] & fluidMask:
+                    for dirIdx in range(1, numDirections):
+                        dx = stencil[dirIdx,0]
+                        dy = stencil[dirIdx,1]
+                        dz = stencil[dirIdx,2]
+                        if flagField[x + dx, y + dy, z + dz] & boundaryMask:
+                            boundaryIndexList.append((x,y,z, dirIdx))
+    return boundaryIndexList
+
+