Skip to content
Snippets Groups Projects
Commit cf5c6230 authored by Vadim Gubaidulin's avatar Vadim Gubaidulin
Browse files

- _compute_chromatic_phase_advance() -> _compute_chromatic_tune_advance() 

- moved the above function to be inside the TransverseSectorMap
- changes phase_advance_* variables to tune_advance_* variables where appropriate
parent 950a07a0
2 merge requests!280.8.0,!16Nonlinear chromaticity
Hide whitespace changes
Inline Side-by-side
mbtrack2/tracking/element.py
+ 49
− 51
View file @ cf5c6230
......@@ -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):
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

The place to be for your source code