From c072e41677c125f7aca1d506e164a57656d026c6 Mon Sep 17 00:00:00 2001
From: Michael Krayer <michael.krayer@posteo.de>
Date: Thu, 24 Mar 2022 21:11:44 +0100
Subject: [PATCH] first version of serial tool: can compute Q, but untested

---
 __init__.py |   1 +
 field.py    |  46 +++++++++++-
 helper.py   |  56 ++++++++++++++-
 serial.py   | 203 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 303 insertions(+), 3 deletions(-)
 create mode 100644 serial.py

diff --git a/__init__.py b/__init__.py
index 0f9973b..35adf69 100644
--- a/__init__.py
+++ b/__init__.py
@@ -4,3 +4,4 @@ from . import field
 from . import ucf
 from . import helper
 from . import parallel
+from . import serial
\ No newline at end of file
diff --git a/field.py b/field.py
index 2619abc..d4464f0 100644
--- a/field.py
+++ b/field.py
@@ -10,7 +10,7 @@ class Field3d:
             self.data = data
         self.origin    = tuple([float(x) for x in origin])
         self.spacing   = tuple([float(x) for x in spacing])
-        self.eps_collapse = 1e-8
+        self.eps_collapse = 1e-4*np.min(spacing)
         self._dim = None
         return
 
@@ -131,6 +131,48 @@ class Field3d:
         if not is_open: f.close()
         return cls(data,origin,spacing)
 
+    @classmethod
+    def from_snapshot(cls,snap,key,procslice=None,dtype=np.float64):
+        ''' 
+        Construct a Field3d from a field saved in an UCF snapshot.
+        Use procslice = (ib,ie,jb,je,kb,je) to load only a block
+        of the full data. 
+        '''
+        from .helper import eval_procslice, key_grid, rank_from_position
+        proc        = snap.procgrid()
+        origin      = snap.origin()
+        spacing     = snap.spacing()
+        periodicity = snap.periodicity()
+        bounds      = snap.bounds()
+        keyg        = key_grid(key)
+        nxp,nyp,nzp = [len(proc[keyg][2*ii]) for ii in range(3)]
+        # Get field dimensions
+        (ipb,ipe,jpb,jpe,kpb,kpe),ori,(nxl,nyl,nzl) = eval_procslice(
+                                    proc,origin,spacing,keyg,procslice)
+        # Create a Field3d: ghost cells will be included!
+        data = np.empty((nxl+2,nyl+2,nzl+2),dtype=dtype)
+        field_ = cls(data,[ori[ii]-spacing[ii] for ii in range(3)],spacing)
+        # Go through each chunk to be read and construct the field
+        for ip_ in range(ipb,ipe+1):
+            for jp_ in range(jpb,jpe+1):
+                for kp_ in range(kpb,kpe+1):
+                    # Determine rank of the chunk to be read
+                    rank_ = rank_from_position(ip_,jp_,kp_,nyp,nzp)
+                    # Compute bounds of this chunk
+                    ib_ = proc[keyg][0][ip_]
+                    ie_ = proc[keyg][1][ip_]
+                    jb_ = proc[keyg][2][jp_]
+                    je_ = proc[keyg][3][jp_]
+                    kb_ = proc[keyg][4][kp_]
+                    ke_ = proc[keyg][5][kp_]
+                    nxl_ = ie_-ib_+1
+                    nyl_ = je_-jb_+1
+                    nzl_ = ke_-kb_+1
+                    # Read data and insert it
+                    subfield = snap.field_chunk(rank_,key)
+                    field_.insert_subfield(subfield)
+        return field_
+
     @classmethod
     def allocate(cls,dim,origin,spacing,fill=None,dtype=np.float64):
         '''Allocates an empty field, or a field filled with 'fill'.'''
@@ -359,7 +401,7 @@ class Field3d:
         assert axis<3, "'axis' must be one of 0,1,2."
         origin = list(self.origin)
         assert shift_origin in ('before','after'), "'shift_origin' must be one of {'before','after'}."
-        if shift_origin=='left':
+        if shift_origin=='before':
             corr_orig = 0
             origin[axis] -= 0.5*self.spacing[axis]           
         else:
