diff --git a/tracking/__init__.py b/tracking/__init__.py
index 41f06b9ee215b768e0f5cd99e1b7308eaf246a22..05db3c85cc9817ec1f16279adae27c34bc6167ff 100644
--- a/tracking/__init__.py
+++ b/tracking/__init__.py
@@ -12,7 +12,8 @@ from mbtrack2.tracking.rf import RFCavity, CavityResonator
from mbtrack2.tracking.parallel import Mpi
from mbtrack2.tracking.optics import Optics, PhyisicalModel
from mbtrack2.tracking.element import (Element, LongitudinalMap,
- TransverseMap, SynchrotronRadiation)
+ TransverseMap, SynchrotronRadiation,
+ SkewQuadrupole)
from mbtrack2.tracking.aperture import (CircularAperture, ElipticalAperture,
RectangularAperture)
from mbtrack2.tracking.wakepotential import WakePotential
diff --git a/tracking/element.py b/tracking/element.py
index 0ef071c4bbde1c90c9be9c9727277017454af4da..2231c04f4a964b3bd753a2510e79d9cb50395ca2 100644
--- a/tracking/element.py
+++ b/tracking/element.py
@@ -128,19 +128,14 @@ class SynchrotronRadiation(Element):
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),2))
- bunch["x"] += self.ring.sigma()[0]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,0]
- bunch["xp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["xp"]
- bunch["xp"] += self.ring.sigma()[1]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,1]
-
+ rand = np.random.normal(size=len(bunch))
+ 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),2))
- bunch["y"] += self.ring.sigma()[2]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,0]
- bunch["yp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["yp"]
- bunch["yp"] += self.ring.sigma()[3]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,1]
-
- # Reset energy change to 0 for next turn
- bunch.energy_change = 0
+ rand = np.random.normal(size=len(bunch))
+ 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):
"""
@@ -205,3 +200,33 @@ class TransverseMap(Element):
bunch["xp"] = xp
bunch["y"] = y
bunch["yp"] = yp
+
+class SkewQuadrupole:
+ """
+ Thin skew quadrupole element used to introduce betatron coupling (the
+ length of the quadrupole is neglected).
+
+ Parameters
+ ----------
+ strength : float
+ Integrated strength of the skew quadrupole [m].
+
+ """
+ def __init__(self, strength):
+ self.strength = strength
+
+ @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
+ """
+
+ bunch['xp'] = bunch['xp'] - self.strength * bunch['y']
+ bunch['yp'] = bunch['yp'] - self.strength * bunch['x']
+
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 6a46e38322048d2005375e45d55911925de4fae3..b682523003d80034a620df6e7317a2e4f6da5eab 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -230,7 +230,8 @@ class Monitor(Element, metaclass=ABCMeta):
class BunchMonitor(Monitor):
"""
- Monitor a single bunch and save attributes (mean, std, emit and current).
+ Monitor a single bunch and save attributes
+ (mean, std, emit, current, and action).
Parameters
----------
@@ -267,9 +268,11 @@ class BunchMonitor(Monitor):
self.bunch_number = bunch_number
group_name = "BunchData_" + str(self.bunch_number)
dict_buffer = {"mean":(6, buffer_size), "std":(6, buffer_size),
- "emit":(3, buffer_size), "current":(buffer_size,)}
+ "emit":(3, buffer_size), "current":(buffer_size,),
+ "cs_invariant":(2, buffer_size)}
dict_file = {"mean":(6, total_size), "std":(6, total_size),
- "emit":(3, total_size), "current":(total_size,)}
+ "emit":(3, total_size), "current":(total_size,),
+ "cs_invariant":(2, total_size)}
self.monitor_init(group_name, save_every, buffer_size, total_size,
dict_buffer, dict_file, file_name, mpi_mode)
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 63ed6bfc6efe4510a10a0c6b45c5c7e09ea6ff89..a457a39865874290f81666bfef337107fb2003a8 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -112,7 +112,7 @@ def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
file.close()
return fig
-def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
+def plot_bunchdata(filename, bunch_number, dataset, dimension="x", x_var="time"):
"""
Plot data recorded by BunchMonitor.
@@ -123,13 +123,14 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
bunch_number : int
Bunch to plot. This has to be identical to 'bunch_number' parameter in
'BunchMonitor' object.
- detaset : {"current","emit","mean","std"}
+ dataset : {"current", "emit", "mean", "std", "cs_invariant"}
HDF5 file's dataset to be plotted.
- option : str, optional
- If dataset is "emit", "mean", or "std", the variable name to be plotted
- needs to be specified :
- for "emit", option = {"x","y","s"}
- for "mean" and "std", option = {"x","xp","y","yp","tau","delta"}
+ dimension : str, optional
+ The dimension of the dataset to plot. Use "None" for "current",
+ otherwise use the following :
+ for "emit", dimension = {"x","y","s"},
+ for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"},
+ for "action", dimension = {"x","y"}.
x_var : {"time", "current"}, optional
Variable to be plotted on the horizontal axis. The default is "time".
@@ -149,19 +150,19 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
label = "current (mA)"
elif dataset == "emit":
- option_dict = {"x":0, "y":1, "s":2}
+ dimension_dict = {"x":0, "y":1, "s":2}
- y_var = file[group][dataset][option_dict[option]]*1e9
+ y_var = file[group][dataset][dimension_dict[dimension]]*1e9
- if option == "x": label = "hor. emittance (nm.rad)"
- elif option == "y": label = "ver. emittance (nm.rad)"
- elif option == "s": label = "long. emittance (nm.rad)"
+ if dimension == "x": label = "hor. emittance (nm.rad)"
+ elif dimension == "y": label = "ver. emittance (nm.rad)"
+ elif dimension == "s": label = "long. emittance (nm.rad)"
- elif dataset == "mean" or "std":
- option_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+ elif dataset == "mean" or dataset == "std":
+ dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
- axis_index = option_dict[option]
+ axis_index = dimension_dict[dimension]
y_var = file[group][dataset][axis_index]*scale[axis_index]
if dataset == "mean":
@@ -174,6 +175,13 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
label = label_list[axis_index]
+ elif dataset == "cs_invariant":
+ dimension_dict = {"x":0, "y":1}
+ axis_index = dimension_dict[dimension]
+ y_var = file[group][dataset][axis_index]
+ label_list = ['$J_x$ (m)', '$J_y$ (m)']
+ label = label_list[axis_index]
+
if x_var == "current":
x_axis = file[group]["current"][:] * 1e3
xlabel = "current (mA)"
@@ -544,5 +552,4 @@ def plot_tunedata(filename, bunch_number):
ax2.set_ylabel("Synchrotron tune")
file.close()
- return (fig1, fig2)
-
\ No newline at end of file
+ return (fig1, fig2)
\ No newline at end of file
diff --git a/tracking/particles.py b/tracking/particles.py
index 87988e833ef880848155d12dd1764e40cb654939..f662f679d21b92a0b303f25379d19cb305108c94 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -233,13 +233,17 @@ class Bunch:
@property
def cs_invariant(self):
- Jx = (self.ring.optics.local_gamma * np.mean(self['x']**2)) + \
- (2*self.ring.optics.local_alpha * np.mean(self['x']*self['xp'])) + \
- (self.ring.optics.local_beta * np.mean(self['xp']**2))
- Jy = (self.ring.optics.local_gamma * np.mean(self['y']**2)) + \
- (2*self.ring.optics.local_alpha * np.mean(self['y']*self['yp'])) + \
- (self.ring.optics.local_beta * np.mean(self['yp']**2))
- return np.array((Jx,Jy))
+ """
+ Return the average Courant-Snyder invariant of each plane.
+
+ """
+ Jx = (self.ring.optics.local_gamma[0] * self['x']**2) + \
+ (2*self.ring.optics.local_alpha[0] * self['x'])*self['xp'] + \
+ (self.ring.optics.local_beta[0] * self['xp']**2)
+ Jy = (self.ring.optics.local_gamma[1] * self['y']**2) + \
+ (2*self.ring.optics.local_alpha[1] * self['y']*self['yp']) + \
+ (self.ring.optics.local_beta[1] * self['yp']**2)
+ return np.array((np.mean(Jx),np.mean(Jy)))
@property
def energy_change(self):
diff --git a/tracking/synchrotron.py b/tracking/synchrotron.py
index f5e2d69a6f5f65e42229ff9f41e5c604a6a24b68..f3842c1e72dc656702f20ca3ad6a346b73259144 100644
--- a/tracking/synchrotron.py
+++ b/tracking/synchrotron.py
@@ -244,11 +244,11 @@ class Synchrotron:
sigma = np.zeros((4,))
sigma[0] = (self.emit[0]*self.optics.local_beta[0] +
self.optics.local_dispersion[0]**2*self.sigma_delta)**0.5
- sigma[1] = (self.emit[0]*self.optics.local_alpha[0] +
+ sigma[1] = (self.emit[0]*self.optics.local_gamma[0] +
self.optics.local_dispersion[1]**2*self.sigma_delta)**0.5
sigma[2] = (self.emit[1]*self.optics.local_beta[1] +
self.optics.local_dispersion[2]**2*self.sigma_delta)**0.5
- sigma[3] = (self.emit[1]*self.optics.local_alpha[1] +
+ sigma[3] = (self.emit[1]*self.optics.local_gamma[1] +
self.optics.local_dispersion[3]**2*self.sigma_delta)**0.5
else:
if isinstance(position, (float, int)):
@@ -258,11 +258,11 @@ class Synchrotron:
sigma = np.zeros((4, n))
sigma[0,:] = (self.emit[0]*self.optics.beta(position)[0] +
self.optics.dispersion(position)[0]**2*self.sigma_delta)**0.5
- sigma[1,:] = (self.emit[0]*self.optics.alpha(position)[0] +
+ sigma[1,:] = (self.emit[0]*self.optics.gamma(position)[0] +
self.optics.dispersion(position)[1]**2*self.sigma_delta)**0.5
sigma[2,:] = (self.emit[1]*self.optics.beta(position)[1] +
self.optics.dispersion(position)[2]**2*self.sigma_delta)**0.5
- sigma[3,:] = (self.emit[1]*self.optics.alpha(position)[1] +
+ sigma[3,:] = (self.emit[1]*self.optics.gamma(position)[1] +
self.optics.dispersion(position)[3]**2*self.sigma_delta)**0.5
return sigma