446 lines
19 KiB
Python
446 lines
19 KiB
Python
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
|
|
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)):
|
|
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].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)
|
|
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"
|
|
assert 'z' in col, "Need column 'z' to initialize 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)
|
|
if select_col is not None:
|
|
for scol in select_col:
|
|
assert scol in col, "Selected column '{}' not found in 'col'.".format(scol)
|
|
attr = {}
|
|
for key in col:
|
|
if (select_col is None or key in select_col or key in ('id','x','y','z')):
|
|
attr[key] = pp[col[key],:,0].squeeze()
|
|
return cls(num,time,attr,period)
|
|
def __getitem__(self,val):
|
|
if isinstance(val,int):
|
|
lo,hi = val,val+1
|
|
hi = hi if hi!=0 else None
|
|
val = slice(lo,hi)
|
|
elif not isinstance(val,slice):
|
|
raise TypeError("Particles can only be sliced by slice objects or integers.")
|
|
return self._slice(val)
|
|
def __str__(self):
|
|
str = '{:d} particles with\n'.format(self.num)
|
|
str+= ' time: {}\n'.format(self.time)
|
|
str += ' attributes:'
|
|
for key in self.attr.keys():
|
|
str += ' {}'.format(key)
|
|
str += '\n'
|
|
str+= ' period: {}\n'.format(self.period)
|
|
str+= ' frame velocity: {}'.format(self.frame_velocity)
|
|
return str
|
|
def _slice(self,slice_part=slice(None)):
|
|
attr = {}
|
|
for key in self.attr:
|
|
attr[key] = self.attr[key][slice_part]
|
|
num = attr['id'].shape[0]
|
|
return Particles(num,self.time,attr,self.period)
|
|
def copy(self):
|
|
attr = {}
|
|
for key in self.attr:
|
|
attr[key] = self.attr[key].copy()
|
|
return Particles(self.num,self.time,attr,period)
|
|
def add_attribute(self,key,val):
|
|
import numpy
|
|
if isinstance(val,(tuple,list,numpy.ndarray)):
|
|
assert len(val)==self.num and val.ndim==1, "Invalid 'val'."
|
|
self.attr[key] = numpy.array(val)
|
|
else:
|
|
self.attr[key] = numpy.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')])
|
|
else:
|
|
assert axis<3, "'axis' must be smaller than 3."
|
|
key = ('x','y','z')[axis]
|
|
return self.attr[key].copy()
|
|
def has_attribute(self,key):
|
|
return key in self.attr
|
|
def translate(self,translation,axis=0):
|
|
'''Translates particles. Periodicity must be enforced manually.'''
|
|
assert axis<3, "'axis' must be smaller than 3."
|
|
key = ('x','y','z')[axis]
|
|
self.attr[key] += translation
|
|
return
|
|
def set_frame_velocity(self,val,axis=0):
|
|
'''Adjust frame of reference by translating particles by
|
|
time*frame_velocity and subtracting frame_velocity from
|
|
particle velocities.'''
|
|
assert axis<3, "'axis' must be smaller than 3."
|
|
valdiff = val-self.frame_velocity[axis]
|
|
self.translate(-self.time*valdiff,axis=axis)
|
|
key = ('u','v','w')[axis]
|
|
self.attr[key] -= valdiff
|
|
self.frame_velocity[axis] = val
|
|
return
|
|
def enforce_periodicity(self,axis=None):
|
|
if axis is None:
|
|
for ii in range(3):
|
|
self.enforce_periodicity(axis=ii)
|
|
else:
|
|
if self.period[axis] is not None:
|
|
key = ('x','y','z')[axis]
|
|
self.attr[key] %= self.period[axis]
|
|
return
|
|
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
|
|
# and the attributes of Particles will be views.
|
|
self.attr = {}
|
|
self.part = []
|
|
self._is_data_list = True
|
|
if copy_particles:
|
|
self.part.append(part.copy())
|
|
else:
|
|
self.part.append(part)
|
|
for key in self.part[0].attr:
|
|
self.attr[key] = [self.part[0].attr[key].view()]
|
|
self.period = self.part[0].period
|
|
self.unraveled = unraveled
|
|
self.frame_velocity = numpy.zeros(3)
|
|
def __str__(self):
|
|
str = 'Trajectories with\n'
|
|
str+= ' time steps: {:d}\n'.format(self.numtime)
|
|
str+= ' particles: {:d}\n'.format(self.numpart)
|
|
str += ' attributes:'
|
|
for key in self.part[0].attr.keys():
|
|
str += ' {}'.format(key)
|
|
str+= '\n'
|
|
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)
|
|
elif isinstance(other,Particles):
|
|
self.add_particles(other)
|
|
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 = []
|
|
for x in val:
|
|
if isinstance(x,int):
|
|
lo,hi = x,x+1
|
|
hi = hi if hi!=0 else None
|
|
sl.append(slice(lo,hi))
|
|
elif isinstance(x,slice):
|
|
sl.append(x)
|
|
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
|
|
pp = load_mat(file,'pp',cast_integer=False) # int casting is expensive and useless on potentially large array
|
|
col,time,ccinfo = load_mat(file,['colpy','time','ccinfo'],cast_integer=True)
|
|
period = [None,None,None]
|
|
if bool(ccinfo['xperiodic']): period[0]=ccinfo['b']
|
|
if bool(ccinfo['yperiodic']): period[1]=ccinfo['d']
|
|
if bool(ccinfo['zperiodic']): period[2]=ccinfo['f']
|
|
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)
|
|
traj = cls(Particles.from_array(pp[:,:,0],col,time[0],period,select_col=select_col),unraveled=unraveled)
|
|
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."
|
|
parts = self.part[slice_time]
|
|
assert len(parts)>0, "Slicing yielded zero time steps."
|
|
for ii,part in enumerate(parts):
|
|
if ii==0:
|
|
traj = Trajectories(part[slice_part],unraveled=self.unraveled)
|
|
else:
|
|
traj += Trajectories(part[slice_part])
|
|
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."
|
|
assert all(part.attr['id']==self.part[0].attr['id']), "Particle IDs differ or a not in the same order."
|
|
assert all([part.period[ii]==self.period[ii] for ii in range(0,3)]), "Period differs!"
|
|
for key in self.attr:
|
|
assert key in part.attr, "Particles to be added are missing attribute '{}'".format(key)
|
|
for key in part.attr:
|
|
assert key in self.attr, "Particles to be added have additional attribute '{}'".format(key)
|
|
# Add new data
|
|
self._make_data_list()
|
|
if self.copy_particles:
|
|
self.part.append(part.copy())
|
|
else:
|
|
self.part.append(part)
|
|
for key in self.attr:
|
|
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."
|
|
assert traj.numpart==self.numpart, "Number of particles is different from previous time steps."
|
|
assert all([traj.period[ii]==self.period[ii] for ii in range(0,3)]), "Period differs!"
|
|
for key in self.attr:
|
|
assert key in traj.attr, "Trajectories to be added are missing attribute '{}'".format(key)
|
|
for key in traj.attr:
|
|
assert key in self.attr, "Trajectories to be added have additional attribute '{}'".format(key)
|
|
# Add new data
|
|
self._make_data_list()
|
|
traj._make_data_list()
|
|
if self.copy_particles:
|
|
self.part += deepcopy(traj.part)
|
|
else:
|
|
self.part += traj.part
|
|
for ii in range(0,traj.numtime):
|
|
for key in self.attr:
|
|
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)
|
|
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
|
|
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([
|
|
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
|
|
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
|
|
self.unravel()
|
|
xyzpath = self.get_trajectories(slice_part=slice_part,slice_time=slice_time)
|
|
npart = xyzpath.shape[1]
|
|
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)
|
|
for axis in range(3):
|
|
if self.period[axis] is not None:
|
|
lastPeriod[axis,:] = numpy.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)
|
|
if itime==ntime:
|
|
hasJumped = numpy.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)
|
|
for ipart in range(npart):
|
|
if hasJumped[ipart]:
|
|
sl = slice(lastJump[ipart],itime)
|
|
segment = xyzpath[:,ipart,sl].copy()
|
|
for axis in range(3):
|
|
if self.period[axis] is not None:
|
|
segment[axis,:] -= lastPeriod[axis,ipart]*self.period[axis]
|
|
out[ipart].append(segment)
|
|
lastJump[ipart] = itime
|
|
lastPeriod = thisPeriod
|
|
# Restore previous ravel state
|
|
if restore_ravel_state and not was_unraveled:
|
|
self.ravel()
|
|
return out
|
|
def slice_time(self,time_beg=None,time_end=None):
|
|
return
|
|
def set_frame_velocity(self,val,axis=0):
|
|
assert axis<3, "'axis' must be smaller than 3."
|
|
for part in self.part:
|
|
part.set_frame_velocity(val,axis=axis)
|
|
self.frame_velocity[axis] = val
|
|
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()
|
|
for axis in range(0,3):
|
|
if self.period[axis] is not None:
|
|
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)
|
|
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 _make_data_array(self):
|
|
'''Transforms self.attr[key] to a numpy 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)
|
|
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
|
|
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)]
|
|
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)
|
|
for itime in range(0,ntime):
|
|
idx = pp[col['id'],:,itime].squeeze().argsort()
|
|
pp[:,:,itime] = pp[:,idx,itime]
|
|
return pp
|
|
|
|
# def mask(pp,col,field,reconstruct=False,cval=np.nan,padding=0.0):
|
|
# from .field import Field3d
|
|
# '''Fills grid points which lie inside of solid phase with values.'''
|
|
# # Loop through local particles
|
|
# for ipart in range(0,self.__npartl):
|
|
# # Slice a box from the field around the particle
|
|
# xp,yp,zp,rp = (self.particles[self.col['x'],ipart],
|
|
# self.particles[self.col['y'],ipart],
|
|
# self.particles[self.col['z'],ipart],
|
|
# self.particles[self.col['r'],ipart])
|
|
# rp += padding*self.__dx[key]
|
|
# # Get coefficients for reconstruction
|
|
# if reconstruct:
|
|
# coeff_lin = self.particles[self.col[key],ipart]
|
|
# if key=='u':
|
|
# coeff_rotx = 0.0
|
|
# coeff_roty = -self.particles[self.col['oz'],ipart]
|
|
# coeff_rotz = self.particles[self.col['oy'],ipart]
|
|
# elif key=='v':
|
|
# coeff_rotx = self.particles[self.col['oz'],ipart]
|
|
# coeff_roty = 0.0
|
|
# coeff_rotz = -self.particles[self.col['ox'],ipart]
|
|
# elif key=='w':
|
|
# coeff_rotx = -self.particles[self.col['oy'],ipart]
|
|
# coeff_roty = self.particles[self.col['ox'],ipart]
|
|
# coeff_rotz = 0.0
|
|
# else:
|
|
# coeff_rotx = 0.0
|
|
# coeff_roty = 0.0
|
|
# coeff_rotz = 0.0
|
|
# # Get bounding box of particle
|
|
# idx_x = np.nonzero((xg>=xp-rp) & (xg<=xp+rp))[0]
|
|
# idx_y = np.nonzero((yg>=yp-rp) & (yg<=yp+rp))[0]
|
|
# idx_z = np.nonzero((zg>=zp-rp) & (zg<=zp+rp))[0]
|
|
# # Triple for loop
|
|
# for ii in range(idx_x[0],idx_x[-1]+1):
|
|
# Dx = xg[ii]-xp
|
|
# for jj in range(idx_y[0],idx_y[-1]+1):
|
|
# Dy = yg[jj]-yp
|
|
# for kk in range(idx_z[0],idx_z[-1]+1):
|
|
# Dz = zg[kk]-zp
|
|
# isInside = Dx*Dx+Dy*Dy+Dz*Dz <= rp*rp
|
|
# if isInside:
|
|
# if reconstruct:
|
|
# self.field[key][ii,jj,kk] = coeff_lin + coeff_rotx*Dx + coeff_roty*Dy + coeff_rotz*Dz
|
|
# else:
|
|
# self.field[key][ii,jj,kk] = cval
|
|
|
|
def slice_columns(pp,col,keys):
|
|
idx_col = []
|
|
col_new = {}
|
|
ii = 0
|
|
for key in keys:
|
|
idx_col.append(col[key])
|
|
col_new[key] = ii
|
|
ii+=1
|
|
return pp[idx_col,:,:], col_new
|
|
|
|
def translate_circular(pp,col,translation,bounds,axis=0):
|
|
'''Translates particles while taking into account
|
|
the bounding box'''
|
|
assert(axis<3)
|
|
L = bounds[2*axis+1]
|
|
keys = ('x','y','z')
|
|
pp[col[keys[axis]],:,:] = (pp[col[keys[axis]],:,:]+translation)%L
|
|
return pp |