some bugfixes for larger datasizes

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'):

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:
+    def get_attribute(self,key):
         return self.attr[key]
-        else:
-            return self.attr[key][id-1]
-    def get_position(self,axis=None,id=None):
+    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
+        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):

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:

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)