diff --git a/helper.py b/helper.py
index 1ead405..c67a744 100644
--- a/helper.py
+++ b/helper.py
@@ -97,4 +97,58 @@ def position_from_rank(rank,nyp,nzp):
     ip = rank//(nyp*nzp)
     jp = (rank//nzp)%nyp
     kp = rank%nzp
-    return (ip,jp,kp)
\ No newline at end of file
+    return (ip,jp,kp)
+
+def key_grid(key):
+    '''Get the key for the grid: scalar fields all share a common grid,
+    so this routine returns 's' if key is 'sX' with X being a (potentially
+    multi-digit) number. Otherwise the key itself is returned.'''
+    if key[0]=='s' and key[1:].isnumeric():
+        return 's'
+    else:
+        return key
+
+def eval_procslice(proc,origin,spacing,key_grid,procslice):
+    # Get number of procs
+    nxp = len(proc[key_grid][0])
+    nyp = len(proc[key_grid][2])
+    nzp = len(proc[key_grid][4])
+    nxyzp = (nxp,nyp,nzp)
+    # Get processor indices for slicing
+    if procslice is not None:
+        ipb = procslice[0] if procslice[0] is not None else 0
+        ipe = procslice[1] if procslice[1] is not None else nxp-1
+        jpb = procslice[2] if procslice[2] is not None else 0
+        jpe = procslice[3] if procslice[3] is not None else nyp-1
+        kpb = procslice[4] if procslice[4] is not None else 0
+        kpe = procslice[5] if procslice[5] is not None else nzp-1
+    else:
+        ipb = 0
+        ipe = nxp-1
+        jpb = 0
+        jpe = nyp-1
+        kpb = 0
+        kpe = nzp-1
+    idxproc = (ipb,ipe,jpb,jpe,kpb,kpe)
+    for ii_ in range(6):
+        assert idxproc[ii_]>=0, \
+            "procslice[{:d}] must be greater than or equal to 0, but is {:d}.".format(
+                ii_,idxproc[ii_])
+        assert idxproc[ii_]<nxyzp[ii_//2], \
+            "procslice[{:d}] must be smaller than {:d}, but is {:d}.".format(
+                ii_,nxyzp[ii_//2],idxproc[ii_])
+    for ii_ in range(3):
+        assert idxproc[2*ii_]<=idxproc[2*ii_+1], \
+            "procslice[{:d}] must be smaller than or equal to procslice[{:d}].".format(
+                2*ii_,2*ii_+1)
+    ib = proc[key_grid][0][ipb]
+    ie = proc[key_grid][1][ipe]
+    jb = proc[key_grid][2][jpb]
+    je = proc[key_grid][3][jpe]
+    kb = proc[key_grid][4][kpb]
+    ke = proc[key_grid][5][kpe]
+    dim = (ie-ib+1,je-jb+1,ke-kb+1)
+    ori = (origin[key_grid][0]+(ib-1)*spacing[0],
+           origin[key_grid][1]+(jb-1)*spacing[1],
+           origin[key_grid][2]+(kb-1)*spacing[2])
+    return idxproc,ori,dim
\ No newline at end of file
diff --git a/serial.py b/serial.py
new file mode 100644
index 0000000..6687f9b
--- /dev/null
+++ b/serial.py
@@ -0,0 +1,203 @@
+import numpy as np
+import math
+from .helper import position_from_rank, rank_from_position, eval_procslice, key_grid
+from .field import Field3d
+
+class SPP:
+    """Serial Python Postprocessor for suspend"""
+    def __init__(self,field,periodicity,bounds,origin,spacing,dim):
+        assert all([key in field for key in ('u','v','w')]), "Required fields: 'u','v','w'."
+        for axis in range(3):
+            if periodicity[axis]:
+                assert bounds[2*axis]==0.0, "Lower bound in periodic direction must equal zero."
+        self.field       = field
+        self.bounds      = tuple(bounds)
+        self.periodicity = tuple(periodicity)
+        self.origin      = tuple(origin)
+        self.spacing     = tuple(spacing)
+        self.dim         = tuple(dim)
+        self.xperiodic,self.yperiodic,self.zperiodic = periodicity
+        return 
+
+    @classmethod
+    def from_snapshot(cls,snap,procslice=None,dtype=np.double,load_pressure=False,verbose=False):
+        # Load fields
+        keys = ('u','v','w','p') if load_pressure else ('u','v','w')
+        field = {}
+        for key_ in keys:
+            if verbose:
+                print('[SPP.from_snapshot] loading {}'.format(key_))
+            field[key_] = Field3d.from_snapshot(snap,key_,procslice,dtype)
+        # As a common grid, the pressure grid is used. Extract its offset 
+        proc      = snap.procgrid()
+        origin    = snap.origin()
+        spacing   = snap.spacing()
+        orip,dimp = eval_procslice(proc,origin,spacing,'p',procslice)[1:]
+        orip = tuple(orip[ii]-spacing[ii] for ii in range(3))
+        dimp = tuple(dimp[ii]+2 for ii in range(3))
+        # Set periodicity and global bounds
+        bounds      = snap.bounds()
+        periodicity = list(snap.periodicity())
+        if procslice is not None: # if not all procs are loaded, set periodicity to false
+            if procslice[0] is not None or procslice[1] is not None:
+                periodicity[0] = False
+            if procslice[2] is not None or procslice[3] is not None:
+                periodicity[1] = False
+            if procslice[4] is not None or procslice[5] is not None:
+                periodicity[2] = False
+        return cls(field,periodicity,bounds,orip,spacing,dimp)
+
+    @classmethod
+    def from_file(cls,file,verbose=False):
+        '''Load saved state from file.'''
+        import h5py
+        if verbose:
+            print('[SPP.from_file] file={:s}'.format(file))
+        f = h5py.File(file,'r')
+        bounds      = tuple(f['bounds'])    
+        periodicity = tuple(f['periodicity'])    
+        origin      = tuple(f['origin'])    
+        spacing     = tuple(f['spacing'])    
+        dim         = tuple(f['dim'])    
+        # Load fields
+        field_names = tuple(x.decode('ascii') for x in f['field_names'])
+        field = {}
+        for key_ in field_names:
+            if verbose: print('[SPP.from_file] loading {:s}'.format(key_))
+            field[key_] = Field3d.from_file(f,key_)
+        f.close()
+        return cls(field,periodicity,bounds,origin,spacing,dim)
+
+    def save(self,file,verbose=False):
+        '''Save current state to file.'''
+        import h5py
+        if verbose:
+            print('[SPP.save] {:s}'.format(file))
+        f = h5py.File(file,'w')
+        f.create_dataset('bounds',data=self.bounds)
+        f.create_dataset('periodicity',data=self.periodicity)
+        f.create_dataset('origin',data=self.origin)
+        f.create_dataset('spacing',data=self.spacing)
+        f.create_dataset('dim',data=self.dim)
+        # Save fields and field names
+        ascii_type = h5py.string_dtype('ascii',32)
+        d = f.create_dataset('field_names',len(self.field),dtype=ascii_type)
+        ii_ = 0
+        for key_ in self.field:
+            self.field[key_].save(f,name=key_)
+            d[ii_] = key_
+            ii_ += 1
+        f.close()
+        return
+
+    def delete(self,key):
+        '''Removes a field from memory.'''
+        import gc
+        assert key in self.field, \
+            "Field {} cannot be deleted, because it does not exist.".format(key)
+        del self.field[key]
+        gc.collect()
+        return
+
+    def differentiate(self,key,axis,key_out=None):
+        assert axis<3, "'axis' must be one of 0,1,2."
+        if key_out is None:
+            key_out = 'd'+key+('dx','dy','dz')[axis]
+        shifting_state = self.shifting_state(key,axis=axis)
+        if shifting_state==-1:
+            shift_origin = 'after'
+        elif shifting_state==0:
+            shift_origin = 'before'
+        elif shifting_state==1:
+            shift_origin = 'before'
+        else:
+            raise ValueError("Invalid shifting state.")
+        self.field[key_out] = self.field[key].derivative(axis,shift_origin=shift_origin)
+        self.exchange_ghost_cells(key_out)
+
+    def qcriterion(self,key_out='Q',from_pressure=False,dtype=np.double):
+        '''Computes Q-criterion'''
+        if from_pressure:
+            assert 'p' in self.field, "Pressure field is not loaded."
+            self.field[key_out] = 0.5*self.field['p'].laplacian(only_keep_interior=False)
+        else:
+            assert 'u' in self.field, "Velocity field is not loaded: u"
+            assert 'v' in self.field, "Velocity field is not loaded: v"
+            assert 'w' in self.field, "Velocity field is not loaded: w"
+            self.field[key_out] = Field3d.allocate(self.dim,self.origin,self.spacing,
+                                                fill=0.0,dtype=dtype)
+            keyx = ('x','y','z')
+            keyu = ('u','v','w')
+            for ii in range(3):
+                key_der = 'd'+keyu[ii]+'d'+keyx[ii]
+                self.differentiate(keyu[ii],ii,key_der)
+                self.shift_to_pressure_grid(key_der)
+                self.field[key_out] += self.field[key_der]**2
+                self.delete(key_der)
+                #
+                iip1 = (ii+1)%3
+                key_der1 = 'd'+keyu[ii]+'d'+keyx[iip1]
+                key_der2 = 'd'+keyu[iip1]+'d'+keyx[ii]
+                self.differentiate(keyu[ii],iip1,key_der1)
+                self.differentiate(keyu[iip1],ii,key_der2)
+                self.shift_to_pressure_grid(key_der1)
+                self.shift_to_pressure_grid(key_der2)
+                self.field[key_out] += 2.0*self.field[key_der1]*self.field[key_der2]
+                self.delete(key_der1)
+                self.delete(key_der2)
+            self.field[key_out] *= -0.5
+        return
+
+    def shift_to_pressure_grid(self,key,key_out=None):
+        from .field import Field3d
+        if key_out is None:
+            if key in ('u','v','w'):
+                raise ValueError("In-place shift of {'u','v','w'} is not permitted.")
+            key_out = key
+        fld_ref   = Field3d.pseudo_field(self.dim,self.origin,self.spacing)
+        rel_shift = self.field[key].relative_shift(fld_ref)
+        self.field[key_out] = self.field[key].shift_origin(rel_shift)
+        self.exchange_ghost_cells(key_out)
+
+    def shifting_state(self,key,axis=None):
+        '''Returns the relative shift with respect to domain
+        boundaries as either 1 (= +0.5h), 0 (= 0h) or -1 (= -0.5h).'''
+        if axis is None:
+            return tuple(self.shifting_state(key,axis=ii) for ii in range(3))
+        assert axis<3, "'axis' must be one of 0,1,2."
+        return int(round(
+            (self.field[key].origin[axis]-self.origin[axis])/(0.5*self.spacing[axis])
+            ))
+
+    def exchange_ghost_cells(self,key):
+        '''Copy ghost cells in periodic directions.'''
+        for axis in range(3):
+            if not self.periodicity[axis]:
+                continue
+            x = self.field[key].grid(axis=axis)
+            sl_in  = 3*[slice(None)]
+            sl_out = 3*[slice(None)]
+            # Loop through grid points below lower bound
+            xgh  = np.mod(x[x<self.bounds[2*axis]],self.bounds[2*axis+1])
+            idx_ = 0
+            for x_ in xgh:
+                ii = self.field[key].nearest_gridpoint(x_,axis=axis)
+                assert math.fabs(x_-x[ii])<self.field[key].eps_collapse, \
+                    "Invalid ghost cell detected: {:f}, {:f}".format(x_,x[ii])
+                sl_in[axis]  = ii
+                sl_out[axis] = idx_
+                self.field[key].data[tuple(sl_out)] = self.field[key].data[tuple(sl_in)]
+                idx_ += 1
+            # Loop through grid points above upper bound
+            xgh  = np.mod(x[x>=self.bounds[2*axis+1]],self.bounds[2*axis+1])
+            idx_ = self.field[key].dim(axis=axis)-len(xgh)
+            for x_ in xgh:
+                ii = self.field[key].nearest_gridpoint(x_,axis=axis)
+                assert math.fabs(x_-x[ii])<self.field[key].eps_collapse, \
+                    "Invalid ghost cell detected:{:f}, {:f}".format(x_,x[ii])
+                sl_in[axis]  = ii
+                sl_out[axis] = idx_
+                self.field[key].data[tuple(sl_out)] = self.field[key].data[tuple(sl_in)]
+                idx_ += 1
+        return
+