Python interface: convenient gather function, slice - CellInterval conversion

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 ]
-
-
 
 
 

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