lbmpy: simplified boundary setup & documentation for creation functions

boundaries/boundaryhandling.py

@@ -3,13 +3,31 @@ import numpy as np
 from pystencils import TypedSymbol, Field
 from pystencils.backends.cbackend import CustomCppCode
 from lbmpy.boundaries.createindexlist import createBoundaryIndexList
+from pystencils.slicing import normalizeSlice
 
 INV_DIR_SYMBOL = TypedSymbol("invDir", "int")
 WEIGHTS_SYMBOL = TypedSymbol("weights", "double")
 
 
 class BoundaryHandling(object):
-    def __init__(self, symbolicPdfField, domainShape, lbMethod, ghostLayers=1, target='cpu'):
+    def __init__(self, pdfField, domainShape, lbMethod, ghostLayers=1, target='cpu'):
+        """
+        Class for managing boundary kernels
+
+        :param pdfField: either pdf numpy array (including ghost layers), or pystencils.Field
+        :param domainShape: domain size without ghost layers
+        :param lbMethod: lattice Boltzmann method
+        :param ghostLayers: number of ghost layers
+        :param target: either 'cpu' or 'gpu'
+        """
+        if isinstance(pdfField, np.ndarray):
+            symbolicPdfField = Field.createFromNumpyArray('pdfs', pdfField, indexDimensions=1)
+            assert pdfField.shape[:-1] == tuple(d + 2*ghostLayers for d in domainShape)
+        elif isinstance(pdfField, Field):
+            symbolicPdfField = pdfField
+        else:
+            raise ValueError("pdfField has to be either a numpy array or a pystencils.Field")
+
         self._symbolicPdfField = symbolicPdfField
         self._shapeWithGhostLayers = [d + 2 * ghostLayers for d in domainShape]
         self._fluidFlag = 2 ** 30
@@ -23,7 +41,48 @@ class BoundaryHandling(object):
         if target not in ('cpu', 'gpu'):
             raise ValueError("Invalid target '%s' . Allowed values: 'cpu' or 'gpu'" % (target,))
 
