Merge branch 'faster2' into develop

mbtrack2/tracking/element.py

@@ -201,27 +201,29 @@ class TransverseMap(Element):
                 self.adts_poly[1](Jx) + self.adts_poly[3](Jy))
 
         # 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
-        matrix = np.zeros((6, 6, len(bunch)))
+        matrix = np.zeros((6, 6, len(bunch)), dtype=np.float64)
 
         # Horizontal
-        matrix[0, 0, :] = np.cos(
-            phase_advance_x) + self.alpha[0] * np.sin(phase_advance_x)
-        matrix[0, 1, :] = self.beta[0] * np.sin(phase_advance_x)
+        c_x = np.cos(phase_advance_x)
+        s_x = np.sin(phase_advance_x)
+
+        matrix[0, 0, :] = c_x + self.alpha[0] * s_x
+        matrix[0, 1, :] = self.beta[0] * s_x
         matrix[0, 2, :] = self.dispersion[0]
-        matrix[1, 0, :] = -1 * self.gamma[0] * np.sin(phase_advance_x)
-        matrix[1, 1, :] = np.cos(
-            phase_advance_x) - self.alpha[0] * np.sin(phase_advance_x)
+        matrix[1, 0, :] = -1 * self.gamma[0] * s_x
+        matrix[1, 1, :] = c_x - self.alpha[0] * s_x
         matrix[1, 2, :] = self.dispersion[1]
         matrix[2, 2, :] = 1
 
         # Vertical
-        matrix[3, 3, :] = np.cos(
-            phase_advance_y) + self.alpha[1] * np.sin(phase_advance_y)
-        matrix[3, 4, :] = self.beta[1] * np.sin(phase_advance_y)
+        c_y = np.cos(phase_advance_y)
+        s_y = np.sin(phase_advance_y)
+
+        matrix[3, 3, :] = c_y + self.alpha[1] * s_y
+        matrix[3, 4, :] = self.beta[1] * s_y
         matrix[3, 5, :] = self.dispersion[2]
-        matrix[4, 3, :] = -1 * self.gamma[1] * np.sin(phase_advance_y)
-        matrix[4, 4, :] = np.cos(
-            phase_advance_y) - self.alpha[1] * np.sin(phase_advance_y)
+        matrix[4, 3, :] = -1 * self.gamma[1] * s_y
+        matrix[4, 4, :] = c_y - self.alpha[1] * s_y
         matrix[4, 5, :] = self.dispersion[3]
         matrix[5, 5, :] = 1
 

mbtrack2/tracking/particles.py

@@ -128,11 +128,14 @@ class Bunch:
             current = 0
         self._mp_number = int(mp_number)
 
-        self.dtype = np.dtype([('x', float), ('xp', float), ('y', float),
-                               ('yp', float), ('tau', float),
-                               ('delta', float)])
-
-        self.particles = np.zeros(self.mp_number, self.dtype)
+        self.particles = {
+            "x": np.zeros(self.mp_number, dtype=np.float64),
+            "xp": np.zeros(self.mp_number, dtype=np.float64),
+            "y": np.zeros(self.mp_number, dtype=np.float64),
+            "yp": np.zeros(self.mp_number, dtype=np.float64),
+            "tau": np.zeros(self.mp_number, dtype=np.float64),
+            "delta": np.zeros(self.mp_number, dtype=np.float64),
+        }
         self.track_alive = track_alive
         self.alive = np.ones((self.mp_number, ), dtype=bool)
         self.current = current
@@ -144,7 +147,7 @@ class Bunch:
 
     def __len__(self):
         """Return the number of alive particles"""
-        return len(self[:])
+        return self.alive.sum()
 
     def __getitem__(self, label):
         """Return the columns label for alive particles"""
@@ -162,11 +165,13 @@ class Bunch:
 
     def __iter__(self):
         """Iterate over labels"""
-        return self.dtype.names.__iter__()
+        return self.particles.keys().__iter__()
 
     def __repr__(self):
         """Return representation of alive particles"""
-        return f'Bunch with macro-particles: \n {pd.DataFrame(self[:])!r}'
+        rep = pd.DataFrame(np.array([self[l] for l in self.__iter__()]).T,
+                           columns=list(self.__iter__()))
+        return f'Bunch with macro-particles: \n {rep!r}'
 
     @property
     def mp_number(self):

mbtrack2/tracking/wakepotential.py

@@ -33,6 +33,11 @@ class WakePotential(Element):
         functions must be uniformly sampled!
     n_bin : int, optional
         Number of bins for constructing the longitudinal bunch profile.
+    interp_on_postion : bool, optional
+        If True, the computed wake potential is interpolated on the exact 
+        particle location. If False, the wake potential is interpolated on the 
+        bin center and each particle of the bin get the same value.
+        Default is True.
         
     Attributes
     ----------
