diff --git a/helper.py b/helper.py index f280941..2c21422 100644 --- a/helper.py +++ b/helper.py @@ -1,8 +1,8 @@ -def load_mat(f,var,squeeze=True): +def load_mat(f,var,squeeze=True,cast_string=True,cast_integer=False): try: return load_matbin(f,var,squeeze=squeeze) except NotImplementedError: - return load_matv73(f,var,squeeze=squeeze) + return load_matv73(f,var,squeeze=squeeze,cast_string=cast_string,cast_integer=cast_integer) def load_matbin(f,var,squeeze=True): '''Loads variables from a matlab binary mat file (not v7.3).''' @@ -24,14 +24,14 @@ def load_matbin(f,var,squeeze=True): else: return parse(data[var]) -def load_matv73(f,var,squeeze=True,cast_string=True,cast_integer=True): +def load_matv73(f,var,squeeze=True,cast_string=True,cast_integer=False): '''Loads variables from a matlab hdf5 mat file (v7.3).''' import numpy import h5py if isinstance(f,str): f = h5py.File(f,'r') if isinstance(var,(list,tuple)): - return tuple(load_matv73(f,x) for x in var) + return tuple(load_matv73(f,x,cast_string=cast_string,cast_integer=cast_integer) for x in var) def parse(val): # String data if cast_string and val.dtype==numpy.dtype('uint16'): diff --git a/particle.py b/particle.py index 835a86e..c10921a 100644 --- a/particle.py +++ b/particle.py @@ -1,38 +1,82 @@ class Particles: - def __init__(self,pp,col,time,period=(None,None,None),select_col=None): + 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." - self.num = pp.shape[1] + 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) - # Copy attributes and sort by ID - self.attr = {} - sortidx = pp[col['id'],:,0].argsort() - idsorted = pp[col['id'],sortidx,0] + attr = {} for key in col: - if key=='id': - self.add_attribute(key,pp[col[key],sortidx,0].astype('int')) - elif (select_col is None or key in select_col or key in ('x','y','z')): - self.add_attribute(key,pp[col[key],sortidx,0]) - self.period = period - self.time = time - def __getitem__(self,key): - return self.get_attribute(key) + 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: {}, {}, {}'.format(*self.period) + 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)): @@ -43,58 +87,74 @@ class Particles: return def del_attribute(self,key): del self.attr[key] - def get_attribute(self,key,id=None): - import numpy - if id is not None: - assert id>=1, "Particle IDs start with 1." - assert id<=self.num, "Requested ID exceeds number of particles." - if id is None: - return self.attr[key] - else: - return self.attr[key][id-1] - def get_position(self,axis=None,id=None): + 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.get_attribute(key,id=id) for key in ('x','y','z')]) + 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.get_attribute(key) + return self.attr[key].copy() def has_attribute(self,key): return key in self.attr def translate(self,translation,axis=0): - '''Translates particles while taking into account periodicity.''' + '''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 - if self.period[axis] is not None: - self.attr[key] %= self.period[axis] 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,copy=False): + def __init__(self,part,unraveled=False,copy_particles=False): import numpy self.numtime = 1 self.numpart = part.num - self.copy = copy + 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: - self.part.append(deepcopy(part)) + 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 = False + self.unraveled = unraveled + self.frame_velocity = numpy.zeros(3) def __str__(self): - str = 'Trajectory with\n' + 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): @@ -117,29 +177,43 @@ class Trajectories: sl.append(x) else: raise TypeError("Trajectories can only be sliced by slice objects or integers.") - return self.get_trajectories(slice_part=sl[0],slice_time=sl[1]) + return self._slice(slice_part=sl[0],slice_time=sl[1]) @classmethod - def from_mat(cls,file,unraveled=False): + def from_mat(cls,file,unraveled=False,select_col=None): from .helper import load_mat - pp,col,time,ccinfo = load_mat(file,['pp','colpy','time','ccinfo']) + 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(pp[:,:,0],col,time[0],period)) + 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(pp[:,:,ii],col,time[ii],period) - traj.unraveled = unraveled + 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['id']==self.part[0].attr['id']), "Particle IDs differ or a not in the same order." + 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) @@ -147,8 +221,8 @@ class Trajectories: assert key in self.attr, "Particles to be added have additional attribute '{}'".format(key) # Add new data self._make_data_list() - if self.copy: - self.part.append(deepcopy(part)) + if self.copy_particles: + self.part.append(part.copy()) else: self.part.append(part) for key in self.attr: @@ -167,7 +241,7 @@ class Trajectories: # Add new data self._make_data_list() traj._make_data_list() - if self.copy: + if self.copy_particles: self.part += deepcopy(traj.part) else: self.part += traj.part @@ -234,6 +308,16 @@ class Trajectories: 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 @@ -249,13 +333,9 @@ class Trajectories: self.attr[key][:,1:] += coeff*self.period[axis] self.unraveled = True return - def ravel(self): - if not self.unraveled: return - raise NotImplementedError('Implemented but untested!') - for axis in range(0,3): - if self.period[axis] is not None: - key = ('x','y','z')[axis] - self.part[ii].attr[key] %= self.period[axis] + def enforce_periodicity(self): + for part in self.part: + part.enforce_periodicity(axis=None) self.unraveled = False return def _make_data_array(self): diff --git a/ucf.py b/ucf.py index 488b9c1..6fcf8e2 100644 --- a/ucf.py +++ b/ucf.py @@ -484,6 +484,22 @@ class UCFTarFile: def read_parameters(self): self._initialize_buffer_params() return self._buffer_params + def bounds(self,axis=None): + params = self.read_parameters() + keys = ('a','b','c','d','e','f') + if axis is None: + return tuple(params.getfloat('geometry',key) for key in keys) + else: + assert axis<3, "'axis' must be smaller than 3" + return tuple(params.getfloat('geometry',key) for key in keys[2*axis:2*(axis+1)]) + def periodicity(self,axis=None): + params = self.read_parameters() + keys = ('xperiodic','yperiodic','zperiodic') + if axis is None: + return tuple(params.getboolean('geometry',key) for key in keys) + else: + assert axis<3, "'axis' must be smaller than 3" + return params.getboolean('geometry',keys[axis]) def period(self,axis=None): if axis is None: return tuple(self.period(axis=ii) for ii in range(0,3)) @@ -569,7 +585,7 @@ class UCFSnapshot(UCFTarFile): '''Returns particle data as Particles object.''' from .particle import Particles pp,col,time = self.read_particles() - return Particles(pp,col,time,self.period(),select_col=select_col) + return Particles.from_array(pp,col,time,self.period(),select_col=select_col) def nproc(self,axis=None): if axis is None: return self._nproc @@ -619,11 +635,10 @@ class UCFAuxiliary(UCFTarFile): def trajectories(self,select_col=None): from .particle import Particles,Trajectories pp,col,time = self.read_particles() - part = Particles(pp[:,:,0],col,time[0],self.period(),select_col=select_col) - traj = Trajectories(part,copy=False) + part = Particles.from_array(pp[:,:,0],col,time[0],self.period(),select_col=select_col) + traj = Trajectories(part,copy_particles=False) for ii in range(1,len(time)): - part = Particles(pp[:,:,ii],col,time[ii],self.period(),select_col=select_col) - traj.add_particles(part) + traj += Particles.from_array(pp[:,:,ii],col,time[ii],self.period(),select_col=select_col) return traj class Ustar: diff --git a/visu.py b/visu.py index a9a24fd..f06872b 100644 --- a/visu.py +++ b/visu.py @@ -32,11 +32,11 @@ def add_particles(plotter,part, import pyvista assert part.has_attribute('r') pts = pyvista.PolyData(part.get_position().transpose()) - pts['r'] = part['r'] + pts['r'] = part.attr['r'] if scalars: assert part.has_attribute(scalars) - pts[scalars] = part[scalars] + pts[scalars] = part.attr[scalars] color=None geom = pyvista.Sphere(radius=1,theta_resolution=theta_resolution,phi_resolution=phi_resolution)