Skip to content
Snippets Groups Projects
Commit 65908062 authored by Vadim GUBAIDULIN's avatar Vadim GUBAIDULIN
Browse files

Merge branch 'feature-nonlinear-chromaticity' into 'develop' 

Nonlinear chromaticity

See merge request !16
parents 9543045d 6233924f
No related branches found
No related tags found
2 merge requests!280.8.0,!16Nonlinear chromaticity
Show whitespace changes
Inline Side-by-side
mbtrack2/tracking/element.py
+ 117
− 204
View file @ 65908062
...@@ -10,6 +10,7 @@ from copy import deepcopy ...@@ -10,6 +10,7 @@ from copy import deepcopy
from functools import wraps from functools import wraps
import numpy as np import numpy as np
from scipy.special import factorial
from mbtrack2.tracking.particles import Beam from mbtrack2.tracking.particles import Beam
...@@ -128,128 +129,26 @@ class SynchrotronRadiation(Element): ...@@ -128,128 +129,26 @@ class SynchrotronRadiation(Element):
---------- ----------
bunch : Bunch or Beam object bunch : Bunch or Beam object
""" """
N = len(bunch)
if self.switch[0] == True: 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[ bunch["delta"] = (1 - 2 * self.ring.T0 / self.ring.tau[2]) * bunch[
"delta"] + 2 * self.ring.sigma_delta * ( "delta"] + 2 * self.ring.sigma_delta * (
self.ring.T0 / self.ring.tau[2])**0.5 * rand self.ring.T0 / self.ring.tau[2])**0.5 * rand
if self.switch[1] == True: 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"] = (1 - 2 * self.ring.T0 / self.ring.tau[0]
) * bunch["xp"] + 2 * self.ring.sigma()[1] * ( ) * bunch["xp"] + 2 * self.ring.sigma()[1] * (
self.ring.T0 / self.ring.tau[0])**0.5 * rand self.ring.T0 / self.ring.tau[0])**0.5 * rand
if self.switch[2] == True: 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"] = (1 - 2 * self.ring.T0 / self.ring.tau[1]
) * bunch["yp"] + 2 * self.ring.sigma()[3] * ( ) * bunch["yp"] + 2 * self.ring.sigma()[3] * (
self.ring.T0 / self.ring.tau[1])**0.5 * rand 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: class SkewQuadrupole:
""" """
Thin skew quadrupole element used to introduce betatron coupling (the Thin skew quadrupole element used to introduce betatron coupling (the
...@@ -346,6 +245,67 @@ class TransverseMapSector(Element): ...@@ -346,6 +245,67 @@ class TransverseMapSector(Element):
else: else:
self.adts_poly = None 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 @Element.parallel
def track(self, bunch): def track(self, bunch):
""" """
...@@ -357,85 +317,45 @@ class TransverseMapSector(Element): ...@@ -357,85 +317,45 @@ class TransverseMapSector(Element):
---------- ----------
bunch : Bunch or Beam object bunch : Bunch or Beam object
""" """
tune_advance_x = self.tune_diff[0]
# Compute phase advance which depends on energy via chromaticity and ADTS tune_advance_y = self.tune_diff[1]
if self.adts_poly is None: # Compute tune advance which depends on energy via chromaticity and ADTS
phase_advance_x = ( if (np.array(self.chro_diff) != 0).any():
2 * np.pi * tune_advance_x_chro, tune_advance_y_chro = self._compute_chromatic_tune_advances(
(self.tune_diff[0] + self.chro_diff[0] * bunch["delta"])) bunch)
phase_advance_y = ( tune_advance_x += tune_advance_x_chro
2 * np.pi * tune_advance_y += tune_advance_y_chro
(self.tune_diff[1] + self.chro_diff[1] * bunch["delta"]))
else: if self.adts_poly is not None:
Jx = ((self.gamma0[0] * bunch["x"]**2) + Jx = ((self.gamma0[0] * bunch["x"]**2) +
(2 * self.alpha0[0] * bunch["x"] * self["xp"]) + (2 * self.alpha0[0] * bunch["x"] * self["xp"]) +
(self.beta0[0] * bunch["xp"]**2)) (self.beta0[0] * bunch["xp"]**2))
Jy = ((self.gamma0[1] * bunch["y"]**2) + Jy = ((self.gamma0[1] * bunch["y"]**2) +
(2 * self.alpha0[1] * bunch["y"] * bunch["yp"]) + (2 * self.alpha0[1] * bunch["y"] * bunch["yp"]) +
(self.beta0[1] * bunch["yp"]**2)) (self.beta0[1] * bunch["yp"]**2))
phase_advance_x = ( tune_advance_x += (self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
2 * np.pi * tune_advance_x += (self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
(self.tune_diff[0] + self.chro_diff[0] * bunch["delta"] +
self.adts_poly[0](Jx) + self.adts_poly[2](Jy))) bunch['x'], bunch['xp'] = self._compute_new_coords(
phase_advance_y = ( bunch, tune_advance_x, 'x')
2 * np.pi * bunch['y'], bunch['yp'] = self._compute_new_coords(
(self.tune_diff[1] + self.chro_diff[1] * bunch["delta"] + bunch, tune_advance_y, 'y')
self.adts_poly[1](Jx) + self.adts_poly[3](Jy)))
# 6x6 matrix corresponding to (x, xp, delta, y, yp, delta) class TransverseMap(TransverseMapSector):
matrix = np.zeros((6, 6, len(bunch))) """
Transverse map for a single turn in the synchrotron.
# Horizontal
matrix[0, 0, :] = np.sqrt(self.beta1[0] / self.beta0[0]) * ( Parameters
np.cos(phase_advance_x) + self.alpha0[0] * np.sin(phase_advance_x)) ----------
matrix[0, 1, :] = np.sqrt( ring : Synchrotron object
self.beta0[0] * self.beta1[0]) * np.sin(phase_advance_x) """
matrix[0, 2, :] = (self.dispersion1[0] -
matrix[0, 0, :] * self.dispersion0[0] - def __init__(self, ring):
matrix[0, 1, :] * self.dispersion0[1]) super().__init__(ring, ring.optics.local_alpha, ring.optics.local_beta,
matrix[1, 0, :] = ( ring.optics.local_dispersion, ring.optics.local_alpha,
(self.alpha0[0] - self.alpha1[0]) * np.cos(phase_advance_x) - ring.optics.local_beta, ring.optics.local_dispersion,
(1 + self.alpha0[0] * self.alpha1[0]) * 2 * np.pi * ring.tune, ring.chro, ring.adts)
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
def transverse_map_sector_generator(ring, positions): def transverse_map_sector_generator(ring, positions):
...@@ -470,36 +390,29 @@ def transverse_map_sector_generator(ring, positions): ...@@ -470,36 +390,29 @@ def transverse_map_sector_generator(ring, positions):
if ring.optics.use_local_values: if ring.optics.use_local_values:
for i in range(N_sec): for i in range(N_sec):
sectors.append( sectors.append(
TransverseMapSector( TransverseMapSector(ring,
ring,
ring.optics.local_alpha, ring.optics.local_alpha,
ring.optics.local_beta, ring.optics.local_beta,
ring.optics.local_dispersion, ring.optics.local_dispersion,
ring.optics.local_alpha, ring.optics.local_alpha,
ring.optics.local_beta, ring.optics.local_beta,
ring.optics.local_dispersion, ring.optics.local_dispersion,
ring.tune / N_sec, 2 * np.pi * ring.tune / N_sec,
ring.chro / N_sec, ring.chro / N_sec,
)) adts=ring.adts /
N_sec if ring.adts else None))
else: else:
import at 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 = deepcopy(ring.optics.lattice)
lat.append(at.Marker("END")) lat.append(at.Marker("END"))
N = len(lat) fit, dpa, tune = at.physics.nonlinear.chromaticity(lat,
refpts = np.arange(N) method='linopt',
(*elem_neg_dp, ) = at.linopt2(lat, refpts=refpts, dp=-dp) dpm=dp,
(*elem_pos_dp, ) = at.linopt2(lat, refpts=refpts, dp=dp) n_points=100,
order=order)
s = elem_neg_dp[2]["s_pos"] Chrox, Chroy = fit[0, 1:], fit[1, 1:]
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
chrox = np.interp(pos, s, Chrox) chrox = np.interp(pos, s, Chrox)
chroy = np.interp(pos, s, Chroy) chroy = np.interp(pos, s, Chroy)
......
mbtrack2/tracking/synchrotron.py
+ 23
− 1
View file @ 65908062
...@@ -109,6 +109,9 @@ class Synchrotron: ...@@ -109,6 +109,9 @@ class Synchrotron:
get_adts() get_adts()
Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar
componenet from AT lattice. componenet from AT lattice.
get_chroma()
Compute chromaticity (linear and nonlinear) from AT lattice and
update the property.
get_mcf_order() get_mcf_order()
Compute momentum compaction factor up to 3rd order from AT lattice. Compute momentum compaction factor up to 3rd order from AT lattice.
get_longitudinal_twiss(V) get_longitudinal_twiss(V)
...@@ -387,6 +390,25 @@ class Synchrotron: ...@@ -387,6 +390,25 @@ class Synchrotron:
coef_yy = np.array([det.A3 / 2, 0]) coef_yy = np.array([det.A3 / 2, 0])
self.adts = [coef_xx, coef_yx, coef_xy, coef_yy] 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): def get_mcf_order(self, add=True, show_fit=False):
""" """
Compute momentum compaction factor up to 3rd order from AT lattice. Compute momentum compaction factor up to 3rd order from AT lattice.
...@@ -408,7 +430,7 @@ class Synchrotron: ...@@ -408,7 +430,7 @@ class Synchrotron:
""" """
import at import at
if self.optics.use_local_values: 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.") " lattice file is loaded.")
deltamin = -1e-4 deltamin = -1e-4
deltamax = 1e-4 deltamax = 1e-4
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment

The place to be for your source code