Documentation & boundary helpers for lbmpy

- documentation for lbmpy creationfunctions
- new tutorial
- boundary handling: various helpers to setup geometry

boundaries/boundaryhandling.py

@@ -5,13 +5,20 @@ from lbmpy.stencils import getStencil
 from pystencils import TypedSymbol, Field
 from pystencils.backends.cbackend import CustomCppCode
 from lbmpy.boundaries.createindexlist import createBoundaryIndexList
-from pystencils.slicing import normalizeSlice
+from pystencils.slicing import normalizeSlice, makeSlice
 
 INV_DIR_SYMBOL = TypedSymbol("invDir", "int")
 WEIGHTS_SYMBOL = TypedSymbol("weights", "double")
 
 
 class BoundaryHandling(object):
+    class BoundaryInfo(object):
+        def __init__(self, name, flag, function, kernel):
+            self.name = name
+            self.flag = flag
+            self.function = function
+            self.kernel = kernel
+
     def __init__(self, pdfField, domainShape, lbMethod, ghostLayers=1, target='cpu'):
         """
         Class for managing boundary kernels
@@ -32,13 +39,16 @@ class BoundaryHandling(object):
 
         self._symbolicPdfField = symbolicPdfField
         self._shapeWithGhostLayers = [d + 2 * ghostLayers for d in domainShape]
-        self._fluidFlag = 2 ** 30
+        self._fluidFlag = 2 ** 0
         self.flagField = np.full(self._shapeWithGhostLayers, self._fluidFlag, dtype=np.int32)
         self._ghostLayers = ghostLayers
         self._lbMethod = lbMethod
-        self._boundaryFunctions = []
-        self._nameToIndex = {}
-        self._boundarySweeps = []
+
+        self._boundaryInfos = []
+        self._nameToBoundary = {}
+        self._periodicityKernels = []
+
+        self._dirty = False
         self._periodicity = [False, False, False]
         self._target = target
         if target not in ('cpu', 'gpu'):
@@ -46,34 +56,59 @@ class BoundaryHandling(object):
 
     @property
     def periodicity(self):
+        """List that indicates for x,y (z) coordinate if domain is periodic in that direction"""
         return self._periodicity
 
+    @property
+    def fluidFlag(self):
+        """Flag that is set where the lattice Boltzmann update should happen"""
+        return self._fluidFlag
+
+    def getFlag(self, name):
+        """Flag that represents the boundary with given name. Raises KeyError if no such boundary exists."""
+        return self._nameToBoundary[name].flag
+
+    def getBoundaryNames(self):
+        """List of names of all registered boundary conditions"""
+        return [b.name for b in self._boundaryInfos]
+
     def setPeriodicity(self, x=False, y=False, z=False):
+        """Enable periodic boundary conditions at the border of the domain"""
         self._periodicity = [x, y, z]
-        self.invalidateIndexCache()
+        self._compilePeriodicityKernels()
+
+    def hasBoundary(self, name):
+        """Returns boolean indicating if a boundary with that name exists"""
+        return name in self._nameToBoundary
 
-    def setBoundary(self, function, indexExpr, maskArr=None, name=None):
+    def setBoundary(self, function, indexExpr=None, maskArr=None, name=None):
         """
         Sets boundary in a rectangular region (slice)
 
-        :param function: boundary
+        :param function: boundary function or the string 'fluid' to remove boundary conditions
         :param indexExpr: slice expression, where boundary should be set, see :mod:`pystencils.slicing`
         :param maskArr: optional boolean (masked) array specifying where the boundary should be set
         :param name: name of the boundary
         """
-        if name is None:
-            if hasattr(function, '__name__'):
-                name = function.__name__
-            elif hasattr(function, 'name'):
-                name = function.name
-            else:
-                raise ValueError("Boundary function has to have a '__name__' or 'name' attribute "
-                                 "if name is not specified")
+        if indexExpr is None:
+            indexExpr = [slice(None, None, None)] * len(self.flagField.shape)
+        if function == 'fluid':
+            flag = self._fluidFlag
+        else:
+            if name is None:
+                if hasattr(function, '__name__'):
+                    name = function.__name__
+                elif hasattr(function, 'name'):
+                    name = function.name
+                else:
+                    raise ValueError("Boundary function has to have a '__name__' or 'name' attribute "
+                                     "if name is not specified")
 
-        if function not in self._boundaryFunctions:
-            self.addBoundary(function, name)
+            if not self.hasBoundary(name):
+                self.addBoundary(function, name)
 
-        flag = self.getFlag(name)
+            flag = self.getFlag(name)
+            assert flag != self._fluidFlag
 
         indexExpr = normalizeSlice(indexExpr, self._shapeWithGhostLayers)
         if maskArr is None:
@@ -81,8 +116,7 @@ class BoundaryHandling(object):
         else:
             flagFieldView = self.flagField[indexExpr]
             flagFieldView[maskArr] = flag
-
-        self.invalidateIndexCache()
+        self._dirty = True
 
     def addBoundary(self, boundaryFunction, name=None):
         """
