implemented particle reading routine

python/ucf/ucf.py

@@ -483,7 +483,7 @@ def readGrid(file,verbosity=False,debug=False):
     obj = UCF(file=file,verbosity=verbosity,debug=debug)
     output = []
     for iset in range(0,obj.NumDataset):
-        (data,params) = obj.readSet(step=1,dset=iset)
+        (data,params) = obj.readSet(step=1,dset=iset+1)
         nx = params[0]
         ny = params[1]
         nz = params[2]
@@ -499,7 +499,7 @@ def readProcgrid(file,verbosity=False,debug=False):
     obj = UCF(file=file,verbosity=verbosity,debug=debug)
     output = []
     for iset in range(0,obj.NumDataset):
-        (data,params) = obj.readSet(step=1,dset=iset)
+        (data,params) = obj.readSet(step=1,dset=iset+1)
         nxp = params[0]
         nyp = params[1]
         nzp = params[2]
@@ -543,4 +543,79 @@ def readFieldChunk(file,step=1,dset=-1,verbosity=False,debug=False):
         tmp['rankijk']= obj.IORank[1:]
         output.append(tmp)
     obj.close()
-    return output
+    return output
+
+def readParticles(file,step=-1,verbosity=False,debug=False):
+    # Check what kind of file was passed: standalone, tar, bytes
+    # TBD: tar is not supported yet
+    obj = UCF(file=file,verbosity=verbosity,debug=debug)
+    if not isinstance(step,list):
+        if step==-1:
+            step = range(1,obj.NumTimestep+1)
+        else:
+            step = [step]
+    # The output will be the following: 
+    # 1) numpy array with dimension (ncol,np,ntime) 
+    # 2) dictionary which specifies the columns
+    # We read the data step by step in a list, which is then converted to a 3D array
+    pp = []
+    for ii in step:
+      (data,params) = obj.readSet(step=ii,dset=1)
+      npart       = params[0]
+      ncol        = params[1]
+      ncol_rank   = params[2]
+      ncol_hybrid = params[3]
+      ncol_dem    = params[4]
+      ncol_scalar = params[5]
+      nscal       = ncol_scalar//2
+      pp.append(data.reshape((ncol,npart),order='F'))
+    # Close UCF obeject
+    obj.close()
+    # Convert list of 2D arrays to 3D array
+    pp = np.stack(pp,axis=2)
+    # Create the dictionary
+    col = colmap_from_flags(ncol_rank,ncol_hybrid,ncol_dem,nscal)
+    # Return result
+    return (pp,col)
+    
+def colmap_from_flags(irank,ihybrid,idem,iscal):
+    '''Creates a dictionary which specifies the columns of a particle array.'''
+    col = {}
+    ioffset = 0
+    if irank>0:
+      col['rank'] = ioffset; ioffset+=1
+    if ihybrid>0:
+      col['id'] = ioffset; ioffset+=1
+      col['x']  = ioffset; ioffset+=1
+      col['y']  = ioffset; ioffset+=1
+      col['z']  = ioffset; ioffset+=1
+      col['r']  = ioffset; ioffset+=1
+      col['rho']= ioffset; ioffset+=1
+      col['ax'] = ioffset; ioffset+=1
+      col['ay'] = ioffset; ioffset+=1
+      col['az'] = ioffset; ioffset+=1
+      col['u']  = ioffset; ioffset+=1
+      col['v']  = ioffset; ioffset+=1
+      col['w']  = ioffset; ioffset+=1
+      col['ox'] = ioffset; ioffset+=1
+      col['oy'] = ioffset; ioffset+=1
+      col['oz'] = ioffset; ioffset+=1
+      col['fx'] = ioffset; ioffset+=1
+      col['fy'] = ioffset; ioffset+=1
+      col['fz'] = ioffset; ioffset+=1
+      col['tx'] = ioffset; ioffset+=1
+      col['ty'] = ioffset; ioffset+=1
+      col['tz'] = ioffset; ioffset+=1
+    if idem>0:
+      col['fxc'] = ioffset; ioffset+=1
+      col['fyc'] = ioffset; ioffset+=1
+      col['fzc'] = ioffset; ioffset+=1
+      col['txc'] = ioffset; ioffset+=1
+      col['tyc'] = ioffset; ioffset+=1
+      col['tzc'] = ioffset; ioffset+=1
+      ioffset = ioffset+6;
+    if iscal>0:
+      for ii in range(0,iscal):
+        col['s'+ii] = ioffset; ioffset+=1
+        col['q'+ii] = ioffset; ioffset+=1
+    return col