From 0e48c9b382709ac9833a9f7e6736b958f79eb0de Mon Sep 17 00:00:00 2001
From: Michael Krayer <michael.krayer@posteo.de>
Date: Wed, 2 Jun 2021 09:35:16 +0200
Subject: [PATCH] added striding to parallel routines in order to gather a
 complete field at rank 0

---
 field.py    |  37 ----------------
 parallel.py | 119 ++++++++++++++++++++++++++++++++++++++++------------
 2 files changed, 93 insertions(+), 63 deletions(-)

diff --git a/field.py b/field.py
index 1e204f4..54705c8 100644
--- a/field.py
+++ b/field.py
@@ -236,7 +236,6 @@ class Field3d:
         self.data = self.data.astype(dtype,copy=False)
         return
 
-
     def derivative(self,axis,only_keep_interior=False,shift_origin='before'):
         '''Computes derivative wrt to direction 'axis' with 2nd order finite differences
         centered between the origin grid points'''
@@ -342,42 +341,6 @@ class Field3d:
             radius.append(int(truncate*sigma_img[ii]+0.5))
         return tuple(radius)
 
-    # def shift_origin(self,rel_shift):
-    #     raise NotImplementedError("Routine has not been verified yet.")
-    #     #TBD: verify this routine
-    #     assert isinstance(rel_shift,(tuple,list,numpy.ndarray)) and len(rel_shift)==3,\
-    #         "'shift' must be tuple/list with length 3."
-    #     assert all([rel_shift[ii]>-1.0 and rel_shift[ii]<1.0 for ii in range(3)]),\
-    #         "'shift' must be in (-1.0,1.0)."
-    #     #data = numpy.full(self.dim,numpy.nan)
-    #     data   = self.data.copy()
-    #     origin = list(self.origin)
-    #     for axis in range(3):
-    #         if abs(rel_shift[axis])<self.eps_collapse:
-    #             continue
-    #         origin[axis] += rel_shift[axis]*self.spacing[axis]
-    #         sl_hi  = 3*[slice(None)]
-    #         sl_lo  = 3*[slice(None)]
-    #         sl_out = 3*[slice(None)]
-    #         sl_clr = 3*[slice(None)]
-    #         sl_hi[axis] = slice(1,None)
-    #         sl_lo[axis] = slice(0,-1)
-    #         if rel_shift[ii]<0.0:
-    #             w = 1.0+rel_shift[axis]
-    #             sl_out[axis] = slice(1,None)
-    #             sl_clr[axis] = slice(0,1)
-    #         else:
-    #             w = rel_shift[axis]
-    #             sl_out[axis] = slice(0,-1)
-    #             sl_clr[axis] = slice(-1,None)
-    #         sl_hi  = tuple(sl_hi)
-    #         sl_lo  = tuple(sl_lo)
-    #         sl_out = tuple(sl_out)
-    #         sl_clr = tuple(sl_clr)
-    #         data[sl_out] = (1.0-w)*data[sl_lo] + w*data[sl_hi]
-    #         data[sl_clr] = numpy.nan
-    #     return Field3d(data,origin,self.spacing)
-
     def shift_origin(self,rel_shift,only_keep_interior=False):
         '''Shifts the origin of a field by multiple of spacing.'''
         from scipy import ndimage
diff --git a/parallel.py b/parallel.py
index b223238..a5ee84c 100644
--- a/parallel.py
+++ b/parallel.py
@@ -491,11 +491,6 @@ class PPP:
         gc.collect()
         return
 
-    def merge(self,key,stride=(1,1,1)):
-        '''Returns the complete (possibly downsampled) field at rank 0.'''
-        return
-
-
     def save_state(self,filename):
         import pickle
         fout = filename+'.{:05d}.pickle'.format(self.rank)
@@ -516,33 +511,105 @@ class PPP:
         with open(fout,"wb") as f:
             pickle.dump(payload,f)
 