@@ -97,45 +131,69 @@ class BoundaryHandling(object):
         if name is None:
             name = boundaryFunction.__name__
 
-        if name in self._nameToIndex:
-            return 2 ** self._nameToIndex[name]
+        if self.hasBoundary(name):
+            return self._boundaryInfos[name].flag
 
-        newIdx = len(self._boundaryFunctions)
-        self._nameToIndex[name] = newIdx
-        self._boundaryFunctions.append(boundaryFunction)
-        return 2 ** newIdx
+        newIdx = len(self._boundaryInfos) + 1  # +1 because 2**0 is reserved for fluid flag
+        boundaryInfo = self.BoundaryInfo(name, 2 ** newIdx, boundaryFunction, None)
+        self._boundaryInfos.append(boundaryInfo)
+        self._nameToBoundary[name] = boundaryInfo
+        self._dirty = True
+        return boundaryInfo.flag
 
-    def invalidateIndexCache(self):
-        self._boundarySweeps = []
+    def indices(self, dx=1.0, includeGhostLayers=False):
+        if not includeGhostLayers:
+            params = [np.arange(start=-1, stop=s-1) * dx for s in self.flagField.shape]
+        else:
+            params = [np.arange(s) * dx for s in self.flagField.shape]
+        return np.meshgrid(*params, indexing='ij')
+
+    def __call__(self, **kwargs):
+        """Run the boundary handling, all keyword args are passed through to the boundary sweeps"""
+        if self._dirty:
+            self.prepare()
+        for boundary in self._boundaryInfos:
+            boundary.kernel(**kwargs)
+        for k in self._periodicityKernels:
+            k(**kwargs)
 
     def clear(self):
+        """Removes all boundaries and fills the domain with fluid"""
         np.fill(self._fluidFlag)
-        self.invalidateIndexCache()
-
-    def getFlag(self, name):
-        return 2 ** self._nameToIndex[name]
+        self._dirty = False
+        self._boundaryInfos = []
+        self._nameToBoundary = {}
 
     def prepare(self):
-        self.invalidateIndexCache()
-        for boundaryIdx, boundaryFunc in enumerate(self._boundaryFunctions):
+        """Fills data structures to speed up the boundary handling and compiles all boundary kernels.
+        This is automatically done when first called. With this function this can be triggered before running."""
+        for boundary in self._boundaryInfos:
+            assert boundary.flag != self._fluidFlag
             idxField = createBoundaryIndexList(self.flagField, self._lbMethod.stencil,
-                                               2 ** boundaryIdx, self._fluidFlag, self._ghostLayers)
-            ast = generateBoundaryHandling(self._symbolicPdfField, idxField, self._lbMethod, boundaryFunc,
+                                               boundary.flag, self._fluidFlag, self._ghostLayers)
+            ast = generateBoundaryHandling(self._symbolicPdfField, idxField, self._lbMethod, boundary.function,
                                            target=self._target)
 
             if self._target == 'cpu':
                 from pystencils.cpu import makePythonFunction as makePythonCpuFunction
-                self._boundarySweeps.append(makePythonCpuFunction(ast, {'indexField': idxField}))
+                boundary.kernel = makePythonCpuFunction(ast, {'indexField': idxField})
             elif self._target == 'gpu':
                 from pystencils.gpucuda import makePythonFunction as makePythonGpuFunction
                 import pycuda.gpuarray as gpuarray
                 idxGpuField = gpuarray.to_gpu(idxField)
-                self._boundarySweeps.append(makePythonGpuFunction(ast, {'indexField': idxGpuField}))
+                boundary.kernel = makePythonGpuFunction(ast, {'indexField': idxGpuField})
             else:
                 assert False
-        self._addPeriodicityHandlers()
+        self._dirty = False
+
+    def invalidateIndexCache(self):
+        """Invalidates the cache for optimization data structures. 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
+        modified."""
+        self._dirty = True
 
-    def _addPeriodicityHandlers(self):
+    def _compilePeriodicityKernels(self):
+        self._periodicityKernels = []
         dim = len(self.flagField.shape)
         if dim == 2:
             stencil = getStencil("D2Q9")
@@ -147,7 +205,7 @@ class BoundaryHandling(object):
         filteredStencil = []
         for direction in stencil:
             useDirection = True
-            if direction == (0,0) or direction == (0,0,0):
+            if direction == (0, 0) or direction == (0, 0, 0):
                 useDirection = False
             for component, periodicity in zip(direction, self._periodicity):
                 if not periodicity and component != 0:
@@ -158,22 +216,16 @@ class BoundaryHandling(object):
         if len(filteredStencil) > 0:
             if self._target == 'cpu':
                 from pystencils.slicing import getPeriodicBoundaryFunctor
-                self._boundarySweeps.append(getPeriodicBoundaryFunctor(filteredStencil, ghostLayers=1))
+                self._periodicityKernels.append(getPeriodicBoundaryFunctor(filteredStencil, ghostLayers=1))
             elif self._target == 'gpu':
                 from pystencils.gpucuda.periodicity import getPeriodicBoundaryFunctor
-                self._boundarySweeps.append(getPeriodicBoundaryFunctor(filteredStencil, self.flagField.shape,
-                                                                       indexDimensions=1,
-                                                                       indexDimShape=len(self._lbMethod.stencil),
-                                                                       ghostLayers=1))
+                self._periodicityKernels.append(getPeriodicBoundaryFunctor(filteredStencil, self.flagField.shape,
+                                                                           indexDimensions=1,
+                                                                           indexDimShape=len(self._lbMethod.stencil),
+                                                                           ghostLayers=1))
             else:
                 assert False
 
-    def __call__(self, **kwargs):
-        if len(self._boundarySweeps) == 0:
-            self.prepare()
-        for boundarySweep in self._boundarySweeps:
-            boundarySweep(**kwargs)
-
 
 # -------------------------------------- Helper Functions --------------------------------------------------------------
 

creationfunctions.py

