new lbm module: creation functions for kernels

createmethodfromdict.py

+"""
+Factory functions for standard LBM methods
+"""
+import sympy as sp
+from copy import copy
+from lbmpy.stencils import getStencil
+from lbmpy.methods.momentbased import createSRT, createTRT, createOrthogonalMRT
+import lbmpy.forcemodels as forceModels
+from lbmpy.simplificationfactory import createSimplificationStrategy
+from lbmpy.updatekernels import createStreamPullKernel, createPdfArray
+from pystencils.cpu.kernelcreation import createKernel as createCpuKernel, addOpenMP
+from pystencils.gpucuda.kernelcreation import createCUDAKernel as createGpuKernel
+from pystencils.cpu import makePythonFunction as makePythonCpuFunction
+from pystencils.gpucuda import makePythonFunction as makePythonGpuFunction
+
+
+def _getParams(params, optParams):
+    defaultMethodDescription = {
+        'target': 'cpu',
+        'stencil': 'D2Q9',
+        'method': 'srt',  # can be srt, trt or mrt
+        'relaxationRates': sp.symbols("omega_:10"),
+        'compressible': False,
+        'order': 2,
+
+        'forceModel': 'none',  # can be 'simple', 'luo' or 'guo'
+        'force': (0, 0, 0),
+    }
+
+    defaultOptimizationDescription = {
+        'doCseInOpposingDirections': False,
+        'doOverallCse': False,
+        'split': True,
+
+        'fieldSize': None,
+        'fieldLayout': 'reverseNumpy',  # can be 'numpy' (='c'), 'reverseNumpy' (='f'), 'fzyx', 'zyxf'
+
+        'openMP': True,
+    }
+    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:
+        raise ValueError("Unknown parameter(s): " + ",".join(unknownParams))
+    if unknownOptParams:
+        raise ValueError("Unknown optimization parameter(s): " + ",".join(unknownOptParams))
+
+    paramsResult = copy(defaultMethodDescription)
+    paramsResult.update(params)
+    optParamsResult = copy(defaultOptimizationDescription)
+    optParamsResult.update(optParams)
+    return paramsResult, optParamsResult
+
+
+def createLatticeBoltzmannFunction(optimizationParams={}, **kwargs):
+    params, optParams = _getParams(kwargs, optimizationParams)
+
+    ast = createLatticeBoltzmannKernel(**params, optimizationParams=optParams)
+    if params['target'] == 'cpu':
+        if 'openMP' in optParams:
+            if isinstance(optParams['openMP'], bool) and optParams['openMP']:
+                addOpenMP(ast)
+            elif isinstance(optParams['openMP'], int):
+                addOpenMP(ast, numThreads=optParams['openMP'])
+        return makePythonCpuFunction(ast)
+    elif params['target'] == 'gpu':
+        return makePythonGpuFunction(ast)
+    else:
+        return ValueError("'target' has to be either 'cpu' or 'gpu'")
+
+
+def createLatticeBoltzmannKernel(optimizationParams={}, **kwargs):
+    params, optParams = _getParams(kwargs, optimizationParams)
+
+    updateRule = createLatticeBoltzmannUpdateRule(**params, optimizationParams=optimizationParams)
+
+    if params['target'] == 'cpu':
+        if 'splitGroups' in updateRule.simplificationHints:
+            splitGroups = updateRule.simplificationHints['splitGroups']
+        else:
+            splitGroups = ()
+        return createCpuKernel(updateRule.allEquations, splitGroups=splitGroups)
+    elif params['target'] == 'gpu':
+        return createGpuKernel(updateRule.allEquations)
+    else:
+        return ValueError("'target' has to be either 'cpu' or 'gpu'")
+
+
+def createLatticeBoltzmannUpdateRule(optimizationParams={}, **kwargs):
+    params, optParams = _getParams(kwargs, optimizationParams)
+    stencil = getStencil(params['stencil'])
+    method = createLatticeBoltzmannCollisionRule(**params)
+
+    splitInnerLoop = 'split' in optParams and optParams['split']
+
+    dirCSE = 'doCseInOpposingDirections'
+    doCseInOpposingDirections = False if dirCSE not in optParams else optParams[dirCSE]
+    doOverallCse = False if 'doOverallCse' not in optParams else optParams['doOverallCse']
+    simplification = createSimplificationStrategy(method, doCseInOpposingDirections, doOverallCse, splitInnerLoop)
+    collisionRule = simplification(method.getCollisionRule())
+
+    if 'fieldSize' in optParams and optParams['fieldSize']:
+        npField = createPdfArray(optParams['fieldSize'], len(stencil), layout=optParams['fieldLayout'])
+        updateRule = createStreamPullKernel(collisionRule, numpyField=npField)
+    else:
+        layoutName = optParams['fieldLayout']
+        if layoutName == 'fzyx' or 'zyxf':
+            dim = len(stencil[0])
+            layoutName = tuple(reversed(range(dim)))
+        updateRule = createStreamPullKernel(collisionRule, genericLayout=layoutName)
+
+    return updateRule
+
+
+def createLatticeBoltzmannCollisionRule(**params):
+    params, _ = _getParams(params, {})
+
+    stencil = getStencil(params['stencil'])
+    dim = len(stencil[0])
+
+    if 'forceModel' in params:
+        forceModelName = params['forceModel']
+        if forceModelName.lower() == 'none':
+            forceModel = None
+        elif forceModelName.lower() == 'simple':
+            forceModel = forceModels.Simple(params['force'][:dim])
+        elif forceModelName.lower() == 'luo':
+            forceModel = forceModels.Luo(params['force'][:dim])
+        elif forceModelName.lower() == 'guo':
+            forceModel = forceModels.Guo(params['force'][:dim])
+        else:
+            raise ValueError("Unknown force model %s" % (forceModelName,))
+    else:
+        forceModel = None
+
+    commonParams = {
+        'compressible': params['compressible'],
+        'equilibriumAccuracyOrder': params['order'],
+        'forceModel': forceModel,
+    }
+    methodName = params['method']
+    relaxationRates = params['relaxationRates']
+
+    if methodName.lower() == 'srt':
+        assert len(relaxationRates) >= 1, "Not enough relaxation rates"
+        method = createSRT(stencil, relaxationRates[0], **commonParams)
+    elif methodName.lower() == 'trt':
+        assert len(relaxationRates) >= 2, "Not enough relaxation rates"
+        method = createTRT(stencil, relaxationRates[0], relaxationRates[1], **commonParams)
+    elif methodName.lower() == 'mrt':
+        nextRelaxationRate = 0
+
+        def relaxationRateGetter(momentGroup):
+            nonlocal nextRelaxationRate
+            res = relaxationRates[nextRelaxationRate]
+            nextRelaxationRate += 1
+            return res
+        method = createOrthogonalMRT(stencil, relaxationRateGetter, **commonParams)
+    else:
+        raise ValueError("Unknown method %s" % (methodName,))
+
+    return method
+

innerloopsplit.py

 import sympy as sp
 from collections import defaultdict
+from pystencils import Field
 
 
 def createLbmSplitGroups(lbmCollisionEqs):
@@ -18,16 +19,27 @@ def createLbmSplitGroups(lbmCollisionEqs):
     sh = lbmCollisionEqs.simplificationHints
     assert 'velocity' in sh, "Needs simplification hint 'velocity': Sequence of velocity symbols"
 
-    pdfSymbols = lbmCollisionEqs.method.preCollisionPdfSymbols
+    preCollisionSymbols = set(lbmCollisionEqs.method.preCollisionPdfSymbols)
+    nonCenterPostCollisionSymbols = set(lbmCollisionEqs.method.postCollisionPdfSymbols[1:])
+    postCollisionSymbols = set(lbmCollisionEqs.method.postCollisionPdfSymbols)
+
     stencil = lbmCollisionEqs.method.stencil
 
     importantSubExpressions = {e.lhs for e in lbmCollisionEqs.subexpressions
-                               if pdfSymbols.intersection(lbmCollisionEqs.getDependentSymbols([e.lhs]))}
-    for eq in lbmCollisionEqs.mainEquations[1:]:
+                               if preCollisionSymbols.intersection(lbmCollisionEqs.getDependentSymbols([e.lhs]))}
+
+    otherWrittenFields = []
+    for eq in lbmCollisionEqs.mainEquations:
+        if eq.lhs not in postCollisionSymbols and isinstance(eq.lhs, Field.Access):
+            otherWrittenFields.append(eq.lhs)
+        if eq.lhs not in nonCenterPostCollisionSymbols:
+            continue
         importantSubExpressions.intersection_update(eq.rhs.atoms(sp.Symbol))
 
