diff --git a/mbtrack2/tracking/element.py b/mbtrack2/tracking/element.py
index a46c6b889f0f9f05660d5826dc548fed8ca3434a..685655f12d3294f187bc1ac6a20d99e161c97ad2 100644
--- a/mbtrack2/tracking/element.py
+++ b/mbtrack2/tracking/element.py
@@ -15,43 +15,6 @@ from scipy.special import factorial
from mbtrack2.tracking.particles import Beam
-def _compute_chromatic_phase_advances(chro, bunch):
- order = len(chro) // 2
- if order == 1:
- phase_advance_x = chro[0] * bunch["delta"]
- phase_advance_y = chro[1] * bunch["delta"]
- elif order == 2:
- phase_advance_x = (chro[0] * bunch["delta"] +
- chro[2] / 2 * bunch["delta"]**2)
- phase_advance_y = (chro[1] * bunch["delta"] +
- chro[3] / 2 * bunch["delta"]**2)
- elif order == 3:
- phase_advance_x = (chro[0] * bunch["delta"] +
- chro[2] / 2 * bunch["delta"]**2 +
- chro[4] / 6 * bunch["delta"]**3)
- phase_advance_y = (chro[1] * bunch["delta"] +
- chro[3] / 2 * bunch["delta"]**2 +
- chro[5] / 6 * bunch["delta"]**3)
- elif order == 4:
- phase_advance_x = (chro[0] * bunch["delta"] +
- chro[2] / 2 * bunch["delta"]**2 +
- chro[4] / 6 * bunch["delta"]**3 +
- chro[6] / 24 * bunch["delta"]**4)
- phase_advance_y = (chro[1] * bunch["delta"] +
- chro[3] / 2 * bunch["delta"]**2 +
- chro[5] / 6 * bunch["delta"]**3 +
- chro[7] / 24 * bunch["delta"]**4)
- else:
- coefs = np.array([1 / factorial(i + 1) for i in range(order + 1)])
- coefs[0] = 0
- chro = np.concatenate(([0, 0], chro))
- phase_advance_x = np.polynomial.polynomial.Polynomial(
- chro[::2] * coefs)(bunch['delta'])
- phase_advance_y = np.polynomial.polynomial.Polynomial(
- chro[1::2] * coefs)(bunch['delta'])
- return phase_advance_x, phase_advance_y
-
-
class Element(metaclass=ABCMeta):
"""
Abstract Element class used for subclass inheritance to define all kinds
@@ -282,15 +245,51 @@ class TransverseMapSector(Element):
else:
self.adts_poly = None
- def _compute_new_coords(self, bunch, phase_advance, plane):
+ def _compute_chromatic_tune_advances(self, chro, bunch):
+ order = len(chro) // 2
+ if order == 1:
+ tune_advance_x = chro[0] * bunch["delta"]
+ tune_advance_y = chro[1] * bunch["delta"]
+ elif order == 2:
+ tune_advance_x = (chro[0] * bunch["delta"] +
+ chro[2] / 2 * bunch["delta"]**2)
+ tune_advance_y = (chro[1] * bunch["delta"] +
+ chro[3] / 2 * bunch["delta"]**2)
+ elif order == 3:
+ tune_advance_x = (chro[0] * bunch["delta"] +
+ chro[2] / 2 * bunch["delta"]**2 +
+ chro[4] / 6 * bunch["delta"]**3)
+ tune_advance_y = (chro[1] * bunch["delta"] +
+ chro[3] / 2 * bunch["delta"]**2 +
+ chro[5] / 6 * bunch["delta"]**3)
+ elif order == 4:
+ tune_advance_x = (chro[0] * bunch["delta"] +
+ chro[2] / 2 * bunch["delta"]**2 +
+ chro[4] / 6 * bunch["delta"]**3 +
+ chro[6] / 24 * bunch["delta"]**4)
+ tune_advance_y = (chro[1] * bunch["delta"] +
+ chro[3] / 2 * bunch["delta"]**2 +
+ chro[5] / 6 * bunch["delta"]**3 +
+ chro[7] / 24 * bunch["delta"]**4)
+ else:
+ coefs = np.array([1 / factorial(i + 1) for i in range(order + 1)])
+ coefs[0] = 0
+ chro = np.concatenate(([0, 0], chro))
+ tune_advance_x = np.polynomial.polynomial.Polynomial(
+ chro[::2] * coefs)(bunch['delta'])
+ tune_advance_y = np.polynomial.polynomial.Polynomial(
+ chro[1::2] * coefs)(bunch['delta'])
+ return tune_advance_x, tune_advance_y
+
+ def _compute_new_coords(self, bunch, tune_advance, plane):
if plane == 'x':
i, j, coord, mom = 0, 0, 'x', 'xp'
elif plane == 'y':
i, j, coord, mom = 1, 2, 'y', 'yp'
else:
raise ValueError('plane should be either x or y')
- c_u = np.cos(2 * np.pi * phase_advance)
- s_u = np.sin(2 * np.pi * phase_advance)
+ c_u = np.cos(2 * np.pi * tune_advance)
+ s_u = np.sin(2 * np.pi * tune_advance)
M00 = np.sqrt(
self.beta1[i] / self.beta0[i]) * (c_u + self.alpha0[i] * s_u)
M01 = np.sqrt(self.beta0[i] * self.beta1[i]) * s_u
@@ -303,7 +302,6 @@ class TransverseMapSector(Element):
self.beta0[i] / self.beta1[i]) * (c_u - self.alpha1[i] * s_u)
M12 = (self.dispersion1[j + 1] - M10 * self.dispersion0[j] -
M11 * self.dispersion0[j + 1])
- M22 = 1
u = (M00 * bunch[coord] + M01 * bunch[mom] + M02 * bunch["delta"])
up = (M10 * bunch[coord] + M11 * bunch[mom] + M12 * bunch["delta"])
return u, up
@@ -319,14 +317,14 @@ class TransverseMapSector(Element):
----------
bunch : Bunch or Beam object
"""
- phase_advance_x = self.tune_diff[0]
- phase_advance_y = self.tune_diff[1]
- # Compute phase advance which depends on energy via chromaticity and ADTS
+ tune_advance_x = self.tune_diff[0]
+ tune_advance_y = self.tune_diff[1]
+ # Compute tune advance which depends on energy via chromaticity and ADTS
if (np.array(self.chro_diff) != 0).any():
- phase_advance_x_chro, phase_advance_y_chro = _compute_chromatic_phase_advances(
+ tune_advance_x_chro, tune_advance_y_chro = _compute_chromatic_tune_advances(
self.chro_diff, bunch)
- phase_advance_x += phase_advance_x_chro
- phase_advance_y += phase_advance_y_chro
+ tune_advance_x += tune_advance_x_chro
+ tune_advance_y += tune_advance_y_chro
if self.adts_poly is not None:
Jx = ((self.gamma0[0] * bunch["x"]**2) +
@@ -335,13 +333,13 @@ class TransverseMapSector(Element):
Jy = ((self.gamma0[1] * bunch["y"]**2) +
(2 * self.alpha0[1] * bunch["y"] * bunch["yp"]) +
(self.beta0[1] * bunch["yp"]**2))
- phase_advance_x += (self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
- phase_advance_x += (self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
+ tune_advance_x += (self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
+ tune_advance_x += (self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
bunch['x'], bunch['xp'] = self._compute_new_coords(
- bunch, phase_advance_x, 'x')
+ bunch, tune_advance_x, 'x')
bunch['y'], bunch['yp'] = self._compute_new_coords(
- bunch, phase_advance_y, 'y')
+ bunch, tune_advance_y, 'y')
class TransverseMap(TransverseMapSector):