diff --git a/python/ucf/ucf.py b/python/ucf/ucf.py index 6776534..e433173 100644 --- a/python/ucf/ucf.py +++ b/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 \ No newline at end of file + 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 \ No newline at end of file