@@ -71,13 +76,14 @@ class WakePotential(Element):
         Reduce wake function samping by an integer factor.
         
     """
-    def __init__(self, ring, wakefield, n_bin=80):
+    def __init__(self, ring, wakefield, n_bin=80, interp_on_postion=True):
         self.wakefield = wakefield
         self.types = self.wakefield.wake_components
         self.n_types = len(self.wakefield.wake_components)
         self.ring = ring
         self.n_bin = n_bin
         self.check_sampling()
+        self.interp_on_postion = interp_on_postion
 
         # Suppress numpy warning for floating-point operations.
         np.seterr(invalid='ignore')
@@ -110,18 +116,12 @@ class WakePotential(Element):
         self.dtau = self.tau[1] - self.tau[0]
 
         # Add N values before and after rho and tau
-        if self.n_bin % 2 == 0:
-            N = int(self.n_bin / 2)
-            self.tau = np.arange(self.tau[0] - self.dtau * N,
-                                 self.tau[-1] + self.dtau * N, self.dtau)
-            self.rho = np.append(self.rho, np.zeros(N))
-            self.rho = np.insert(self.rho, 0, np.zeros(N))
-        else:
-            N = int(np.floor(self.n_bin / 2))
-            self.tau = np.arange(self.tau[0] - self.dtau * N,
-                                 self.tau[-1] + self.dtau * (N+1), self.dtau)
-            self.rho = np.append(self.rho, np.zeros(N))
-            self.rho = np.insert(self.rho, 0, np.zeros(N + 1))
+        N = int(self.n_bin / 2)
+        self.tau = np.arange(self.tau[0] - self.dtau * N,
+                             self.tau[-1] + self.dtau * N, self.dtau)
+        self.rho = np.pad(self.rho, (N, N + self.n_bin % 2),
+                          mode="constant",
+                          constant_values=(0, 0))
 
         if len(self.tau) != len(self.rho):
             self.tau = np.append(self.tau, self.tau[-1] + self.dtau)
@@ -148,7 +148,7 @@ class WakePotential(Element):
             Dipole moment of the bunch.
 
         """
-        dipole = np.zeros((self.n_bin - 1, ))
+        dipole = np.empty((self.n_bin - 1, ))
         for i in range(self.n_bin - 1):
             dipole[i] = bunch[plane][self.sorted_index == i].sum()
         dipole = dipole / self.profile
@@ -219,14 +219,14 @@ class WakePotential(Element):
                 tau0 = np.arange(tau0[0] - dtau0*n_assym, tau0[-1] + dtau0,
                                  dtau0)
                 n_to_add = len(tau0) - len(W0)
-                W0 = np.insert(W0, 0, np.zeros(n_to_add))
+                W0 = np.concatenate((np.zeros(n_to_add), W0))
 
             # add at tail
             elif np.abs(tau0[0]) > np.abs(tau0[-1]):
                 tau0 = np.arange(tau0[0], tau0[-1] + dtau0 * (n_assym+1),
                                  dtau0)
                 n_to_add = len(tau0) - len(W0)
-                W0 = np.insert(W0, 0, np.zeros(n_to_add))
+                W0 = np.concatenate((np.zeros(n_to_add), W0))
 
         # Check is the wf is shorter than rho then add zeros
         if (tau0[0] > tau[0]) or (tau0[-1] < tau[-1]):
@@ -237,7 +237,9 @@ class WakePotential(Element):
                              dtau0)
             W0 = np.insert(W0, 0, np.zeros(n))
             n_to_add = len(tau0) - len(W0)
-            W0 = np.insert(W0, len(W0), np.zeros(n_to_add))
+            W0 = np.pad(W0, (0, n_to_add),
+                        mode="constant",
+                        constant_values=(0, 0))
 
         if save_data:
             setattr(self, "tau0_" + wake_type, tau0)
@@ -300,8 +302,13 @@ class WakePotential(Element):
             self.charge_density(bunch)
             for wake_type in self.types:
                 tau0, Wp = self.get_wakepotential(bunch, wake_type)
-                Wp_interp = np.interp(bunch["tau"], tau0 + self.tau_mean, Wp,
-                                      0, 0)
+                if self.interp_on_postion:
+                    Wp_interp = np.interp(bunch["tau"], tau0 + self.tau_mean, Wp,
+                                          0, 0)
+                else:
+                    Wp_interp = np.interp(self.center, tau0 + self.tau_mean, Wp,
+                                          0, 0)
+                    Wp_interp= Wp_interp[self.sorted_index]
                 if wake_type == "Wlong":
                     bunch["delta"] += Wp_interp * bunch.charge / self.ring.E0
                 elif wake_type == "Wxdip":