diff --git a/boundaries/boundaryhandling.py b/boundaries/boundaryhandling.py index 3f68a725c7edd45d5b89a4f058e3e7796dc30c57..c4b9cc684740ae782171f5bde61080d08241c70a 100644 --- a/boundaries/boundaryhandling.py +++ b/boundaries/boundaryhandling.py @@ -2,9 +2,7 @@ import sympy as sp import numpy as np from pystencils import TypedSymbol, Field from pystencils.backends.cbackend import CustomCppCode -from pystencils.astnodes import Block, SympyAssignment, LoopOverCoordinate, KernelFunction -from pystencils.transformations import moveConstantsBeforeLoop, resolveFieldAccesses, typingFromSympyInspection, \ - typeAllEquations +from pystencils.cpu.kernelcreation import createIndexedKernel from lbmpy.boundaries.createindexlist import createBoundaryIndexList INV_DIR_SYMBOL = TypedSymbol("invDir", "int") @@ -74,6 +72,7 @@ class BoundaryHandling(object): for boundaryIdx, boundaryFunc in enumerate(self._boundaryFunctions): idxField = createBoundaryIndexList(self.flagField, self._lbMethod.stencil, 2 ** boundaryIdx, self._fluidFlag, self._ghostLayers) + idxField = transformIndexListToStruct(idxField) ast = generateBoundaryHandling(self._symbolicPdfField, idxField, self._lbMethod, boundaryFunc) if self._target == 'cpu': @@ -135,47 +134,28 @@ class LbmMethodInfo(CustomCppCode): super(LbmMethodInfo, self).__init__(code, symbolsRead=set(), symbolsDefined=symbolsDefined) -def generateBoundaryHandling(pdfField, indexArr, lbMethod, boundaryFunctor): - dim = lbMethod.dim - - cellLoopBody = Block([]) - cellLoop = LoopOverCoordinate(cellLoopBody, coordinateToLoopOver=0, start=0, stop=indexArr.shape[0]) +def transformIndexListToStruct(arr): + #TODO create in correct form right away + dim = arr.shape[-1] -1 + coordinateNames = ['x', 'y', 'z'][:dim] + dataTypeInfo = [(name, np.int) for name in coordinateNames] + [('dir', np.int)] + indexArrStruct = np.empty((arr.shape[0]), dtype=np.dtype(dataTypeInfo)) + for idx, name in enumerate(coordinateNames): + indexArrStruct[name] = arr[:, idx] + indexArrStruct['dir'] = arr[:, -1] + return indexArrStruct - indexField = Field.createFromNumpyArray("indexField", indexArr, indexDimensions=1) - indexField.isIndexField = True - coordinateSymbols = [TypedSymbol(name, "int") for name in ['x', 'y', 'z']] - for d in range(dim): - cellLoopBody.append(SympyAssignment(coordinateSymbols[d], indexField[0](d))) - dirSymbol = TypedSymbol("dir", "int") - cellLoopBody.append(SympyAssignment(dirSymbol, indexField[0](dim))) +def generateBoundaryHandling(pdfField, indexArr, lbMethod, boundaryFunctor): + indexField = Field.createFromNumpyArray("indexField", indexArr) - boundaryEqList = boundaryFunctor(pdfField, dirSymbol, lbMethod) + elements = [LbmMethodInfo(lbMethod)] + boundaryEqList = boundaryFunctor(pdfField, indexField[0]('dir'), lbMethod) if type(boundaryEqList) is tuple: boundaryEqList, additionalNodes = boundaryEqList + elements += boundaryEqList + elements += additionalNodes else: - additionalNodes = [] - - typeInfos = typingFromSympyInspection(boundaryEqList, pdfField.dtype) - fieldsRead, fieldsWritten, assignments = typeAllEquations(boundaryEqList, typeInfos) - fieldsAccessed = fieldsRead.union(fieldsWritten) - set([indexField]) - - for be in assignments: - cellLoopBody.append(be) - - functionBody = Block([cellLoop]) - ast = KernelFunction(functionBody, [indexField] + list(fieldsAccessed)) - - if len(additionalNodes) > 0: - loops = ast.atoms(LoopOverCoordinate) - assert len(loops) == 1 - loop = list(loops)[0] - for node in additionalNodes: - loop.body.append(node) - - functionBody.insertFront(LbmMethodInfo(lbMethod)) + elements += boundaryEqList + return createIndexedKernel(elements, [indexField]) - fixedCoordinateMapping = {f.name: coordinateSymbols[:dim] for f in fieldsAccessed} - resolveFieldAccesses(ast, set(['indexField']), fieldToFixedCoordinates=fixedCoordinateMapping) - moveConstantsBeforeLoop(ast) - return ast