diff --git a/mbtrack2/tracking/element.py b/mbtrack2/tracking/element.py
index 7ccfa4d8caae64fab0190b9414693df4887b52cd..7fab92b56d9c37dabb6f3bdd2f110e32dfd490fe 100644
--- a/mbtrack2/tracking/element.py
+++ b/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
diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py
index 1c3f20b12b53e40837e7c99922db83fa0cc6f096..f334be881a7fd0a559eadb6826fb30e9b8835b4b 100644
--- a/mbtrack2/tracking/particles.py
+++ b/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):
diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py
index 5881a2d706b5c041910b3913706e7004d582c847..46abfbd835b7bbd8e9fee32516c0e22010aea4f1 100644
--- a/mbtrack2/tracking/wakepotential.py
+++ b/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":