+    def setBoundary(self, function, indexExpr, maskArr=None, name=None):
+        """
+        Sets boundary in a rectangular region (slice)
+
+        :param function: boundary
+        :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 function not in self._boundaryFunctions:
+            self.addBoundary(function, name)
+
+        flag = self.getFlag(name)
+
+        indexExpr = normalizeSlice(indexExpr, self._shapeWithGhostLayers)
+        if maskArr is None:
+            self.flagField[indexExpr] = flag
+        else:
+            flagFieldView = self.flagField[indexExpr]
+            flagFieldView[maskArr] = flag
+
+        self.invalidateIndexCache()
+
     def addBoundary(self, boundaryFunction, name=None):
+        """
+        Adds a boundary condition, i.e. reserves a flog in the flag field and returns that flag
+        If a boundary with that name already exists, the existing flag is returned.
+        This flag can be logicalled or'ed to the boundaryHandling.flagField
+
+        :param boundaryFunction: boundary condition function, see :mod:`lbmpy.boundaries.boundaryconditions`
+        :param name: boundaries with different name are considered different. If not given
+                     ```boundaryFunction.__name__`` is used
+        """
         if name is None:
             name = boundaryFunction.__name__
 
@@ -45,27 +104,6 @@ class BoundaryHandling(object):
     def getFlag(self, name):
         return 2 ** self._nameToIndex[name]
 
-    def setBoundary(self, function, indexExpr, maskArr=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 function not in self._boundaryFunctions:
-            self.addBoundary(function, name)
-
-        flag = self.getFlag(name)
-
-        if maskArr is None:
-            self.flagField[indexExpr] = flag
-        else:
-            flagFieldView = self.flagField[indexExpr]
-            flagFieldView[maskArr] = flag
-
-        self.invalidateIndexCache()
-
     def prepare(self):
         self.invalidateIndexCache()
         for boundaryIdx, boundaryFunc in enumerate(self._boundaryFunctions):
@@ -86,6 +124,8 @@ class BoundaryHandling(object):
                 assert False
 
     def __call__(self, **kwargs):
+        if len(self._boundaryFunctions) == 0:
+            return
         if len(self._boundarySweeps) == 0:
             self.prepare()
         for boundarySweep in self._boundarySweeps:

creationfunctions.py

-"""
-Factory functions for standard LBM methods
+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
+----------
+
+The following list describes common parameters for the creation functions. They have to be passed as keyword parameters.
+
+
+Method parameters
+^^^^^^^^^^^^^^^^^
+
+General:
+
+- ``stencil='D2Q9'``: stencil name e.g. 'D2Q9', 'D3Q19'. See :func:`pystencils.stencils.getStencil` for details
+- ``method='srt'``: name of lattice Boltzmann method. This determines the selection and relaxation pattern of
+  moments/cumulants, i.e. which moment/cumulant basis is chosen, and which of the basis vectors are relaxed together
+    - ``srt``: single relaxation time (:func:`lbmpy.methods.createSRT`)
+    - ``trt``: two relaxation time, first relaxation rate is for even moments and determines the viscosity (as in SRT),
+      the second relaxation rate is used for relaxing odd moments, and controls the bulk viscosity.
+      (:func:`lbmpy.methods.createTRT`)
+    - ``mrt``: orthogonal multi relaxation time model, number of relaxation rates depends on the stencil
+      (:func:`lbmpy.methods.createOrthogonalMRT`)
+    - ``mrt3``: three relaxation time method, where shear moments are relaxed with first relaxation rate (and therefore
+      determine viscosity, second rate relaxes the shear tensor trace (determines bulk viscosity) and last rate relaxes
+      all other, higher order moments. If two relaxation rates are chosen the same this is equivalent to a KBC type
+      relaxation (:func:`lbmpy.methods.createThreeRelaxationRateMRT`)
+    - ``mrt_raw``: non-orthogonal MRT where all relaxation rates can be specified independently i.e. there are as many
+      relaxation rates as stencil entries. Look at the generated method in Jupyter to see which moment<->relaxation rate
+      mapping (:func:`lbmpy.methods.createRawMRT`)
+    - ``trt-kbc-n<N>`` where <N> is 1,2,3 or 4. Special two-relaxation method. This is not the entropic method
+      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.
+- ``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.
+- ``equilibriumAccuracyOrder=2``: order in velocity, at which the equilibrium moment/cumulant approximation is
+  truncated. Order 2 is sufficient to approximate Navier-Stokes
+- ``forceModel=None``: possible values: ``None``, ``'simple'``, ``'luo'``, ``'guo'``. For details see
+  :mod:`lbmpy.forcemodels`
+- ``useContinuousMaxwellianEquilibrium=True``: way to compute equilibrium moments/cumulants, if False the standard
+  discretized LBM equilibrium is used, otherwise the equilibrium moments are computed from the continuous Maxwellian.
+  This makes only a difference if sparse stencils are used e.g. D2Q9 and D3Q27 are not affected, D319 and DQ15 are
+- ``cumulant=False``: use cumulants instead of moments
+- ``initialVelocity=None``: initial velocity in domain, can either be a tuple (x,y,z) velocity to set a constant
+  velocity everywhere, or a numpy array with the same size of the domain, with a last coordinate of shape dim to set
+  velocities on cell level
+
+Entropic methods:
+
+- ``entropic=False``: In case there are two distinct relaxation rate in a method, one of them (usually the one, not
+  determining the viscosity) can be automatically chosen w.r.t an entropy condition. For details see
+  :mod:`lbmpy.methods.entropic`
+- ``entropicNewtonIterations=None``: For moment methods the entropy optimum can be calculated in closed form.
+  For cumulant methods this is not possible, in that case it is computed using Newton iterations. This parameter can be
+  used to force Newton iterations and specify how many should be done
+- ``omegaOutputField=None``: you can pass a pystencils Field here, where the calculated free relaxation
+  rate is written to
+
+
+Optimization Parameters
+^^^^^^^^^^^^^^^^^^^^^^^
+
+Simplifications / Transformations:
+
+- ``doCseInOpposingDirections=False``: run common subexpression elimination for opposing stencil directions
+- ``doOverallCse=False``: run common subexpression elimination after all other simplifications have been executed
+- ``split=False``: split innermost loop, to handle only 2 directions per loop. This reduces the number of parallel
+  load/store streams and thus speeds up the kernel on most architectures
+
+
+Field size information:
+
+- ``pdfArr``: pass a numpy array here to create kernels with fixed size and create the loop nest according to layout
+  of this array
+- ``fieldSize``: create kernel for fixed field size
+- ``fieldLayout='c'``:   ``'c'`` or ``'numpy'`` for standard numpy layout, ``'reverseNumpy'`` or ``'f'`` for fortran
+  layout, this does not apply when pdfArr was given, then the same layout as pdfArr is used
+
+GPU:
+
+- ``target='cpu'``: ``'cpu'`` or ``'gpu'``, last option requires a CUDA enabled graphics card
+  and installed *pycuda* package
+- ``gpuIndexing='block'``: determines mapping of CUDA threads to cells. Can be either ``'block'`` or ``'line'``
+- ``gpuIndexingParams='block'``: parameters passed to init function of gpu indexing.
+  For ``'block'`` indexing one can e.g. specify the block size ``{'blockSize' : (128, 4, 1)}``
+
+Other:
+
+- ``openMP=True``: only applicable for cpu simulations. Can be a boolean to turn multi threading on/off, or an integer
+  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
 """
 import sympy as sp
 from copy import copy
 from functools import partial
 
-from lbmpy.methods.creationfunctions import createKBCTypeTRT, createRawMRT, createThreeRelaxationRateMRT
+from lbmpy.methods import createSRT, createTRT, createOrthogonalMRT, createKBCTypeTRT, \
+    createRawMRT, createThreeRelaxationRateMRT
 from lbmpy.methods.entropic import addIterativeEntropyCondition, addEntropyCondition
 from lbmpy.stencils import getStencil
-from lbmpy.methods import createSRT, createTRT, createOrthogonalMRT
 import lbmpy.forcemodels as forceModels
 from lbmpy.simplificationfactory import createSimplificationStrategy
 from lbmpy.updatekernels import createStreamPullKernel, createPdfArray
@@ -22,15 +152,15 @@ def updateWithDefaultParameters(params, optParams):
         'compressible': False,
         'equilibriumAccuracyOrder': 2,
 
-        'entropic': False,
-        'entropicNewtonIterations': None,
-        'omegaOutputField': None,
-
-        'useContinuousMaxwellianEquilibrium': False,
-        'cumulant': False,
         'forceModel': 'none',  # can be 'simple', 'luo' or 'guo'
         'force': (0, 0, 0),
+        'useContinuousMaxwellianEquilibrium': True,
+        'cumulant': False,
         'initialVelocity': None,
+
+        'entropic': False,
+        'entropicNewtonIterations': None,
+        'omegaOutputField': None,
     }
 
     defaultOptimizationDescription = {
@@ -39,7 +169,7 @@ def updateWithDefaultParameters(params, optParams):
         'split': False,
 
         'fieldSize': None,
-        'fieldLayout': 'reverseNumpy',  # can be 'numpy' (='c'), 'reverseNumpy' (='f'), 'fzyx', 'zyxf'
+        'fieldLayout': 'c',  # can be 'numpy' (='c'), 'reverseNumpy' (='f'), 'fzyx', 'zyxf'
 
         'target': 'cpu',
         'openMP': True,
@@ -140,7 +270,7 @@ def createLatticeBoltzmannUpdateRule(lbMethod=None, optimizationParams={}, **kwa
 
     if params['entropic']:
         if params['entropicNewtonIterations']:
-            if isinstance(params['entropicNewtonIterations'], bool):
+            if isinstance(params['entropicNewtonIterations'], bool) or params['cumulant']:
                 iterations = 3
             else:
                 iterations = params['entropicNewtonIterations']

methods/__init__.py

 from lbmpy.methods.abstractlbmethod import AbstractLbMethod
 from lbmpy.methods.momentbased import MomentBasedLbMethod, RelaxationInfo
 from lbmpy.methods.creationfunctions import createSRT, createTRT, createTRTWithMagicNumber, createOrthogonalMRT, \
-    createWithContinuousMaxwellianEqMoments, createWithDiscreteMaxwellianEqMoments
+    createWithContinuousMaxwellianEqMoments, createWithDiscreteMaxwellianEqMoments, createKBCTypeTRT, createRawMRT, \
+    createThreeRelaxationRateMRT
 from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation, DensityVelocityComputation

plot2d.py

@@ -8,8 +8,10 @@ def vectorField(field, step=2, **kwargs):
 
 def vectorFieldMagnitude(field, **kwargs):
     from numpy.linalg import norm
-    field = norm(field, axis=2, ord=2)
-    return scalarField(field, **kwargs)
+    norm = norm(field, axis=2, ord=2)
+    if hasattr(field, 'mask'):
+        norm = np.ma.masked_array(norm, mask=field.mask[:,:,0])
+    return scalarField(norm, **kwargs)
 
 
 def scalarField(field, **kwargs):
@@ -32,9 +34,11 @@ def vectorFieldMagnitudeAnimation(runFunction, plotSetupFunction=lambda: None,
 
     def updatefig(*args):
         f = runFunction()
-        f = norm(f, axis=2, ord=2)
-        f = np.swapaxes(f, 0, 1)
-        im.set_array(f)
+        normed = norm(f, axis=2, ord=2)
+        if hasattr(f, 'mask'):
+            normed = np.ma.masked_array(normed, mask=f.mask[:, :, 0])
+        normed = np.swapaxes(normed, 0, 1)
+        im.set_array(normed)
         plotUpdateFunction()
         return im,
 

scenarios.py

+import numpy as np
+from functools import partial
+from pystencils.field import getLayoutOfArray, createNumpyArrayWithLayout
+from pystencils.slicing import sliceFromDirection, addGhostLayers, getPeriodicBoundaryFunctor, removeGhostLayers, \
+    normalizeSlice, makeSlice
+from lbmpy.creationfunctions import createLatticeBoltzmannFunction, updateWithDefaultParameters
+from lbmpy.macroscopic_value_kernels import compileMacroscopicValuesGetter, compileMacroscopicValuesSetter
+from lbmpy.boundaries import BoundaryHandling, noSlip, ubb, fixedDensity
+from lbmpy.stencils import getStencil
+from lbmpy.updatekernels import createPdfArray
+
+
+class Scenario(object):
+    def __init__(self, domainSize, methodParameters, optimizationParams, lbmKernel=None,
+                 initialVelocity=None, preUpdateFunctions=[], kernelParams={}):
+        """
+        Constructor for generic scenarios. You probably want to use one of the simpler scenario factory functions
+        of this file.
+
+        :param domainSize: tuple, defining the domain size without ghost layers
+        :param methodParameters: dict with method parameters, as documented in :mod:`lbmpy.creationfunctions`,
+                                 passed to :func:`lbmpy.creationfunction.createLatticeBoltzmannFunction`
+        :param optimizationParams: dict with optimization parameters, as documented in :mod:`lbmpy.creationfunctions`,
+                                   passed to :func:`lbmpy.creationfunction.createLatticeBoltzmannFunction`
+        :param lbmKernel: a lattice boltzmann function can be passed here, if None it is created with the parameters
+                          specified above
+        :param initialVelocity: tuple with initial velocity of the domain, can either be a constant or a numpy array
+                                with first axes shaped like the domain, and the last dimension of size #dimensions
+        :param preUpdateFunctions: list of functions that are called before the LBM kernel. They get the pdf array as
+                                   only argument. Can be used for custom boundary conditions, periodicity, etc.
+        :param kernelParams: dict which is passed to the kernel as additional parameters
+        """
+
+        ghostLayers = 1
+        domainSizeWithGhostLayer = tuple([s + 2 * ghostLayers for s in domainSize])
+        D = len(domainSize)
+        if 'stencil' not in methodParameters:
+            methodParameters['stencil'] = 'D2Q9' if D == 2 else 'D3Q27'
+
+        if isinstance(initialVelocity, np.ndarray):
+            initialVelocity = addGhostLayers(initialVelocity, indexDimensions=1, ghostLayers=1)
+
+        methodParameters, optimizationParams = updateWithDefaultParameters(methodParameters, optimizationParams)
+
+        Q = len(getStencil(methodParameters['stencil']))
+        self._pdfArrays = [createPdfArray(domainSize, Q, layout=optimizationParams['fieldLayout']),
+                           createPdfArray(domainSize, Q, layout=optimizationParams['fieldLayout'])]
+
+        # Create kernel
+        if lbmKernel is None:
+            optimizationParams['pdfArr'] = self._pdfArrays[0]
+            methodParameters['optimizationParams'] = optimizationParams
+            self._lbmKernel = createLatticeBoltzmannFunction(**methodParameters)
+        else:
+            self._lbmKernel = lbmKernel
+        assert D == self._lbmKernel.method.dim, "Domain size and stencil do not match"
+
+        # Macroscopic value input/output
+        pdfArrLayout = getLayoutOfArray(self._pdfArrays[0])
+        pdfArrLayoutNoIdx = getLayoutOfArray(self._pdfArrays[0], indexDimensionIds=[D])
+        self._density = createNumpyArrayWithLayout(domainSizeWithGhostLayer, layout=pdfArrLayoutNoIdx)
+        self._velocity = createNumpyArrayWithLayout(list(domainSizeWithGhostLayer) + [D], layout=pdfArrLayout)
+        self._getMacroscopic = compileMacroscopicValuesGetter(self._lbmKernel.method, ['density', 'velocity'],
+                                                              pdfArr=self._pdfArrays[0], target='cpu')
+
+        self._boundaryHandling = BoundaryHandling(self._pdfArrays[0], domainSize, self.method,
+                                                  target=optimizationParams['target'])
+
+        self._preUpdateFunctions = preUpdateFunctions
+        self._kernelParams = kernelParams
+        self._pdfGpuArrays = []
+
+        if initialVelocity is None:
+            initialVelocity = [0] * D
+
+        setMacroscopic = compileMacroscopicValuesSetter(self.method, {'density': 1.0, 'velocity': initialVelocity},
+                                                        pdfArr=self._pdfArrays[0], target='cpu')
+        setMacroscopic(pdfs=self._pdfArrays[0])
+
+        if optimizationParams['target'] == 'gpu':
+            import pycuda.gpuarray as gpuarray
+            self._pdfGpuArrays = [gpuarray.to_gpu(a) for a in self._pdfArrays]
+        else:
+            self._pdfGpuArrays = []
+
+    def run(self, timeSteps=1):
+        """Run the scenario for the given amount of time steps"""
+        if len(self._pdfGpuArrays) > 0:
+            self._gpuTimeloop(timeSteps)
+        else:
+            self._cpuTimeloop(timeSteps)
+
+    @property
+    def velocity(self):
+        """Velocity as numpy array"""
+        mask = np.logical_not(np.bitwise_and(self.boundaryHandling.flagField, self.boundaryHandling._fluidFlag))
+        mask = np.repeat(mask[..., np.newaxis], self.dim, axis=2)
+        return removeGhostLayers(np.ma.masked_array(self._velocity, mask), indexDimensions=1)
+
+    @property
+    def density(self):
+        """Density as numpy array"""
+        mask = np.logical_not(np.bitwise_and(self.boundaryHandling.flagField, self.boundaryHandling._fluidFlag))
+        return removeGhostLayers(np.ma.masked_array(self._density, mask))
+
+    @property
+    def pdfs(self):
+        """Particle distribution functions as numpy array"""
+        return removeGhostLayers(self._pdfArrays[0], indexDimensions=1)
+
+    @property
+    def boundaryHandling(self):
+        """Boundary handling instance of the scenario. Use this to change the boundary setup"""
+        return self._boundaryHandling
+
+    @property
+    def method(self):
+        """Lattice boltzmann method description"""
+        return self._lbmKernel.method
+
+    @property
+    def dim(self):
+        """Dimension of the domain"""
+        return self.method.dim
+
+    @property
+    def ast(self):
+        """Returns abstract syntax tree of the kernel"""
+        return self._lbmKernel.ast
+
+    @property
+    def updateRule(self):
+        """Equation collection defining the LBM update rule (already in simplified form)"""
+        return self._lbmKernel.updateRule
+
+    def animateVelocity(self, steps=10, **kwargs):
+        import lbmpy.plot2d as plt
+
+        def runFunction():
+            self.run(steps)
+            return self.velocity
+        return plt.vectorFieldMagnitudeAnimation(runFunction, **kwargs)
+
+    def plotVelocity(self, **kwargs):
+        import lbmpy.plot2d as plt
+        if self.dim == 2:
+            plt.vectorFieldMagnitude(self.velocity, **kwargs)
+            plt.title("Velocity Magnitude")
+            plt.colorbar()
+            plt.axis('equal')
+        elif self.dim == 3:
+            idxSlice = normalizeSlice(makeSlice[:, :, 0.5], self.velocity.shape[:3])
+            plt.vectorFieldMagnitude(self.velocity[idxSlice], **kwargs)
+            plt.title("Velocity Magnitude in x-y (z at half of domain)")
+            plt.colorbar()
+            plt.axis('equal')
+        else:
+            raise NotImplementedError("Can only plot 2D and 3D scenarios")
+
+    def plotDensity(self, **kwargs):
+        import lbmpy.plot2d as plt
+        if self.dim == 2:
+            plt.scalarField(self.density, **kwargs)
+            plt.title("Density")
+            plt.colorbar()
+            plt.axis('equal')
+        elif self.dim == 3:
+            idxSlice = normalizeSlice(makeSlice[:, :, 0.5], self.density.shape)
+            plt.scalarField(self.density[idxSlice], **kwargs)
+            plt.title("Density in x-y (z at half of domain)")
+            plt.colorbar()
+            plt.axis('equal')
+        else:
+            raise NotImplementedError("Can only plot 2D and 3D scenarios")
+
+    def _timeloop(self, pdfArrays, timeSteps):
+        for t in range(timeSteps):
+            for f in self._preUpdateFunctions:
+                f(pdfArrays[0])
+            if self._boundaryHandling is not None:
+                self._boundaryHandling(pdfs=pdfArrays[0])
+            self._lbmKernel(src=pdfArrays[0], dst=pdfArrays[1], **self._kernelParams)
+
+            pdfArrays[0], pdfArrays[1] = pdfArrays[1], pdfArrays[0]  # swap
+
+    def _cpuTimeloop(self, timeSteps):
+        self._timeloop(self._pdfArrays, timeSteps)
+        self._getMacroscopic(pdfs=self._pdfArrays[0], density=self._density, velocity=self._velocity)
+
+    def _gpuTimeloop(self, timeSteps):
+        # Transfer data to gpu
+        for cpuArr, gpuArr in zip(self._pdfArrays, self._pdfGpuArrays):
+            gpuArr.set(cpuArr)
+
+        self._timeloop(self._pdfGpuArrays, timeSteps)
+
+        # Transfer data from gpu to cpu
+        for cpuArr, gpuArr in zip(self._pdfArrays, self._pdfGpuArrays):
+            gpuArr.get(cpuArr)
+
+        self._getMacroscopic(pdfs=self._pdfArrays[0], density=self._density, velocity=self._velocity)
+
+
+# ---------------------------------------- Example Scenarios -----------------------------------------------------------
+
+
+def createFullyPeriodicFlow(initialVelocity, optimizationParams={}, lbmKernel=None, kernelParams={}, **kwargs):
+    """
+    Creates a fully periodic setup with prescribed velocity field
+    """
+    domainSize = initialVelocity.shape
+
+    stencil = getStencil(kwargs['stencil'])
+    periodicity = getPeriodicBoundaryFunctor(stencil)
+
+    return Scenario(domainSize, kwargs, optimizationParams, lbmKernel, initialVelocity, kernelParams, [periodicity])
+
+
+def createLidDrivenCavity(domainSize, lidVelocity=0.005, optimizationParams={}, lbmKernel=None,
+                          kernelParams={}, **kwargs):
+    """
+    Creates a lid driven cavity scenario with prescribed lid velocity in lattice coordinates
+    """
+    scenario = Scenario(domainSize, kwargs, optimizationParams, lbmKernel=lbmKernel, kernelParams=kernelParams)
+
+    myUbb = partial(ubb, velocity=[lidVelocity, 0, 0][:scenario.method.dim])
+    myUbb.name = 'ubb'
+    dim = scenario.method.dim
+    scenario.boundaryHandling.setBoundary(myUbb, sliceFromDirection('N', dim))
+    for direction in ('W', 'E', 'S') if scenario.method.dim == 2 else ('W', 'E', 'S', 'T', 'B'):
+        scenario.boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, dim))
+
+    return scenario
+
+
+def createPressureGradientDrivenChannel(dim, pressureDifference, domainSize=None, radius=None, length=None,
+                                        lbmKernel=None, optimizationParams={}, initialVelocity=None,
+                                        boundarySetupFunctions=[], kernelParams={}, **kwargs):
+    """
+    Creates a channel flow in x direction, which is driven by two pressure boundaries.
+    """
+    assert dim in (2, 3)
+
+    if radius is not None:
+        assert length is not None
+        if dim == 3:
+            domainSize = (length, 2 * radius + 1, 2 * radius + 1)
+            roundChannel = True
+        else:
+            if domainSize is None:
+                domainSize = (length, 2 * radius)
+    else:
+        roundChannel = False
+
+    if domainSize is None:
+        raise ValueError("Missing domainSize or radius and length parameters!")
+
+    scenario = Scenario(domainSize, kwargs, optimizationParams, lbmKernel=lbmKernel, kernelParams=kernelParams,
+                        initialVelocity=initialVelocity)
+    boundaryHandling = scenario.boundaryHandling
+    pressureBoundaryInflow = partial(fixedDensity, density=1.0 + pressureDifference)
+    pressureBoundaryInflow.__name__ = "Inflow"
+
+    pressureBoundaryOutflow = partial(fixedDensity, density=1.0)
+    pressureBoundaryOutflow.__name__ = "Outflow"
+    boundaryHandling.setBoundary(pressureBoundaryInflow, sliceFromDirection('W', dim))
+    boundaryHandling.setBoundary(pressureBoundaryOutflow, sliceFromDirection('E', dim))
+
+    if dim == 2:
+        for direction in ('N', 'S'):
+            boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, dim))
+    elif dim == 3:
+        if roundChannel:
+            noSlipIdx = boundaryHandling.addBoundary(noSlip)
+            ff = boundaryHandling.flagField
+            yMid = ff.shape[1] // 2
+            zMid = ff.shape[2] // 2
+            y, z = np.meshgrid(range(ff.shape[1]), range(ff.shape[2]))
+            ff[(y - yMid) ** 2 + (z - zMid) ** 2 > radius ** 2] = noSlipIdx
+        else:
+            for direction in ('N', 'S', 'T', 'B'):
+                boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, dim))
+
+    for userFunction in boundarySetupFunctions:
+        userFunction(boundaryHandling, scenario.method, domainSize)
+
+    return scenario
+
+
+def createForceDrivenChannel(dim, force, domainSize=None, radius=None, length=None, lbmKernel=None,
+                             optimizationParams={}, initialVelocity=None, boundarySetupFunctions=[],
+                             kernelParams={}, **kwargs):
+    """
+    Creates a channel flow in x direction, which is driven by a constant force along the x axis
+    """
+    assert dim in (2, 3)
+    kwargs['force'] = tuple([force, 0, 0][:dim])
+
+    if radius is not None:
+        assert length is not None
+        if dim == 3:
+            domainSize = (length, 2 * radius + 1, 2 * radius + 1)
+            roundChannel = True
+        else:
+            if domainSize is None:
+                domainSize = (length, 2 * radius)
+    else:
+        roundChannel = False
+
+    def periodicity(pdfArr):
+        pdfArr[0, :, :] = pdfArr[-2, :, :]
+        pdfArr[-1, :, :] = pdfArr[1, :, :]
+
+    scenario = Scenario(domainSize, kwargs, optimizationParams, lbmKernel=lbmKernel,
+                        initialVelocity=initialVelocity, preUpdateFunctions=[periodicity],
+                        kernelParams=kernelParams)
+
+    boundaryHandling = scenario.boundaryHandling
+
+    if dim == 2:
+        for direction in ('N', 'S'):
+            boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, dim))
+    elif dim == 3:
+        if roundChannel:
+            noSlipIdx = boundaryHandling.addBoundary(noSlip)
+            ff = boundaryHandling.flagField
+            yMid = ff.shape[1] // 2
+            zMid = ff.shape[2] // 2
+            y, z = np.meshgrid(range(ff.shape[1]), range(ff.shape[2]))
+            ff[(y - yMid) ** 2 + (z - zMid) ** 2 > radius ** 2] = noSlipIdx
+        else:
+            for direction in ('N', 'S', 'T', 'B'):
+                boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, dim))
+    for userFunction in boundarySetupFunctions:
+        userFunction(boundaryHandling, scenario.method, domainSize)
+
+    assert domainSize is not None
+    if 'forceModel' not in kwargs:
+        kwargs['forceModel'] = 'guo'
+
+    return scenario

serialscenario.py

-from functools import partial
-import numpy as np
-from pystencils import Field
-from pystencils.field import getLayoutOfArray, createNumpyArrayWithLayout
-from pystencils.slicing import sliceFromDirection, addGhostLayers, getPeriodicBoundaryFunctor
-from lbmpy.creationfunctions import createLatticeBoltzmannFunction, updateWithDefaultParameters
-from lbmpy.macroscopic_value_kernels import compileMacroscopicValuesGetter, compileMacroscopicValuesSetter
-from lbmpy.boundaries import BoundaryHandling, noSlip, ubb, fixedDensity
-from lbmpy.stencils import getStencil
-from lbmpy.updatekernels import createPdfArray
-
-
-def createScenario(domainSize, boundarySetupFunction, methodParameters, optimizationParams, lbmKernel=None,
-                   initialVelocity=None, preUpdateFunctions=[], kernelParams={}):
-    ghostLayers = 1
-    domainSizeWithGhostLayer = tuple([s + 2 * ghostLayers for s in domainSize])
-    D = len(domainSize)
-    if 'stencil' not in methodParameters:
-        methodParameters['stencil'] = 'D2Q9' if D == 2 else 'D3Q27'
-
-    methodParameters, optimizationParams = updateWithDefaultParameters(methodParameters, optimizationParams)
-
-    Q = len(getStencil(methodParameters['stencil']))
-    pdfArrays = [createPdfArray(domainSize, Q, layout=optimizationParams['fieldLayout']),
-                 createPdfArray(domainSize, Q, layout=optimizationParams['fieldLayout'])]
-
-    # Create kernel
-    if lbmKernel is None:
-        methodParameters['optimizationParams'] = optimizationParams
-        lbmKernel = createLatticeBoltzmannFunction(**methodParameters)
-    method = lbmKernel.method
-
-    assert D == method.dim, "Domain size and stencil do not match"
-
-    # Boundary setup
-    if boundarySetupFunction is not None:
-        symPdfField = Field.createFromNumpyArray('pdfs', pdfArrays[0], indexDimensions=1)
-        boundaryHandling = BoundaryHandling(symPdfField, domainSize, lbmKernel.method,
-                                            target=optimizationParams['target'])
-        boundarySetupFunction(boundaryHandling=boundaryHandling, method=method)
-        boundaryHandling.prepare()
-    else:
-        boundaryHandling = None
-
-    # Macroscopic value input/output
-    pdfArrLayout = getLayoutOfArray(pdfArrays[0])
-    pdfArrLayoutNoIdx = getLayoutOfArray(pdfArrays[0], indexDimensionIds=[D])
-    densityArr = [createNumpyArrayWithLayout(domainSizeWithGhostLayer, layout=pdfArrLayoutNoIdx)]
-    velocityArr = [createNumpyArrayWithLayout(list(domainSizeWithGhostLayer) + [D], layout=pdfArrLayout)]
-    getMacroscopic = compileMacroscopicValuesGetter(method, ['density', 'velocity'], pdfArr=pdfArrays[0], target='cpu')
-
-    if initialVelocity is None:
-        initialVelocity = [0] * D
-    setMacroscopic = compileMacroscopicValuesSetter(method, {'density': 1.0, 'velocity': initialVelocity},
-                                                    pdfArr=pdfArrays[0], target='cpu')
-    setMacroscopic(pdfs=pdfArrays[0])
-
-    if optimizationParams['target'] == 'gpu':
-        import pycuda.gpuarray as gpuarray
-        pdfGpuArrays = [gpuarray.to_gpu(a) for a in pdfArrays]
-    else:
-        pdfGpuArrays = []
-
-    def cpuTimeLoop(timeSteps):
-        for t in range(timeSteps):
-            for f in preUpdateFunctions:
-                f(pdfArrays[0])
-            if boundaryHandling is not None:
-                boundaryHandling(pdfs=pdfArrays[0])
-            lbmKernel(src=pdfArrays[0], dst=pdfArrays[1], **kernelParams)
-
-            pdfArrays[0], pdfArrays[1] = pdfArrays[1], pdfArrays[0]
-        getMacroscopic(pdfs=pdfArrays[0], density=densityArr[0], velocity=velocityArr[0])
-        return pdfArrays[0], densityArr[0], velocityArr[0]
-
-    def gpuTimeLoop(timeSteps):
-        # Transfer data to gpu
-        for cpuArr, gpuArr in zip(pdfArrays, pdfGpuArrays):
-            gpuArr.set(cpuArr)
-
-        for t in range(timeSteps):
-            for f in preUpdateFunctions:
-                f(pdfGpuArrays[0])
-            if boundaryHandling is not None:
-                boundaryHandling(pdfs=pdfGpuArrays[0])
-            lbmKernel(src=pdfGpuArrays[0], dst=pdfGpuArrays[1], **kernelParams)
-
-            pdfGpuArrays[0], pdfGpuArrays[1] = pdfGpuArrays[1], pdfGpuArrays[0]
-
-        # Transfer data from gpu to cpu
-        for cpuArr, gpuArr in zip(pdfArrays, pdfGpuArrays):
-            gpuArr.get(cpuArr)
-
-        getMacroscopic(pdfs=pdfArrays[0], density=densityArr[0], velocity=velocityArr[0])
-        return pdfArrays[0], densityArr[0], velocityArr[0]
-
-    cpuTimeLoop.kernel = lbmKernel
-    gpuTimeLoop.kernel = lbmKernel
-
-    return gpuTimeLoop if optimizationParams['target'] == 'gpu' else cpuTimeLoop
-
-
-def createFullyPeriodicFlow(initialVelocity, optimizationParams={}, lbmKernel=None, kernelParams={}, **kwargs):
-    domainSize = initialVelocity.shape
-    # extend velocity with ghost layer
-    initialVelocityWithGl = addGhostLayers(initialVelocity, indexDimensions=1, ghostLayers=1)
-
-    stencil = getStencil(kwargs['stencil'])
-    periodicity = getPeriodicBoundaryFunctor(stencil)
-
-    return createScenario(domainSize, None, kwargs, optimizationParams, lbmKernel=lbmKernel,
-                          initialVelocity=initialVelocityWithGl, kernelParams=kernelParams,
-                          preUpdateFunctions=[periodicity])
-
-
-def createLidDrivenCavity(domainSize, lidVelocity=0.005, optimizationParams={}, lbmKernel=None,
-                          kernelParams={}, **kwargs):
-    def boundarySetupFunction(boundaryHandling, method):
-        myUbb = partial(ubb, velocity=[lidVelocity, 0, 0][:method.dim])
-        myUbb.name = 'ubb'
-        boundaryHandling.setBoundary(myUbb, sliceFromDirection('N', method.dim))
-        for direction in ('W', 'E', 'S') if method.dim == 2 else ('W', 'E', 'S', 'T', 'B'):
-            boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, method.dim))
-
-    return createScenario(domainSize, boundarySetupFunction, kwargs, optimizationParams, lbmKernel=lbmKernel,
-                          kernelParams=kernelParams)
-
-
-def createPressureGradientDrivenChannel(dim, pressureDifference, domainSize=None, radius=None, length=None,
-                                        lbmKernel=None, optimizationParams={}, boundarySetupFunctions=[],
-                                        kernelParams={}, **kwargs):
-    assert dim in (2, 3)
-
-    if radius is not None:
-        assert length is not None
-        if dim == 3:
-            domainSize = (length, 2 * radius + 1, 2 * radius + 1)
-            roundChannel = True
-        else:
-            if domainSize is None:
-                domainSize = (length, 2 * radius)
-    else:
-        roundChannel = False
-
-    def boundarySetupFunction(boundaryHandling, method):
-        pressureBoundaryInflow = partial(fixedDensity, density=1.0 + pressureDifference)
-        pressureBoundaryInflow.__name__ = "Inflow"
-
-        pressureBoundaryOutflow = partial(fixedDensity, density=1.0)
-        pressureBoundaryOutflow.__name__ = "Outflow"
-        boundaryHandling.setBoundary(pressureBoundaryInflow, sliceFromDirection('W', method.dim))
-        boundaryHandling.setBoundary(pressureBoundaryOutflow, sliceFromDirection('E', method.dim))
-
-        if method.dim == 2:
-            for direction in ('N', 'S'):
-                boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, method.dim))
-        elif method.dim == 3:
-            if roundChannel:
-                noSlipIdx = boundaryHandling.addBoundary(noSlip)
-                ff = boundaryHandling.flagField
-                yMid = ff.shape[1] // 2
-                zMid = ff.shape[2] // 2
-                y, z = np.meshgrid(range(ff.shape[1]), range(ff.shape[2]))
-                ff[(y - yMid) ** 2 + (z - zMid) ** 2 > radius ** 2] = noSlipIdx
-            else:
-                for direction in ('N', 'S', 'T', 'B'):
-                    boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, method.dim))
-
-        for userFunction in boundarySetupFunctions:
-            userFunction(boundaryHandling, method, domainSize)
-
-    assert domainSize is not None
-    return createScenario(domainSize, boundarySetupFunction, kwargs, optimizationParams, lbmKernel=lbmKernel,
-                          kernelParams=kernelParams)
-
-
-def createForceDrivenChannel(dim, force, domainSize=None, radius=None, length=None, lbmKernel=None,
-                             optimizationParams={}, initialVelocity=None, boundarySetupFunctions=[],
-                             kernelParams={}, **kwargs):
-    assert dim in (2, 3)
-    kwargs['force'] = tuple([force, 0, 0][:dim])
-
-    if radius is not None:
-        assert length is not None
-        if dim == 3:
-            domainSize = (length, 2 * radius + 1, 2 * radius + 1)
-            roundChannel = True
-        else:
-            if domainSize is None:
-                domainSize = (length, 2 * radius)
-    else:
-        roundChannel = False
-
-    def boundarySetupFunction(boundaryHandling, method):
-        if method.dim == 2:
-            for direction in ('N', 'S'):
-                boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, method.dim))
-        elif method.dim == 3:
-            if roundChannel:
-                noSlipIdx = boundaryHandling.addBoundary(noSlip)
-                ff = boundaryHandling.flagField
-                yMid = ff.shape[1] // 2
-                zMid = ff.shape[2] // 2
-                y, z = np.meshgrid(range(ff.shape[1]), range(ff.shape[2]))
-                ff[(y - yMid) ** 2 + (z - zMid) ** 2 > radius ** 2] = noSlipIdx
-            else:
-                for direction in ('N', 'S', 'T', 'B'):
-                    boundaryHandling.setBoundary(noSlip, sliceFromDirection(direction, method.dim))
-        for userFunction in boundarySetupFunctions:
-            userFunction(boundaryHandling, method, domainSize)
-
-    def periodicity(pdfArr):
-        pdfArr[0, :, :] = pdfArr[-2, :, :]
-        pdfArr[-1, :, :] = pdfArr[1, :, :]
-
-    assert domainSize is not None
-    if 'forceModel' not in kwargs:
-        kwargs['forceModel'] = 'guo'
-
-    return createScenario(domainSize, boundarySetupFunction, kwargs, optimizationParams, lbmKernel=lbmKernel,
-                          initialVelocity=initialVelocity, preUpdateFunctions=[periodicity],
-                          kernelParams=kernelParams)
-