Use numpy array instead of panda DataFrame to store particles

1st test, not complete

tracking/particles.py

@@ -97,7 +97,8 @@ class Bunch:
     Springer, Eq.(8.39) of p224.
     """
     
-    def __init__(self, ring, mp_number=1e3, current=1e-3, alive=True):
+    def __init__(self, ring, mp_number=1e3, current=1e-3, alive=True, 
+                 track_alive=False):
         
         self.ring = ring
         if not alive:
@@ -105,20 +106,22 @@ class Bunch:
             current = 0
         self._mp_number = int(mp_number)
         
-        particles = {"x":np.zeros((self.mp_number,)),
-                     "xp":np.zeros((self.mp_number,)),
-                     "y":np.zeros((self.mp_number,)),
-                     "yp":np.zeros((self.mp_number,)),
-                     "tau":np.zeros((self.mp_number,)),
-                     "delta":np.zeros((self.mp_number,)),
-                     }
-        self.particles = pd.DataFrame(particles)
-        self.alive = pd.Series(np.ones((self.mp_number,),dtype=bool))
+        self.dtype = np.dtype([('x',np.float),
+                       ('xp',np.float),
+                       ('y',np.float),
+                       ('yp',np.float),
+                       ('tau',np.float),
+                       ('delta',np.float)])
+        
+        self.particles = np.zeros(self.mp_number, self.dtype)
+        self.track_alive = track_alive
+        
+        self.alive = np.ones((self.mp_number,),dtype=bool)
         self.current = current
         if not alive:
-            self.alive = pd.Series(np.zeros((self.mp_number,),dtype=bool))  
+            self.alive = np.zeros((self.mp_number,),dtype=bool)
             
-        self._energy_change = pd.Series(np.zeros((self.mp_number,),dtype=float))  
+        self._energy_change = np.zeros((self.mp_number,),dtype=float)
         
     def __len__(self):
         """Return the number of alive particles"""
@@ -126,15 +129,21 @@ class Bunch:
         
     def __getitem__(self, label):
         """Return the columns label for alive particles"""
-        return self.particles.loc[self.alive, label]
+        if self.track_alive is True:
+            return self.particles[label][self.alive]
+        else:
+            return self.particles[label]
     
     def __setitem__(self, label, value):
         """Set value to the columns label for alive particles"""
-        self.particles.loc[self.alive, label] = value
+        if self.track_alive is True:
+            self.particles[label][self.alive] = value
+        else:
+            self.particles[label] = value
     
     def __iter__(self):
         """Iterate over labels"""
-        return self[:].__iter__()
+        return self.dtype.names.__iter__()
     
     def __repr__(self):
         """Return representation of alive particles"""
@@ -222,11 +231,17 @@ class Bunch:
         the SynchrotronRadiation class to compute radiation damping. To be
         changed by Element objects which change the energy, for example RF
         cavities."""
-        return self._energy_change.loc[self.alive]
+        if self.track_alive is True:
+            return self._energy_change[self.alive]
+        else:
+            return self._energy_change
     
     @energy_change.setter
     def energy_change(self, value):
-        self._energy_change.loc[self.alive] = value 
+        if self.track_alive is True:
+            self._energy_change[self.alive] = value
+        else:
+            self._energy_change = value
         
     def init_gaussian(self, cov=None, mean=None, **kwargs):
         """
@@ -297,14 +312,18 @@ class Bunch:
             Number of marco-particles in each bin.
 
         """
-        bins = pd.interval_range(start=min(self[dimension]), 
-                                 end=max(self[dimension]), 
-                                 periods=n_bin,
-                                 name=dimension)
-        sorted_index = pd.cut(self[dimension], bins)
-        sorted_index = sorted_index.reindex(self.alive.index)
-        profile = sorted_index.value_counts().sort_index()
-        return (bins, sorted_index, profile)
+        bin_min = self[dimension].min()
+        bin_min = min(bin_min*0.99, bin_min*1.01)
+        bin_max = self[dimension].max()
+        bin_max = max(bin_max*0.99, bin_max*1.01)
+        
+        bins = np.linspace(bin_min, bin_max, n_bin)
+        center = (bins[1:] + bins[:-1])/2
+        sorted_index = np.searchsorted(bins, self[dimension], side='left')
+        sorted_index -= 1
+        profile = np.bincount(sorted_index, minlength=n_bin-1)
+        
+        return (bins, sorted_index, profile, center)
     
     def plot_profile(self, dimension="tau", n_bin=75):
         """
@@ -318,10 +337,10 @@ class Bunch:
             Number of bins. The default is 75.
 
         """
-        bins, sorted_index, profile = self.binning(dimension, n_bin)
+        bins, sorted_index, profile, center = self.binning(dimension, n_bin)
         fig = plt.figure()
         ax = fig.gca()
-        ax.plot(bins.mid, profile)
+        ax.plot(center, profile)
         
     def plot_phasespace(self, x_var="tau", y_var="delta", plot_type="j"):
         """

tracking/wakepotential.py

@@ -72,18 +72,19 @@ class WakePotential(Element):
         """
         
         # Get binning data
-        a, b, c = bunch.binning(n_bin=self.n_bin)
+        a, b, c, d = bunch.binning(n_bin=self.n_bin)
         self.bins = a
         self.sorted_index = b
         self.profile = c
-        self.bin_size = self.bins.length
+        self.center = d
+        self.bin_size = self.bins[1] - self.bins[0]
         
         # Compute charge density
         self.rho = bunch.charge_per_mp*self.profile/(self.bin_size*bunch.charge)
         self.rho = np.array(self.rho)
         
         # Compute time array
-        self.tau = np.array(self.bins.mid)
+        self.tau = np.array(self.center)
         self.dtau = self.tau[1] - self.tau[0]
         
         # Add N values before and after rho and tau