From 3366816adfeb9966127eac5d78537bb6c86e2ee9 Mon Sep 17 00:00:00 2001 From: Michael Krayer Date: Tue, 1 Jun 2021 12:16:10 +0200 Subject: [PATCH] first production version for filtered plotting --- field.py | 343 ++++++++++++++++++-------- parallel.py | 683 ++++++++++++++++++++++++++++++++++------------------ particle.py | 111 +++++---- 3 files changed, 759 insertions(+), 378 deletions(-) diff --git a/field.py b/field.py index 360c95f..1e204f4 100644 --- a/field.py +++ b/field.py @@ -4,19 +4,74 @@ class Field3d: assert len(origin)==3, "'origin' must be of length 3" assert len(spacing)==3, "'spacing' must be of length 3" self.data = numpy.array(data) - self.numpoints = self.data.shape self.origin = tuple([float(x) for x in origin]) self.spacing = tuple([float(x) for x in spacing]) self.eps_collapse = 1e-12 + self._dim = None return def __str__(self): str = 'Field3d with\n' - str+= ' dimension: {}, {}, {}\n'.format(*self.numpoints) + str+= ' dimension: {}, {}, {}\n'.format(*self.dim()) str+= ' origin: {:G}, {:G}, {:G}\n'.format(*self.origin) - str+= ' spacing: {:G}, {:G}, {:G}'.format(*self.spacing) + str+= ' spacing: {:G}, {:G}, {:G}\n'.format(*self.spacing) + str+= ' datatype: {}'.format(self.dtype()) return str + def __add__(self,other): + if isinstance(other,Field3d): + assert self.has_same_grid(other), "Grid mismatch." + return Field3d(self.data+other.data,self.origin,self.spacing) + else: + return Field3d(self.data+other,self.origin,self.spacing) + + def __sub__(self,other): + if isinstance(other,Field3d): + assert self.has_same_grid(other), "Grid mismatch." + return Field3d(self.data-other.data,self.origin,self.spacing) + else: + return Field3d(self.data-other,self.origin,self.spacing) + + def __mul__(self,other): + if isinstance(other,Field3d): + assert self.has_same_grid(other), "Grid mismatch." + return Field3d(self.data*other.data,self.origin,self.spacing) + else: + return Field3d(self.data*other,self.origin,self.spacing) + + def __radd__(self,other): + return Field3d(other+self.data,self.origin,self.spacing) + + def __rmul__(self,other): + return Field3d(other*self.data,self.origin,self.spacing) + + def __pow__(self,other): + return Field3d(self.data**other,self.origin,self.spacing) + + def __iadd__(self,other): + if isinstance(other,Field3d): + assert self.has_same_grid(other), "Grid mismatch." + self.data += other.data + else: + self.data += other + return self + + def __isub__(self,other): + if isinstance(other,Field3d): + assert self.has_same_grid(other), "Grid mismatch." + self.data -= other.data + else: + self.data -= other + return self + + def __imul__(self,other): + if isinstance(other,Field3d): + assert self.has_same_grid(other), "Grid mismatch." + self.data *= other.data + else: + self.data *= other + return self + # TBD: this should return another Field3d object # def __getitem__(self,val): # assert isinstance(val,tuple) and len(val)==3, "Field3d must be indexed by [ii,jj,kk]." @@ -45,6 +100,30 @@ class Field3d: assert (chunk['nzl']+2*chunk['ighost'])==nz, "Invalid chunk data: nzl != chunk['data'].shape[2]" return cls(chunk['data'],origin=(xo,yo,zo),spacing=(dx,dy,dz)) + @classmethod + def allocate(cls,dim,origin,spacing,fill=None,dtype=numpy.float64,pseudo=False): + '''Allocates an empty field, or a field filled with 'fill'.''' + assert isinstance(dim,(tuple,list,numpy.ndarray)) and len(dim)==3,\ + "'dim' must be a tuple/list of length 3." + assert isinstance(origin,(tuple,list,numpy.ndarray)) and len(origin)==3,\ + "'origin' must be a tuple/list of length 3." + assert isinstance(spacing,(tuple,list,numpy.ndarray)) and len(spacing)==3,\ + "'spacing' must be a tuple/list of length 3." + if pseudo: + data = numpy.empty((0,0,0)) + r = cls(data,origin,spacing) + r._dim = dim + return r + else: + if fill is None: + data = numpy.empty(dim,dtype=dtype) + else: + data = numpy.full(dim,fill,dtype=dtype) + return cls(data,origin,spacing) + + def copy(self): + return Field3d(self.data.copy(),self.origin,self.spacing) + def insert_subfield(self,subfield): assert all([abs(subfield.spacing[ii]-self.spacing[ii])=0 and idx_origin[ii]=0 and idx_origin[ii]0 and idx_nearest0 and idx_nearest-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])-1.0 and rel_shift[ii]<1.0 for ii in range(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.numpoints,numpy.nan) - data = self.data.copy() + data = self.data.copy() origin = list(self.origin) + sl = 3*[slice(None)] for axis in range(3): if abs(rel_shift[axis])0: + weights = (1.0-rel_shift[axis],rel_shift[axis]) + data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=numpy.nan,origin=-1) + origin[axis] += rel_shift[axis]*self.spacing[axis] + if only_keep_interior: + sl[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 + weights = (-rel_shift[axis],1.0+rel_shift[axis]) + data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=numpy.nan,origin=0) + origin[axis] += rel_shift[axis]*self.spacing[axis] + if only_keep_interior: + sl[axis] = slice(1,None) + origin[axis] += self.spacing[axis] + if only_keep_interior: + data = data[sl] return Field3d(data,origin,self.spacing) - def change_grid(self,origin,spacing,numpoints,padding=None,numpad=1): + def relative_shift(self,field): + '''Compute the relative shift (in terms of spacing) to shift self onto field.''' + assert self.has_same_spacing(field), "spacing differs." + rel_shift = [0.0,0.0,0.0] + for axis in range(3): + dist = field.origin[axis]-self.origin[axis] + if abs(dist)>self.eps_collapse: + rel_shift[axis] = dist/self.spacing[axis] + return tuple(rel_shift) + + def change_grid(self,origin,spacing,dim,padding=None,numpad=1): assert all([origin[ii]>=self.origin[ii] for ii in range(0,3)]), "New origin is out of bounds." - endpoint = [origin[ii]+(numpoints[ii]-1)*spacing[ii] for ii in range(0,3)] + endpoint = [origin[ii]+(dim[ii]-1)*spacing[ii] for ii in range(0,3)] assert all([endpoint[ii]<=self.endpoint(ii) for ii in range(0,3)]), "New end point is out of bounds." # Allocate (possibly padded array) - origin_pad,numpoints_pad,sl_out = padding(origin,spacing,numpoints,padding,numpad): - data = numpy.zeros(numpoints_pad,dtype=self.data.dtype) + origin_pad,dim_pad,sl_out = padding(origin,spacing,dim,padding,numpad) + data = numpy.zeros(dim_pad,dtype=self.data.dtype) # Trilinear interpolation if numpy.allclose(spacing,self.spacing): # spacing is the same: we can construct universal weights for the stencil @@ -295,14 +437,14 @@ class Field3d: for kk in range(0,2): if c[ii,jj,kk]>self.eps_collapse: data[sl_out] += c[ii,jj,kk]*self.data[ - i0+ii:i0+ii+numpoints[0], - j0+jj:j0+jj+numpoints[1], - k0+kk:k0+kk+numpoints[2]] + i0+ii:i0+ii+dim[0], + j0+jj:j0+jj+dim[1], + k0+kk:k0+kk+dim[2]] else: data_ref = data[sl_out] - for ii in range(0,numpoints[0]): - for jj in range(0,numpoints[1]): - for kk in range(0,numpoints[2]): + for ii in range(0,dim[0]): + for jj in range(0,dim[1]): + for kk in range(0,dim[2]): coord = ( origin[0]+ii*spacing[0], origin[1]+jj*spacing[1], @@ -326,36 +468,36 @@ class Field3d: return val @staticmethod - def padding(origin,spacing,numpoints,padding,numpad): + def padding(origin,spacing,dim,padding,numpad): if isinstance(numpad,int): numpad = numpy.fill(3,numpad,dtype=int) else: numpad = numpy.array(numpad,dtype=int) assert len(numpad)==3, "'numpad' must be either an integer or tuple/list of length 3." origin_pad = numpy.array(origin) - numpoints_pad = numpy.array(numpoints) + dim_pad = numpy.array(dim) sl_out = [slice(None),slice(None),slice(None)] if padding is not None: if padding=='before': - numpoints_pad += numpad + dim_pad += numpad origin_pad -= numpad*spacing for axis in range(3): sl_out[axis] = slice(numpad[axis],None) elif padding=='after': - numpoints_pad += numpad + dim_pad += numpad for axis in range(3): sl_out[axis] = slice(0,-numpad[axis]) elif padding=='both': - numpoints_pad += 2*numpad + dim_pad += 2*numpad origin_pad -= numpad*spacing for axis in range(3): sl_out[axis] = slice(numpad[axis],-numpad[axis]) else: raise ValueError("'padding' must either be None or one of {'before','after','both'}.") - sl_out = tuple(sl_out) - origin_pad = tuple(origin_pad) - numpoints_pad = tuple(numpoints_pad) - return (origin_pad,numpoints_pad,sl_out) + sl_out = tuple(sl_out) + origin_pad = tuple(origin_pad) + dim_pad = tuple(dim_pad) + return (origin_pad,dim_pad,sl_out) def weights_trilinear(self,rel_dist): assert len(rel_dist)==3, "len(rel_dist) must be 3." @@ -408,6 +550,19 @@ class Field3d: "'origin' must be tuple of length 3." return self.to_vtk(deep=deep).slice(normal=normal,origin=origin) +def gaussian_filter_umean_channel(array,spacing,sigma,truncate=4.0): + '''Applies a Gaussian filter to a numpy array of dimension (1,ny,1) which + contains the mean streamwise velocity of the channel (or possibly sth else). + Since yu[0] = c and yu[-1] = d, we can use scipy's mirror settings and don't + need ghost cells.''' + from scipy import ndimage + assert array.ndim==3, "Expected an array with three dimensions/axes." + assert array.shape[0]==1 and array.shape[1]>1 and array.shape[2]==1,\ + "Expected an array with shape (1,ny,1)." + sigma_img = sigma/spacing + array = ndimage.gaussian_filter1d(array,sigma_img,axis=1,truncate=truncate,mode='mirror') + return array + class ChunkIterator: '''Iterates through all chunks. 'snapshot' must be an instance of a class which returns a Field3d from the method call diff --git a/parallel.py b/parallel.py index 5d9757a..b223238 100644 --- a/parallel.py +++ b/parallel.py @@ -1,22 +1,43 @@ +import numpy as np class PPP: """Parallel Python Postprocessor for suspend""" - def __init__(self,comm,func_load,num_ghost,precision,origin,spacing,periodicity,bounds,proc,chunks_per_proc): - '''Constructor: except for comm, only rank 0 needs initialized data.''' - self.comm = comm - self.rank = comm.Get_rank() - self.func_load = func_load - self.init_settings(num_ghost,precision) - self.init_domain(origin,spacing,periodicity,bounds) - self.init_procgrid(proc,chunks_per_proc) - self.field = {} - self.symmetries = {} - return - + def __init__(self,comm,func_load, + num_ghost,chunks_per_proc,origin,spacing,periodicity,bounds, + proc_grid,nxp,nyp,nzp, + proc_grid_ext,nxp_ext,nyp_ext,nzp_ext, + nghbr,field,symmetries): + '''Class constructor: it is expected that all communcation is done in the corresponding + class methods, such that every rank already holds the required information.''' + # Initialize object with passed arguments + self.comm = comm + self.func_load = func_load + self.num_ghost = num_ghost + self.chunks_per_proc = chunks_per_proc + self.origin = origin + self.spacing = spacing + self.periodicity = periodicity + self.bounds = bounds + self.proc_grid = proc_grid + self.nxp,self.nyp,self.nzp = nxp,nyp,nzp + self.proc_grid_ext = proc_grid_ext + self.nxp_ext,self.nyp_ext,self.nzp_ext = nxp_ext,nyp_ext,nzp_ext + self.nghbr = nghbr + self.field = field + self.symmetries = symmetries + # Initialize some deduced variables + self.rank = comm.Get_rank() + self.nproc = nxp*nyp*nzp + self.ip,self.jp,self.kp = self.position_from_rank(self.rank,external=False) + self.nghx,self.nghy,self.nghz = num_ghost + self.xperiodic,self.yperiodic,self.zperiodic = periodicity + return + @classmethod - def from_snapshot(cls,snap,chunks_per_proc=(1,1,1),num_ghost=(1,1,1),precision='float64'): + def from_snapshot(cls,snap,chunks_per_proc=(1,1,1),num_ghost=(1,1,1)): from mpi4py import MPI comm = MPI.COMM_WORLD rank = comm.Get_rank() + # Let rank 0 read the processor grid as well as some info and then broadcast it if rank==0: proc = snap.procgrid() origin = snap.origin() @@ -24,182 +45,53 @@ class PPP: periodicity = snap.periodicity() bounds = snap.bounds() else: - proc = None - origin = None - spacing = None - periodicity = None - bounds = None + proc,origin,spacing,periodicity,bounds = 5*[None] + proc = comm.bcast(proc ,root=0) + origin = comm.bcast(origin ,root=0) + spacing = comm.bcast(spacing ,root=0) + periodicity = comm.bcast(periodicity,root=0) + bounds = comm.bcast(bounds ,root=0) + # Make sure all processors got the same arguments + chunks_per_proc = comm.bcast(chunks_per_proc,root=0) + num_ghost = comm.bcast(num_ghost,root=0) + # Set data loading function func_load = snap.field_chunk - return cls(comm,func_load,num_ghost,precision,origin,spacing,periodicity,bounds,proc,chunks_per_proc) + # Construct new processor grid + (proc_grid,nxp,nyp,nzp, + proc_grid_ext,nxp_ext,nyp_ext,nzp_ext, + nghbr) = PPP._construct_procgrid(comm,proc,chunks_per_proc,num_ghost,periodicity) + # Initially, no field is loaded + field = {} + symmetries = {} + return cls(comm,func_load, + num_ghost,chunks_per_proc,origin,spacing,periodicity,bounds, + proc_grid,nxp,nyp,nzp, + proc_grid_ext,nxp_ext,nyp_ext,nzp_ext, + nghbr,field,symmetries) - def init_settings(self,num_ghost,precision): - '''Initializes PPP settings for all processors.''' - # TBD: some assertions - self.num_ghost = self.comm.bcast(num_ghost,root=0) - self.precision = self.comm.bcast(precision,root=0) - # Some shortcuts - self.nghx,self.nghy,self.nghz = self.num_ghost + @classmethod + def from_state(cls,filename): + import pickle + from mpi4py import MPI + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + fin = filename+'.{:05d}.pickle'.format(rank) + with open(fin,"rb") as f: + payload = pickle.load(f) + (num_ghost,chunks_per_proc,origin,spacing,periodicity,bounds, + proc_grid,nxp,nyp,nzp,nghbr,field,symmetries) = payload + assert nxp*nyp*nzp==comm.Get_size(), "The loaded state requires {} processors, but "\ + "we are currently running with {}.".format(nxp*nyp*nzp,comm.Get_size()) + func_load = None + proc_grid_ext = None + nxp_ext,nyp_ext,nzp_ext = 3*[None] + return cls(comm,func_load, + num_ghost,chunks_per_proc,origin,spacing,periodicity,bounds, + proc_grid,nxp,nyp,nzp, + proc_grid_ext,nxp_ext,nyp_ext,nzp_ext, + nghbr,field,symmetries) - def init_symmetries(self,key,mirror=(True,True)): - '''Sets the symmetries for ghost cells behind the wall''' - # Example: wall in y - # No-slip boundary (no mirror) free-slip boundary (mirror) - # u -> -u u -> u - # v -> -v v -> -v - # w -> -w w -> w - # p -> p p -> p - import numpy as np - self.symmetries[key] = np.zeros((3,3,3),dtype='i') - if not self.xperiodic: - if key=='u': - self.symmetries[key][0,1,1] = -1 - self.symmetries[key][2,1,1] = -1 - else: - self.symmetries[key][0,1,1] = 1 if mirror[0] else -1 - self.symmetries[key][2,1,1] = 1 if mirror[1] else -1 - if not self.yperiodic: - if key=='v': - self.symmetries[key][1,0,1] = -1 - self.symmetries[key][1,2,1] = -1 - else: - self.symmetries[key][1,0,1] = 1 if mirror[0] else -1 - self.symmetries[key][1,2,1] = 1 if mirror[1] else -1 - if not self.zperiodic: - if key=='w': - self.symmetries[key][1,1,0] = -1 - self.symmetries[key][1,1,2] = -1 - else: - self.symmetries[key][1,1,0] = 1 if mirror[0] else -1 - self.symmetries[key][1,1,2] = 1 if mirror[1] else -1 - - def init_domain(self,origin,spacing,periodicity,bounds): - '''Sets up domain information for all processors''' - # TBD: some assertions - self.origin = self.comm.bcast(origin,root=0) - self.spacing = self.comm.bcast(spacing,root=0) - self.periodicity = self.comm.bcast(periodicity,root=0) - self.bounds = self.comm.bcast(bounds,root=0) - # Some shortcuts - self.xperiodic,self.yperiodic,self.zperiodic = self.periodicity - return - - def init_procgrid(self,proc,chunks_per_proc): - # Note: requires nghx, xperiodic to be set - '''Read input processor grid, compute processor grid for workers''' - import numpy as np - self.chunks_per_proc = self.comm.bcast(chunks_per_proc,root=0) - self.nxcpp = chunks_per_proc[0] - self.nycpp = chunks_per_proc[1] - self.nzcpp = chunks_per_proc[2] - if self.rank==0: - # Assert proc and add it to class - assert all(k in proc for k in ('u','v','w','p','s')), "'proc' must be a dictionary with "\ - "keys 'u','v','w','p','s'" - for key in proc: - assert len(proc[key])==6, "Entries of 'proc' must have length of 6." - proc_grid_ext = proc - # Initialize chunks per processor - # Verify that this processor grid can be redistributed onto the requested processor layout - nxp_ext = len(proc_grid_ext['u'][0]) - nyp_ext = len(proc_grid_ext['u'][2]) - nzp_ext = len(proc_grid_ext['u'][4]) - nproc_ext = nxp_ext*nyp_ext*nzp_ext - # - assert nxp_ext%self.nxcpp==0, "Number of processors must be divisible by the number "\ - "of processors per process. (nxp_ext={}, nxcpp={})".format( - nxp_ext,self.nxcpp) - assert nyp_ext%self.nycpp==0, "Number of processors must be divisible by the number "\ - "of processors per process. (nyp_ext={}, nycpp={})".format( - nyp_ext,self.nycpp) - assert nzp_ext%self.nzcpp==0, "Number of processors must be divisible by the number "\ - "of processors per process. (nzp_ext={}, nzcpp={})".format( - nzp_ext,self.nzcpp) - # Determine the new processor layout and verify total number of MPI processes - nxp = nxp_ext//self.nxcpp - nyp = nyp_ext//self.nycpp - nzp = nzp_ext//self.nzcpp - nproc = nxp*nyp*nzp - assert nproc==self.comm.Get_size(), "Number of MPI processes does not match the requested "\ - "processor layout. (MPI procs: {}, required procs: {})".format( - self.comm.Get_size(),nproc) - # Construct internal processor grid - proc_grid = {} - for key in proc_grid_ext: - proc_grid[key] = [None]*6 - proc_grid[key][0] = proc_grid_ext[key][0][0::self.nxcpp] - proc_grid[key][1] = proc_grid_ext[key][1][self.nxcpp-1::self.nxcpp] - proc_grid[key][2] = proc_grid_ext[key][2][0::self.nycpp] - proc_grid[key][3] = proc_grid_ext[key][3][self.nycpp-1::self.nycpp] - proc_grid[key][4] = proc_grid_ext[key][4][0::self.nzcpp] - proc_grid[key][5] = proc_grid_ext[key][5][self.nzcpp-1::self.nzcpp] - else: - proc_grid_ext = None - proc_grid = None - nxp_ext,nyp_ext,nzp_ext,nproc_ext = None,None,None,None - nxp,nyp,nzp,nproc = None,None,None,None - # Broadcast the data - self.proc_grid_ext = self.comm.bcast(proc_grid_ext,root=0) - self.proc_grid = self.comm.bcast(proc_grid,root=0) - self.nxp_ext,self.nyp_ext,\ - self.nzp_ext,self.nproc_ext = self.comm.bcast((nxp_ext,nyp_ext,nzp_ext,nproc_ext),root=0) - self.nxp,self.nyp,\ - self.nzp,self.nproc = self.comm.bcast((nxp,nyp,nzp,nproc),root=0) - # Get position in processor grid - self.ip,self.jp,self.kp = self.position_from_rank(self.rank,external=False) - # Determine local grid indices and size - for key in self.proc_grid: - nxl = self.proc_grid[key][1][self.ip]-self.proc_grid[key][0][self.ip]+1 - nyl = self.proc_grid[key][3][self.jp]-self.proc_grid[key][2][self.jp]+1 - nzl = self.proc_grid[key][5][self.kp]-self.proc_grid[key][4][self.kp]+1 - # Verify that local grid size is not smaller than ghost cell size - assert (nxl>=self.nghx and nyl>=self.nghy and nzl>=self.nghz),\ - "Local grid size must be greater than number "\ - "of ghost cells in each direction!" - # Initialize neighbor array - nghbr = np.empty((3,3,3),dtype='int') - # wrap-around x - ipl = self.ip-1 - if ipl<0: - if self.xperiodic: ipl = self.nxp-1 - else: ipl = -1 - ipr = self.ip+1 - if ipr>self.nxp-1: - if self.xperiodic: ipr = 0 - else: ipr = -1 - inghbr = (ipl,self.ip,ipr) - # wrap-around y - jpl = self.jp-1 - if jpl<0: - if self.yperiodic: jpl = self.nyp-1 - else: jpl = -1 - jpr = self.jp+1 - if jpr>self.nyp-1: - if self.yperiodic: jpr = 0 - else: jpr = -1 - jnghbr = (jpl,self.jp,jpr) - # wrap-around z - kpl = self.kp-1 - if kpl<0: - if self.zperiodic: kpl = self.nzp-1 - else: kpl = -1 - kpr = self.kp+1 - if kpr>self.nzp-1: - if self.zperiodic: kpr = 0 - else: kpr = -1 - knghbr = (kpl,self.kp,kpr) - # Construct array of neighbors - for ip in range(3): - for jp in range(3): - for kp in range(3): - # Assign rank to neighbor array - if inghbr[ip]<0 or jnghbr[jp]<0 or knghbr[kp]<0: - nghbr[ip,jp,kp] = -1 - else: - nghbr[ip,jp,kp] = self.rank_from_position(inghbr[ip],jnghbr[jp],knghbr[kp],external=False) - # Save neighbors as class variable - self.nghbr = nghbr - - def load_field(self,key,io_limit=None,barrier=False): + def load_field(self,key,io_limit=None,barrier=False,dtype=np.float64): '''Loads the required chunks from file''' from .field import Field3d import numpy as np @@ -226,7 +118,7 @@ class PPP: data = np.empty( (nxl+2*self.nghx, nyl+2*self.nghy, - nzl+2*self.nghz),dtype=self.precision) + nzl+2*self.nghz),dtype=dtype) # Compute origin of subfield origin = (self.origin[key][0]+(ib-1-self.nghx)*self.spacing[0], self.origin[key][1]+(jb-1-self.nghy)*self.spacing[1], @@ -257,34 +149,32 @@ class PPP: # Exchange ghost cells self.exchange_ghost_cells(key) # Initialize symmetries and impose BC - self.init_symmetries(key) + self.init_symmetry(key) self.impose_boundary_conditions(key) # Syncronize processes if requested if barrier: self.comm.Barrier() - def differentiate(self,key,axis,key_out=None): + def differentiate(self,key,axis,key_out=None,on_pressure_grid=False): assert axis<3, "'axis' must be one of 0,1,2." if key_out is None: - key_out = key+('x','y','z')[axis] - origin = list(self.origin) + key_out = 'd'+key+('/dx','/dy','/dz')[axis] + origin = list(self.origin[key]) shifting_state = self.shifting_state(key,axis=axis) if shifting_state==-1: - padding = 'after' + shift_origin = 'after' origin[axis] += 0.5*self.spacing[axis] elif shifting_state==0: - padding = 'before' + shift_origin = 'before' origin[axis] -= 0.5*self.spacing[axis] elif shifting_state==1: - padding = 'before' + shift_origin = 'before' origin[axis] -= 0.5*self.spacing[axis] else: raise ValueError("Invalid shifting state.") - self.field[key_out] = self.field[key].derivative(axis,padding=padding) - self.origin[key_out] = tuple(origin) - self.spacing[key_out] = self.spacing[key].copy() - self.symmetries[key_out] = self.symmetries[key].copy() - self.proc_grid[key_out] = self.proc_grid[key].copy() - self.proc_grid_ext[key_out] = self.proc_grid_ext[key].copy() + self.field[key_out] = self.field[key].derivative(axis,only_keep_interior=False,shift_origin=shift_origin) + self.origin[key_out] = tuple(origin) + self.symmetries[key_out] = self.symmetries[key].copy() + self.proc_grid[key_out] = self.proc_grid[key].copy() if axis==0: self.symmetries[key_out][0,1,1] = -self.symmetries[key_out][0,1,1] self.symmetries[key_out][2,1,1] = -self.symmetries[key_out][2,1,1] @@ -297,15 +187,90 @@ class PPP: # Exchange ghost cells and set boundary conditions self.exchange_ghost_cells(key_out) self.impose_boundary_conditions(key_out) + # Shift to pressure grid if desired + if on_pressure_grid: + self.shift_to_pressure_grid(key_out,key_out=key_out) + + def allocate(self,key_grid,fill=None,pseudo=False,dtype=np.float64): + from .field import Field3d + ib = self.proc_grid[key_grid][0][self.ip] + ie = self.proc_grid[key_grid][1][self.ip] + jb = self.proc_grid[key_grid][2][self.jp] + je = self.proc_grid[key_grid][3][self.jp] + kb = self.proc_grid[key_grid][4][self.kp] + ke = self.proc_grid[key_grid][5][self.kp] + origin = (self.origin[key_grid][0]+(ib-1-self.nghx)*self.spacing[0], + self.origin[key_grid][1]+(jb-1-self.nghy)*self.spacing[1], + self.origin[key_grid][2]+(kb-1-self.nghz)*self.spacing[2]) + dim = (ie-ib+1+2*self.nghx, + je-jb+1+2*self.nghy, + ke-kb+1+2*self.nghz) + return Field3d.allocate(dim,origin,self.spacing,fill=fill,pseudo=pseudo,dtype=dtype) + + def qcriterion(self,key_out=None,from_pressure=False,keep_derivatives=False): + '''Computes Q-criterion''' + from .field import Field3d + if key_out is None: + key_out = 'Q' + if from_pressure: + assert 'p' in self.field, "Pressure field is not loaded." + self.copy('p',key_out,skip_data=True) + 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.copy('p',key_out,skip_data=True,skip_symmetries=True) + self.init_symmetry(key_out) + self.field[key_out] = self.allocate('p',fill=0.0,dtype=self.field['u'].dtype()) + keyx = ('x','y','z') + keyu = ('u','v','w') + for ii in range(3): + key_tmp = 'd'+keyu[ii]+'/d'+keyx[ii] + parprint(key_tmp) + self.differentiate(keyu[ii],axis=ii,key_out=key_tmp,on_pressure_grid=True) + self.field[key_out] += self.field[key_tmp]**2 + if not keep_derivatives: + self.delete(key_tmp) + # + iip1 = (ii+1)%3 + key_tmp1 = 'd'+keyu[ii]+'/d'+keyx[iip1] + key_tmp2 = 'd'+keyu[iip1]+'/d'+keyx[ii] + parprint(key_tmp1,key_tmp2) + self.differentiate(keyu[ii],axis=iip1,key_out=key_tmp1,on_pressure_grid=True) + self.differentiate(keyu[iip1],axis=ii,key_out=key_tmp2,on_pressure_grid=True) + self.field[key_out] += 2.0*self.field[key_tmp1]*self.field[key_tmp2] + if not keep_derivatives: + self.delete(key_tmp1) + self.delete(key_tmp2) + self.field[key_out] *= -0.5 + # Exchange ghost cells and set boundary conditions + self.exchange_ghost_cells(key_out) + self.impose_boundary_conditions(key_out) + return + + #def vorticity(self,axis=None,key_out=None,keep_derivatives=False,on_pressure_grid=True): + + def multiply(self,key1,key2,key_out=None,on_pressure_grid=True): + if key_out is None: + key_out = key1 + if not self.field[key1].has_same_grid(self.field[key2]) or on_pressure_grid: + self.shift_to_pressure_grid(key1,key_out) + self.shift_to_pressure_grid(key2,'tmp') + self.field[key_out] *= self.field['tmp'] + self.delete('tmp') + else: + self.copy(key1,key_out) + self.field[key_out] *= self.field[key2] + return def gaussian_filter(self,key,sigma,truncate=4.0,key_out=None,iterate=False): '''Applies a gaussian filter to a field as in-place operation. Sigma is the std of the filter in terms of grid width.''' import numpy as np if key_out is None: key_out = key - else: - self.origin[key_out] = self.origin[key].copy() - self.spacing[key_out] = self.spacing[key].copy() + if key!=key_out: + self.copy(key,key_out,skip_data=True) # Compute radius of Gaussian filter radius = self.field[key].gaussian_filter_radius(sigma,truncate=truncate) if not iterate: @@ -324,23 +289,27 @@ class PPP: assert all([radius[ii]>0 if sigma[ii]>0.0 else True for ii in range(3)]),\ "Iterative procedure leads to invalid stencil radius: "\ "increase number of ghost cells. {}".format(radius) - print('Iterations: {}, stencil radius: {}'.format(niter,radius)) + parprint('Gaussian filter: iterations={}, stencil radius={}'.format(niter,radius)) for iiter in range(niter): + parprint('Iter #{:d}'.format(iiter)) # Filter field: if key_out is None, perform operation inplace self.field[key_out] = self.field[key].gaussian_filter(sigma, - truncate=truncate,border='constant',const_val=0.0) + truncate=truncate,only_keep_interior=False) # Exchange ghost cells and set boundary conditions self.exchange_ghost_cells(key_out) self.impose_boundary_conditions(key_out) # Iterate inplace from now on key = key_out - def broadcast(self,key,arg,operation): - '''Broadcasts an inplace operation involving a scalar or matrix on + def broadcast(self,key,operation,arg,key_out=None): + '''Broadcasts an operation involving a scalar or matrix on the entire grid. If 'arg' is a matrix, it must be three-dimensional and its axes must be singular or of length nx/ny/nz.''' import numpy as np import operator + # Broadcast argument + arg = self.comm.bcast(arg,root=0) + # Select operator if operation in ('add','+'): op = operator.iadd elif operation in ('subtract','sub','-'): @@ -358,7 +327,7 @@ class PPP: for axis in range(3): if arg.shape[axis]==1: continue - elif arg.shape[axis]==self.numpoints(key,axis=axis): + elif arg.shape[axis]==self.dim(key,axis=axis): pos = self.position_from_rank(self.rank,external=False)[axis] sl_arg[axis] = slice( self.proc_grid[key][2*axis][pos]-1, @@ -366,16 +335,24 @@ class PPP: else: raise ValueError("'arg' must either be singular or match global "\ "grid dimension. (axis={}: got {:d}, expected {:d}".format( - axis,arg.shape[axis],self.numpoints(key,axis=axis))) + axis,arg.shape[axis],self.dim(key,axis=axis))) # Only operate on interior and communcate ghosts later sl_int = tuple(slice(self.num_ghost[ii],-self.num_ghost[ii]) for ii in range(3)) sl_arg = tuple(sl_arg) - op(self.field[key].data[sl_int],arg[sl_arg]) + if key_out is None: + op(self.field[key].data[sl_int],arg[sl_arg]) + else: + self.copy(key,key_out) + op(self.field[key_out].data[sl_int],arg[sl_arg]) # Exchange ghost cells and set boundary conditions self.exchange_ghost_cells(key) self.impose_boundary_conditions(key) elif isinstance(arg,(int,float)): - op(self.field[key].data,arg) + if key_out is None: + op(self.field[key].data,arg) + else: + self.copy(key,key_out) + op(self.field[key_out].data,arg) return def integrate(self,key,integrate_axis,ufunc=None,ignore_nan=False,average=False): @@ -396,7 +373,7 @@ class PPP: weights_local = np.array(weights_local) # Simple implementation: all processor communicate directly to root if self.rank==0: - dim_final = tuple(1 if integrate_axis[axis] else self.numpoints(key,axis=axis) for axis in range(3)) + dim_final = tuple(1 if integrate_axis[axis] else self.dim(key,axis=axis) for axis in range(3)) # Allocate a nice spot of memory integ = np.zeros(dim_final,dtype=integ_local.dtype) if average: @@ -432,47 +409,169 @@ class PPP: self.comm.send(weights_local,dest=0,tag=2) return None + 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 + # If pressure is loaded, use field information, else create a + # fake Field3d object so that we can use its methods. + field_ref = self.allocate('p',pseudo=True) + # Compute the relative shift and shift it + rel_shift = self.field[key].relative_shift(field_ref) + self.field[key_out] = self.field[key].shift_origin(rel_shift) + # Match number of gridpoints + sl = 3*[slice(None)] + for axis in range(3): + if not self.field[key_out].has_same_dim(field_ref,axis=axis): + sl[axis] = slice(0,field_ref.dim(axis=axis)) + self.field[key_out].data = self.field[key_out].data[sl] + # Copy metadata of pressure grid + self.copy('p',key_out,skip_data=True,skip_symmetries=True) + self.symmetries[key_out] = self.symmetries[key].copy() + # Exchange ghost cells and set boundary conditions + self.exchange_ghost_cells(key_out) + self.impose_boundary_conditions(key_out) + + def init_symmetry(self,key,mirror=(True,True)): + '''Sets the symmetries for ghost cells behind the wall''' + # Example: wall in y + # No-slip boundary (no mirror) free-slip boundary (mirror) + # u -> -u u -> u + # v -> -v v -> -v + # w -> -w w -> w + # p -> p p -> p + import numpy as np + self.symmetries[key] = np.zeros((3,3,3),dtype='i') + if not self.xperiodic: + if key=='u': + self.symmetries[key][0,1,1] = -1 + self.symmetries[key][2,1,1] = -1 + else: + self.symmetries[key][0,1,1] = 1 if mirror[0] else -1 + self.symmetries[key][2,1,1] = 1 if mirror[1] else -1 + if not self.yperiodic: + if key=='v': + self.symmetries[key][1,0,1] = -1 + self.symmetries[key][1,2,1] = -1 + else: + self.symmetries[key][1,0,1] = 1 if mirror[0] else -1 + self.symmetries[key][1,2,1] = 1 if mirror[1] else -1 + if not self.zperiodic: + if key=='w': + self.symmetries[key][1,1,0] = -1 + self.symmetries[key][1,1,2] = -1 + else: + self.symmetries[key][1,1,0] = 1 if mirror[0] else -1 + self.symmetries[key][1,1,2] = 1 if mirror[1] else -1 + + def copy(self,key,key_out,skip_data=False,skip_symmetries=False): + '''Copies a field to a new key.''' + if key==key_out: + return + self.origin[key_out] = tuple(self.origin[key]) + self.proc_grid[key_out] = self.proc_grid[key].copy() + if not skip_symmetries: + self.symmetries[key_out] = self.symmetries[key].copy() + if not skip_data: + self.field[key_out] = self.field[key].copy() + return + + def delete(self,key): + '''Removes a field from memory.''' + assert key in self.field, "Field {} cannot be deleted, because it "\ + "does not exist.".format(key) + import gc + del self.field[key] + del self.symmetries[key] + if key not in ('u','v','w','p','s'): + del self.origin[key] + del self.proc_grid[key] + 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) + payload = ( + self.num_ghost, + self.chunks_per_proc, + self.origin, + self.spacing, + self.periodicity, + self.bounds, + self.proc_grid, + self.nxp, + self.nyp, + self.nzp, + self.nghbr, + self.field, + self.symmetries) + with open(fout,"wb") as f: + pickle.dump(payload,f) + + def to_vtk(self,key,no_ghost=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() + def vtk_contour(self,key,val): '''Compute isocontour for chunks.''' - if any([self.num_ghost[ii]>1 for ii in range(3)]): - idx_origin = tuple(self.num_ghost[ii]-1 for ii in range(3)) - numpoints = self.chunk_size(axis=None) - return self.field[key].extract_subfield( - idx_origin,numpoints).vtk_contour(val) - else: - return self.field[key].vtk_contour(val) - return + return self._subfield_for_vtk(key,False).vtk_contour(val) def vtk_slice(self,key,normal,origin): '''Extracts a plane from field.''' - if any([self.num_ghost[ii]>1 for ii in range(3)]): - idx_origin = tuple(self.num_ghost[ii]-1 for ii in range(3)) - numpoints = tuple(self.field[key].numpoints[ii]-2*(self.num_ghost[ii]-1) - for ii in range(3)) - return self.field[key].extract_subfield( - idx_origin,numpoints).vtk_slice(normal,origin) - else: - return self.field[key].vtk_slice(normal,origin) - return + return self._subfield_for_vtk(key,False).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 rank_from_position(self,ip,jp,kp,external=False): if external: nyp,nzp = self.nyp_ext,self.nzp_ext else: nyp,nzp = self.nyp,self.nzp - return ip*nyp*nzp+jp*nzp+kp + return self._rank_from_position(ip,jp,kp,nyp,nzp) def position_from_rank(self,rank,external=False): if external: nyp,nzp = self.nyp_ext,self.nzp_ext else: nyp,nzp = self.nyp,self.nzp + return self._position_from_rank(rank,nyp,nzp) + + @staticmethod + def _rank_from_position(ip,jp,kp,nyp,nzp): + return ip*nyp*nzp+jp*nzp+kp + + @staticmethod + def _position_from_rank(rank,nyp,nzp): ip = rank//(nyp*nzp) jp = (rank//nzp)%nyp kp = rank%nzp return (ip,jp,kp) def shifting_state(self,key,axis=None): + '''Returns the relative shift of first global grid point 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." @@ -483,13 +582,13 @@ class PPP: if axis is None: return tuple(self.chunk_size(key,axis=ii) for ii in range(3)) assert axis<3, "'axis' must be one of 0,1,2." - return self.field[key].numpoints[axis]-2*self.num_ghost[axis] + return self.field[key].dim(axis=axis)-2*self.num_ghost[axis] - def numpoints(self,key,axis=None): + def dim(self,key,axis=None): '''Returns the total number of gridpoints across all processors without ghost cells.''' if axis is None: - return tuple(self.numpoints(key,axis=ii) for ii in range(3)) + return tuple(self.dim(key,axis=ii) for ii in range(3)) assert axis<3, "'axis' must be one of 0,1,2." return self.proc_grid[key][2*axis+1][-1] @@ -600,7 +699,7 @@ class PPP: recvbuf = np.zeros( (ii_dst.stop-ii_dst.start, jj_dst.stop-jj_dst.start, - kk_dst.stop-kk_dst.start),dtype=self.precision) + kk_dst.stop-kk_dst.start),dtype=self.field[key].data.dtype) reqrecv = self.comm.Irecv(recvbuf,source=rank_src,tag=tag) # [recv] wait for data to be received if reqrecv is not None: @@ -625,7 +724,7 @@ class PPP: for axis in range(3): if positionBd[axis]==-1: # lower boundary # index of first point within the domain - idx = self.field[key].nearest_gridpoint(self.bounds[2*axis],axis=axis,lower=True)+1 + idx = self.field[key].nearest_gridpoint(self.bounds[2*axis],axis=axis,lower=True)+1 # distance first point to wall: should be either 0 or dx/2 dist = np.abs(self.field[key].coordinate(idx,axis=axis)-self.bounds[2*axis]) if dist>=0.25*self.spacing[axis]: # no point on boundary @@ -637,19 +736,123 @@ class PPP: break elif positionBd[axis]==1: # upper boundary # index of last point within the domain - idx = self.field[key].nearest_gridpoint(self.bounds[2*axis+1],axis=axis,lower=True) + idx = self.field[key].nearest_gridpoint(self.bounds[2*axis+1],axis=axis,lower=True) # distance last point to wall: should be either 0 or -dx/2 dist = np.abs(self.field[key].coordinate(idx,axis=axis)-self.bounds[2*axis+1]) if dist>=0.25*self.spacing[axis]: # no point on boundary - sl_dst[axis] = slice(idx+1,self.field[key].numpoints[axis],1) - sl_src[axis] = slice(idx,2*idx+1-self.field[key].numpoints[axis],-1) + sl_dst[axis] = slice(idx+1,self.field[key].dim(axis=axis),1) + sl_src[axis] = slice(idx,2*idx+1-self.field[key].dim(axis=axis),-1) else: # point on boundary - sl_dst[axis] = slice(idx+1,self.field[key].numpoints[axis],1) - sl_src[axis] = slice(idx-1,2*idx-self.field[key].numpoints[axis],-1) + sl_dst[axis] = slice(idx+1,self.field[key].dim(axis=axis),1) + sl_src[axis] = slice(idx-1,2*idx-self.field[key].dim(axis=axis),-1) break + # print(sl_src,sl_dst,self.field[key].dim()) self.field[key].data[tuple(sl_dst)] = self.symmetries[key][inghbr,jnghbr,knghbr]*\ self.field[key].data[tuple(sl_src)] + @staticmethod + def _construct_procgrid(comm,proc,chunks_per_proc,num_ghost,periodicity): + import numpy as np + rank = comm.Get_rank() + nxcpp,nycpp,nzcpp = chunks_per_proc + nghx,nghy,nghz = num_ghost + xperiodic,yperiodic,zperiodic = periodicity + # Assert proc + assert all(k in proc for k in ('u','v','w','p','s')), "'proc' must be a dictionary with "\ + "keys 'u','v','w','p','s'" + for key in proc: + assert len(proc[key])==6, "Entries of 'proc' must have length of 6." + proc_grid_ext = proc + # Initialize chunks per processor + # Verify that this processor grid can be redistributed onto the requested processor layout + nxp_ext = len(proc_grid_ext['u'][0]) + nyp_ext = len(proc_grid_ext['u'][2]) + nzp_ext = len(proc_grid_ext['u'][4]) + nproc_ext = nxp_ext*nyp_ext*nzp_ext + # + assert nxp_ext%nxcpp==0, "Number of processors must be divisible by the number "\ + "of processors per process. (nxp_ext={}, nxcpp={})".format( + nxp_ext,nxcpp) + assert nyp_ext%nycpp==0, "Number of processors must be divisible by the number "\ + "of processors per process. (nyp_ext={}, nycpp={})".format( + nyp_ext,nycpp) + assert nzp_ext%nzcpp==0, "Number of processors must be divisible by the number "\ + "of processors per process. (nzp_ext={}, nzcpp={})".format( + nzp_ext,nzcpp) + # Determine the new processor layout and verify total number of MPI processes + nxp = nxp_ext//nxcpp + nyp = nyp_ext//nycpp + nzp = nzp_ext//nzcpp + nproc = nxp*nyp*nzp + assert nproc==comm.Get_size(), "Number of MPI processes does not match the requested "\ + "processor layout. (MPI procs: {}, required procs: {})".format( + comm.Get_size(),nproc) + # Construct internal processor grid + proc_grid = {} + for key in proc_grid_ext: + proc_grid[key] = [None]*6 + proc_grid[key][0] = proc_grid_ext[key][0][0::nxcpp] + proc_grid[key][1] = proc_grid_ext[key][1][nxcpp-1::nxcpp] + proc_grid[key][2] = proc_grid_ext[key][2][0::nycpp] + proc_grid[key][3] = proc_grid_ext[key][3][nycpp-1::nycpp] + proc_grid[key][4] = proc_grid_ext[key][4][0::nzcpp] + proc_grid[key][5] = proc_grid_ext[key][5][nzcpp-1::nzcpp] + # Get position in processor grid + ip,jp,kp = PPP._position_from_rank(rank,nyp,nzp) + # Determine local grid indices and size + for key in proc_grid: + nxl = proc_grid[key][1][ip]-proc_grid[key][0][ip]+1 + nyl = proc_grid[key][3][jp]-proc_grid[key][2][jp]+1 + nzl = proc_grid[key][5][kp]-proc_grid[key][4][kp]+1 + # Verify that local grid size is not smaller than ghost cell size + assert (nxl>=nghx and nyl>=nghy and nzl>=nghz),\ + "Local grid size must be greater than number "\ + "of ghost cells in each direction!" + # Initialize neighbor array + nghbr = np.empty((3,3,3),dtype='int') + # wrap-around x + ipl = ip-1 + if ipl<0: + if xperiodic: ipl = nxp-1 + else: ipl = -1 + ipr = ip+1 + if ipr>nxp-1: + if xperiodic: ipr = 0 + else: ipr = -1 + inghbr = (ipl,ip,ipr) + # wrap-around y + jpl = jp-1 + if jpl<0: + if yperiodic: jpl = nyp-1 + else: jpl = -1 + jpr = jp+1 + if jpr>nyp-1: + if yperiodic: jpr = 0 + else: jpr = -1 + jnghbr = (jpl,jp,jpr) + # wrap-around z + kpl = kp-1 + if kpl<0: + if zperiodic: kpl = nzp-1 + else: kpl = -1 + kpr = kp+1 + if kpr>nzp-1: + if zperiodic: kpr = 0 + else: kpr = -1 + knghbr = (kpl,kp,kpr) + # Construct array of neighbors + for ip in range(3): + for jp in range(3): + for kp in range(3): + # Assign rank to neighbor array + if inghbr[ip]<0 or jnghbr[jp]<0 or knghbr[kp]<0: + nghbr[ip,jp,kp] = -1 + else: + nghbr[ip,jp,kp] = PPP._rank_from_position(inghbr[ip],jnghbr[jp],knghbr[kp],nyp,nzp) + return (proc_grid,nxp,nyp,nzp, + proc_grid_ext,nxp_ext,nyp_ext,nzp_ext, + nghbr) + def _baton_wait(self,batch_size,tag=420): '''Block execution until an empty message from rank-batch_wait is received (issued by _baton_pass)''' diff --git a/particle.py b/particle.py index a23f427..d7bd0f3 100644 --- a/particle.py +++ b/particle.py @@ -1,37 +1,36 @@ +import numpy as np class Particles: def __init__(self,num,time,attr,period): '''Initialize a Particles object. num number of particles time time value - attr dictionary of numpy arrays with length equal to num + attr dictionary of np arrays with length equal to num period the domain periods ''' - import numpy assert 'id' in attr, "Need attribute 'id' to initialize Particles" assert 'x' in attr, "Need attribute 'x' to initialize Particles" assert 'y' in attr, "Need attribute 'y' to initialize Particles" assert 'z' in attr, "Need attribute 'z' to initialize Particles" assert isinstance(period,(tuple,list)) and len(period)==3, "'period' must be tuple/list with length 3" self.num = num - if isinstance(time,(tuple,list,numpy.ndarray)): + if isinstance(time,(tuple,list,np.ndarray)): assert len(time)==1, "'time' must be a scalar value." time = time[0] self.time = time self.attr = {} sortidx = attr['id'].argsort() for key in attr: - assert attr[key].ndim==1, "Attributes must be a one-dimensional numpy array." + assert attr[key].ndim==1, "Attributes must be a one-dimensional np array." assert attr[key].shape[0]==num, "Attributes must be of length 'num'." if key=='id': self.add_attribute(key,attr[key][sortidx].astype('int')) else: self.add_attribute(key,attr[key][sortidx]) self.period = tuple(period) - self.frame_velocity = numpy.zeros(3) + self.frame_velocity = np.zeros(3) return @classmethod def from_array(cls,pp,col,time,period,select_col=None): - import numpy assert 'id' in col, "Need column 'id' to initialize Particles" assert 'x' in col, "Need column 'x' to initialize Particles" assert 'y' in col, "Need column 'y' to initialize Particles" @@ -39,7 +38,7 @@ class Particles: assert pp.ndim==2 or pp.shape[2]==1, "'pp' cannot be a time series." num = pp.shape[1] if pp.ndim==2: - pp = numpy.expand_dims(pp,axis=2) + pp = np.expand_dims(pp,axis=2) if select_col is not None: for scol in select_col: assert scol in col, "Selected column '{}' not found in 'col'.".format(scol) @@ -78,21 +77,19 @@ class Particles: attr[key] = self.attr[key].copy() return Particles(self.num,self.time,attr,self.period) def add_attribute(self,key,val): - import numpy - if isinstance(val,(tuple,list,numpy.ndarray)): + if isinstance(val,(tuple,list,np.ndarray)): assert len(val)==self.num and val.ndim==1, "Invalid 'val'." - self.attr[key] = numpy.array(val) + self.attr[key] = np.array(val) else: - self.attr[key] = numpy.full(self.num,val) + self.attr[key] = np.full(self.num,val) return def del_attribute(self,key): del self.attr[key] def get_attribute(self,key): return self.attr[key] def get_position(self,axis=None): - import numpy if axis is None: - return numpy.vstack([self.attr[key].copy() for key in ('x','y','z')]) + return np.vstack([self.attr[key].copy() for key in ('x','y','z')]) else: assert axis<3, "'axis' must be smaller than 3." key = ('x','y','z')[axis] @@ -113,7 +110,8 @@ class Particles: valdiff = val-self.frame_velocity[axis] self.translate(-self.time*valdiff,axis=axis) key = ('u','v','w')[axis] - self.attr[key] -= valdiff + if self.has_attribute(key): + self.attr[key] -= valdiff self.frame_velocity[axis] = val return def enforce_periodicity(self,axis=None): @@ -127,7 +125,6 @@ class Particles: return def to_vtk(self,deep=False): import pyvista as pv - import numpy as np position = np.vstack([self.attr[key] for key in ('x','y','z')]).transpose() mesh = pv.PolyData(position,deep=deep) for key in self.attr: @@ -143,12 +140,11 @@ class Particles: class Trajectories: def __init__(self,part,unraveled=False,copy_particles=False): - import numpy self.numtime = 1 self.numpart = part.num self.copy_particles = copy_particles # Attributes are first stored as a list of view to instances of - # Particles. Later they will be concatenated into numpy array + # Particles. Later they will be concatenated into np array # and the attributes of Particles will be views. self.attr = {} self.part = [] @@ -161,7 +157,8 @@ class Trajectories: self.attr[key] = [self.part[0].attr[key].view()] self.period = self.part[0].period self.unraveled = unraveled - self.frame_velocity = numpy.zeros(3) + self.frame_velocity = np.zeros(3) + def __str__(self): str = 'Trajectories with\n' str+= ' time steps: {:d}\n'.format(self.numtime) @@ -173,6 +170,7 @@ class Trajectories: str+= ' frame velocity: {}\n'.format(self.frame_velocity) str+= ' unraveled: {}'.format(self.unraveled) return str + def __iadd__(self,other): if isinstance(other,Trajectories): self.add_trajectories(other) @@ -181,6 +179,7 @@ class Trajectories: else: raise TypeError('Cannot add type {}'.format(type(other))) return self + def __getitem__(self,val): assert isinstance(val,tuple) and len(val)==2, "Trajectories must be indexed by [part,time]." sl = [] @@ -194,6 +193,7 @@ class Trajectories: else: raise TypeError("Trajectories can only be sliced by slice objects or integers.") return self._slice(slice_part=sl[0],slice_time=sl[1]) + @classmethod def from_mat(cls,file,unraveled=False,select_col=None): from .helper import load_mat @@ -206,6 +206,7 @@ class Trajectories: for key in col: col[key]-=1 return cls.from_array(pp,col,time,period,unraveled=unraveled,select_col=select_col) + @classmethod def from_array(cls,pp,col,time,period,unraveled=False,select_col=None): ntime = len(time) @@ -213,6 +214,7 @@ class Trajectories: for ii in range(1,ntime): traj += Particles.from_array(pp[:,:,ii],col,time[ii],period,select_col=select_col) return traj + def _slice(self,slice_part=slice(None),slice_time=slice(None)): assert isinstance(slice_part,slice), "'slice_part' must be a slice object." assert isinstance(slice_time,slice), "'slice_time' must be a slice object." @@ -225,8 +227,8 @@ class Trajectories: traj += Trajectories(part[slice_part]) traj.frame_velocity = self.frame_velocity return traj + def add_particles(self,part): - import numpy # Verify part assert part.time>self.part[-1].time, "Time steps must be added in monotonically increasing order." assert part.num==self.numpart, "Number of particles is different from previous time steps." @@ -246,6 +248,7 @@ class Trajectories: self.attr[key].append(self.part[-1].attr[key].view()) self.numtime += 1 return + def add_trajectories(self,traj): # Verify traj assert traj.part[0].time>self.part[-1].time, "Time steps must be added in monotonically increasing order." @@ -267,26 +270,26 @@ class Trajectories: self.attr[key].append(self.part[-traj.numtime+ii].attr[key].view()) self.numtime += traj.numtime return + def get_time(self): - import numpy - time = numpy.zeros(self.numtime) + time = np.zeros(self.numtime) for ii,part in enumerate(self.part): time[ii] = part.time return time + def get_trajectories(self,slice_part=slice(None),slice_time=slice(None)): - '''Get (x,y,z) trajectories in numpy array of dimension (3,npart,ntime).''' - import numpy + '''Get (x,y,z) trajectories in np array of dimension (3,npart,ntime).''' assert isinstance(slice_part,slice), "'slice_part' must be a slice." assert isinstance(slice_time,slice), "'slice_time' must be a slice." self._make_data_array() - return numpy.stack([ + return np.stack([ self.attr['x'][slice_part,slice_time], self.attr['y'][slice_part,slice_time], self.attr['z'][slice_part,slice_time]],axis=0) + def get_trajectories_segmented(self,slice_part=slice(None),slice_time=slice(None),restore_ravel_state=True): '''Get (x,y,z) segments of trajectories as a tuple (len: npart) of - lists (len: nsegments) of numpy arrays (shape: 3 X ntime of the segment)''' - import numpy + lists (len: nsegments) of np arrays (shape: 3 X ntime of the segment)''' assert isinstance(slice_part,slice), "'slice_part' must be a slice." assert isinstance(slice_time,slice), "'slice_time' must be a slice." was_unraveled = self.unraveled @@ -296,21 +299,21 @@ class Trajectories: ntime = xyzpath.shape[2] # Initialize output and helper arrays out = tuple([] for ii in range(npart)) - lastJump = numpy.zeros(npart,dtype=numpy.uint) - lastPeriod = numpy.zeros((3,npart),dtype=numpy.int) + lastJump = np.zeros(npart,dtype=np.uint) + lastPeriod = np.zeros((3,npart),dtype=np.int) for axis in range(3): if self.period[axis] is not None: - lastPeriod[axis,:] = numpy.floor_divide(xyzpath[axis,:,0],self.period[axis]) + lastPeriod[axis,:] = np.floor_divide(xyzpath[axis,:,0],self.period[axis]) # Construct output tuple for itime in range(1,ntime+1): - thisPeriod = numpy.zeros((3,npart),dtype=int) + thisPeriod = np.zeros((3,npart),dtype=int) if itime==ntime: - hasJumped = numpy.ones(npart,dtype=bool) + hasJumped = np.ones(npart,dtype=bool) else: for axis in range(3): if self.period[axis] is not None: - thisPeriod[axis,:] = numpy.floor_divide(xyzpath[axis,:,itime],self.period[axis]) - hasJumped = numpy.any(thisPeriod!=lastPeriod,axis=0) + thisPeriod[axis,:] = np.floor_divide(xyzpath[axis,:,itime],self.period[axis]) + hasJumped = np.any(thisPeriod!=lastPeriod,axis=0) for ipart in range(npart): if hasJumped[ipart]: sl = slice(lastJump[ipart],itime) @@ -325,8 +328,20 @@ class Trajectories: if restore_ravel_state and not was_unraveled: self.enforce_periodicity() return out + def slice_time(self,time_beg=None,time_end=None): - return + time = self.get_time() + if time_beg is not None: + idx_beg = np.searchsorted(time,time_beg,side='left') + else: + idx_beg = None + if time_end is not None: + idx_end = np.searchsorted(time,time_end,side='left') + else: + idx_end = None + sl = slice(idx_beg,idx_end+1) + return self._slice(slice_time=sl) + def set_frame_velocity(self,val,axis=0): assert axis<3, "'axis' must be smaller than 3." for part in self.part: @@ -335,8 +350,8 @@ class Trajectories: if not self.unraveled: self.enforce_periodicity() return + def unravel(self): - import numpy if self.unraveled: return #raise NotImplementedError('Implemented but untested!') self._make_data_array() @@ -345,48 +360,56 @@ class Trajectories: key = ('x','y','z')[axis] #posdiff = (self.part[ii].attr[key]-self.part[ii-1].attr[key]).squeeze() posdiff = self.attr[key][:,1:]-self.attr[key][:,0:-1] - coeff = -numpy.sign(posdiff)*(numpy.abs(posdiff)>0.5*self.period[axis]) - coeff = numpy.cumsum(coeff,axis=1) + coeff = -np.sign(posdiff)*(np.abs(posdiff)>0.5*self.period[axis]) + coeff = np.cumsum(coeff,axis=1) self.attr[key][:,1:] += coeff*self.period[axis] self.unraveled = True return + def enforce_periodicity(self): for part in self.part: part.enforce_periodicity(axis=None) self.unraveled = False return + + def to_vtk(self): + import pyvista as pv + mesh = pv.PolyData() + for part in self.get_trajectories_segmented(): + for seg in part: + mesh += pv.helpers.lines_from_points(seg.transpose()) + return mesh + def _make_data_array(self): - '''Transforms self.attr[key] to a numpy array of dimension npart X ntime + '''Transforms self.attr[key] to a np array of dimension npart X ntime and updates self.part[itime].attr[key] to views on self.attr[key][:,itime].''' - import numpy #print('DEBUG: _make_data_array') if not self._is_data_list: return #print([x.shape for x in self.attr['x']]) for key in self.attr: - self.attr[key] = numpy.stack(self.attr[key],axis=1) + self.attr[key] = np.stack(self.attr[key],axis=1) for ii in range(0,self.numtime): self.part[ii].attr[key] = self.attr[key][:,ii] #print(self.attr['x'].shape) self._is_data_list = False return + def _make_data_list(self): - '''Transforms self.attr[key] to a list of length ntime of numpy arrays of + '''Transforms self.attr[key] to a list of length ntime of np arrays of dimension npart and updates self.part[itime].attr[key] to views on self.attr[key][itime].''' - import numpy #print('DEBUG: _make_data_list') if self._is_data_list: return #print(self.attr['x'].shape) for key in self.attr: - self.attr[key] = [numpy.squeeze(x) for x in numpy.split(self.attr[key],self.numtime,axis=1)] + self.attr[key] = [np.squeeze(x) for x in np.split(self.attr[key],self.numtime,axis=1)] for ii in range(0,self.numtime): self.part[ii].attr[key] = self.attr[key][ii].view() #print([x.shape for x in self.attr['x']]) self._is_data_list = True return - def sort(pp,col): ncol,npart,ntime = pp.shape assert('id' in col)