diff --git a/mbtrack2/tracking/element.py b/mbtrack2/tracking/element.py index 4385b57d2f44e87d82648d2ba97831c22051f581..049f7006f60c3f844e7c50c391204b9816a0c8ca 100644 --- a/mbtrack2/tracking/element.py +++ b/mbtrack2/tracking/element.py @@ -10,6 +10,7 @@ from copy import deepcopy from functools import wraps import numpy as np +from scipy.special import factorial from mbtrack2.tracking.particles import Beam @@ -128,128 +129,26 @@ class SynchrotronRadiation(Element): ---------- bunch : Bunch or Beam object """ + N = len(bunch) if self.switch[0] == True: - rand = np.random.normal(size=len(bunch)) + rand = np.random.standard_normal(size=N) bunch["delta"] = (1 - 2 * self.ring.T0 / self.ring.tau[2]) * bunch[ "delta"] + 2 * self.ring.sigma_delta * ( self.ring.T0 / self.ring.tau[2])**0.5 * rand if self.switch[1] == True: - rand = np.random.normal(size=len(bunch)) + rand = np.random.standard_normal(size=N) bunch["xp"] = (1 - 2 * self.ring.T0 / self.ring.tau[0] ) * bunch["xp"] + 2 * self.ring.sigma()[1] * ( self.ring.T0 / self.ring.tau[0])**0.5 * rand if self.switch[2] == True: - rand = np.random.normal(size=len(bunch)) + rand = np.random.standard_normal(size=N) bunch["yp"] = (1 - 2 * self.ring.T0 / self.ring.tau[1] ) * bunch["yp"] + 2 * self.ring.sigma()[3] * ( self.ring.T0 / self.ring.tau[1])**0.5 * rand -class TransverseMap(Element): - """ - Transverse map for a single turn in the synchrotron. - - Parameters - ---------- - ring : Synchrotron object - """ - - def __init__(self, ring): - self.ring = ring - self.alpha = self.ring.optics.local_alpha - self.beta = self.ring.optics.local_beta - self.gamma = self.ring.optics.local_gamma - self.dispersion = self.ring.optics.local_dispersion - if self.ring.adts is not None: - self.adts_poly = [ - np.poly1d(self.ring.adts[0]), - np.poly1d(self.ring.adts[1]), - np.poly1d(self.ring.adts[2]), - np.poly1d(self.ring.adts[3]), - ] - - @Element.parallel - def track(self, bunch): - """ - Tracking method for the element. - No bunch to bunch interaction, so written for Bunch objects and - @Element.parallel is used to handle Beam objects. - - Parameters - ---------- - bunch : Bunch or Beam object - """ - - # Compute phase advance which depends on energy via chromaticity and ADTS - if self.ring.adts is None: - phase_advance_x = ( - 2 * np.pi * - (self.ring.tune[0] + self.ring.chro[0] * bunch["delta"])) - phase_advance_y = ( - 2 * np.pi * - (self.ring.tune[1] + self.ring.chro[1] * bunch["delta"])) - else: - Jx = ((self.ring.optics.local_gamma[0] * bunch["x"]**2) + - (2 * self.ring.optics.local_alpha[0] * bunch["x"] * - bunch["xp"]) + - (self.ring.optics.local_beta[0] * bunch["xp"]**2)) - Jy = ((self.ring.optics.local_gamma[1] * bunch["y"]**2) + - (2 * self.ring.optics.local_alpha[1] * bunch["y"] * - bunch["yp"]) + - (self.ring.optics.local_beta[1] * bunch["yp"]**2)) - phase_advance_x = ( - 2 * np.pi * - (self.ring.tune[0] + self.ring.chro[0] * bunch["delta"] + - self.adts_poly[0](Jx) + self.adts_poly[2](Jy))) - phase_advance_y = ( - 2 * np.pi * - (self.ring.tune[1] + self.ring.chro[1] * bunch["delta"] + - 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)), dtype=np.float64) - - # Horizontal - 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] * s_x - matrix[1, 1, :] = c_x - self.alpha[0] * s_x - matrix[1, 2, :] = self.dispersion[1] - matrix[2, 2, :] = 1 - - # Vertical - 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] * s_y - matrix[4, 4, :] = c_y - self.alpha[1] * s_y - matrix[4, 5, :] = self.dispersion[3] - matrix[5, 5, :] = 1 - - x = (matrix[0, 0] * bunch["x"] + matrix[0, 1] * bunch["xp"] + - matrix[0, 2] * bunch["delta"]) - xp = (matrix[1, 0] * bunch["x"] + matrix[1, 1] * bunch["xp"] + - matrix[1, 2] * bunch["delta"]) - y = (matrix[3, 3] * bunch["y"] + matrix[3, 4] * bunch["yp"] + - matrix[3, 5] * bunch["delta"]) - yp = (matrix[4, 3] * bunch["y"] + matrix[4, 4] * bunch["yp"] + - matrix[4, 5] * bunch["delta"]) - - bunch["x"] = x - bunch["xp"] = xp - bunch["y"] = y - bunch["yp"] = yp - - class SkewQuadrupole: """ Thin skew quadrupole element used to introduce betatron coupling (the @@ -346,6 +245,67 @@ class TransverseMapSector(Element): else: self.adts_poly = None + def _compute_chromatic_tune_advances(self, bunch): + order = len(self.chro_diff) // 2 + if order == 1: + tune_advance_x = self.chro_diff[0] * bunch["delta"] + tune_advance_y = self.chro_diff[1] * bunch["delta"] + elif order == 2: + tune_advance_x = (self.chro_diff[0] * bunch["delta"] + + self.chro_diff[2] / 2 * bunch["delta"]**2) + tune_advance_y = (self.chro_diff[1] * bunch["delta"] + + self.chro_diff[3] / 2 * bunch["delta"]**2) + elif order == 3: + tune_advance_x = (self.chro_diff[0] * bunch["delta"] + + self.chro_diff[2] / 2 * bunch["delta"]**2 + + self.chro_diff[4] / 6 * bunch["delta"]**3) + tune_advance_y = (self.chro_diff[1] * bunch["delta"] + + self.chro_diff[3] / 2 * bunch["delta"]**2 + + self.chro_diff[5] / 6 * bunch["delta"]**3) + elif order == 4: + tune_advance_x = (self.chro_diff[0] * bunch["delta"] + + self.chro_diff[2] / 2 * bunch["delta"]**2 + + self.chro_diff[4] / 6 * bunch["delta"]**3 + + self.chro_diff[6] / 24 * bunch["delta"]**4) + tune_advance_y = (self.chro_diff[1] * bunch["delta"] + + self.chro_diff[3] / 2 * bunch["delta"]**2 + + self.chro_diff[5] / 6 * bunch["delta"]**3 + + self.chro_diff[7] / 24 * bunch["delta"]**4) + else: + coefs = np.array([1 / factorial(i + 1) for i in range(order + 1)]) + coefs[0] = 0 + self.chro_diff = np.concatenate(([0, 0], self.chro_diff)) + tune_advance_x = np.polynomial.polynomial.Polynomial( + self.chro_diff[::2] * coefs)(bunch['delta']) + tune_advance_y = np.polynomial.polynomial.Polynomial( + self.chro_diff[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 * 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 + M02 = (self.dispersion1[j] - M00 * self.dispersion0[j] - + M01 * self.dispersion0[j + 1]) + M10 = ((self.alpha0[i] - self.alpha1[i]) * c_u - + (1 + self.alpha0[i] * self.alpha1[i]) * s_u) / np.sqrt( + self.beta0[i] * self.beta1[i]) + M11 = np.sqrt( + 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]) + u = (M00 * bunch[coord] + M01 * bunch[mom] + M02 * bunch["delta"]) + up = (M10 * bunch[coord] + M11 * bunch[mom] + M12 * bunch["delta"]) + return u, up + @Element.parallel def track(self, bunch): """ @@ -357,85 +317,45 @@ class TransverseMapSector(Element): ---------- bunch : Bunch or Beam object """ - - # Compute phase advance which depends on energy via chromaticity and ADTS - if self.adts_poly is None: - phase_advance_x = ( - 2 * np.pi * - (self.tune_diff[0] + self.chro_diff[0] * bunch["delta"])) - phase_advance_y = ( - 2 * np.pi * - (self.tune_diff[1] + self.chro_diff[1] * bunch["delta"])) - else: + 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(): + tune_advance_x_chro, tune_advance_y_chro = self._compute_chromatic_tune_advances( + bunch) + 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) + (2 * self.alpha0[0] * bunch["x"] * self["xp"]) + (self.beta0[0] * bunch["xp"]**2)) Jy = ((self.gamma0[1] * bunch["y"]**2) + (2 * self.alpha0[1] * bunch["y"] * bunch["yp"]) + (self.beta0[1] * bunch["yp"]**2)) - phase_advance_x = ( - 2 * np.pi * - (self.tune_diff[0] + self.chro_diff[0] * bunch["delta"] + - self.adts_poly[0](Jx) + self.adts_poly[2](Jy))) - phase_advance_y = ( - 2 * np.pi * - (self.tune_diff[1] + self.chro_diff[1] * bunch["delta"] + - 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))) - - # Horizontal - matrix[0, 0, :] = np.sqrt(self.beta1[0] / self.beta0[0]) * ( - np.cos(phase_advance_x) + self.alpha0[0] * np.sin(phase_advance_x)) - matrix[0, 1, :] = np.sqrt( - self.beta0[0] * self.beta1[0]) * np.sin(phase_advance_x) - matrix[0, 2, :] = (self.dispersion1[0] - - matrix[0, 0, :] * self.dispersion0[0] - - matrix[0, 1, :] * self.dispersion0[1]) - matrix[1, 0, :] = ( - (self.alpha0[0] - self.alpha1[0]) * np.cos(phase_advance_x) - - (1 + self.alpha0[0] * self.alpha1[0]) * - np.sin(phase_advance_x)) / np.sqrt(self.beta0[0] * self.beta1[0]) - matrix[1, 1, :] = np.sqrt(self.beta0[0] / self.beta1[0]) * ( - np.cos(phase_advance_x) - self.alpha1[0] * np.sin(phase_advance_x)) - matrix[1, 2, :] = (self.dispersion1[1] - - matrix[1, 0, :] * self.dispersion0[0] - - matrix[1, 1, :] * self.dispersion0[1]) - matrix[2, 2, :] = 1 - - # Vertical - matrix[3, 3, :] = np.sqrt(self.beta1[1] / self.beta0[1]) * ( - np.cos(phase_advance_y) + self.alpha0[1] * np.sin(phase_advance_y)) - matrix[3, 4, :] = np.sqrt( - self.beta0[1] * self.beta1[1]) * np.sin(phase_advance_y) - matrix[3, 5, :] = (self.dispersion1[2] - - matrix[3, 3, :] * self.dispersion0[2] - - matrix[3, 4, :] * self.dispersion0[3]) - matrix[4, 3, :] = ( - (self.alpha0[1] - self.alpha1[1]) * np.cos(phase_advance_y) - - (1 + self.alpha0[1] * self.alpha1[1]) * - np.sin(phase_advance_y)) / np.sqrt(self.beta0[1] * self.beta1[1]) - matrix[4, 4, :] = np.sqrt(self.beta0[1] / self.beta1[1]) * ( - np.cos(phase_advance_y) - self.alpha1[1] * np.sin(phase_advance_y)) - matrix[4, 5, :] = (self.dispersion1[3] - - matrix[4, 3, :] * self.dispersion0[2] - - matrix[4, 4, :] * self.dispersion0[3]) - matrix[5, 5, :] = 1 - - x = (matrix[0, 0, :] * bunch["x"] + matrix[0, 1, :] * bunch["xp"] + - matrix[0, 2, :] * bunch["delta"]) - xp = (matrix[1, 0, :] * bunch["x"] + matrix[1, 1, :] * bunch["xp"] + - matrix[1, 2, :] * bunch["delta"]) - y = (matrix[3, 3, :] * bunch["y"] + matrix[3, 4, :] * bunch["yp"] + - matrix[3, 5, :] * bunch["delta"]) - yp = (matrix[4, 3, :] * bunch["y"] + matrix[4, 4, :] * bunch["yp"] + - matrix[4, 5, :] * bunch["delta"]) - - bunch["x"] = x - bunch["xp"] = xp - bunch["y"] = y - bunch["yp"] = yp + 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, tune_advance_x, 'x') + bunch['y'], bunch['yp'] = self._compute_new_coords( + bunch, tune_advance_y, 'y') + + +class TransverseMap(TransverseMapSector): + """ + Transverse map for a single turn in the synchrotron. + + Parameters + ---------- + ring : Synchrotron object + """ + + def __init__(self, ring): + super().__init__(ring, ring.optics.local_alpha, ring.optics.local_beta, + ring.optics.local_dispersion, ring.optics.local_alpha, + ring.optics.local_beta, ring.optics.local_dispersion, + 2 * np.pi * ring.tune, ring.chro, ring.adts) def transverse_map_sector_generator(ring, positions): @@ -470,36 +390,29 @@ def transverse_map_sector_generator(ring, positions): if ring.optics.use_local_values: for i in range(N_sec): sectors.append( - TransverseMapSector( - ring, - ring.optics.local_alpha, - ring.optics.local_beta, - ring.optics.local_dispersion, - ring.optics.local_alpha, - ring.optics.local_beta, - ring.optics.local_dispersion, - ring.tune / N_sec, - ring.chro / N_sec, - )) + TransverseMapSector(ring, + ring.optics.local_alpha, + ring.optics.local_beta, + ring.optics.local_dispersion, + ring.optics.local_alpha, + ring.optics.local_beta, + ring.optics.local_dispersion, + 2 * np.pi * ring.tune / N_sec, + ring.chro / N_sec, + adts=ring.adts / + N_sec if ring.adts else None)) else: import at - def _compute_chro(ring, pos, dp=1e-4): + def _compute_chro(ring, pos, dp=1e-2, order=4): lat = deepcopy(ring.optics.lattice) lat.append(at.Marker("END")) - N = len(lat) - refpts = np.arange(N) - (*elem_neg_dp, ) = at.linopt2(lat, refpts=refpts, dp=-dp) - (*elem_pos_dp, ) = at.linopt2(lat, refpts=refpts, dp=dp) - - s = elem_neg_dp[2]["s_pos"] - mux0 = elem_neg_dp[2]["mu"][:, 0] - mux1 = elem_pos_dp[2]["mu"][:, 0] - muy0 = elem_neg_dp[2]["mu"][:, 1] - muy1 = elem_pos_dp[2]["mu"][:, 1] - - Chrox = (mux1-mux0) / (2*dp) / 2 / np.pi - Chroy = (muy1-muy0) / (2*dp) / 2 / np.pi + fit, dpa, tune = at.physics.nonlinear.chromaticity(lat, + method='linopt', + dpm=dp, + n_points=100, + order=order) + Chrox, Chroy = fit[0, 1:], fit[1, 1:] chrox = np.interp(pos, s, Chrox) chroy = np.interp(pos, s, Chroy) diff --git a/mbtrack2/tracking/synchrotron.py b/mbtrack2/tracking/synchrotron.py index 15235ecd1b60f38a4984074ab71c87c1fba54a85..41458eff38988e909c1f3c90b2d4863c0a3e12ae 100644 --- a/mbtrack2/tracking/synchrotron.py +++ b/mbtrack2/tracking/synchrotron.py @@ -109,6 +109,9 @@ class Synchrotron: get_adts() Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar componenet from AT lattice. + get_chroma() + Compute chromaticity (linear and nonlinear) from AT lattice and + update the property. get_mcf_order() Compute momentum compaction factor up to 3rd order from AT lattice. get_longitudinal_twiss(V) @@ -387,6 +390,25 @@ class Synchrotron: coef_yy = np.array([det.A3 / 2, 0]) self.adts = [coef_xx, coef_yx, coef_xy, coef_yy] + def get_chroma(self, order=4, dpm=0.02, n_points=100): + """ + Compute chromaticity (linear and nonlinear) from AT lattice and update the property. + """ + import at + if self.optics.use_local_values: + raise ValueError("Value needs to be provided manualy as no AT" + + " lattice file is loaded.") + fit, dpa, tune = at.physics.nonlinear.chromaticity(self.optics.lattice, + method='linopt', + dpm=dpm, + n_points=n_points, + order=order) + chrox, chroy = fit + self.chro = [ + elem for pair in zip(chrox[1:], chroy[1:]) for elem in pair + ] + return chrox[1:], chroy[1:] + def get_mcf_order(self, add=True, show_fit=False): """ Compute momentum compaction factor up to 3rd order from AT lattice. @@ -408,7 +430,7 @@ class Synchrotron: """ import at if self.optics.use_local_values: - raise ValueError("ADTS needs to be provided manualy as no AT" + + raise ValueError("Value needs to be provided manualy as no AT" + " lattice file is loaded.") deltamin = -1e-4 deltamax = 1e-4