-    subexpressionsToPreCompute = list(sh['velocity']) + list(importantSubExpressions)
-    splitGroups = [subexpressionsToPreCompute, ]
+    importantSubExpressions.update(sh['velocity'])
+
+    subexpressionsToPreCompute = list(importantSubExpressions)
+    splitGroups = [subexpressionsToPreCompute + otherWrittenFields, ]
 
     directionGroups = defaultdict(list)
     dim = len(stencil[0])
@@ -45,4 +57,5 @@ def createLbmSplitGroups(lbmCollisionEqs):
             directionGroups[direction].append(eq.lhs)
     splitGroups += directionGroups.values()
 
-    return splitGroups
+    lbmCollisionEqs.simplificationHints['splitGroups'] = splitGroups
+    return lbmCollisionEqs

methods/momentbased.py

@@ -380,20 +380,20 @@ def createOrthogonalMRT(stencil, relaxationRateGetter=None, compressible=False,
         sq = x ** 2 + y ** 2 + z ** 2
         nestedMoments = [
             [one, x, y, z],  # [0, 3, 5, 7]
-            [3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z],  # [9, 10, 11, 12, 13]
             [sq - 1],  # [1]
             [3 * sq ** 2 - 6 * sq + 1],  # [2]
             [(3 * sq - 5) * x, (3 * sq - 5) * y, (3 * sq - 5) * z],  # [4, 6, 8]
+            [3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z],  # [9, 10, 11, 12, 13]
             [x * y * z]
         ]
     elif stencilsHaveSameEntries(stencil, getStencil("D3Q19")):
         sq = x ** 2 + y ** 2 + z ** 2
         nestedMoments = [
             [one, x, y, z],  # [0, 3, 5, 7]
-            [3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z],  # [9, 11, 13, 14, 15]
             [sq - 1],  # [1]
             [3 * sq ** 2 - 6 * sq + 1],  # [2]
             [(3 * sq - 5) * x, (3 * sq - 5) * y, (3 * sq - 5) * z],  # [4, 6, 8]
+            [3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z],  # [9, 11, 13, 14, 15]
             [(2 * sq - 3) * (3 * x ** 2 - sq), (2 * sq - 3) * (y ** 2 - z ** 2)],  # [10, 12]
             [(y ** 2 - z ** 2) * x, (z ** 2 - x ** 2) * y, (x ** 2 - y ** 2) * z]  # [16, 17, 18]
         ]

simplificationfactory.py

 from functools import partial
 import sympy as sp
+
+from lbmpy.innerloopsplit import createLbmSplitGroups
 from pystencils.equationcollection.simplifications import applyOnAllEquations, \
     subexpressionSubstitutionInMainEquations, sympyCSE, addSubexpressionsForDivisions
 
 
-def createSimplificationStrategy(lbmMethod, doCseInOpposingDirections=False, doOverallCse=False):
+def createSimplificationStrategy(lbmMethod, doCseInOpposingDirections=False, doOverallCse=False, splitInnerLoop=False):
     from pystencils.equationcollection import SimplificationStrategy
     from lbmpy.methods import MomentBasedLbmMethod
     from lbmpy.methods.momentbasedsimplifications import replaceSecondOrderVelocityProducts, \
@@ -25,6 +27,8 @@ def createSimplificationStrategy(lbmMethod, doCseInOpposingDirections=False, doO
         s.add(replaceCommonQuadraticAndConstantTerm)
         s.add(factorDensityAfterFactoringRelaxationTimes)
         s.add(subexpressionSubstitutionInMainEquations)
+        if splitInnerLoop:
+            s.add(createLbmSplitGroups)
 
     if doCseInOpposingDirections:
         s.add(cseInOpposingDirections)
@@ -34,16 +38,3 @@ def createSimplificationStrategy(lbmMethod, doCseInOpposingDirections=False, doO
     s.add(addSubexpressionsForDivisions)
 
     return s
-
-
-
-if __name__ == '__main__':
-    from lbmpy.stencils import getStencil
-    from lbmpy.methods.momentbased import createOrthogonalMRT
-
-    stencil = getStencil("D2Q9")
-    m = createOrthogonalMRT(stencil, compressible=True)
-    cr = m.getCollisionRule()
-
-    simp = createSimplificationStrategy(m)
-    simp(cr)

updatekernels.py

 import numpy as np
 from pystencils import Field
+from pystencils.sympyextensions import fastSubs
 from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor
 
 
@@ -18,7 +19,7 @@ def createLBMKernel(collisionRule, inputField, outputField, accessor):
     :return: LbmCollisionRule where pre- and post collision symbols have been replaced
     """
     method = collisionRule.method
-    preCollisionSymbols = method.preCollisionSymbols
+    preCollisionSymbols = method.preCollisionPdfSymbols
     postCollisionSymbols = method.postCollisionPdfSymbols
     substitutions = {}
 
@@ -29,7 +30,15 @@ def createLBMKernel(collisionRule, inputField, outputField, accessor):
         substitutions[preCollisionSymbols[idx]] = inputAccess
         substitutions[postCollisionSymbols[idx]] = outputAccess
 
-    return collisionRule.copyWithSubstitutionsApplied(substitutions)
+    result = collisionRule.copyWithSubstitutionsApplied(substitutions)
+
+    if 'splitGroups' in result.simplificationHints:
+        newSplitGroups = []
+        for splitGroup in result.simplificationHints['splitGroups']:
+            newSplitGroups.append([fastSubs(e, substitutions) for e in splitGroup])
+        result.simplificationHints['splitGroups'] = newSplitGroups
+
+    return result
 
 
 def createStreamPullKernel(collisionRule, numpyField=None, srcFieldName="src", dstFieldName="dst",
@@ -60,6 +69,36 @@ def createStreamPullKernel(collisionRule, numpyField=None, srcFieldName="src", d
     return createLBMKernel(collisionRule, src, dst, StreamPullTwoFieldsAccessor)
 
 
+# ---------------------------------- Pdf array creation for various layouts --------------------------------------------
+
+def createPdfArray(size, numDirections, ghostLayers=1, layout='fzyx'):
+    """
+    Creates an empy numpy array for a pdf field with the specified memory layout.
+
+    Examples:
+        >>> createPdfArray((3, 4, 5), 9, layout='zyxf', ghostLayers=0).shape
+        (3, 4, 5, 9)
+        >>> createPdfArray((3, 4, 5), 9, layout='zyxf', ghostLayers=0).strides
+        (72, 216, 864, 8)
+        >>> createPdfArray((3, 4), 9, layout='zyxf', ghostLayers=1).shape
+        (5, 6, 9)
+        >>> createPdfArray((3, 4), 9, layout='zyxf', ghostLayers=1).strides
+        (72, 360, 8)
+    """
+    sizeWithGl = [s + 2 * ghostLayers for s in size]
+    if layout == "fzyx" or layout == 'f' or layout == 'reverseNumpy':
+        return np.empty(sizeWithGl + [numDirections], order='f')
+    elif layout == 'c' or layout == 'numpy':
+        return np.empty(sizeWithGl + [numDirections], order='c')
+    elif layout == 'zyxf':
+        res = np.empty(list(reversed(sizeWithGl)) + [numDirections], order='c')
+        res = res.swapaxes(0, 1)
+        if len(size) == 3:
+            res = res.swapaxes(1, 2)
+            res = res.swapaxes(0, 1)
+        return res
+
+
 # ------------------------------------------- Add output fields to kernel ----------------------------------------------