diff --git a/python/waLBerla/core_extension.py b/python/waLBerla/core_extension.py index 4450868e0ca08351e59db27457942b5464b5faa8..8720e3faf254636785a8461811630b9dc3edd306 100644 --- a/python/waLBerla/core_extension.py +++ b/python/waLBerla/core_extension.py @@ -57,19 +57,14 @@ def sliceToCellInterval( s ): return walberla_cpp.CellInterval( newMin,newMax) -def cellIntervalToSlice( cellInterval ): +def cellIntervalToSlice(cellInterval, collapseExtentOne=True): slices = [] for i in range(3): - if cellInterval.min[i] == cellInterval.max[i]: + if collapseExtentOne and cellInterval.min[i] == cellInterval.max[i]: slices.append( cellInterval.min[i] ) else: slices.append( slice(cellInterval.min[i], cellInterval.max[i]+1,None ) ) return slices - #return makeSlice[ cellInterval.min[0]:cellInterval.max[0]+1, - # cellInterval.min[1]:cellInterval.max[1]+1, - # cellInterval.min[2]:cellInterval.max[2]+1 ] - - diff --git a/python/waLBerla/field_extension.py b/python/waLBerla/field_extension.py index 000a4126363bef2993a4dbce79b8ad09887a77f9..7856c4ed97d42ebb3a5eef5fa2f077645d677617 100644 --- a/python/waLBerla/field_extension.py +++ b/python/waLBerla/field_extension.py @@ -3,9 +3,11 @@ try: from . import walberla_cpp except ImportError: import walberla_cpp + + # ----------------------------- Python functions to extend the C++ field module --------------------------------- -def npArrayFromWaLBerlaField( field, withGhostLayers=False ): +def npArrayFromWaLBerlaField(field, withGhostLayers=False): """ Creates a numpy array view on the waLBerla field data @field: the waLBerla field @withGhostLayers: Possible values: @@ -50,11 +52,11 @@ def npArrayFromWaLBerlaField( field, withGhostLayers=False ): return view -def arrayFromWaLBerlaAdaptor( field, withGhostLayers=False ): +def arrayFromWaLBerlaAdaptor(field, withGhostLayers=False): return npArrayFromWaLBerlaField( field.copyToField(), withGhostLayers ) -def copyArrayToField( dstField, srcArray, slice=[ slice(None,None,None) ]*3, withGhostLayers=False ): +def copyArrayToField(dstField, srcArray, slice=[slice(None,None,None) ]*3, withGhostLayers=False): """ Copies a numpy array into (part of) a waLBerla field Usually no copying has to take place between waLBerla fields and numpy arrays, since an array view can be @@ -69,11 +71,19 @@ def copyArrayToField( dstField, srcArray, slice=[ slice(None,None,None) ]*3, wit """ dstAsArray = npArrayFromWaLBerlaField(dstField, withGhostLayers) numpy.copyto( dstAsArray[slice], srcArray ) - - + + def extend(cppFieldModule): + def gatherGenerator(blocks, blockDataName, sliceObj, allGather=False): + field = cppFieldModule.gather(blocks, blockDataName, sliceObj, targetRank=-1 if allGather else 0) + if field is not None: + field = npArrayFromWaLBerlaField(field) + field.flags.writeable = False + yield field + + cppFieldModule.toArray = npArrayFromWaLBerlaField cppFieldModule.adaptorToArray = arrayFromWaLBerlaAdaptor cppFieldModule.copyArrayToField = copyArrayToField - + cppFieldModule.gatherGenerator = gatherGenerator