@@ -2,42 +2,6 @@ r"""
 Creating LBM kernels
 ====================
 
-
-Terminology and creation pipeline
----------------------------------
-
-Kernel functions are created in three steps:
-
-1. *Method*:
-         the method defines the collision process. Currently there are two big categories:
-         moment and cumulant based methods. A method defines how each moment or cumulant is relaxed by
-         storing the equilibrium value and the relaxation rate for each moment/cumulant.
-2. *Collision/Update Rule*:
-         Methods can generate a "collision rule" which is an equation collection that define the
-         post collision values as a function of the pre-collision values. On these equation collection
-         simplifications are applied to reduce the number of floating point operations.
-         At this stage an entropic optimization step can also be added to determine one relaxation rate by an
-         entropy condition.
-         Then a streaming rule is added which transforms the collision rule into an update rule.
-         The streaming step depends on the pdf storage (source/destination, AABB pattern, EsoTwist).
-         Currently only the simple source/destination  pattern is supported.
-3. *AST*:
-        The abstract syntax tree describes the structure of the kernel, including loops and conditionals.
-        The ast can be modified e.g. to add OpenMP pragmas, reorder loops or apply other optimizations.
-4. *Function*:
-        This step compiles the AST into an executable function, either for CPU or GPUs. This function
-        behaves like a normal Python function and runs one LBM time step.
-
-The function :func:`createLatticeBoltzmannFunction` runs the whole pipeline, the other functions in this module
-execute this pipeline only up to a certain step. Each function optionally also takes the result of the previous step.
-
-For example, to modify the AST one can run::
-
-    ast = createLatticeBoltzmannAst(...)
-    # modify ast here
-    func = createLatticeBoltzmannFunction(ast=ast, ...)
-
-
 Parameters
 ----------
 
@@ -69,7 +33,8 @@ General:
       yet, only the relaxation pattern. To get the entropic method, see parameters below!
       (:func:`lbmpy.methods.createKBCTypeTRT`)
 - ``relaxationRates``: sequence of relaxation rates, number depends on selected method. If you specify more rates than
-  method needs, the additional rates are ignored.
+  method needs, the additional rates are ignored. For SRT and TRT models it is possible ot define a single
+  ``relaxationRate`` instead of a list, the second rate for TRT is then determined via magic number.
 - ``compressible=False``: affects the selection of equilibrium moments. Both options approximate the *incompressible*
   Navier Stokes Equations. However when chosen as False, the approximation is better, the standard LBM derivation is
   compressible.
@@ -130,6 +95,44 @@ Other:
   specifying the number of threads. If True is specified OpenMP chooses the number of threads
 - ``doublePrecision=True``:  by default simulations run with double precision floating point numbers, by setting this
   parameter to False, single precision is used, which is much faster, especially on GPUs
+
+
+
+
+Terminology and creation pipeline
+---------------------------------
+
+Kernel functions are created in three steps:
+
+1. *Method*:
+         the method defines the collision process. Currently there are two big categories:
+         moment and cumulant based methods. A method defines how each moment or cumulant is relaxed by
+         storing the equilibrium value and the relaxation rate for each moment/cumulant.
+2. *Collision/Update Rule*:
+         Methods can generate a "collision rule" which is an equation collection that define the
+         post collision values as a function of the pre-collision values. On these equation collection
+         simplifications are applied to reduce the number of floating point operations.
+         At this stage an entropic optimization step can also be added to determine one relaxation rate by an
+         entropy condition.
+         Then a streaming rule is added which transforms the collision rule into an update rule.
+         The streaming step depends on the pdf storage (source/destination, AABB pattern, EsoTwist).
+         Currently only the simple source/destination  pattern is supported.
+3. *AST*:
+        The abstract syntax tree describes the structure of the kernel, including loops and conditionals.
+        The ast can be modified e.g. to add OpenMP pragmas, reorder loops or apply other optimizations.
+4. *Function*:
+        This step compiles the AST into an executable function, either for CPU or GPUs. This function
+        behaves like a normal Python function and runs one LBM time step.
+
+The function :func:`createLatticeBoltzmannFunction` runs the whole pipeline, the other functions in this module
+execute this pipeline only up to a certain step. Each function optionally also takes the result of the previous step.
+
+For example, to modify the AST one can run::
+
+    ast = createLatticeBoltzmannAst(...)
+    # modify ast here
+    func = createLatticeBoltzmannFunction(ast=ast, ...)
+
 """
 import sympy as sp
 from copy import copy
@@ -138,6 +141,7 @@ from functools import partial
 from lbmpy.methods import createSRT, createTRT, createOrthogonalMRT, createKBCTypeTRT, \
     createRawMRT, createThreeRelaxationRateMRT
 from lbmpy.methods.entropic import addIterativeEntropyCondition, addEntropyCondition
+from lbmpy.methods.relaxationrates import relaxationRateFromMagicNumber
 from lbmpy.stencils import getStencil
 import lbmpy.forcemodels as forceModels
 from lbmpy.simplificationfactory import createSimplificationStrategy
@@ -179,6 +183,12 @@ def updateWithDefaultParameters(params, optParams):
         'gpuIndexing': 'block',
         'gpuIndexingParams': {},
     }
+    if 'relaxationRate' in params:
+        if 'relaxationRates' not in params:
+            params['relaxationRates'] = [params['relaxationRate'],
+                                         relaxationRateFromMagicNumber(params['relaxationRate'])]
+            del params['relaxationRate']
+
     unknownParams = [k for k in params.keys() if k not in defaultMethodDescription]
     unknownOptParams = [k for k in optParams.keys() if k not in defaultOptimizationDescription]
     if unknownParams:

plot2d.py

@@ -2,28 +2,127 @@ from matplotlib.pyplot import *
 
 
 def vectorField(field, step=2, **kwargs):
+    """
+    Plot given vector field as quiver (arrow) plot.
+
+    :param field: numpy array with 3 dimensions, first two are spatial x,y coordinate, the last
+                  coordinate should have shape 2 and stores the 2 velocity components
+    :param step: plots only every steps's cell
+    :param kwargs: keyword arguments passed to :func:`matplotlib.pyplot.quiver`
+    """
     veln = field.swapaxes(0, 1)
