diff --git a/boundaries/boundaryhandling.py b/boundaries/boundaryhandling.py
index 92269a161390acd8c2bef9c694b95d1504fe5e1f..ce481c2d012d115af3e6a2d5cf1c296fa77f7e6a 100644
--- a/boundaries/boundaryhandling.py
+++ b/boundaries/boundaryhandling.py
@@ -15,9 +15,9 @@ class BoundaryHandling(GenericBoundaryHandling, PeriodicityHandling):
def setBoundary(self, boundaryObject, indexExpr=None, maskCallback=None, sliceNormalizationGhostLayers=1):
mask = None
if maskCallback is not None:
- meshgridParams = [np.arange(start=-sliceNormalizationGhostLayers, stop=s-sliceNormalizationGhostLayers)
- for s in self.flagField.shape]
- indices = np.meshgrid(*meshgridParams, indexing='ij')
+ gridParams = [np.arange(start=-sliceNormalizationGhostLayers, stop=s - sliceNormalizationGhostLayers) + 0.5
+ for s in self.flagField.shape]
+ indices = np.meshgrid(*gridParams, indexing='ij')
mask = maskCallback(*indices)
return GenericBoundaryHandling.setBoundary(self, boundaryObject, indexExpr, mask)
diff --git a/boundaries/handlinginterface.py b/boundaries/handlinginterface.py
index 7de44c7239a04eee6d57d28d71fa59466599700e..55db0f6ba738e27ea8f5ce4f308fca566c65e17a 100644
--- a/boundaries/handlinginterface.py
+++ b/boundaries/handlinginterface.py
@@ -83,7 +83,8 @@ class GenericBoundaryHandling(object):
if self._dirty:
self.prepare()
for boundary in self._boundaryInfos.values():
- boundary.kernel(**kwargs)
+ if boundary.kernel:
+ boundary.kernel(**kwargs)
def getBoundaryIndexArray(self, boundaryObject):
return self._boundaryInfos[boundaryObject].indexArray
@@ -111,6 +112,8 @@ class GenericBoundaryHandling(object):
boundaryFlag = self._flagFieldInterface.getFlag(boundaryObject)
idxArray = createBoundaryIndexList(self.flagField, self._lbMethod.stencil,
boundaryFlag, fluidFlag, self.ghostLayers)
+ if len(idxArray) == 0:
+ continue
if boundaryObject.additionalData:
coordinateNames = ["x", "y", "z"][:dim]
@@ -149,6 +152,9 @@ class GenericBoundaryHandling(object):
self._dirty = False
+ def reserveFlag(self, boundaryObject):
+ self._flagFieldInterface.getFlag(boundaryObject)
+
def setBoundary(self, boundaryObject, indexExpr=None, maskArr=None):
"""
Sets boundary in a rectangular region (slice)
diff --git a/parallel/blockiteration.py b/parallel/blockiteration.py
new file mode 100644
index 0000000000000000000000000000000000000000..18e388552f6a6b88ed2e6286489b0cef0453389a
--- /dev/null
+++ b/parallel/blockiteration.py
@@ -0,0 +1,30 @@
+import numpy as np
+import waLBerla as wlb
+from pystencils.slicing import normalizeSlice
+
+
+def slicedBlockIteration(blocks, indexExpr=None, ghostLayers=1, sliceNormalizationGhostLayers=1):
+ if indexExpr is None:
+ indexExpr = [slice(None, None, None)] * 3
+
+ domainCellBB = blocks.getDomainCellBB()
+ domainExtent = [s + 2 * sliceNormalizationGhostLayers for s in domainCellBB.size]
+ indexExpr = normalizeSlice(indexExpr, domainExtent)
+ targetCellBB = wlb.CellInterval.fromSlice(indexExpr)
+ targetCellBB.shift(*[a - sliceNormalizationGhostLayers for a in domainCellBB.min])
+
+ for block in blocks:
+ intersection = blocks.getBlockCellBB(block).getExpanded(ghostLayers)
+ intersection.intersect(targetCellBB)
+ if intersection.empty():
+ continue
+
+ meshGridParams = [offset + 0.5 + np.arange(width)
+ for offset, width in zip(intersection.min, intersection.size)]
+
+ indexArrays = np.meshgrid(*meshGridParams, indexing='ij')
+
+ localTargetBB = blocks.transformGlobalToLocal(block, intersection)
+ localTargetBB.shift(ghostLayers, ghostLayers, ghostLayers)
+ localSlice = localTargetBB.toSlice()
+ yield block, indexArrays, localSlice
diff --git a/parallel/boundaryhandling.py b/parallel/boundaryhandling.py
index 1ae128a21aaeb6d5e320b73e35f86e3adac6ac95..92ec8465bc1a0bb565b2af1e5d63daf9b923bdd1 100644
--- a/parallel/boundaryhandling.py
+++ b/parallel/boundaryhandling.py
@@ -1,11 +1,11 @@
import numpy as np
import waLBerla as wlb
from lbmpy.boundaries.handlinginterface import GenericBoundaryHandling, FlagFieldInterface
+from lbmpy.parallel.blockiteration import slicedBlockIteration
from pystencils.slicing import normalizeSlice
class BoundaryHandling(object):
-
def __init__(self, blocks, lbMethod, pdfFieldId, flagFieldId, boundaryId='boundary', target='cpu', openMP=True):
def addBoundaryHandling(block, *args, **kwargs):
@@ -16,12 +16,14 @@ class BoundaryHandling(object):
target=target, openMP=openMP)
self._boundaryId = boundaryId
+ self._pdfFieldId = pdfFieldId
self._blocks = blocks
self.dim = lbMethod.dim
blocks.addBlockData(boundaryId, addBoundaryHandling)
def __call__(self, **kwargs):
for block in self._blocks:
+ kwargs['pdfs'] = wlb.field.toArray(block[self._pdfFieldId], withGhostLayers=True)
block[self._boundaryId](**kwargs)
def prepare(self):
@@ -32,35 +34,46 @@ class BoundaryHandling(object):
if indexExpr is None:
indexExpr = [slice(None, None, None)] * self.dim
- domainCellBB = self._blocks.getDomainCellBB()
- domainExtent = [s + 2 * sliceNormalizationGhostLayers for s in domainCellBB.size]
- indexExpr = normalizeSlice(indexExpr, domainExtent)
- targetCellBB = wlb.CellInterval.fromSlice(indexExpr)
- targetCellBB.shift(*[a - sliceNormalizationGhostLayers for a in domainCellBB.min])
+ if len(self._blocks) == 0:
+ return
+ gl = self._blocks[0][self._boundaryId].ghostLayers
+ ngl = sliceNormalizationGhostLayers
for block in self._blocks:
- boundaryHandling = block[self._boundaryId]
- ghostLayers = boundaryHandling.ghostLayers
- intersection = self._blocks.getBlockCellBB(block).getExpanded(ghostLayers)
- intersection.intersect(targetCellBB)
- if not intersection.empty():
- if maskCallback is not None:
- meshGridParams = [offset + np.arange(width)
- for offset, width in zip(intersection.min, intersection.size)]
- indexArr = np.meshgrid(*meshGridParams, indexing='ij')
- mask = maskCallback(*indexArr)
- else:
- mask = None
- localTargetBB = self._blocks.transformGlobalToLocal(block, intersection)
- localTargetBB.shift(ghostLayers, ghostLayers, ghostLayers)
- block[self._boundaryId].setBoundary(boundaryObject, indexExpr=localTargetBB.toSlice(), maskArr=mask)
+ block[self._boundaryId].reserveFlag(boundaryObject)
+
+ for block, indexArrs, localSlice in slicedBlockIteration(self._blocks, indexExpr, gl, ngl):
+ mask = maskCallback(*indexArrs) if maskCallback else None
+ block[self._boundaryId].setBoundary(boundaryObject, indexExpr=localSlice, maskArr=mask)
+
+ #domainCellBB = self._blocks.getDomainCellBB()
+ #domainExtent = [s + 2 * sliceNormalizationGhostLayers for s in domainCellBB.size]
+ #indexExpr = normalizeSlice(indexExpr, domainExtent)
+ #targetCellBB = wlb.CellInterval.fromSlice(indexExpr)
+ #targetCellBB.shift(*[a - sliceNormalizationGhostLayers for a in domainCellBB.min])
+ #
+ #for block in self._blocks:
+ # boundaryHandling = block[self._boundaryId]
+ # ghostLayers = boundaryHandling.ghostLayers
+ # intersection = self._blocks.getBlockCellBB(block).getExpanded(ghostLayers)
+ # intersection.intersect(targetCellBB)
+ # if not intersection.empty():
+ # if maskCallback is not None:
+ # meshGridParams = [offset + 0.5 + np.arange(width)
+ # for offset, width in zip(intersection.min, intersection.size)]
+ # indexArr = np.meshgrid(*meshGridParams, indexing='ij')
+ # mask = maskCallback(*indexArr)
+ # else:
+ # mask = None
+ # localTargetBB = self._blocks.transformGlobalToLocal(block, intersection)
+ # localTargetBB.shift(ghostLayers, ghostLayers, ghostLayers)
+ # block[self._boundaryId].setBoundary(boundaryObject, indexExpr=localTargetBB.toSlice(), maskArr=mask)
# ----------------------------------------------------------------------------------------------------------------------
class WalberlaFlagFieldInterface(FlagFieldInterface):
-
def __init__(self, block, flagFieldId):
self._block = block
self._flagFieldId = flagFieldId
@@ -68,10 +81,13 @@ class WalberlaFlagFieldInterface(FlagFieldInterface):
assert self._flagArray.shape[3] == 1
self._flagArray = self._flagArray[..., 0]
- fluidFlag = self._block[self._flagFieldId].registerFlag('fluid')
- self._boundaryObjectToName = {'fluid': fluidFlag}
+ self._fluidFlag = self._block[self._flagFieldId].registerFlag('fluid')
+ self._flagArray.fill(self._fluidFlag)
+ self._boundaryObjectToName = {}
def getFlag(self, boundaryObject):
+ if boundaryObject == 'fluid':
+ return self._fluidFlag
if boundaryObject not in self._boundaryObjectToName:
name = self._makeBoundaryName(boundaryObject, self._boundaryObjectToName.values())
self._boundaryObjectToName[boundaryObject] = name
@@ -91,14 +107,16 @@ class WalberlaFlagFieldInterface(FlagFieldInterface):
return self._boundaryObjectToName.keys()
def clear(self):
+ self._flagArray.fill(self._fluidFlag)
raise NotImplementedError()
+
if __name__ == '__main__':
from lbmpy.creationfunctions import createLatticeBoltzmannMethod
from lbmpy.boundaries.boundaryconditions import NoSlip
from pystencils.slicing import makeSlice
- blocks = wlb.createUniformBlockGrid(cellsPerBlock=(3, 3, 3), blocks=(2, 2, 2), oneBlockPerProcess=False)
+ blocks = wlb.createUniformBlockGrid(cellsPerBlock=(4, 4, 4), blocks=(2, 2, 2), oneBlockPerProcess=False)
lbMethod = createLatticeBoltzmannMethod(stencil='D3Q19', method='srt', relaxationRate=1.8)
@@ -109,6 +127,17 @@ if __name__ == '__main__':
def maskCallback(x, y, z):
- return x > -100
+ midPoint = (4, 4, 4)
+ radius = 2.0
+
+ resultMask = (x - midPoint[0]) ** 2 + \
+ (y - midPoint[1]) ** 2 + \
+ (z - midPoint[2]) ** 2 <= radius ** 2
+ return resultMask
+
- bh.setBoundary(NoSlip(), makeSlice[0, :, :], maskCallback)
+ bh.setBoundary(NoSlip(), makeSlice[:, :, :], maskCallback, sliceNormalizationGhostLayers=0)
+ vtkOutput = wlb.vtk.makeOutput(blocks, "vtkOutput")
+ flagFieldWriter = wlb.field.createVTKWriter(blocks, "flagField")
+ vtkOutput.addCellDataWriter(flagFieldWriter)
+ vtkOutput.write(1)
diff --git a/parallel/scenarios.py b/parallel/scenarios.py
index 7e909e832e1fd016d7bb346a36a3d708c5225b74..acc7bc125777489fd85a4768a59c3bf31f5f8d92 100644
--- a/parallel/scenarios.py
+++ b/parallel/scenarios.py
@@ -1,18 +1,37 @@
import sympy as sp
+import numpy as np
+import numbers
import waLBerla as wlb
from lbmpy.creationfunctions import createLatticeBoltzmannFunction, updateWithDefaultParameters, \
switchToSymbolicRelaxationRatesForEntropicMethods
+from lbmpy.macroscopic_value_kernels import compileMacroscopicValuesGetter, compileMacroscopicValuesSetter
from lbmpy.parallel.boundaryhandling import BoundaryHandling
+from lbmpy.parallel.blockiteration import slicedBlockIteration
+from pystencils.field import createNumpyArrayWithLayout, getLayoutOfArray
class Scenario(object):
def __init__(self, blocks, methodParameters, optimizationParams, lbmKernel=None,
- initialVelocityCallback=None, preUpdateFunctions=[], kernelParams={},
- blockDataPrefix='', directCommunication=False):
+ preUpdateFunctions=[], kernelParams={}, blockDataPrefix='', directCommunication=False):
+
+ self.blocks = blocks
+ self._blockDataPrefix = blockDataPrefix
+ self.kernelParams = kernelParams
methodParameters, optimizationParams = updateWithDefaultParameters(methodParameters, optimizationParams)
- target = optimizationParams['target']
+ self._target = optimizationParams['target']
+ if optimizationParams['fieldLayout'] in ('f', 'reverseNumpy'):
+ optimizationParams['fieldLayout'] = 'fzyx'
+ if optimizationParams['fieldLayout'] not in ('fzyx', 'zyxf'):
+ raise ValueError("In parallel scenarios only layouts 'fxyz' and 'zyxf' are supported")
+
+ wlbLayoutMap = {
+ 'fzyx': wlb.field.Layout.fzyx,
+ 'zyxf': wlb.field.Layout.zyxf,
+ }
+ wlbLayout = wlbLayoutMap[optimizationParams['fieldLayout']]
+ self._fieldLayout = optimizationParams['fieldLayout']
# ----- Create LBM kernel
if lbmKernel is None:
@@ -22,14 +41,12 @@ class Scenario(object):
else:
self._lbmKernel = lbmKernel
- self._blockDataPrefix = blockDataPrefix
-
# ----- Add fields
- numPdfs = len(self._lbmKernel.method.stencils)
- wlb.field.addToStorage(blocks, blockDataPrefix + 'pdfs', float, fSize=numPdfs, ghostLayers=1)
- wlb.field.addFlagFieldToStorage(blocks, blockDataPrefix + "flags", nrOfBits=16, ghostLayers=1)
+ numPdfs = len(self._lbmKernel.method.stencil)
+ wlb.field.addToStorage(blocks, self.pdfFieldId, float, fSize=numPdfs, ghostLayers=1, layout=wlbLayout)
+ wlb.field.addFlagFieldToStorage(blocks, self.flagFieldId, nrOfBits=16, ghostLayers=1)
- if target == 'gpu':
+ if self._target == 'gpu':
wlb.cuda.addGpuFieldToStorage(blocks, blockDataPrefix + "gpuPdfs", float, fSize=numPdfs,
usePitchedMem=False)
@@ -39,39 +56,192 @@ class Scenario(object):
communicationFunctions = {
('cpu', True): (wlb.createUniformDirectScheme, wlb.field.createMPIDatatypeInfo),
('cpu', False): (wlb.createUniformBufferedScheme, wlb.field.createPackInfo),
- ('gpu', True): (wlb.createUniformDirectScheme, wlb.cuda.createMPIDatatypeInfo),
- ('gpu', False): (wlb.createUniformBufferedScheme, wlb.cuda.createPackInfo),
+ #('gpu', True): (wlb.createUniformDirectScheme, wlb.cuda.createMPIDatatypeInfo),
+ #('gpu', False): (wlb.createUniformBufferedScheme, wlb.cuda.createPackInfo),
}
- createScheme, createCommunicationData = communicationFunctions[(target, directCommunication)]
+ createScheme, createCommunicationData = communicationFunctions[(self._target, directCommunication)]
self._communicationScheme = createScheme(blocks, communicationStencil)
- self._communicationScheme.addDataToCommunicate(createCommunicationData(blocks, blockDataPrefix + 'pdfs'))
+ self._communicationScheme.addDataToCommunicate(createCommunicationData(blocks, self.pdfFieldId))
# ----- Create boundary handling
- self._boundaryHandling = BoundaryHandling(blocks, self._lbmKernel.method, blockDataPrefix + 'pdfs',
- blockDataPrefix + "flags", target=target)
+ self.boundaryHandling = BoundaryHandling(blocks, self._lbmKernel.method, self.pdfFieldId, self.flagFieldId,
+ target=self._target)
+
+ # ----- Macroscopic value input/output
+
+ wlb.field.addToStorage(blocks, self.densityFieldId, float, fSize=1, ghostLayers=1, layout=wlbLayout)
+ wlb.field.addToStorage(blocks, self.velocityFieldId, float, fSize=3, ghostLayers=1, layout=wlbLayout)
+
+ self._getMacroscopicKernel = compileMacroscopicValuesGetter(self._lbmKernel.method, ['density', 'velocity'],
+ fieldLayout=optimizationParams['fieldLayout'],
+ target='cpu')
+
+ self._macroscopicValueSetter = None
+ self._macroscopicValueGetter = None
+
+ self._vtkOutput = wlb.vtk.makeOutput(blocks, "vtk")
+ self._vtkOutput.addCellDataWriter(wlb.field.createVTKWriter(self.blocks, self.flagFieldId))
+ self._vtkOutput.addCellDataWriter(wlb.field.createVTKWriter(self.blocks, self.densityFieldId))
+ self._vtkOutput.addCellDataWriter(wlb.field.createVTKWriter(self.blocks, self.velocityFieldId))
+
+ self.timeStepsRun = 0
+ self._dstFieldCache = dict()
+
+ def run(self, timeSteps):
+ if self._target == 'cpu':
+ for t in range(timeSteps):
+ self._communicationScheme()
+ self.boundaryHandling()
+ for block in self.blocks:
+ srcField = block[self.pdfFieldId]
+ pdfArr = wlb.field.toArray(srcField, withGhostLayers=True)
+ swapIdx = (srcField.size, srcField.layout)
+ if swapIdx not in self._dstFieldCache:
+ self._dstFieldCache[swapIdx] = srcField.cloneUninitialized()
+ dstField = self._dstFieldCache[swapIdx]
+ dstArr = wlb.field.toArray(dstField, withGhostLayers=True)
+ self._lbmKernel(src=pdfArr, dst=dstArr, **self.kernelParams)
+ srcField.swapDataPointers(dstField)
+
+ self._getMacroscopicValues()
+ self.timeStepsRun += timeSteps
+
+ def writeVTK(self):
+ self._vtkOutput.write(self.timeStepsRun)
+
+ def setMacroscopicValue(self, density=None, velocity=None, indexExpr=None):
+
+ if density is None and velocity is None:
+ raise ValueError("Please specify either density or velocity")
+
+ # Check velocity
+ velocityCallback = None
+ velocityValue = [0, 0, 0]
+ if isinstance(velocity, numbers.Number):
+ velocityValue = [velocity] * 3
+ elif hasattr(velocity, "__len__"):
+ assert len(velocity) == 3
+ velocityValue = velocity
+ elif hasattr(velocity, "__call__"):
+ velocityCallback = velocity
+ elif velocity is None:
+ pass
+ else:
+ raise ValueError("velocity has be a number, sequence of length 3 or a callback function")
+
+ # Check density
+ densityCallback = None
+ densityValue = 1.0
+ if isinstance(density, numbers.Number):
+ densityValue = density
+ elif hasattr(density, "__call__"):
+ densityCallback = density
+ elif density is None:
+ pass
+ else:
+ raise ValueError("density has to be a number or a callback function")
+
+ for block, (x, y, z), localSlice in slicedBlockIteration(self.blocks, indexExpr, 1, 1):
+ if density:
+ densityArr = wlb.field.toArray(block[self.densityFieldId], withGhostLayers=True)
+ densityArr[localSlice] = densityCallback(x, y, z) if densityCallback else densityValue
+ if velocity:
+ velArr = wlb.field.toArray(block[self.velocityFieldId], withGhostLayers=True)
+ velArr[localSlice, :] = velocityCallback(x, y, z) if velocityCallback else velocityValue
+
+ self._setMacroscpicValues()
+
+ @property
+ def pdfFieldId(self):
+ return self._blockDataPrefix + "pdfs"
+
+ @property
+ def flagFieldId(self):
+ return self._blockDataPrefix + "flags"
+
+ @property
+ def densityFieldId(self):
+ return self._blockDataPrefix + "density"
+
+ @property
+ def velocityFieldId(self):
+ return self._blockDataPrefix + "velocity"
+
+ @property
+ def method(self):
+ """Lattice boltzmann method description"""
+ return self._lbmKernel.method
+
+ def _getMacroscopicValues(self):
+ """Takes values from pdf field and writes them to density and velocity field"""
+ if len(self.blocks) == 0:
+ return
+
+ if self._macroscopicValueGetter is None:
+ shape = wlb.field.toArray(self.blocks[0][self.flagFieldId], withGhostLayers=True).shape
+ self._macroscopicValueGetter = compileMacroscopicValuesGetter(self.method, ['density', 'velocity'], target='cpu',
+ fieldLayout=self._fieldLayout)
+ for block in self.blocks:
+ densityArr = wlb.field.toArray(block[self.densityFieldId], withGhostLayers=True)
+ velArr = wlb.field.toArray(block[self.velocityFieldId], withGhostLayers=True)
+ pdfArr = wlb.field.toArray(block[self.pdfFieldId], withGhostLayers=True)
+ self._macroscopicValueGetter(pdfs=pdfArr, velocity=velArr, density=densityArr)
+
+ def _setMacroscpicValues(self):
+ """Takes values from density / velocity field and writes them to pdf field"""
+ if len(self.blocks) == 0:
+ return
+
+ if self._macroscopicValueSetter is None:
+ shape = wlb.field.toArray(self.blocks[0][self.flagFieldId], withGhostLayers=True).shape
+ vals = {
+ 'density': np.ones(shape),
+ 'velocity': np.zeros(shape + (3,)),
+ }
+ self._macroscopicValueSetter = compileMacroscopicValuesSetter(self.method, vals, target='cpu',
+ fieldLayout=self._fieldLayout,)
+
+ for block in self.blocks:
+ densityArr = wlb.field.toArray(block[self.densityFieldId], withGhostLayers=True)
+ velArr = wlb.field.toArray(block[self.velocityFieldId], withGhostLayers=True)
+ pdfArr = wlb.field.toArray(block[self.pdfFieldId], withGhostLayers=True)
+ self._macroscopicValueSetter(pdfArr=pdfArr, velocity=velArr, density=densityArr)
+
if __name__ == '__main__':
- class A(object):
- def __call__(self, call_from):
- print("foo from A, call from %s" % call_from)
+ from lbmpy.boundaries import NoSlip, FixedDensity
+ from pystencils.slicing import sliceFromDirection
+
+ blocks = wlb.createUniformBlockGrid(cellsPerBlock=(64, 20, 20), blocks=(2, 2, 2), oneBlockPerProcess=False,
+ periodic=(1, 0, 0))
+
+ scenario = Scenario(blocks, methodParameters={'stencil': 'D3Q19', 'relaxationRate': 1.92, 'force': (1e-5, 0, 0)},
+ optimizationParams={})
+
+ wallBoundary = NoSlip()
+ for direction in ('N', 'S', 'T', 'B'):
+ scenario.boundaryHandling.setBoundary(wallBoundary, sliceFromDirection(direction, 3))
+ #scenario.boundaryHandling.setBoundary(FixedDensity(1.02), sliceFromDirection('W', 3, normalOffset=1))
+ #scenario.boundaryHandling.setBoundary(FixedDensity(1.00), sliceFromDirection('E', 3, normalOffset=1))
- class B(object):
- def __call__(self, call_from):
- print( "foo from B, call from %s" % call_from)
+ def maskCallback(x, y, z):
+ midPoint = (32, 10, 10)
+ radius = 5.0
+ resultMask = (x - midPoint[0]) ** 2 + \
+ (y - midPoint[1]) ** 2 + \
+ (z - midPoint[2]) ** 2 <= radius ** 2
+ return resultMask
- class C(object):
- def __call__(self, call_from):
- print( "foo from C, call from %s" % call_from )
+ scenario.boundaryHandling.setBoundary(wallBoundary, maskCallback=maskCallback)
+ scenario.run(1)
+ for i in range(50):
+ print(i)
+ scenario.writeVTK()
+ scenario.run(200)
- class D(A, B, C):
- def foo(self):
- for cls in D.__bases__:
- cls.__call__(self, "D")
- d = D()
- d.foo()
\ No newline at end of file
diff --git a/updatekernels.py b/updatekernels.py
index 7bba6051c94ed2e4e9943468c6cfb58b766360ad..092e72af1ee71bd8b5d6eb3c9a846dd102f3afc8 100644
--- a/updatekernels.py
+++ b/updatekernels.py
@@ -1,7 +1,7 @@
import numpy as np
import sympy as sp
from pystencils import Field
-from pystencils.field import createNumpyArrayWithLayout
+from pystencils.field import createNumpyArrayWithLayout, layoutStringToTuple
from pystencils.sympyextensions import fastSubs
from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, Pseudo2DTwoFieldsAccessor, PeriodicTwoFieldsAccessor
@@ -97,12 +97,8 @@ def createPdfArray(size, numDirections, ghostLayers=1, layout='fzyx'):
"""
sizeWithGl = [s + 2 * ghostLayers for s in size]
dim = len(size)
- if layout == "fzyx" or layout == 'f' or layout == 'reverseNumpy':
- layout = tuple(reversed(range(dim+1)))
- elif layout == 'c' or layout == 'numpy':
- layout = tuple(range(dim+1))
- elif layout == 'zyxf':
- layout = tuple(reversed(range(dim))) + (dim,)
+ if isinstance(layout, str):
+ layout = layoutStringToTuple(dim+1)
return createNumpyArrayWithLayout(sizeWithGl + [numDirections], layout)