-    def to_vtk(self,key,no_ghost=False):
+    def to_vtk(self,key,stride=(1,1,1),merge_at_root=False):
         '''Returns the field (only its interior + some ghost cells for plotting)
-        as a vtk mesh.'''
-        return self._subfield_for_vtk(key,no_ghost).to_vtk()
+        as a vtk mesh. If other ghost cells are required, use one of the _subfield
+        methods and apply .to_vtk() to the result.'''
+        from .field import Field3d
+        if merge_at_root:
+            if self.rank==0:
+                stride = np.array(stride)
+                # Allocate the full output field
+                origin = np.array(self.origin[key])
+                spacing = stride*np.array(self.spacing)
+                dim = (np.array(self.dim(key,axis=None))/stride+1).astype(int)
+                # the following seems necessary and seems to work, but i didn't think much about it
+                for axis in range(3):
+                    if self.dim(key,axis=axis)%stride[axis]!=0:
+                        dim[axis]+=1
+                output = Field3d.allocate(dim,origin,spacing,dtype=self.field[key].dtype())
+                # Recieve subfields and insert them
+                output.insert_subfield(self._subfield_for_vtk_merge(key,stride=stride))
+                for rank_src in range(1,self.nproc):
+                    output.insert_subfield(self.comm.recv(source=rank_src))
+                return output.to_vtk()
+            else:
+                self.comm.send(self._subfield_for_vtk_merge(key,stride=stride),dest=0)
+                return None
+        else:
+            return self._subfield_for_vtk(key,stride=stride).to_vtk()
 
-    def vtk_contour(self,key,val):
+    def vtk_contour(self,key,val,stride=(1,1,1)):
         '''Compute isocontour for chunks.'''
-        return self._subfield_for_vtk(key,False).vtk_contour(val)
+        return self._subfield_for_vtk(key,stride=stride).vtk_contour(val)
 
-    def vtk_slice(self,key,normal,origin):
+    def vtk_slice(self,key,normal,origin,stride=(1,1,1)):
         '''Extracts a plane from field.'''
-        return self._subfield_for_vtk(key,False).vtk_slice(normal,origin)
+        return self._subfield_for_vtk(key,stride=stride).vtk_slice(normal,origin)
 
-    def _subfield_for_vtk(self,key,no_ghost):
-        num_ghost = np.array(self.num_ghost)
-        idx_origin = num_ghost
-        dim  = np.array(self.field[key].dim(axis=None))
-        dim -= 2*num_ghost
-        if not no_ghost:
-            idx_origin -= 1
-            dim += 1
-            if self.ip==self.nxp-1: dim[0]+=1
-            if self.jp==self.nyp-1: dim[1]+=1
-            if self.kp==self.nzp-1: dim[2]+=1
-        idx_origin = tuple(idx_origin)
-        dim = tuple(dim)
-        return self.field[key].extract_subfield(idx_origin,dim)
+    def _subfield_for_vtk(self,key,stride=(1,1,1)):
+        '''Extracts a subfield from chunk which has one leading and no trailing
+        ghost cell except when there is no trailing neighbor. This guarantees 
+        non-overlapping contours.'''
+        no_upper_ghost = [True,True,True]
+        if self.ip==self.nxp-1: no_upper_ghost[0]=False
+        if self.jp==self.nyp-1: no_upper_ghost[1]=False
+        if self.kp==self.nzp-1: no_upper_ghost[2]=False
+        return self._subfield(key,stride,(1,1,1),no_upper_ghost=no_upper_ghost)
+
+    def _subfield_for_vtk_merge(self,key,stride=(1,1,1)):
+        '''Extracts a subfield from chunk which has one leading and no trailing
+        ghost cell except when there is no trailing neighbor. This guarantees 
+        non-overlapping contours.'''
+        no_lower_ghost = [True,True,True]
+        no_upper_ghost = [True,True,True]
+        if self.ip==self.nxp-1: no_upper_ghost[0]=False
+        if self.jp==self.nyp-1: no_upper_ghost[1]=False
+        if self.kp==self.nzp-1: no_upper_ghost[2]=False
+        return self._subfield(key,stride,(1,1,1),no_lower_ghost=no_lower_ghost,
+                                                 no_upper_ghost=no_upper_ghost)
+
+    def _subfield_interior(self,key,stride=(1,1,1)):
+        return self._subfield(key,stride,(0,0,0))
+
+    def _subfield(self,key,stride,num_ghost,no_lower_ghost=(False,False,False),
+                                            no_upper_ghost=(False,False,False)):
+        '''Returns the field with a stride applied.'''
+        stride    = np.array(stride,dtype=int)
+        num_ghost = np.array(num_ghost,dtype=int)
+        no_lower_ghost = np.array(no_lower_ghost,dtype=bool)
+        no_upper_ghost = np.array(no_upper_ghost,dtype=bool)
+        assert len(stride)==3, "'stride' needs to be of length 3."
+        assert len(num_ghost)==3, "'num_ghost' needs to be of length 3."
+        assert len(no_lower_ghost)==3, "'no_lower_ghost' needs to be of length 3."
+        assert len(no_upper_ghost)==3, "'no_upper_ghost' needs to be of length 3."
+        assert all(np.array(self.num_ghost)>=stride*num_ghost), "At least stride*num_ghost "\
+            "ghost cells are required in the original field. Have {}, need {}.".format(
+                tuple(self.num_ghost),tuple(stride*num_ghost))
+        ib = self.proc_grid[key][0][self.ip]-1
+        ie = self.proc_grid[key][1][self.ip]-1
+        jb = self.proc_grid[key][2][self.jp]-1
+        je = self.proc_grid[key][3][self.jp]-1
+        kb = self.proc_grid[key][4][self.kp]-1
+        ke = self.proc_grid[key][5][self.kp]-1
+        ijkb = np.array((ib,jb,kb))
+        ijke = np.array((ie,je,ke))
+        # Determine the first point to consider in this chunk for the case of output
+        # without ghost cells
+        idx_origin = (np.ceil(ijkb/stride)*stride).astype(int)-ijkb+self.num_ghost
+        # Now the last point to consider under the same circumstances
+        idx_endpoint = (np.floor(ijke/stride)*stride).astype(int)-ijkb+self.num_ghost
+        # If we want the subfield to have ghost cells, shift the origin and endpoint
+        idx_origin   -= (~no_lower_ghost).astype(int)*stride*num_ghost
+        idx_endpoint += (~no_upper_ghost).astype(int)*stride*num_ghost
+        # Compute the number of points
+        num_points = ((idx_endpoint-idx_origin)/stride+1).astype(int)
+        # Some last assertions which should never fail
+        # print(self.rank,idx_origin,idx_endpoint,num_points,self.field[key].dim(),self.field[key].dim())
+        assert all(idx_origin>=0)
+        assert all(idx_endpoint<np.array(self.field[key].dim()))
+        # Return the subfield
+        return self.field[key].extract_subfield(idx_origin,num_points,stride=stride)
 
     def rank_from_position(self,ip,jp,kp,external=False):
         if external: