Merge branch 'numpy_bunch'

tracking/particles.py

@@ -8,15 +8,15 @@ Module where particles, bunches and beams are described as objects.
 """
 
 import numpy as np
-import pandas as pd
 import matplotlib.pyplot as plt
-from mbtrack2.tracking.parallel import Mpi
 import seaborn as sns
+import pandas as pd
+from mbtrack2.tracking.parallel import Mpi
 from scipy.constants import c, m_e, m_p, e
 
 class Particle:
     """
-    Define a particle object
+    Define a particle object.
 
     Attributes
     ----------
@@ -47,35 +47,41 @@ class Proton(Particle):
 
 class Bunch:
     """
-    Define a bunch
+    Define a bunch object.
     
     Parameters
     ----------
     ring : Synchrotron object
     mp_number : float, optional
-        Macro-particle number
+        Macro-particle number.
     current : float, optional
-        Bunch current in [A]
+        Bunch current in [A].
+    track_alive : bool, optional
+        If False, the code no longer take into account alive/dead particles.
+        Should be set to True if element such as apertures are used.
+        Can be set to False to gain a speed increase.
     alive : bool, optional
-        If False, the bunch is defined as empty 
+        If False, the bunch is defined as empty.
         
     Attributes
     ----------
     mp_number : int
-        Macro-particle number
+        Macro-particle number.
+    alive : array of bool of shape (mp_number,)
+        Array used to monitor which particle is dead or alive.
     charge : float
-        Bunch charge in [C]
+        Bunch charge in [C].
     charge_per_mp : float
-        Charge per macro-particle in [C]
+        Charge per macro-particle in [C].
     particle_number : int
-        Number of particles in the bunch
+        Number of particles in the bunch.
     current : float
-        Bunch current in [A]
+        Bunch current in [A].
     mean : array of shape (6,)
-        Mean position of alive particles for each coordinates
+        Mean position of alive particles for each coordinates.
     std : array of shape (6,)
         Standard deviation of the position of alive particles for each 
-        coordinates
+        coordinates.
     emit : array of shape (3,)
         Bunch emittance for each plane [1]. !!! -> Correct for long ?
     energy_change : series of shape (alive,) 
@@ -97,7 +103,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, track_alive=True,
+                 alive=True):
         
         self.ring = ring
         if not alive:
@@ -105,20 +112,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,19 +135,25 @@ 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"""
-        return f'{self[:]!r}'
+        return f'{pd.DataFrame(self[:])!r}'
         
     @property
     def mp_number(self):
@@ -222,11 +237,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):
         """
@@ -289,22 +310,28 @@ class Bunch:
 
         Returns
         -------
-        bins : IntervalIndex of shape (n_bin,)
+        bins : array of shape (n_bin,)
             Bins where the particles are sorted.
-        sorted_index : Series of shape (self.mp_number,)
+        sorted_index : array of shape (self.mp_number,)
             Bin number of each macro-particles.
-        profile : Series of shape (n_bin,)
+        profile : array of shape (n_bin - 1,)
             Number of marco-particles in each bin.
+        center : array of shape (n_bin - 1,)
+            Center of 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 +345,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"):
         """
@@ -476,7 +503,7 @@ class Beam:
         return distance
         
     def init_beam(self, filling_pattern, current_per_bunch=1e-3, 
-                  mp_per_bunch=1e3):
+                  mp_per_bunch=1e3, track_alive=True):
         """
         Initialize beam with a given filling pattern and marco-particle number 
         per bunch. Then initialize the different bunches with a 6D gaussian
@@ -496,6 +523,10 @@ class Beam:
             Current per bunch in [A]
         mp_per_bunch : float, optional
             Macro-particle number per bunch
+        track_alive : bool, optional
+            If False, the code no longer take into account alive/dead particles.
+            Should be set to True if element such as apertures are used.
+            Can be set to False to gain a speed increase.
         """
         
         if (len(filling_pattern) != self.ring.h):
@@ -507,13 +538,15 @@ class Beam:
         if filling_pattern.dtype == np.dtype("bool"):
             for value in filling_pattern:
                 if value == True:
-                    bunch_list.append(Bunch(self.ring, mp_per_bunch, current_per_bunch))
+                    bunch_list.append(Bunch(self.ring, mp_per_bunch, 
+                                            current_per_bunch, track_alive))
                 elif value == False:
                     bunch_list.append(Bunch(self.ring, alive=False))
         elif filling_pattern.dtype == np.dtype("float64"):
             for current in filling_pattern:
                 if current != 0:
-                    bunch_list.append(Bunch(self.ring, mp_per_bunch, current))
+                    bunch_list.append(Bunch(self.ring, mp_per_bunch, 
+                                            current_per_bunch, track_alive))
                 elif current == 0:
                     bunch_list.append(Bunch(self.ring, alive=False))
         else:

tracking/wakepotential.py

@@ -85,18 +85,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