-    quiver(veln[::step, ::step, 0], veln[::step, ::step, 1], **kwargs)
+    res = quiver(veln[::step, ::step, 0], veln[::step, ::step, 1], **kwargs)
+    axis('equal')
+    return res
 
 
 def vectorFieldMagnitude(field, **kwargs):
+    """
+    Plots the magnitude of a vector field as colormap
+    :param field: numpy array with 3 dimensions, first two are spatial x,y coordinate, the last
+                  coordinate should have shape 2 and stores the 2 velocity components
+    :param kwargs: keyword arguments passed to :func:`matplotlib.pyplot.imshow`
+    """
     from numpy.linalg import norm
     norm = norm(field, axis=2, ord=2)
     if hasattr(field, 'mask'):
-        norm = np.ma.masked_array(norm, mask=field.mask[:,:,0])
+        norm = np.ma.masked_array(norm, mask=field.mask[:, :, 0])
     return scalarField(norm, **kwargs)
 
 
 def scalarField(field, **kwargs):
+    """
+    Plots field values as colormap
+
+    :param field: two dimensional numpy array
+    :param kwargs: keyword arguments passed to :func:`matplotlib.pyplot.imshow`
+    """
     import numpy as np
     field = np.swapaxes(field, 0, 1)
-    return imshow(field, origin='lower', **kwargs)
+    res = imshow(field, origin='lower', **kwargs)
+    axis('equal')
+    return res
+
+
+def plotBoundaryHandling(boundaryHandling, boundaryNameToColor=None):
+    """
+    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
+
+    if len(boundaryHandling.flagField.shape) != 2:
+        raise NotImplementedError("Only implemented for 2D boundary handlings")
+
+    if boundaryNameToColor:
+        fixedColors = boundaryNameToColor
+    else:
+        fixedColors = {
+            'fluid': '#1f77ff11',
+            'noSlip': '#000000'
+        }
+
+    boundaryNames = ['fluid'] + boundaryHandling.getBoundaryNames()
+    flagValues = [boundaryHandling.fluidFlag] + [boundaryHandling.getFlag(n)
+                                                 for n in boundaryHandling.getBoundaryNames()]
+
+    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.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')
+    plt.legend(handles=patches, bbox_to_anchor=(1.02, 0.5), loc=2, borderaxespad=0.)
+
+# ------------------------------------------- Animations ---------------------------------------------------------------
+
+
+def vectorFieldAnimation(runFunction, step=2, rescale=True, plotSetupFunction=lambda: None,
+                         plotUpdateFunction=lambda: None, interval=30, frames=180, **kwargs):
+    import matplotlib.animation as animation
+    from numpy.linalg import norm
+
+    fig = gcf()
+    im = None
+    field = runFunction()
+    if rescale:
+        maxNorm = np.max(norm(field, axis=2, ord=2))
+        field /= maxNorm
+        if 'scale' not in kwargs:
+            kwargs['scale'] = 1.0
+
+    quiverPlot = vectorField(field, step=step, **kwargs)
+    plotSetupFunction()
+
+    def updatefig(*args):
+        f = runFunction()
+        f = np.swapaxes(f, 0, 1)
+        if rescale:
+            maxNorm = np.max(norm(f, axis=2, ord=2))
+            f /= maxNorm
+        u, v = f[::step, ::step, 0], f[::step, ::step, 1]
+        quiverPlot.set_UVC(u, v)
+        plotUpdateFunction()
+        return im,
+
+    return animation.FuncAnimation(fig, updatefig, interval=interval, frames=frames)
 
 
 def vectorFieldMagnitudeAnimation(runFunction, plotSetupFunction=lambda: None,
                                   plotUpdateFunction=lambda: None, interval=30, frames=180, **kwargs):
     import matplotlib.animation as animation
-    import numpy as np
     from numpy.linalg import norm
 
     fig = gcf()
@@ -42,4 +141,4 @@ def vectorFieldMagnitudeAnimation(runFunction, plotSetupFunction=lambda: None,
         plotUpdateFunction()
         return im,
 
-    return animation.FuncAnimation(fig, updatefig, interval=interval, frames=frames)
+    return animation.FuncAnimation(fig, updatefig, interval=interval, frames=frames)
\ No newline at end of file