From 28282fa264a23b2c0f64c89ebeb21ee1da96bb40 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Thu, 31 Oct 2024 10:21:30 +0100 Subject: [PATCH 01/21] Survey testing with MadPoint --- examples/survey/000_dev.py | 58 +++++++++++++++ examples/survey/_madpoint.py | 137 +++++++++++++++++++++++++++++++++++ 2 files changed, 195 insertions(+) create mode 100644 examples/survey/000_dev.py create mode 100644 examples/survey/_madpoint.py diff --git a/examples/survey/000_dev.py b/examples/survey/000_dev.py new file mode 100644 index 000000000..4c1104dc4 --- /dev/null +++ b/examples/survey/000_dev.py @@ -0,0 +1,58 @@ +import xtrack as xt +import numpy as np +import matplotlib.pyplot as plt + +from _madpoint import MadPoint + +env = xt.Environment() + +env.particle_ref = xt.Particles(p0c = 1E9) +env['k0'] = 0.00 +env['h'] = 1E-3 + +line = env.new_line(name = 'line', components=[ + env.new('bend', xt.Bend, k0 = 'k0', h = 'h', length = 0.5, at=2), + env.new('xyshift', xt.XYShift, dx = 0.1, dy=0.2), + env.new('end', xt.Marker, at = 5)]) + +line.configure_bend_model(core = 'bend-kick-bend', edge = 'suppressed') +line.cut_at_s(np.linspace(0, line.get_length(), 101)) + +sv = line.survey() +tw = line.twiss4d(betx = 1, bety =1) + +madpoints = [] +xx = [] +yy = [] +zz = [] +for nn in tw.name: + madpoints.append( + MadPoint(name = nn, xsuite_twiss = tw, xsuite_survey = sv)) + xx.append(madpoints[-1].p[0]) + yy.append(madpoints[-1].p[1]) + zz.append(madpoints[-1].p[2]) + +xx = np.array(xx) +yy = np.array(yy) +zz = np.array(zz) + +sv['xx'] = xx +sv['yy'] = yy +sv['zz'] = zz + + +fig, axs = plt.subplots(3, 1, sharex=True) +axs[0].plot(sv.Z, sv.X) +axs[1].plot(tw.s, tw.x) +axs[2].plot(sv.s, sv.xx) + +axs[0].set_xlabel('Z [m]') +axs[0].set_ylabel('X [m]') + +axs[1].set_xlabel('s [m]') +axs[1].set_ylabel('x [m]') + +axs[2].set_xlabel('s [m]') +axs[2].set_ylabel('xx [m]') + +plt.show() \ No newline at end of file diff --git a/examples/survey/_madpoint.py b/examples/survey/_madpoint.py new file mode 100644 index 000000000..56437c86a --- /dev/null +++ b/examples/survey/_madpoint.py @@ -0,0 +1,137 @@ +import numpy as np + +class MadPoint(object): + @classmethod + def from_survey(cls, name, mad=None, xsuite_survey=None): + return cls(name, mad=mad, use_twiss=False, use_survey=True, + xsuite_survey=xsuite_survey) + + @classmethod + def from_twiss(cls, name, mad): + return cls(name, mad, use_twiss=True, use_survey=False) + + def __init__(self, name, mad=None, use_twiss=True, use_survey=True, + xsuite_twiss=None, xsuite_survey=None): + + self.use_twiss = use_twiss + self.use_survey = use_survey + + if not (use_survey) and not (use_twiss): + raise ValueError( + "use_survey and use_twiss cannot be False at the same time" + ) + + self.tx = None + self.ty = None + self.tpx = None + self.tpy = None + + self.sx = None + self.sy = None + self.sz = None + self.sp = None + theta = 0.0 + phi = 0.0 + psi = 0.0 + + if mad is not None: + + self.name = name + if use_twiss: + assert xsuite_survey is None + twiss = mad.table.twiss + names = twiss.name + if use_survey: + assert xsuite_twiss is None + survey = mad.table.survey + names = survey.name + # patch for this issue https://github.com/hibtc/cpymad/issues/91 + for ii, nn in enumerate(names): + if not nn.endswith(':1'): + names[ii] = nn+':1' + + idx = np.where(names == name)[0][0] + + if use_twiss: + self.tx = twiss.x[idx] + self.ty = twiss.y[idx] + self.tpx = twiss.px[idx] + self.tpy = twiss.py[idx] + + if use_survey: + self.sx = survey.x[idx] + self.sy = survey.y[idx] + self.sz = survey.z[idx] + self.sp = np.array([self.sx, self.sy, self.sz]) + theta = survey.theta[idx] + phi = survey.phi[idx] + psi = survey.psi[idx] + else: + + + if use_twiss: + assert xsuite_twiss is not None + idx = np.where(np.array(xsuite_twiss['name']) == name)[0][0] + self.tx = xsuite_twiss.x[idx] + self.ty = xsuite_twiss.y[idx] + self.tpx = xsuite_twiss.px[idx] + self.tpy = xsuite_twiss.py[idx] + + if use_survey: + assert xsuite_survey is not None + idx = np.where(np.array(xsuite_survey['name']) == name)[0][0] + self.sx = xsuite_survey.X[idx] + self.sy = xsuite_survey.Y[idx] + self.sz = xsuite_survey.Z[idx] + self.sp = np.array([self.sx, self.sy, self.sz]) + theta = xsuite_survey.theta[idx] + phi = xsuite_survey.phi[idx] + psi = xsuite_survey.psi[idx] + + thetam = np.array( + [ + [np.cos(theta), 0, np.sin(theta)], + [0, 1, 0], + [-np.sin(theta), 0, np.cos(theta)], + ] + ) + phim = np.array( + [ + [1, 0, 0], + [0, np.cos(phi), np.sin(phi)], + [0, -np.sin(phi), np.cos(phi)], + ] + ) + psim = np.array( + [ + [np.cos(psi), -np.sin(psi), 0], + [np.sin(psi), np.cos(psi), 0], + [0, 0, 1], + ] + ) + wm = np.dot(thetam, np.dot(phim, psim)) + self.ex = np.dot(wm, np.array([1, 0, 0])) + self.ey = np.dot(wm, np.array([0, 1, 0])) + self.ez = np.dot(wm, np.array([0, 0, 1])) + + self.p = np.array([0.0, 0.0, 0.0]) + + if use_twiss: + self.p += self.ex * self.tx + self.ey * self.ty + + if use_survey: + self.p += self.sp + + def shift_survey(self, delta): + self.sx -= delta[0] + self.sy -= delta[1] + self.sz -= delta[2] + self.sp -= delta + self.p -= delta + + def dist(self, other): + return np.sqrt(np.sum((self.p - other.p) ** 2)) + + def distxy(self, other): + dd = self.p - other.p + return np.dot(dd, self.ex), np.dot(dd, self.ey) From 77ae7cac8023584b804cf1e8bdcacefa48027075 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Thu, 31 Oct 2024 10:21:42 +0100 Subject: [PATCH 02/21] Addition of XYShift to allowed element types in new --- xtrack/line.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index 4e68a5d2c..cb052c28c 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -56,13 +56,14 @@ _ALLOWED_ELEMENT_TYPES_IN_NEW = [xt.Drift, xt.Bend, xt.Quadrupole, xt.Sextupole, xt.Octupole, xt.Cavity, xt.Multipole, xt.Solenoid, - xt.Marker, xt.Replica] + xt.Marker, xt.Replica, xt.XYShift] _ALLOWED_ELEMENT_TYPES_DICT = {'Drift': xt.Drift, 'Bend': xt.Bend, 'Quadrupole': xt.Quadrupole, 'Sextupole': xt.Sextupole, 'Octupole': xt.Octupole, 'Cavity': xt.Cavity, 'Multipole': xt.Multipole, 'Solenoid': xt.Solenoid, - 'Marker': xt.Marker, 'Replica': xt.Replica} + 'Marker': xt.Marker, 'Replica': xt.Replica, + 'XYShift': xt.XYShift} _STR_ALLOWED_ELEMENT_TYPES_IN_NEW = ', '.join([tt.__name__ for tt in _ALLOWED_ELEMENT_TYPES_IN_NEW]) From 11037bebf6d70d49f40dc2c92d72597f7d21bafb Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Thu, 31 Oct 2024 11:10:15 +0100 Subject: [PATCH 03/21] Added handling of XY shifts --- xtrack/survey.py | 42 +++++++++++++++++++++++++++++++++++------- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/xtrack/survey.py b/xtrack/survey.py index 7d9e0d06c..cec43a386 100644 --- a/xtrack/survey.py +++ b/xtrack/survey.py @@ -71,8 +71,16 @@ def advance_drift(v, w, R): return np.dot(w, R) + v, w -def advance_element(v, w, length=0, angle=0, tilt=0): +def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0): """Computing the advance element-by-element. See MAD-X manual for generation of R and S""" + if dx != 0 or dy != 0: + assert angle == 0, "dx and dy are only supported for angle=0" + assert tilt == 0, "dx and dy are only supported for tilt=0" + assert length == 0, "dx and dy are only supported for length=0" + + v = v + np.array([dx, dy, 0]) + return v, w + if angle == 0: R = np.array([0, 0, length]) return advance_drift(v, w, R) @@ -133,6 +141,8 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, out_drift_length = list(self.drift_length[:-1][::-1]) out_angle = list(-self.angle[:-1][::-1]) out_tilt = list(-self.tilt[:-1][::-1]) + out_dx = list(-self.dx[:-1][::-1]) + out_dy = list(-self.dy[:-1][::-1]) out_name = list(self.name[:-1][::-1]) if type(element0) is str: @@ -141,7 +151,7 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, X, Y, Z, theta, phi, psi = compute_survey( X0, Y0, Z0, theta0, phi0, psi0, out_drift_length, out_angle, out_tilt, - element0=element0) + out_dx, out_dy, element0=element0) # Initializing dictionary out_columns = {} @@ -158,6 +168,8 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, out_columns['drift_length'] = np.array(out_drift_length + [0.]) out_columns['angle'] = np.array(out_angle + [0.]) out_columns['tilt'] = np.array(out_tilt + [0.]) + out_columns["dx"] = np.array(out_dx + [0.]) + out_columns["dy"] = np.array(out_dy + [0.]) out_scalars = {} out_scalars["element0"] = element0 @@ -218,12 +230,20 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, drift_length = tt.length drift_length[~tt.isthick] = 0 + # Extract xy shifts from elements + dx = tt._own_dx + dy = tt._own_dy + + # Handling of rotation elements + sin_angle = tt._own_sin_angle + print(sin_angle) + if type(element0) == str: element0 = line.element_names.index(element0) X, Y, Z, theta, phi, psi = compute_survey( X0, Y0, Z0, theta0, phi0, psi0, drift_length[:-1], angle[:-1], tilt[:-1], - element0=element0) + dx[:-1], dy[:-1], element0=element0) # Initializing dictionary out_columns = {} @@ -240,6 +260,8 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, out_columns['drift_length'] = drift_length out_columns['angle'] = angle out_columns['tilt'] = tilt + out_columns['dx'] = dx + out_columns['dy'] = dy out_scalars['element0'] = element0 @@ -251,23 +273,29 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, - element0=0, reverse_xs=False): + dx, dy, element0=0, reverse_xs=False): if element0 != 0: assert not(reverse_xs), "Not implemented yet" drift_forward = drift_length[element0:] angle_forward = angle[element0:] tilt_forward = tilt[element0:] + dx_forward = dx[element0:] + dy_forward = dy[element0:] (X_forward, Y_forward, Z_forward, theta_forward, phi_forward, psi_forward) = compute_survey(X0, Y0, Z0, theta0, phi0, psi0, - drift_forward, angle_forward, tilt_forward) + drift_forward, angle_forward, tilt_forward, + dx_forward, dy_forward) drift_backward = drift_length[:element0][::-1] angle_backward = -np.array(angle[:element0][::-1]) tilt_backward = -np.array(tilt[:element0][::-1]) + dx_backward = -np.array(dx[:element0][::-1]) + dy_backward = -np.array(dy[:element0][::-1]) (X_backward, Y_backward, Z_backward, theta_backward, phi_backward, psi_backward) = compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_backward, angle_backward, tilt_backward, + dx_backward, dy_backward, reverse_xs=True) X = np.array(X_backward[::-1][:-1] + X_forward) @@ -288,7 +316,7 @@ def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, w = get_w_from_angles(theta=theta0, phi=phi0, psi=psi0, reverse_xs=reverse_xs) # Advancing element by element - for ll, aa, tt in zip(drift_length, angle, tilt): + for ll, aa, tt, xx, yy in zip(drift_length, angle, tilt, dx, dy): th, ph, ps = get_angles_from_w(w, reverse_xs=reverse_xs) @@ -300,7 +328,7 @@ def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, psi.append(ps) # Advancing - v, w = advance_element(v, w, length=ll, angle=aa, tilt=tt) + v, w = advance_element(v, w, length=ll, angle=aa, tilt=tt, dx=xx, dy=yy) # Last marker th, ph, ps = get_angles_from_w(w, reverse_xs=reverse_xs) From c333053c807ecea542b8d849e768b7bdf78f92b1 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Thu, 31 Oct 2024 11:10:29 +0100 Subject: [PATCH 04/21] Allow use of XRotation and YRotation in new line --- xtrack/line.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index cb052c28c..3aac4dc34 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -56,14 +56,15 @@ _ALLOWED_ELEMENT_TYPES_IN_NEW = [xt.Drift, xt.Bend, xt.Quadrupole, xt.Sextupole, xt.Octupole, xt.Cavity, xt.Multipole, xt.Solenoid, - xt.Marker, xt.Replica, xt.XYShift] + xt.Marker, xt.Replica, xt.XYShift, xt.XRotation, xt. YRotation] _ALLOWED_ELEMENT_TYPES_DICT = {'Drift': xt.Drift, 'Bend': xt.Bend, 'Quadrupole': xt.Quadrupole, 'Sextupole': xt.Sextupole, 'Octupole': xt.Octupole, 'Cavity': xt.Cavity, 'Multipole': xt.Multipole, 'Solenoid': xt.Solenoid, 'Marker': xt.Marker, 'Replica': xt.Replica, - 'XYShift': xt.XYShift} + 'XYShift': xt.XYShift, + 'XRotation': xt.XRotation, 'YRotation': xt.YRotation} _STR_ALLOWED_ELEMENT_TYPES_IN_NEW = ', '.join([tt.__name__ for tt in _ALLOWED_ELEMENT_TYPES_IN_NEW]) @@ -4151,6 +4152,12 @@ def _get_attr_cache(self): '_own_k4sl': ('ksl', 4), '_own_k5sl': ('ksl', 5), + '_own_dx': 'dx', + '_own_dy': 'dy', + + '_own_sin_angle': 'sin_angle', + '_own_cos_angle': 'cos_angle', + '_parent_length': (('_parent', 'length'), None), '_parent_sin_rot_s': (('_parent', '_sin_rot_s'), None), '_parent_cos_rot_s': (('_parent', '_cos_rot_s'), None), @@ -4189,6 +4196,12 @@ def _get_attr_cache(self): '_parent_k4sl': (('_parent', 'ksl'), 4), '_parent_k5sl': (('_parent', 'ksl'), 5), + '_parent_dx': (('_parent', 'dx'), None), + '_parent_dy': (('_parent', 'dy'), None), + + '_parent_sin_angle': (('_parent', 'sin_angle'), None), + '_parent_cos_angle': (('_parent', 'cos_angle'), None), + }, derived_fields={ 'length': lambda attr: @@ -4266,6 +4279,10 @@ def _get_attr_cache(self): + attr['_parent_k5s'] * attr['_parent_length'] * attr['weight'] * attr._inherit_strengths), 'hkick': lambda attr: attr["angle_rad"] - attr["k0l"], 'vkick': lambda attr: attr["k0sl"], + 'dx': lambda attr: attr['_own_dx'] + attr['_parent_dx'], + 'dy': lambda attr: attr['_own_dy'] + attr['_parent_dy'], + 'sin_angle': lambda attr: attr['_own_sin_angle'] + attr['_parent_sin_angle'], + 'cos_angle': lambda attr: attr['_own_cos_angle'] + attr['_parent_cos_angle'], } ) return cache From 48a54494fea43f7f3b3c6e981f557a329d7be01a Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Thu, 31 Oct 2024 17:22:58 +0100 Subject: [PATCH 05/21] Helper functions Twiss and survey with madpoint Comparison plotting --- examples/survey/_helpers.py | 76 +++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 examples/survey/_helpers.py diff --git a/examples/survey/_helpers.py b/examples/survey/_helpers.py new file mode 100644 index 000000000..9b878b0a4 --- /dev/null +++ b/examples/survey/_helpers.py @@ -0,0 +1,76 @@ +""" +Helpers for testing the MadPoint class +""" +import numpy as np +import matplotlib.pyplot as plt +from _madpoint import MadPoint + +def survey_test(line): + """ + Produce twiss and survey for a line, including MadPoint + """ + survey = line.survey() + twiss = line.twiss4d(_continue_if_lost = True, betx = 1, bety = 1) + + madpoints = [] + xx = [] + yy = [] + zz = [] + for nn in twiss.name: + madpoints.append( + MadPoint(name = nn, xsuite_twiss = twiss, xsuite_survey = survey)) + xx.append(madpoints[-1].p[0]) + yy.append(madpoints[-1].p[1]) + zz.append(madpoints[-1].p[2]) + + survey['xx'] = np.array(xx) + survey['yy'] = np.array(yy) + survey['zz'] = np.array(zz) + + return survey, twiss + +def summary_plot(survey, twiss, title, zero_tol = 1E-12): + """ + Summary plot comparing survey, twiss and MadPoint + """ + def zero_small_values(arr, tol = zero_tol): + return np.where(np.abs(arr) < tol, 0, arr) + + fix, axs = plt.subplots(3, 2, figsize=(16, 8)) + axs[0, 0].plot( + zero_small_values(survey.Z), + zero_small_values(survey.X)) + axs[1, 0].plot( + zero_small_values(twiss.s), + zero_small_values(twiss.x)) + axs[2, 0].plot( + zero_small_values(survey.s), + zero_small_values(survey.xx)) + + axs[0, 1].plot( + zero_small_values(survey.Z), + zero_small_values(survey.Y)) + axs[1, 1].plot( + zero_small_values(twiss.s), + zero_small_values(twiss.y)) + axs[2, 1].plot( + zero_small_values(survey.s), + zero_small_values(survey.yy)) + + axs[0, 0].set_xlabel('Z [m]') + axs[0, 0].set_ylabel('X [m]') + axs[1, 0].set_xlabel('s [m]') + axs[1, 0].set_ylabel('x [m]') + axs[2, 0].set_xlabel('s [m]') + axs[2, 0].set_ylabel('xx [m]') + + axs[0, 1].set_xlabel('Z [m]') + axs[0, 1].set_ylabel('Y [m]') + axs[1, 1].set_xlabel('s [m]') + axs[1, 1].set_ylabel('y [m]') + axs[2, 1].set_xlabel('s [m]') + axs[2, 1].set_ylabel('yy [m]') + + fix.suptitle(title) + + plt.tight_layout() From 08e3bc1f97f72a5ff330543df825f293d49716fb Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Thu, 31 Oct 2024 17:23:11 +0100 Subject: [PATCH 06/21] Survey test examples --- examples/survey/001_bend.py | 217 +++++++++++++++++++++++++++++++ examples/survey/002_multipole.py | 108 +++++++++++++++ examples/survey/003_xyshift.py | 161 +++++++++++++++++++++++ examples/survey/004_xrotation.py | 77 +++++++++++ examples/survey/005_yrotation.py | 78 +++++++++++ 5 files changed, 641 insertions(+) create mode 100644 examples/survey/001_bend.py create mode 100644 examples/survey/002_multipole.py create mode 100644 examples/survey/003_xyshift.py create mode 100644 examples/survey/004_xrotation.py create mode 100644 examples/survey/005_yrotation.py diff --git a/examples/survey/001_bend.py b/examples/survey/001_bend.py new file mode 100644 index 000000000..caafab6b9 --- /dev/null +++ b/examples/survey/001_bend.py @@ -0,0 +1,217 @@ +""" +Test the survey on Bends +""" +################################################################################ +# Packages +################################################################################ +import numpy as np +import matplotlib.pyplot as plt + +from _helpers import survey_test, summary_plot +import xtrack as xt + +################################################################################ +# User variables +################################################################################ +PLOT_COMPARISONS = True +TOLERANCE = 1E-12 +REL_TOLERANCE = 1E-3 + +bend_h = 1E-3 + +################################################################################ +# Create line +################################################################################ +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('bend', xt.Bend, k0 = 0, h = 0, length = 0.5, at = 1), + env.new('end', xt.Marker, at = 2)]) +line.configure_bend_model(edge = 'suppressed') +line.config.XTRACK_USE_EXACT_DRIFTS = True + +################################################################################ +# Horizontal Bend +################################################################################ + +######################################## +# h = k0 = 0 +######################################## +line['bend'].k0 = 0 +line['bend'].h = 0 +line['bend'].rot_s_rad = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Horizontal Bend: h = k0 = 0') + +# Everything zero here +assert np.allclose(sv.X[-1], 0, atol=TOLERANCE) +assert np.allclose(sv.xx[-1], 0, atol=TOLERANCE) +assert np.allclose(tw.x[-1], 0, atol=TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# h = 0, k0 != 0 +######################################## +line['bend'].k0 = bend_h +line['bend'].h = 0 +line['bend'].rot_s_rad = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Horizontal Bend: h = 0, k0 != 0') + +# No survey with h = 0 +assert np.allclose(sv.X[-1], 0, atol=TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# h != 0, k0 = 0 +######################################## +line['bend'].k0 = 0 +line['bend'].h = bend_h +line['bend'].rot_s_rad = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Horizontal Bend: h != 0, k0 = 0') + +# Survey negative of Twiss for h != 0, k0 = 0 +assert np.allclose(sv.X[-1], -tw.x[-1], rtol=REL_TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# h != k0 !=0 +######################################## +line['bend'].k0 = bend_h * 2 +line['bend'].h = bend_h +line['bend'].rot_s_rad = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Horizontal Bend: h = k0 /2') + +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# h = k0 != 0 +######################################## +line['bend'].k0 = bend_h +line['bend'].h = bend_h +line['bend'].rot_s_rad = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Horizontal Bend: h = k0 != 0') + +# No orbit with h=k0 +assert np.allclose(tw.x[-1], 0, atol=TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +################################################################################ +# Vertical Bend +################################################################################ + +######################################## +# h = k0 = 0 +######################################## +line['bend'].k0 = 0 +line['bend'].h = 0 +line['bend'].rot_s_rad = np.pi / 2 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Vertical Bend: h = k0 = 0') + +# Everything zero here +assert np.allclose(sv.Y[-1], 0, atol=TOLERANCE) +assert np.allclose(sv.yy[-1], 0, atol=TOLERANCE) +assert np.allclose(tw.y[-1], 0, atol=TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) + +######################################## +# h = 0, k0 != 0 +######################################## +line['bend'].k0 = bend_h +line['bend'].h = 0 +line['bend'].rot_s_rad = np.pi / 2 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Vertical Bend: h = 0, k0 != 0') + +assert np.allclose(sv.Y[-1], 0, atol=TOLERANCE) +assert np.allclose(sv.yy[-1], tw.y[-1], rtol=REL_TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) + +######################################## +# h != 0, k0 = 0 +######################################## +line['bend'].k0 = 0 +line['bend'].h = bend_h +line['bend'].rot_s_rad = np.pi / 2 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Vertical Bend: h != 0, k0 = 0') + +# Survey negative of Twiss +assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=REL_TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) + +######################################## +# h != k0 !=0 +######################################## +line['bend'].k0 = bend_h * 2 +line['bend'].h = bend_h +line['bend'].rot_s_rad = np.pi / 2 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Vertical Bend: h = k0 /2') + +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) + +######################################## +# h = k0 != 0 +######################################## +line['bend'].k0 = bend_h +line['bend'].h = bend_h +line['bend'].rot_s_rad = np.pi / 2 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Vertical Bend: h = k0 != 0') + +assert np.allclose(tw.y[-1], 0, atol=TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) + +######################################## +# Show plots +######################################## +if PLOT_COMPARISONS: + plt.show() diff --git a/examples/survey/002_multipole.py b/examples/survey/002_multipole.py new file mode 100644 index 000000000..f3144289d --- /dev/null +++ b/examples/survey/002_multipole.py @@ -0,0 +1,108 @@ +""" +Test the survey on Bends +""" +################################################################################ +# Packages +################################################################################ +import numpy as np +import matplotlib.pyplot as plt + +from _helpers import survey_test, summary_plot +import xtrack as xt + +################################################################################ +# User variables +################################################################################ +PLOT_COMPARISONS = True +TOLERANCE = 1E-12 +REL_TOLERANCE = 1E-3 + +hxl = 1E-3 + +################################################################################ +# Create line +################################################################################ +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('mult', xt.Multipole, hxl = 0, length = 0.5, at = 1), + env.new('end', xt.Marker, at = 2)]) +line.configure_bend_model(edge = 'suppressed') +line.config.XTRACK_USE_EXACT_DRIFTS = True + +################################################################################ +# Horizontal +################################################################################ + +######################################## +# hxl = 0 +######################################## +line['mult'].hxl = 0 +line['mult'].rot_s_rad = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'hxl = 0') + +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# hxl != 0 +######################################## +line['mult'].hxl = hxl +line['mult'].rot_s_rad = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'hxl') + +# Survey negative of Twiss for h != 0, k0 = 0 +assert np.allclose(sv.X[-1], -tw.x[-1], rtol=REL_TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +################################################################################ +# Vertical +################################################################################ + +######################################## +# hyl = 0 +######################################## +line['mult'].hxl = 0 +line['mult'].rot_s_rad = np.pi / 2 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'hyl = 0') + +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# hyl != 0 +######################################## +line['mult'].hxl = hxl +line['mult'].rot_s_rad = np.pi / 2 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'hyl') + +# Survey negative of Twiss for h != 0, k0 = 0 +assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=REL_TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) + +######################################## +# Show plots +######################################## +if PLOT_COMPARISONS: + plt.show() diff --git a/examples/survey/003_xyshift.py b/examples/survey/003_xyshift.py new file mode 100644 index 000000000..a1e7a1e19 --- /dev/null +++ b/examples/survey/003_xyshift.py @@ -0,0 +1,161 @@ +""" +Test the survey on XY Shifts +""" +################################################################################ +# Packages +################################################################################ +import numpy as np +import matplotlib.pyplot as plt + +from _helpers import survey_test, summary_plot +import xtrack as xt + +################################################################################ +# User variables +################################################################################ +PLOT_COMPARISONS = True +TOLERANCE = 1E-12 +REL_TOLERANCE = 1E-3 + +################################################################################ +# Create line +################################################################################ +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('bend', xt.Bend, k0 = 0, h = 0, length = 0.5, at = 1), + env.new('xyshift', xt.XYShift, dx = 0, dy = 0, at = 2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +line.config.XTRACK_USE_EXACT_DRIFTS = True + +################################################################################ +# X Shift +################################################################################ + +######################################## +# Shift only +######################################## +line['bend'].k0 = 0 +line['bend'].h = 0 +line['bend'].rot_s_rad = 0 + +line['xyshift'].dx = 0.1 +line['xyshift'].dy = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Shift X, bend off') + +# Survey negative of Twiss for h != 0, k0 = 0 +assert np.allclose(sv.X[-1], -tw.x[-1], rtol=REL_TOLERANCE) +# MadPoint sum of twiss and survey +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Shift + horizontal bend +######################################## +line['bend'].k0 = 1 +line['bend'].h = 1 +line['bend'].rot_s_rad = 0 + +line['xyshift'].dx = 0.1 +line['xyshift'].dy = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Shift X, bend off') + +# MadPoint sum of twiss and survey +# TODO: This is broken? +# assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Shift + vertical bend +######################################## +line['bend'].k0 = 1 +line['bend'].h = 1 +line['bend'].rot_s_rad = np.pi / 2 + +line['xyshift'].dx = 0.1 +line['xyshift'].dy = 0 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Shift X, bend off') + +# MadPoint sum of twiss and survey +# TODO: This is broken? +# assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) + +################################################################################ +# Y Shift +################################################################################ + +######################################## +# Shift only +######################################## +line['bend'].k0 = 0 +line['bend'].h = 0 +line['bend'].rot_s_rad = 0 + +line['xyshift'].dx = 0 +line['xyshift'].dy = 0.1 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Shift X, bend off') + +# Survey negative of Twiss for h != 0, k0 = 0 +assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=REL_TOLERANCE) +# MadPoint sum of twiss and survey +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) + +######################################## +# Shift + horizontal bend +######################################## +line['bend'].k0 = 1 +line['bend'].h = 1 +line['bend'].rot_s_rad = 0 + +line['xyshift'].dx = 0 +line['xyshift'].dy = 0.1 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Shift X, bend off') + +# MadPoint sum of twiss and survey +# TODO: This is broken? +# assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) + +######################################## +# Shift + vertical bend +######################################## +line['bend'].k0 = 1 +line['bend'].h = 1 +line['bend'].rot_s_rad = np.pi / 2 + +line['xyshift'].dx = 0 +line['xyshift'].dy = 0.1 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Shift X, bend off') + +# MadPoint sum of twiss and survey +# TODO: This is broken? +# assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) + +######################################## +# Show plots +######################################## +if PLOT_COMPARISONS: + plt.show() diff --git a/examples/survey/004_xrotation.py b/examples/survey/004_xrotation.py new file mode 100644 index 000000000..c9a1107a9 --- /dev/null +++ b/examples/survey/004_xrotation.py @@ -0,0 +1,77 @@ +""" +Test the survey on XRotation +""" +################################################################################ +# Packages +################################################################################ +import numpy as np +import matplotlib.pyplot as plt + +from _helpers import survey_test, summary_plot +import xtrack as xt + +################################################################################ +# User variables +################################################################################ +PLOT_COMPARISONS = True +TOLERANCE = 1E-12 +REL_TOLERANCE = 1E-3 + +hxl = 1E-3 + +################################################################################ +# Create line +################################################################################ +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('mult', xt.Multipole, hxl = 0, length = 0.5, at = 1), + env.new('xrotation', xt.XRotation, angle = 0, at = 2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +line.config.XTRACK_USE_EXACT_DRIFTS = True + +################################################################################ +# Tests +################################################################################ + +######################################## +# Rotation only +######################################## +line['mult'].hxl = 0 +line['mult'].rot_s_rad = 0 + +line['xrotation'].angle = np.rad2deg(hxl) + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Rotation X, bend off') + +######################################## +# Shift + horizontal bend +######################################## +line['mult'].hxl = hxl +line['mult'].rot_s_rad = 0 + +line['xrotation'].angle = np.rad2deg(hxl) + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Rotation X, Horiztonal Mult') + +######################################## +# Shift + vertical bend +######################################## +line['mult'].hxl = hxl +line['mult'].rot_s_rad = np.pi / 2 + +line['xrotation'].angle = np.rad2deg(hxl) + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Rotation X, Vertical Mult') + +######################################## +# Show plots +######################################## +if PLOT_COMPARISONS: + plt.show() diff --git a/examples/survey/005_yrotation.py b/examples/survey/005_yrotation.py new file mode 100644 index 000000000..4c1c8a2e6 --- /dev/null +++ b/examples/survey/005_yrotation.py @@ -0,0 +1,78 @@ +""" +Test the survey on YRotation +""" +################################################################################ +# Packages +################################################################################ +import numpy as np +import matplotlib.pyplot as plt + +from _helpers import survey_test, summary_plot +import xtrack as xt + +################################################################################ +# User variables +################################################################################ +PLOT_COMPARISONS = True +TOLERANCE = 1E-12 +REL_TOLERANCE = 1E-3 + +################################################################################ +# Create line +################################################################################ +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('bend', xt.Bend, k0 = 0, h = 0, length = 0.5, at = 1), + env.new('yrotation', xt.YRotation, angle = 0, at = 2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +line.config.XTRACK_USE_EXACT_DRIFTS = True + +################################################################################ +# Tests +################################################################################ + +######################################## +# Rotation only +######################################## +line['bend'].k0 = 0 +line['bend'].h = 0 +line['bend'].rot_s_rad = 0 + +line['yrotation'].angle = 10 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Rotation Y, bend off') + +######################################## +# Shift + horizontal bend +######################################## +line['bend'].k0 = 1 +line['bend'].h = 1 +line['bend'].rot_s_rad = 0 + +line['yrotation'].angle = 10 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Rotation Y, Horiztonal Bend') + +######################################## +# Shift + vertical bend +######################################## +line['bend'].k0 = 1 +line['bend'].h = 1 +line['bend'].rot_s_rad = np.pi / 2 + +line['yrotation'].angle = 10 + +sv, tw = survey_test(line) +if PLOT_COMPARISONS: + summary_plot(sv, tw, 'Rotation Y, Vertical Bend') + +######################################## +# Show plots +######################################## +if PLOT_COMPARISONS: + plt.show() From a63da05045ee400d6c47c137e1c803012f4455d4 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Fri, 1 Nov 2024 16:19:27 +0100 Subject: [PATCH 07/21] Allowed S rotations Parameter name change to 'transform' to avoid confusion with element rotations --- xtrack/line.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index 3aac4dc34..6ea8dad46 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -56,15 +56,16 @@ _ALLOWED_ELEMENT_TYPES_IN_NEW = [xt.Drift, xt.Bend, xt.Quadrupole, xt.Sextupole, xt.Octupole, xt.Cavity, xt.Multipole, xt.Solenoid, - xt.Marker, xt.Replica, xt.XYShift, xt.XRotation, xt. YRotation] + xt.Marker, xt.Replica, xt.XYShift, xt.XRotation, + xt.YRotation, xt.SRotation] _ALLOWED_ELEMENT_TYPES_DICT = {'Drift': xt.Drift, 'Bend': xt.Bend, 'Quadrupole': xt.Quadrupole, 'Sextupole': xt.Sextupole, 'Octupole': xt.Octupole, 'Cavity': xt.Cavity, 'Multipole': xt.Multipole, 'Solenoid': xt.Solenoid, 'Marker': xt.Marker, 'Replica': xt.Replica, - 'XYShift': xt.XYShift, - 'XRotation': xt.XRotation, 'YRotation': xt.YRotation} + 'XYShift': xt.XYShift, 'XRotation': xt.XRotation, + 'YRotation': xt.YRotation, 'SRotation': xt.SRotation} _STR_ALLOWED_ELEMENT_TYPES_IN_NEW = ', '.join([tt.__name__ for tt in _ALLOWED_ELEMENT_TYPES_IN_NEW]) @@ -4158,6 +4159,10 @@ def _get_attr_cache(self): '_own_sin_angle': 'sin_angle', '_own_cos_angle': 'cos_angle', + # TODO: This should be corrected at the element level + '_own_sin_z': 'sin_z', + '_own_cos_z': 'cos_z', + '_parent_length': (('_parent', 'length'), None), '_parent_sin_rot_s': (('_parent', '_sin_rot_s'), None), '_parent_cos_rot_s': (('_parent', '_cos_rot_s'), None), @@ -4202,6 +4207,10 @@ def _get_attr_cache(self): '_parent_sin_angle': (('_parent', 'sin_angle'), None), '_parent_cos_angle': (('_parent', 'cos_angle'), None), + # TODO: This should be corrected at the element level + '_parent_sin_z': (('_parent', 'sin_z'), None), + '_parent_cos_z': (('_parent', 'cos_z'), None), + }, derived_fields={ 'length': lambda attr: @@ -4281,8 +4290,11 @@ def _get_attr_cache(self): 'vkick': lambda attr: attr["k0sl"], 'dx': lambda attr: attr['_own_dx'] + attr['_parent_dx'], 'dy': lambda attr: attr['_own_dy'] + attr['_parent_dy'], - 'sin_angle': lambda attr: attr['_own_sin_angle'] + attr['_parent_sin_angle'], - 'cos_angle': lambda attr: attr['_own_cos_angle'] + attr['_parent_cos_angle'], + 'transform_angle_rad': lambda attr: np.arctan2( + attr['_own_sin_angle'] + attr['_parent_sin_angle'] +\ + attr['_own_sin_z'] + attr['_parent_sin_z'], + attr['_own_cos_angle'] + attr['_parent_cos_angle'] +\ + attr['_own_cos_z'] + attr['_parent_cos_z']), } ) return cache From 350307f8d1302aa4faf5956a94b5430a2173f0b3 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Fri, 1 Nov 2024 16:23:59 +0100 Subject: [PATCH 08/21] Added explicit advance rotation function Handling of XYSRotations Renaming of variables --- xtrack/survey.py | 122 ++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 100 insertions(+), 22 deletions(-) diff --git a/xtrack/survey.py b/xtrack/survey.py index cec43a386..01b07aef1 100644 --- a/xtrack/survey.py +++ b/xtrack/survey.py @@ -64,6 +64,10 @@ def advance_bend(v, w, R, S): v2 = w1*R + v1 | w2 = w1*S""" return np.dot(w, R) + v, np.dot(w, S) +def advance_rotation(v, w, S): + """Advancing through rotation element: + Rotate w matrix according to transformation matrix S""" + return v, np.dot(w, S) def advance_drift(v, w, R): """Advancing through drift element, see MAD-X manual: @@ -71,19 +75,65 @@ def advance_drift(v, w, R): return np.dot(w, R) + v, w -def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0): - """Computing the advance element-by-element. See MAD-X manual for generation of R and S""" +def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, + transf_x_rad=0, transf_y_rad=0, transf_s_rad=0): + """Computing the advance element-by-element. + See MAD-X manual for generation of R and S""" + # XYShift Handling if dx != 0 or dy != 0: - assert angle == 0, "dx and dy are only supported for angle=0" - assert tilt == 0, "dx and dy are only supported for tilt=0" - assert length == 0, "dx and dy are only supported for length=0" + assert angle == 0, "dx and dy are only supported for angle = 0" + assert tilt == 0, "dx and dy are only supported for tilt = 0" + assert length == 0, "dx and dy are only supported for length = 0" - v = v + np.array([dx, dy, 0]) - return v, w + R = np.array([dx, dy, 0]) + # XYShift tarnsforms as a drift + return advance_drift(v, w, R) + + # XRotation Handling + if transf_x_rad != 0: + assert angle == 0, "rot_x_rad is only supported for angle = 0" + assert tilt == 0, "rot_x_rad is only supported for tilt = 0" + assert length == 0, "rot_x_rad is only supported for length = 0" + + # Relevant sine/cosine + cr = np.cos(-transf_x_rad) + sr = np.sin(-transf_x_rad) + # ------ + S = np.array([[1, 0, 0], [0, cr, sr], [0, -sr, cr]]) # x rotation matrix + return advance_rotation(v, w, S) + + # YRotation Handling + if transf_y_rad != 0: + assert angle == 0, "rot_y_rad is only supported for angle = 0" + assert tilt == 0, "rot_y_rad is only supported for tilt = 0" + assert length == 0, "rot_y_rad is only supported for length = 0" + + # Relevant sine/cosine + cr = np.cos(transf_y_rad) + sr = np.sin(transf_y_rad) + # ------ + S = np.array([[cr, 0, -sr], [0, 1, 0], [sr, 0, cr]]) # y rotation matrix + return advance_rotation(v, w, S) + # SRotation Handling + if transf_s_rad != 0: + assert angle == 0, "rot_s_rad is only supported for angle = 0" + assert tilt == 0, "rot_s_rad is only supported for tilt = 0" + assert length == 0, "rot_s_rad is only supported for length = 0" + + # Relevant sine/cosine + cr = np.cos(transf_s_rad) + sr = np.sin(transf_s_rad) + # ------ + S = np.array([[cr, sr, 0], [-sr, cr, 0], [0, 0, 1]]) # z rotation matrix + return advance_rotation(v, w, S) + + # Non bending elements if angle == 0: R = np.array([0, 0, length]) return advance_drift(v, w, R) + + # Horizontal bends elif tilt == 0: # Relevant sine/cosine ca = np.cos(angle) @@ -94,6 +144,7 @@ def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0): S = np.array([[ca, 0, -sa], [0, 1, 0], [sa, 0, ca]]) return advance_bend(v, w, R, S) + # Tilted bends else: # Relevant sine/cosine ca = np.cos(angle) @@ -143,6 +194,9 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, out_tilt = list(-self.tilt[:-1][::-1]) out_dx = list(-self.dx[:-1][::-1]) out_dy = list(-self.dy[:-1][::-1]) + out_rot_x_rad = list(-self.rot_x_rad[:-1][::-1]) + out_rot_y_rad = list(-self.rot_y_rad[:-1][::-1]) + out_rot_s_rad = list(-self.rot_s_rad[:-1][::-1]) out_name = list(self.name[:-1][::-1]) if type(element0) is str: @@ -151,7 +205,8 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, X, Y, Z, theta, phi, psi = compute_survey( X0, Y0, Z0, theta0, phi0, psi0, out_drift_length, out_angle, out_tilt, - out_dx, out_dy, element0=element0) + out_dx, out_dy, out_rot_x_rad, + out_rot_y_rad, out_rot_s_rad, element0=element0) # Initializing dictionary out_columns = {} @@ -166,11 +221,17 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, out_columns["s"] = self.s[-1] - self.s[::-1] out_columns['drift_length'] = np.array(out_drift_length + [0.]) + out_columns['angle'] = np.array(out_angle + [0.]) out_columns['tilt'] = np.array(out_tilt + [0.]) + out_columns["dx"] = np.array(out_dx + [0.]) out_columns["dy"] = np.array(out_dy + [0.]) + out_columns["rot_x_rad"] = np.array(out_rot_x_rad + [0.]) + out_columns["rot_y_rad"] = np.array(out_rot_y_rad + [0.]) + out_columns["rot_s_rad"] = np.array(out_rot_s_rad + [0.]) + out_scalars = {} out_scalars["element0"] = element0 @@ -224,26 +285,28 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, assert not values_at_element_exit, "Not implemented yet" # Extract angle and tilt from elements - tt = line.get_table(attr = True) - angle = tt.angle_rad - tilt = tt.rot_s_rad + tt = line.get_table(attr = True) + angle = tt.angle_rad + tilt = tt.rot_s_rad drift_length = tt.length drift_length[~tt.isthick] = 0 # Extract xy shifts from elements - dx = tt._own_dx - dy = tt._own_dy + dx = tt.dx + dy = tt.dy - # Handling of rotation elements - sin_angle = tt._own_sin_angle - print(sin_angle) + # Handling of XYSRotation elements + transf_angle_rad = tt.transform_angle_rad + transf_x_rad = transf_angle_rad * np.array(tt.element_type == 'XRotation') + transf_y_rad = transf_angle_rad * np.array(tt.element_type == 'YRotation') + transf_s_rad = transf_angle_rad * np.array(tt.element_type == 'SRotation') if type(element0) == str: element0 = line.element_names.index(element0) X, Y, Z, theta, phi, psi = compute_survey( X0, Y0, Z0, theta0, phi0, psi0, drift_length[:-1], angle[:-1], tilt[:-1], - dx[:-1], dy[:-1], element0=element0) + dx[:-1], dy[:-1], transf_x_rad, transf_y_rad, transf_s_rad, element0=element0) # Initializing dictionary out_columns = {} @@ -260,9 +323,14 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, out_columns['drift_length'] = drift_length out_columns['angle'] = angle out_columns['tilt'] = tilt + out_columns['dx'] = dx out_columns['dy'] = dy + out_columns['transf_x_rad'] = transf_x_rad + out_columns['transf_y_rad'] = transf_y_rad + out_columns['transf_s_rad'] = transf_s_rad + out_scalars['element0'] = element0 out = SurveyTable(data={**out_columns, **out_scalars}, # this is a merge @@ -273,7 +341,7 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, - dx, dy, element0=0, reverse_xs=False): + dx, dy, transf_x_rad, transf_y_rad, transf_s_rad, element0=0, reverse_xs=False): if element0 != 0: assert not(reverse_xs), "Not implemented yet" @@ -282,20 +350,28 @@ def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, tilt_forward = tilt[element0:] dx_forward = dx[element0:] dy_forward = dy[element0:] + transf_x_rad_forward = transf_x_rad[element0:] + transf_y_rad_forward = transf_y_rad[element0:] + transf_s_rad_forward = transf_s_rad[element0:] (X_forward, Y_forward, Z_forward, theta_forward, phi_forward, psi_forward) = compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_forward, angle_forward, tilt_forward, - dx_forward, dy_forward) + dx_forward, dy_forward, transf_x_rad_forward, + transf_y_rad_forward, transf_s_rad_forward) drift_backward = drift_length[:element0][::-1] angle_backward = -np.array(angle[:element0][::-1]) tilt_backward = -np.array(tilt[:element0][::-1]) dx_backward = -np.array(dx[:element0][::-1]) dy_backward = -np.array(dy[:element0][::-1]) + transf_x_rad_backward = -np.array(transf_x_rad[:element0][::-1]) + transf_y_rad_backward = -np.array(transf_y_rad[:element0][::-1]) + transf_s_rad_backward = -np.array(transf_s_rad[:element0][::-1]) (X_backward, Y_backward, Z_backward, theta_backward, phi_backward, psi_backward) = compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_backward, angle_backward, tilt_backward, - dx_backward, dy_backward, + dx_backward, dy_backward, transf_x_rad_backward, + transf_y_rad_backward, transf_s_rad_backward, reverse_xs=True) X = np.array(X_backward[::-1][:-1] + X_forward) @@ -316,7 +392,8 @@ def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, w = get_w_from_angles(theta=theta0, phi=phi0, psi=psi0, reverse_xs=reverse_xs) # Advancing element by element - for ll, aa, tt, xx, yy in zip(drift_length, angle, tilt, dx, dy): + for ll, aa, tt, xx, yy, tx, ty, ts, in zip(drift_length, angle, tilt, dx, dy, + transf_x_rad, transf_x_rad, transf_x_rad): th, ph, ps = get_angles_from_w(w, reverse_xs=reverse_xs) @@ -328,7 +405,8 @@ def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, psi.append(ps) # Advancing - v, w = advance_element(v, w, length=ll, angle=aa, tilt=tt, dx=xx, dy=yy) + v, w = advance_element(v, w, length=ll, angle=aa, tilt=tt, dx=xx, dy=yy, + transf_x_rad=tx, transf_y_rad=ty, transf_s_rad=ts) # Last marker th, ph, ps = get_angles_from_w(w, reverse_xs=reverse_xs) From d20ed39318ecb149e49bd2eef95c1a1fd134d139 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Fri, 1 Nov 2024 17:29:28 +0100 Subject: [PATCH 09/21] Working on examples --- examples/survey/000_dev.py | 218 ++++++++++++++++++++++++++++------- examples/survey/001_bend.py | 219 +++++++++++++++++++++++++----------- examples/survey/_helpers.py | 138 ++++++++++++++++------- 3 files changed, 432 insertions(+), 143 deletions(-) diff --git a/examples/survey/000_dev.py b/examples/survey/000_dev.py index 4c1104dc4..3015a4783 100644 --- a/examples/survey/000_dev.py +++ b/examples/survey/000_dev.py @@ -1,58 +1,200 @@ +""" +Test the survey on Bends +""" +################################################################################ +# Packages +################################################################################ import xtrack as xt import numpy as np import matplotlib.pyplot as plt -from _madpoint import MadPoint +from _helpers import madpoint_twiss_survey, summary_plot -env = xt.Environment() +################################################################################ +# Setup +################################################################################ -env.particle_ref = xt.Particles(p0c = 1E9) -env['k0'] = 0.00 -env['h'] = 1E-3 +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +env['k0'] = 1E-3 +env['h'] = 1E-3 +env['dx'] = 2E-6 +env['dy'] = 3E-6 +# TODO: These don't work with the env.new function +env['angle_x'] = 2E-4 +env['angle_y'] = 3E-4 + +line = env.new_line(name = 'line', components=[ + env.new('bend1', xt.Bend, k0 = 0, h = 0, length = 0.5, at=2), + env.new('xyshift', xt.XYShift, dx = 'dx', dy = 'dy', at = 3,), + env.new('xyshift2', 'xyshift', mode = 'replica', at = 4), + env.new('yrotation', xt.YRotation, angle = env['angle_x'], at = 5), + env.new('xrotation', xt.XRotation, angle = env['angle_y'], at = 6), + env.new('xyshift3', 'xyshift', mode = 'replica', at = 7), + # env.new('bend2', xt.Bend, k0 = 'k0', h = 'h', length = 0.5, at=8), + # env.new('xyshift3', 'xyshift', mode = 'replica', at = 9), + env.new('end', xt.Marker, at = 10)]) + +# line.cut_at_s(np.linspace(0, line.get_length(), 1001)) + +################################################################################ +# Horizontal Mult vs YRotation +################################################################################ + +######################################## +# Rotation +######################################## +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('yrotation', xt.YRotation, angle = 0, at = 1), + env.new('xyshift', xt.XYShift, dx = 1E-6, dy = 1E-6, at = 2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +# line.config.XTRACK_USE_EXACT_DRIFTS = True + +# NB Minus sign here as rotation works opposite direction +line['yrotation'].angle = -1E-3 + +sv, tw = madpoint_twiss_survey(line) +summary_plot(sv, tw, 'Rotation Y') + +######################################## +# Bend +######################################## +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('mult', xt.Multipole, hxl = 0, at = 1), + env.new('xyshift', xt.XYShift, dx = 1E-6, dy = 1E-6, at = 2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +# line.config.XTRACK_USE_EXACT_DRIFTS = True + +line['mult'].hxl = np.deg2rad(1E-3) + +sv, tw = madpoint_twiss_survey(line) +summary_plot(sv, tw, 'Horiztonal Mult') + +################################################################################ +# Vertical Bend vs XRotation +################################################################################ + +######################################## +# Rotation +######################################## +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) line = env.new_line(name = 'line', components=[ - env.new('bend', xt.Bend, k0 = 'k0', h = 'h', length = 0.5, at=2), - env.new('xyshift', xt.XYShift, dx = 0.1, dy=0.2), - env.new('end', xt.Marker, at = 5)]) + env.new('xrotation', xt.XRotation, angle = 0, at = 1), + env.new('xyshift', xt.XYShift, dx = 1E-6, dy = 1E-6, at = 2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +# line.config.XTRACK_USE_EXACT_DRIFTS = True -line.configure_bend_model(core = 'bend-kick-bend', edge = 'suppressed') -line.cut_at_s(np.linspace(0, line.get_length(), 101)) +line['xrotation'].angle = 1E-3 -sv = line.survey() -tw = line.twiss4d(betx = 1, bety =1) +sv, tw = madpoint_twiss_survey(line) +summary_plot(sv, tw, 'Rotation X') -madpoints = [] -xx = [] -yy = [] -zz = [] -for nn in tw.name: - madpoints.append( - MadPoint(name = nn, xsuite_twiss = tw, xsuite_survey = sv)) - xx.append(madpoints[-1].p[0]) - yy.append(madpoints[-1].p[1]) - zz.append(madpoints[-1].p[2]) +######################################## +# Bend +######################################## +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('mult', xt.Multipole, hxl = 0, at = 1, rot_s_rad = np.pi / 2), + env.new('xyshift', xt.XYShift, dx = 1E-6, dy = 1E-6, at = 2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +# line.config.XTRACK_USE_EXACT_DRIFTS = True + +line['mult'].hxl = np.deg2rad(1E-3) + +sv, tw = madpoint_twiss_survey(line) +summary_plot(sv, tw, 'Vertical Mult') + +################################################################################ +# Vertical shift vs SRotation + Horizontal shift +################################################################################ + +######################################## +# Vertical Shift +######################################## +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('xyshift', xt.XYShift, dy = -1E-6, at = 2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +# line.config.XTRACK_USE_EXACT_DRIFTS = True + +sv, tw = madpoint_twiss_survey(line) +summary_plot(sv, tw, 'Vertical Shift') -xx = np.array(xx) -yy = np.array(yy) -zz = np.array(zz) +######################################## +# Bend +######################################## +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('srot', xt.SRotation, angle = 90, at=1), + env.new('xyshift', xt.XYShift, dx = 1E-6, at = 2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +# line.config.XTRACK_USE_EXACT_DRIFTS = True + +sv, tw = madpoint_twiss_survey(line) +summary_plot(sv, tw, 'SRot + Horizontal Shift') -sv['xx'] = xx -sv['yy'] = yy -sv['zz'] = zz +################################################################################ +# SRotation then Horizontal Bend +################################################################################ + +######################################## +# Vertical Shift +######################################## +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('bend', xt.Bend, k0 = 1E-3, h = 1E-3, length = 0.5, at=2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +# line.config.XTRACK_USE_EXACT_DRIFTS = True + +sv, tw = madpoint_twiss_survey(line) +summary_plot(sv, tw, 'Horizontal Bend') + +######################################## +# Bend +######################################## +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line(name = 'line', components=[ + env.new('srot', xt.SRotation, angle=90, at=1), + env.new('bend', xt.Bend, k0 = 1E-3, h = 1E-3, length = 0.5, at=2), + env.new('end', xt.Marker, at = 3)]) +line.configure_bend_model(edge = 'suppressed') +# line.config.XTRACK_USE_EXACT_DRIFTS = True +sv, tw = madpoint_twiss_survey(line) +summary_plot(sv, tw, 'SRot + Horizontal Bend') -fig, axs = plt.subplots(3, 1, sharex=True) -axs[0].plot(sv.Z, sv.X) -axs[1].plot(tw.s, tw.x) -axs[2].plot(sv.s, sv.xx) +################################################################################ +# Multiple +################################################################################ -axs[0].set_xlabel('Z [m]') -axs[0].set_ylabel('X [m]') +######################################## +# Rotation +######################################## +env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +line = env.new_line( + name = 'line', + components= [env.new('drift', xt.Drift, length = 1E-3)] +\ + [env.new('xrotation', xt.XRotation, angle = 1E-2), + env.new('drift', xt.Drift, length = 1E-3)]*1000 +\ + [env.new('xyshift', xt.XYShift, dy = 1E-2)] +\ + [env.new('xrotation', xt.XRotation, angle = 1E-2), + env.new('drift', xt.Drift, length = 1E-3)]*1000 +\ + [env.new('end', xt.Marker)]) +line.configure_bend_model(edge = 'suppressed') +sv, tw = madpoint_twiss_survey(line) +summary_plot(sv, tw, 'Check that XYShift is suitably rotated') -axs[1].set_xlabel('s [m]') -axs[1].set_ylabel('x [m]') +# TODO: Replica +# TODO: Reverse -axs[2].set_xlabel('s [m]') -axs[2].set_ylabel('xx [m]') plt.show() \ No newline at end of file diff --git a/examples/survey/001_bend.py b/examples/survey/001_bend.py index caafab6b9..20ab17cff 100644 --- a/examples/survey/001_bend.py +++ b/examples/survey/001_bend.py @@ -6,8 +6,9 @@ ################################################################################ import numpy as np import matplotlib.pyplot as plt +from matplotlib.lines import Line2D -from _helpers import survey_test, summary_plot +from _helpers import madpoint_twiss_survey, summary_plot, zero_small_values import xtrack as xt ################################################################################ @@ -15,17 +16,25 @@ ################################################################################ PLOT_COMPARISONS = True TOLERANCE = 1E-12 -REL_TOLERANCE = 1E-3 -bend_h = 1E-3 +# Small angle needed to paraxial approximation tests +BEND_ANGLE = 1E-3 +BEND_LENGTH = 0.5 ################################################################################ # Create line ################################################################################ -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line(name = 'line', components=[ - env.new('bend', xt.Bend, k0 = 0, h = 0, length = 0.5, at = 1), - env.new('end', xt.Marker, at = 2)]) +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) + +line = env.new_line( + name = 'line', + components = [ + env.new('bend', xt.Bend, k0 = 0, h = 0, length = BEND_LENGTH, at = 1), + env.new('end', xt.Marker, at = 2)]) + +######################################## +# Configure Bend and Drift Model +######################################## line.configure_bend_model(edge = 'suppressed') line.config.XTRACK_USE_EXACT_DRIFTS = True @@ -33,6 +42,42 @@ # Horizontal Bend ################################################################################ +######################################## +# Plot setup +######################################## +fig = plt.figure(figsize=(16, 8)) +gs = fig.add_gridspec(5, 3, hspace = 0, wspace=0.5) +axs = gs.subplots(sharex=True) + +def add_to_plot(axes, survey, twiss, index, tol = 1E-12): + axes[index, 0].plot( + zero_small_values(survey.Z, tol), + zero_small_values(survey.X, tol), + c = 'k') + axes[index, 0].plot( + zero_small_values(survey.Z, tol), + zero_small_values(survey.Y, tol), + c = 'r') + + axes[index, 1].plot( + zero_small_values(twiss.s, tol), + zero_small_values(twiss.x, tol), + c = 'k') + axes[index, 1].plot( + zero_small_values(twiss.s, tol), + zero_small_values(twiss.y, tol), + c = 'r') + + axes[index, 2].plot( + zero_small_values(survey.s, tol), + zero_small_values(survey.xx, tol), + c = 'k') + axes[index, 2].plot( + zero_small_values(survey.s, tol), + zero_small_values(survey.yy, tol), + c = 'r') + + ######################################## # h = k0 = 0 ######################################## @@ -40,86 +85,130 @@ line['bend'].h = 0 line['bend'].rot_s_rad = 0 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Horizontal Bend: h = k0 = 0') + summary_plot(sv, tw, 'Horizontal Bend: h = 0, k0 = 0') + add_to_plot(axs, sv, tw, 0) # Everything zero here -assert np.allclose(sv.X[-1], 0, atol=TOLERANCE) -assert np.allclose(sv.xx[-1], 0, atol=TOLERANCE) -assert np.allclose(tw.x[-1], 0, atol=TOLERANCE) +assert np.allclose(sv.X[-1], 0, atol = TOLERANCE) +assert np.allclose(sv.xx[-1], 0, atol = TOLERANCE) +assert np.allclose(tw.x[-1], 0, atol = TOLERANCE) # All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) ######################################## # h = 0, k0 != 0 ######################################## -line['bend'].k0 = bend_h +line['bend'].k0 = BEND_ANGLE / BEND_LENGTH line['bend'].h = 0 line['bend'].rot_s_rad = 0 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: summary_plot(sv, tw, 'Horizontal Bend: h = 0, k0 != 0') + add_to_plot(axs, sv, tw, 1) # No survey with h = 0 -assert np.allclose(sv.X[-1], 0, atol=TOLERANCE) +assert np.allclose(sv.X[-1], 0, atol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) # All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) ######################################## # h != 0, k0 = 0 ######################################## line['bend'].k0 = 0 -line['bend'].h = bend_h +line['bend'].h = BEND_ANGLE / BEND_LENGTH line['bend'].rot_s_rad = 0 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: summary_plot(sv, tw, 'Horizontal Bend: h != 0, k0 = 0') + add_to_plot(axs, sv, tw, 2) # Survey negative of Twiss for h != 0, k0 = 0 -assert np.allclose(sv.X[-1], -tw.x[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.X[-1], -tw.x[-1], rtol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) # All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) ######################################## -# h != k0 !=0 +# h != k0 !=0 ######################################## -line['bend'].k0 = bend_h * 2 -line['bend'].h = bend_h +line['bend'].k0 = BEND_ANGLE / BEND_LENGTH * 2 +line['bend'].h = BEND_ANGLE / BEND_LENGTH line['bend'].rot_s_rad = 0 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Horizontal Bend: h = k0 /2') + summary_plot(sv, tw, 'Horizontal Bend: h = k0 / 2') + add_to_plot(axs, sv, tw, 3) # MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) # All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) ######################################## # h = k0 != 0 ######################################## -line['bend'].k0 = bend_h -line['bend'].h = bend_h +line['bend'].k0 = BEND_ANGLE / BEND_LENGTH +line['bend'].h = BEND_ANGLE / BEND_LENGTH line['bend'].rot_s_rad = 0 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: summary_plot(sv, tw, 'Horizontal Bend: h = k0 != 0') + add_to_plot(axs, sv, tw, 4) -# No orbit with h=k0 -assert np.allclose(tw.x[-1], 0, atol=TOLERANCE) +# No orbit with h = k0 +assert np.allclose(tw.x[-1], 0, atol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) # All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) + +######################################## +# Plot +######################################## +# Titles +axs[0, 0].set_title('Survey') +axs[0, 1].set_title('Twiss') +axs[0, 2].set_title('MadPoint') + +# x labels +axs[4, 0].set_xlabel('Z [m]') +axs[4, 1].set_xlabel('s [m]') +axs[4, 2].set_xlabel('s [m]') + +# y labels +axs[2, 0].set_ylabel('X [m]') +axs[2, 1].set_ylabel('x [m]') +axs[2, 2].set_ylabel('x [m]') + +# Row labels +for j, label in enumerate([ + 'h = 0, k0 = 0', + 'h = 0, k0 != 0', + 'h != 0, k0 = 0', + 'h = k0 / 2 != 0', + 'h = k0 != 0']): + fig.text(0.05, 0.2 * j, label, va='center', ha='center', fontsize=12) + +legend_elements = [ + Line2D([0], [0], color='black', lw=2, label='x'), + Line2D([0], [0], color='red', lw=2, label='y')] + +# Place the custom legend below the x-axis labels +fig.legend(handles=legend_elements, loc='lower center', ncol=2, frameon=False, bbox_to_anchor=(0.5, -0.05)) + +fig.suptitle('Horizontal Bend') + +plt.tight_layout() +plt.show() ################################################################################ # Vertical Bend @@ -132,83 +221,83 @@ line['bend'].h = 0 line['bend'].rot_s_rad = np.pi / 2 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Vertical Bend: h = k0 = 0') + summary_plot(sv, tw, 'Vertical Bend: h = 0, k0 = 0') # Everything zero here -assert np.allclose(sv.Y[-1], 0, atol=TOLERANCE) -assert np.allclose(sv.yy[-1], 0, atol=TOLERANCE) -assert np.allclose(tw.y[-1], 0, atol=TOLERANCE) +assert np.allclose(sv.Y[-1], 0, atol = TOLERANCE) +assert np.allclose(sv.yy[-1], 0, atol = TOLERANCE) +assert np.allclose(tw.y[-1], 0, atol = TOLERANCE) # All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) ######################################## # h = 0, k0 != 0 ######################################## -line['bend'].k0 = bend_h +line['bend'].k0 = BEND_ANGLE / BEND_LENGTH line['bend'].h = 0 line['bend'].rot_s_rad = np.pi / 2 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: summary_plot(sv, tw, 'Vertical Bend: h = 0, k0 != 0') -assert np.allclose(sv.Y[-1], 0, atol=TOLERANCE) -assert np.allclose(sv.yy[-1], tw.y[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.Y[-1], 0, atol = TOLERANCE) +assert np.allclose(sv.yy[-1], tw.y[-1], rtol = TOLERANCE) # All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) ######################################## # h != 0, k0 = 0 ######################################## line['bend'].k0 = 0 -line['bend'].h = bend_h +line['bend'].h = BEND_ANGLE / BEND_LENGTH line['bend'].rot_s_rad = np.pi / 2 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: summary_plot(sv, tw, 'Vertical Bend: h != 0, k0 = 0') # Survey negative of Twiss -assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.Y[-1], -tw.y[-1], rtol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) # All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) ######################################## # h != k0 !=0 ######################################## -line['bend'].k0 = bend_h * 2 -line['bend'].h = bend_h +line['bend'].k0 = BEND_ANGLE / BEND_LENGTH * 2 +line['bend'].h = BEND_ANGLE / BEND_LENGTH line['bend'].rot_s_rad = np.pi / 2 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Vertical Bend: h = k0 /2') + summary_plot(sv, tw, 'Vertical Bend: h = k0 / 2') # MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) # All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) ######################################## # h = k0 != 0 ######################################## -line['bend'].k0 = bend_h -line['bend'].h = bend_h +line['bend'].k0 = BEND_ANGLE / BEND_LENGTH +line['bend'].h = BEND_ANGLE / BEND_LENGTH line['bend'].rot_s_rad = np.pi / 2 -sv, tw = survey_test(line) +sv, tw = madpoint_twiss_survey(line) if PLOT_COMPARISONS: summary_plot(sv, tw, 'Vertical Bend: h = k0 != 0') -assert np.allclose(tw.y[-1], 0, atol=TOLERANCE) +assert np.allclose(tw.y[-1], 0, atol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) # All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) ######################################## # Show plots diff --git a/examples/survey/_helpers.py b/examples/survey/_helpers.py index 9b878b0a4..d1a97d366 100644 --- a/examples/survey/_helpers.py +++ b/examples/survey/_helpers.py @@ -5,7 +5,13 @@ import matplotlib.pyplot as plt from _madpoint import MadPoint -def survey_test(line): +def zero_small_values(arr, tol): + """ + Set values within a tolerance of 0, to 0 + """ + return np.where(np.abs(arr) < tol, 0, arr) + +def madpoint_twiss_survey(line): """ Produce twiss and survey for a line, including MadPoint """ @@ -29,48 +35,100 @@ def survey_test(line): return survey, twiss +# def summary_plot(survey, twiss, title, zero_tol = 1E-12): +# """ +# Summary plot comparing survey, twiss and MadPoint +# """ +# def zero_small_values(arr, tol = zero_tol): +# return np.where(np.abs(arr) < tol, 0, arr) + +# fig, axs = plt.subplots(3, 2, figsize=(16, 8)) +# axs[0, 0].plot( +# zero_small_values(survey.Z), +# zero_small_values(survey.X)) +# axs[1, 0].plot( +# zero_small_values(twiss.s), +# zero_small_values(twiss.x)) +# axs[2, 0].plot( +# zero_small_values(survey.s), +# zero_small_values(survey.xx)) + +# axs[0, 1].plot( +# zero_small_values(survey.Z), +# zero_small_values(survey.Y)) +# axs[1, 1].plot( +# zero_small_values(twiss.s), +# zero_small_values(twiss.y)) +# axs[2, 1].plot( +# zero_small_values(survey.s), +# zero_small_values(survey.yy)) + +# axs[0, 0].set_xlabel('Z [m]') +# axs[0, 0].set_ylabel('X [m]') +# axs[1, 0].set_xlabel('s [m]') +# axs[1, 0].set_ylabel('x [m]') +# axs[2, 0].set_xlabel('s [m]') +# axs[2, 0].set_ylabel('xx [m]') + +# axs[0, 1].set_xlabel('Z [m]') +# axs[0, 1].set_ylabel('Y [m]') +# axs[1, 1].set_xlabel('s [m]') +# axs[1, 1].set_ylabel('y [m]') +# axs[2, 1].set_xlabel('s [m]') +# axs[2, 1].set_ylabel('yy [m]') + +# fig.suptitle(title) + +# plt.tight_layout() + def summary_plot(survey, twiss, title, zero_tol = 1E-12): """ Summary plot comparing survey, twiss and MadPoint """ - def zero_small_values(arr, tol = zero_tol): - return np.where(np.abs(arr) < tol, 0, arr) - - fix, axs = plt.subplots(3, 2, figsize=(16, 8)) - axs[0, 0].plot( - zero_small_values(survey.Z), - zero_small_values(survey.X)) - axs[1, 0].plot( - zero_small_values(twiss.s), - zero_small_values(twiss.x)) - axs[2, 0].plot( - zero_small_values(survey.s), - zero_small_values(survey.xx)) - - axs[0, 1].plot( - zero_small_values(survey.Z), - zero_small_values(survey.Y)) - axs[1, 1].plot( - zero_small_values(twiss.s), - zero_small_values(twiss.y)) - axs[2, 1].plot( - zero_small_values(survey.s), - zero_small_values(survey.yy)) - - axs[0, 0].set_xlabel('Z [m]') - axs[0, 0].set_ylabel('X [m]') - axs[1, 0].set_xlabel('s [m]') - axs[1, 0].set_ylabel('x [m]') - axs[2, 0].set_xlabel('s [m]') - axs[2, 0].set_ylabel('xx [m]') - - axs[0, 1].set_xlabel('Z [m]') - axs[0, 1].set_ylabel('Y [m]') - axs[1, 1].set_xlabel('s [m]') - axs[1, 1].set_ylabel('y [m]') - axs[2, 1].set_xlabel('s [m]') - axs[2, 1].set_ylabel('yy [m]') - - fix.suptitle(title) + + fig, axs = plt.subplots(3, 1, figsize=(16, 8)) + axs[0].plot( + zero_small_values(survey.Z, zero_tol), + zero_small_values(survey.X, zero_tol), + c = 'k') + axs[0].plot( + zero_small_values(survey.Z, zero_tol), + zero_small_values(survey.Y, zero_tol), + c = 'r') + + axs[1].plot( + zero_small_values(twiss.s, zero_tol), + zero_small_values(twiss.x, zero_tol), + c = 'k') + axs[1].plot( + zero_small_values(twiss.s, zero_tol), + zero_small_values(twiss.y, zero_tol), + c = 'r') + + axs[2].plot( + zero_small_values(survey.s, zero_tol), + zero_small_values(survey.xx, zero_tol), + c = 'k') + axs[2].plot( + zero_small_values(survey.s, zero_tol), + zero_small_values(survey.yy, zero_tol), + c = 'r') + + axs[0].set_xlabel('Z [m]') + axs[0].set_ylabel('X [m]') + axs[0].set_title('Survey') + + axs[1].set_xlabel('s [m]') + axs[1].set_ylabel('x [m]') + axs[1].set_title('Twiss') + + axs[2].set_xlabel('s [m]') + axs[2].set_ylabel('xx [m]') + axs[2].set_title('MadPoint') + + axs[2].legend(loc='upper center', bbox_to_anchor=(0.5, -0.05), + fancybox=True, ncol=2) + + fig.suptitle(title) plt.tight_layout() From 0d048789008adc5690679b467437d46c24a0973a Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Mon, 4 Nov 2024 17:13:10 +0100 Subject: [PATCH 10/21] Corrected variable naming Cleanup of structure for troubleshooting ease --- xtrack/survey.py | 311 ++++++++++++++++++++++++++++++----------------- 1 file changed, 202 insertions(+), 109 deletions(-) diff --git a/xtrack/survey.py b/xtrack/survey.py index 01b07aef1..30d91ad2d 100644 --- a/xtrack/survey.py +++ b/xtrack/survey.py @@ -64,11 +64,13 @@ def advance_bend(v, w, R, S): v2 = w1*R + v1 | w2 = w1*S""" return np.dot(w, R) + v, np.dot(w, S) + def advance_rotation(v, w, S): """Advancing through rotation element: Rotate w matrix according to transformation matrix S""" return v, np.dot(w, S) + def advance_drift(v, w, R): """Advancing through drift element, see MAD-X manual: v2 = w1*R + v1 | w2 = w1*S -> S is unity""" @@ -132,8 +134,8 @@ def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, if angle == 0: R = np.array([0, 0, length]) return advance_drift(v, w, R) - - # Horizontal bends + + # Horizontal bending elements elif tilt == 0: # Relevant sine/cosine ca = np.cos(angle) @@ -144,7 +146,7 @@ def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, S = np.array([[ca, 0, -sa], [0, 1, 0], [sa, 0, ca]]) return advance_bend(v, w, R, S) - # Tilted bends + # Tilted bending elements else: # Relevant sine/cosine ca = np.cos(angle) @@ -157,13 +159,16 @@ def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, S = np.array([[ca, 0, -sa], [0, 1, 0], [sa, 0, ca]]) # Orthogonal rotation matrix for tilt - T = np.array([[ct, -st, 0], [st, ct, 0], [0, 0, 1]]) - Tinv = np.array([[ct, st, 0], [-st, ct, 0], [0, 0, 1]]) + T = np.array([[ct, -st, 0], [st, ct, 0], [0, 0, 1]]) + Tinv = np.array([[ct, st, 0], [-st, ct, 0], [0, 0, 1]]) return advance_bend(v, w, np.dot(T, R), np.dot(T, np.dot(S, Tinv))) class SurveyTable(Table): + """ + Table for survey data. + """ _error_on_row_not_found = True @@ -194,45 +199,54 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, out_tilt = list(-self.tilt[:-1][::-1]) out_dx = list(-self.dx[:-1][::-1]) out_dy = list(-self.dy[:-1][::-1]) - out_rot_x_rad = list(-self.rot_x_rad[:-1][::-1]) - out_rot_y_rad = list(-self.rot_y_rad[:-1][::-1]) - out_rot_s_rad = list(-self.rot_s_rad[:-1][::-1]) + out_transf_x_rad = list(-self.transf_x_rad[:-1][::-1]) + out_transf_y_rad = list(-self.transf_y_rad[:-1][::-1]) + out_transf_s_rad = list(-self.transf_s_rad[:-1][::-1]) out_name = list(self.name[:-1][::-1]) - if type(element0) is str: + if isinstance(element0, str): element0 = out_name.index(element0) X, Y, Z, theta, phi, psi = compute_survey( - X0, Y0, Z0, theta0, phi0, psi0, - out_drift_length, out_angle, out_tilt, - out_dx, out_dy, out_rot_x_rad, - out_rot_y_rad, out_rot_s_rad, element0=element0) + X0 = X0, + Y0 = Y0, + Z0 = Z0, + theta0 = theta0, + phi0 = phi0, + psi0 = psi0, + drift_length = out_drift_length, + angle = out_angle, + tilt = out_tilt, + dx = out_dx, + dy = out_dy, + transf_x_rad = out_transf_x_rad, + transf_y_rad = out_transf_y_rad, + transf_s_rad = out_transf_s_rad, + element0 = element0, + reverse_xs = False) # Initializing dictionary out_columns = {} - out_columns["X"] = np.array(X) - out_columns["Y"] = np.array(Y) - out_columns["Z"] = np.array(Z) - out_columns["theta"] = np.unwrap(theta) - out_columns["phi"] = np.unwrap(phi) - out_columns["psi"] = np.unwrap(psi) - - out_columns["name"] = np.array(list(out_name) + ["_end_point"]) - out_columns["s"] = self.s[-1] - self.s[::-1] + out_scalars = {} + # Fill survey data + out_columns["X"] = np.array(X) + out_columns["Y"] = np.array(Y) + out_columns["Z"] = np.array(Z) + out_columns["theta"] = np.unwrap(theta) + out_columns["phi"] = np.unwrap(phi) + out_columns["psi"] = np.unwrap(psi) + out_columns["name"] = np.array(list(out_name) + ["_end_point"]) + out_columns["s"] = self.s[-1] - self.s[::-1] out_columns['drift_length'] = np.array(out_drift_length + [0.]) + out_columns['angle'] = np.array(out_angle + [0.]) + out_columns['tilt'] = np.array(out_tilt + [0.]) + out_columns["dx"] = np.array(out_dx + [0.]) + out_columns["dy"] = np.array(out_dy + [0.]) + out_columns["transf_x_rad"] = np.array(out_transf_x_rad + [0.]) + out_columns["transf_y_rad"] = np.array(out_transf_y_rad + [0.]) + out_columns["transf_s_rad"] = np.array(out_transf_s_rad + [0.]) - out_columns['angle'] = np.array(out_angle + [0.]) - out_columns['tilt'] = np.array(out_tilt + [0.]) - - out_columns["dx"] = np.array(out_dx + [0.]) - out_columns["dy"] = np.array(out_dy + [0.]) - - out_columns["rot_x_rad"] = np.array(out_rot_x_rad + [0.]) - out_columns["rot_y_rad"] = np.array(out_rot_y_rad + [0.]) - out_columns["rot_s_rad"] = np.array(out_rot_s_rad + [0.]) - - out_scalars = {} out_scalars["element0"] = element0 out = SurveyTable(data=(out_columns | out_scalars), @@ -241,6 +255,9 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, return out def plot(self, element_width=None, legend=True, **kwargs): + """ + Plot the survey using xplt.FloorPlot + """ # Shallow copy of self out_sv_table = SurveyTable.__new__(SurveyTable) out_sv_table.__dict__.update(self.__dict__) @@ -284,10 +301,14 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, assert not values_at_element_exit, "Not implemented yet" - # Extract angle and tilt from elements + # Get line table to extract attributes tt = line.get_table(attr = True) + + # Extract angle and tilt from elements angle = tt.angle_rad tilt = tt.rot_s_rad + + # Extract drift lengths drift_length = tt.length drift_length[~tt.isthick] = 0 @@ -296,107 +317,170 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, dy = tt.dy # Handling of XYSRotation elements - transf_angle_rad = tt.transform_angle_rad - transf_x_rad = transf_angle_rad * np.array(tt.element_type == 'XRotation') - transf_y_rad = transf_angle_rad * np.array(tt.element_type == 'YRotation') - transf_s_rad = transf_angle_rad * np.array(tt.element_type == 'SRotation') + transf_angle_rad = tt.transform_angle_rad + transf_x_rad = transf_angle_rad * np.array(tt.element_type == 'XRotation') + transf_y_rad = transf_angle_rad * np.array(tt.element_type == 'YRotation') + transf_s_rad = transf_angle_rad * np.array(tt.element_type == 'SRotation') - if type(element0) == str: + if isinstance(element0, str): element0 = line.element_names.index(element0) X, Y, Z, theta, phi, psi = compute_survey( - X0, Y0, Z0, theta0, phi0, psi0, drift_length[:-1], angle[:-1], tilt[:-1], - dx[:-1], dy[:-1], transf_x_rad, transf_y_rad, transf_s_rad, element0=element0) + X0 = X0, + Y0 = Y0, + Z0 = Z0, + theta0 = theta0, + phi0 = phi0, + psi0 = psi0, + drift_length = drift_length[:-1], + angle = angle[:-1], + tilt = tilt[:-1], + dx = dx[:-1], + dy = dy[:-1], + transf_x_rad = transf_x_rad[:-1], + transf_y_rad = transf_y_rad[:-1], + transf_s_rad = transf_s_rad[:-1], + element0 = element0, + reverse_xs = False) # Initializing dictionary out_columns = {} out_scalars = {} - out_columns["X"] = np.array(X) - out_columns["Y"] = np.array(Y) - out_columns["Z"] = np.array(Z) - out_columns["theta"] = np.unwrap(theta) - out_columns["phi"] = np.unwrap(phi) - out_columns["psi"] = np.unwrap(psi) - - out_columns["name"] = tt.name - out_columns["s"] = tt.s - out_columns['drift_length'] = drift_length - out_columns['angle'] = angle - out_columns['tilt'] = tilt - - out_columns['dx'] = dx - out_columns['dy'] = dy + # Fill survey data + out_columns["X"] = np.array(X) + out_columns["Y"] = np.array(Y) + out_columns["Z"] = np.array(Z) + out_columns["theta"] = np.unwrap(theta) + out_columns["phi"] = np.unwrap(phi) + out_columns["psi"] = np.unwrap(psi) + out_columns["name"] = tt.name + out_columns["s"] = tt.s + out_columns['drift_length'] = drift_length + out_columns['angle'] = angle + out_columns['tilt'] = tilt + out_columns['dx'] = dx + out_columns['dy'] = dy out_columns['transf_x_rad'] = transf_x_rad out_columns['transf_y_rad'] = transf_y_rad out_columns['transf_s_rad'] = transf_s_rad - out_scalars['element0'] = element0 + out_scalars['element0'] = element0 - out = SurveyTable(data={**out_columns, **out_scalars}, # this is a merge - col_names=out_columns.keys()) + out = SurveyTable( + data = {**out_columns, **out_scalars}, # this is a merge + col_names = out_columns.keys()) out._data['line'] = line return out -def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, - dx, dy, transf_x_rad, transf_y_rad, transf_s_rad, element0=0, reverse_xs=False): +def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, + drift_length, angle, tilt, + dx, dy, transf_x_rad, transf_y_rad, transf_s_rad, + element0 = 0, reverse_xs = False): + """ + Compute survey from initial position and orientation. + """ + # If element0 is not the first element, split the survey if element0 != 0: + # Assert that reverse_xs is not implemented yet assert not(reverse_xs), "Not implemented yet" - drift_forward = drift_length[element0:] - angle_forward = angle[element0:] - tilt_forward = tilt[element0:] - dx_forward = dx[element0:] - dy_forward = dy[element0:] - transf_x_rad_forward = transf_x_rad[element0:] - transf_y_rad_forward = transf_y_rad[element0:] - transf_s_rad_forward = transf_s_rad[element0:] + + # Forward section of survey + drift_forward = drift_length[element0:] + angle_forward = angle[element0:] + tilt_forward = tilt[element0:] + dx_forward = dx[element0:] + dy_forward = dy[element0:] + transf_x_rad_forward = transf_x_rad[element0:] + transf_y_rad_forward = transf_y_rad[element0:] + transf_s_rad_forward = transf_s_rad[element0:] + + # Evaluate forward survey (X_forward, Y_forward, Z_forward, theta_forward, phi_forward, - psi_forward) = compute_survey(X0, Y0, Z0, theta0, phi0, psi0, - drift_forward, angle_forward, tilt_forward, - dx_forward, dy_forward, transf_x_rad_forward, - transf_y_rad_forward, transf_s_rad_forward) - - drift_backward = drift_length[:element0][::-1] - angle_backward = -np.array(angle[:element0][::-1]) - tilt_backward = -np.array(tilt[:element0][::-1]) - dx_backward = -np.array(dx[:element0][::-1]) - dy_backward = -np.array(dy[:element0][::-1]) - transf_x_rad_backward = -np.array(transf_x_rad[:element0][::-1]) - transf_y_rad_backward = -np.array(transf_y_rad[:element0][::-1]) - transf_s_rad_backward = -np.array(transf_s_rad[:element0][::-1]) + psi_forward) = compute_survey( + X0 = X0, + Y0 = Y0, + Z0 = Z0, + theta0 = theta0, + phi0 = phi0, + psi0 = psi0, + drift_length = drift_forward, + angle = angle_forward, + tilt = tilt_forward, + dx = dx_forward, + dy = dy_forward, + transf_x_rad = transf_x_rad_forward, + transf_y_rad = transf_y_rad_forward, + transf_s_rad = transf_s_rad_forward, + element0 = 0, + reverse_xs = False) + + # Backward section of survey + drift_backward = drift_length[:element0][::-1] + angle_backward = -np.array(angle[:element0][::-1]) + tilt_backward = -np.array(tilt[:element0][::-1]) + dx_backward = -np.array(dx[:element0][::-1]) + dy_backward = -np.array(dy[:element0][::-1]) + transf_x_rad_backward = -np.array(transf_x_rad[:element0][::-1]) + transf_y_rad_backward = -np.array(transf_y_rad[:element0][::-1]) + transf_s_rad_backward = -np.array(transf_s_rad[:element0][::-1]) + + # Evaluate backward survey (X_backward, Y_backward, Z_backward, theta_backward, phi_backward, - psi_backward) = compute_survey(X0, Y0, Z0, theta0, phi0, psi0, - drift_backward, angle_backward, tilt_backward, - dx_backward, dy_backward, transf_x_rad_backward, - transf_y_rad_backward, transf_s_rad_backward, - reverse_xs=True) - - X = np.array(X_backward[::-1][:-1] + X_forward) - Y = np.array(Y_backward[::-1][:-1] + Y_forward) - Z = np.array(Z_backward[::-1][:-1] + Z_forward) - theta = np.array(theta_backward[::-1][:-1] + theta_forward) - phi = np.array(phi_backward[::-1][:-1] + phi_forward) - psi = np.array(psi_backward[::-1][:-1]+ psi_forward) + psi_backward) = compute_survey( + X0 = X0, + Y0 = Y0, + Z0 = Z0, + theta0 = theta0, + phi0 = phi0, + psi0 = psi0, + drift_length = drift_backward, + angle = angle_backward, + tilt = tilt_backward, + dx = dx_backward, + dy = dy_backward, + transf_x_rad = transf_x_rad_backward, + transf_y_rad = transf_y_rad_backward, + transf_s_rad = transf_s_rad_backward, + element0 = 0, + reverse_xs = True) + + # Concatenate forward and backward + X = np.array(X_backward[::-1][:-1] + X_forward) + Y = np.array(Y_backward[::-1][:-1] + Y_forward) + Z = np.array(Z_backward[::-1][:-1] + Z_forward) + theta = np.array(theta_backward[::-1][:-1] + theta_forward) + phi = np.array(phi_backward[::-1][:-1] + phi_forward) + psi = np.array(psi_backward[::-1][:-1]+ psi_forward) return X, Y, Z, theta, phi, psi - X = [] - Y = [] - Z = [] - theta = [] - phi = [] - psi = [] - v = np.array([X0, Y0, Z0]) - w = get_w_from_angles(theta=theta0, phi=phi0, psi=psi0, - reverse_xs=reverse_xs) + # Initialise lists for storing the survey + X = [] + Y = [] + Z = [] + theta = [] + phi = [] + psi = [] + + # Initial position and orientation + v = np.array([X0, Y0, Z0]) + w = get_w_from_angles( + theta = theta0, + phi = phi0, + psi = psi0, + reverse_xs = reverse_xs) + # Advancing element by element for ll, aa, tt, xx, yy, tx, ty, ts, in zip(drift_length, angle, tilt, dx, dy, - transf_x_rad, transf_x_rad, transf_x_rad): + transf_x_rad, transf_y_rad, transf_s_rad): - th, ph, ps = get_angles_from_w(w, reverse_xs=reverse_xs) + # Get angles from w matrix after previous element + th, ph, ps = get_angles_from_w(w, reverse_xs = reverse_xs) + # Store position and orientation at element entrance X.append(v[0]) Y.append(v[1]) Z.append(v[2]) @@ -405,11 +489,20 @@ def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, psi.append(ps) # Advancing - v, w = advance_element(v, w, length=ll, angle=aa, tilt=tt, dx=xx, dy=yy, - transf_x_rad=tx, transf_y_rad=ty, transf_s_rad=ts) + v, w = advance_element( + v = v, + w = w, + length = ll, + angle = aa, + tilt = tt, + dx = xx, + dy = yy, + transf_x_rad = tx, + transf_y_rad = ty, + transf_s_rad = ts) # Last marker - th, ph, ps = get_angles_from_w(w, reverse_xs=reverse_xs) + th, ph, ps = get_angles_from_w(w, reverse_xs = reverse_xs) X.append(v[0]) Y.append(v[1]) Z.append(v[2]) @@ -417,5 +510,5 @@ def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, phi.append(ph) psi.append(ps) - # Returns as SurveyTable object + # Return data for SurveyTable object return X, Y, Z, theta, phi, psi From 455b6bd52346077898bf94bb33968cdbe9083db0 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Thu, 7 Nov 2024 18:18:13 +0100 Subject: [PATCH 11/21] Changes to formatting Increased clarity now that there are many more parameters to keep track of --- xtrack/survey.py | 94 +++++++++++++++++++++++++++++------------------- 1 file changed, 57 insertions(+), 37 deletions(-) diff --git a/xtrack/survey.py b/xtrack/survey.py index 30d91ad2d..3ae893b88 100644 --- a/xtrack/survey.py +++ b/xtrack/survey.py @@ -14,7 +14,7 @@ # Required functions # ================================================== -def get_w_from_angles(theta, phi, psi, reverse_xs=False): +def get_w_from_angles(theta, phi, psi, reverse_xs = False): """W matrix, see MAD-X manual""" costhe = np.cos(theta) cosphi = np.cos(phi) @@ -40,7 +40,7 @@ def get_w_from_angles(theta, phi, psi, reverse_xs=False): return w -def get_angles_from_w(w, reverse_xs=False): +def get_angles_from_w(w, reverse_xs = False): """Inverse function of get_w_from_angles()""" # w[0, 2]/w[2, 2] = (sinthe * cosphi)/(costhe * cosphi) # w[1, 0]/w[1, 1] = (cosphi * sinpsi)/(cosphi * cospsi) @@ -55,7 +55,8 @@ def get_angles_from_w(w, reverse_xs=False): psi = np.arctan2(w[1, 0], w[1, 1]) phi = np.arctan2(w[1, 2], w[1, 1] / np.cos(psi)) - # TODO: arctan2 returns angle between [-pi,pi]. Hence theta ends up not at 2pi after a full survey + # TODO: arctan2 returns angle between [-pi,pi]. + # Hence theta ends up not at 2pi after a full survey return theta, phi, psi @@ -77,8 +78,9 @@ def advance_drift(v, w, R): return np.dot(w, R) + v, w -def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, - transf_x_rad=0, transf_y_rad=0, transf_s_rad=0): +def advance_element( + v, w, length = 0, angle = 0, tilt = 0, + dx = 0, dy = 0, transf_x_rad = 0, transf_y_rad = 0, transf_s_rad = 0): """Computing the advance element-by-element. See MAD-X manual for generation of R and S""" # XYShift Handling @@ -97,7 +99,7 @@ def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, assert tilt == 0, "rot_x_rad is only supported for tilt = 0" assert length == 0, "rot_x_rad is only supported for length = 0" - # Relevant sine/cosine + # Rotation sine/cosine cr = np.cos(-transf_x_rad) sr = np.sin(-transf_x_rad) # ------ @@ -110,7 +112,7 @@ def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, assert tilt == 0, "rot_y_rad is only supported for tilt = 0" assert length == 0, "rot_y_rad is only supported for length = 0" - # Relevant sine/cosine + # Rotation sine/cosine cr = np.cos(transf_y_rad) sr = np.sin(transf_y_rad) # ------ @@ -123,7 +125,7 @@ def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, assert tilt == 0, "rot_s_rad is only supported for tilt = 0" assert length == 0, "rot_s_rad is only supported for length = 0" - # Relevant sine/cosine + # Rotation sine/cosine cr = np.cos(transf_s_rad) sr = np.sin(transf_s_rad) # ------ @@ -137,7 +139,7 @@ def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, # Horizontal bending elements elif tilt == 0: - # Relevant sine/cosine + # Angle sine/cosine ca = np.cos(angle) sa = np.sin(angle) # ------ @@ -148,9 +150,10 @@ def advance_element(v, w, length=0, angle=0, tilt=0, dx=0, dy=0, # Tilted bending elements else: - # Relevant sine/cosine + # Angle sine/cosine ca = np.cos(angle) sa = np.sin(angle) + # Tilt sine/cosine ct = np.cos(tilt) st = np.sin(tilt) # ------ @@ -172,8 +175,13 @@ class SurveyTable(Table): _error_on_row_not_found = True - def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, - phi0=None, psi0=None, element0=None): + def reverse( + self, + X0 = None, Y0 = None, Z0 = None, + theta0 = None, phi0 = None, psi0 = None, element0 = None): + """ + Reverse the survey. + """ if element0 is None: element0 = len(self.name) - self.element0 - 1 @@ -194,15 +202,15 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, psi0 = self.psi[self.element0] # We cut away the last marker (added by survey) and reverse the order - out_drift_length = list(self.drift_length[:-1][::-1]) - out_angle = list(-self.angle[:-1][::-1]) - out_tilt = list(-self.tilt[:-1][::-1]) - out_dx = list(-self.dx[:-1][::-1]) - out_dy = list(-self.dy[:-1][::-1]) - out_transf_x_rad = list(-self.transf_x_rad[:-1][::-1]) - out_transf_y_rad = list(-self.transf_y_rad[:-1][::-1]) - out_transf_s_rad = list(-self.transf_s_rad[:-1][::-1]) - out_name = list(self.name[:-1][::-1]) + out_drift_length = list(self.drift_length[:-1][::-1]) + out_angle = list(-self.angle[:-1][::-1]) + out_tilt = list(-self.tilt[:-1][::-1]) + out_dx = list(-self.dx[:-1][::-1]) + out_dy = list(-self.dy[:-1][::-1]) + out_transf_x_rad = list(-self.transf_x_rad[:-1][::-1]) + out_transf_y_rad = list(-self.transf_y_rad[:-1][::-1]) + out_transf_s_rad = list(-self.transf_s_rad[:-1][::-1]) + out_name = list(self.name[:-1][::-1]) if isinstance(element0, str): element0 = out_name.index(element0) @@ -249,15 +257,20 @@ def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, out_scalars["element0"] = element0 - out = SurveyTable(data=(out_columns | out_scalars), - col_names=out_columns.keys()) + out = SurveyTable( + data=(out_columns | out_scalars), + col_names=out_columns.keys()) return out - def plot(self, element_width=None, legend=True, **kwargs): + def plot(self, element_width = None, legend = True, **kwargs): """ Plot the survey using xplt.FloorPlot """ + # Import the xplt module here + # (Not at the top as not default installation with xsuite) + import xplt + # Shallow copy of self out_sv_table = SurveyTable.__new__(SurveyTable) out_sv_table.__dict__.update(self.__dict__) @@ -266,13 +279,19 @@ def plot(self, element_width=None, legend=True, **kwargs): # Removing the count for repeated elements out_sv_table.name = np.array([nn.split('::')[0] for nn in out_sv_table.name]) + # Setting element width for plotting if element_width is None: x_range = max(self.X) - min(self.X) y_range = max(self.Y) - min(self.Y) z_range = max(self.Z) - min(self.Z) - element_width = max([x_range, y_range, z_range]) * 0.03 - import xplt - xplt.FloorPlot(out_sv_table, self.line, element_width=element_width, **kwargs) + element_width = max([x_range, y_range, z_range]) * 0.03 + + xplt.FloorPlot( + survey = out_sv_table, + line = self.line, + element_width = element_width, + **kwargs) + if legend: import matplotlib.pyplot as plt plt.legend() @@ -282,10 +301,10 @@ def plot(self, element_width=None, legend=True, **kwargs): # Main function # ================================================== -def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, - element0=0, - values_at_element_exit=False, - reverse=True): +def survey_from_line( + line, + X0 = 0, Y0 = 0, Z0 = 0, theta0 = 0, phi0 = 0, psi0 = 0, + element0 = 0, values_at_element_exit = False, reverse = True): """Execute SURVEY command. Based on MADX equivalent. Attributes, must be given in this order in the dictionary: X0 (float) Initial X position in meters. @@ -375,10 +394,11 @@ def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, return out -def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, - drift_length, angle, tilt, - dx, dy, transf_x_rad, transf_y_rad, transf_s_rad, - element0 = 0, reverse_xs = False): +def compute_survey( + X0, Y0, Z0, theta0, phi0, psi0, + drift_length, angle, tilt, + dx, dy, transf_x_rad, transf_y_rad, transf_s_rad, + element0 = 0, reverse_xs = False): """ Compute survey from initial position and orientation. """ @@ -474,8 +494,8 @@ def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, reverse_xs = reverse_xs) # Advancing element by element - for ll, aa, tt, xx, yy, tx, ty, ts, in zip(drift_length, angle, tilt, dx, dy, - transf_x_rad, transf_y_rad, transf_s_rad): + for ll, aa, tt, xx, yy, tx, ty, ts, in zip( + drift_length, angle, tilt, dx, dy, transf_x_rad, transf_y_rad, transf_s_rad): # Get angles from w matrix after previous element th, ph, ps = get_angles_from_w(w, reverse_xs = reverse_xs) From 0be854e64865b2fb8355ecc3c3e8e9dca33a761d Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Thu, 7 Nov 2024 18:18:23 +0100 Subject: [PATCH 12/21] Test plots --- examples/survey/000_dev.py | 200 ----------- examples/survey/001_bend.py | 277 ++++++++------- examples/survey/002_multipole.py | 295 +++++++++++++--- examples/survey/003_xyshift.py | 201 ++++++----- examples/survey/004_xrotation.py | 564 ++++++++++++++++++++++++++++-- examples/survey/005_yrotation.py | 567 ++++++++++++++++++++++++++++-- examples/survey/006_srotation.py | 575 +++++++++++++++++++++++++++++++ examples/survey/_helpers.py | 117 ++----- 8 files changed, 2163 insertions(+), 633 deletions(-) delete mode 100644 examples/survey/000_dev.py create mode 100644 examples/survey/006_srotation.py diff --git a/examples/survey/000_dev.py b/examples/survey/000_dev.py deleted file mode 100644 index 3015a4783..000000000 --- a/examples/survey/000_dev.py +++ /dev/null @@ -1,200 +0,0 @@ -""" -Test the survey on Bends -""" -################################################################################ -# Packages -################################################################################ -import xtrack as xt -import numpy as np -import matplotlib.pyplot as plt - -from _helpers import madpoint_twiss_survey, summary_plot - -################################################################################ -# Setup -################################################################################ - -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -env['k0'] = 1E-3 -env['h'] = 1E-3 -env['dx'] = 2E-6 -env['dy'] = 3E-6 - -# TODO: These don't work with the env.new function -env['angle_x'] = 2E-4 -env['angle_y'] = 3E-4 - -line = env.new_line(name = 'line', components=[ - env.new('bend1', xt.Bend, k0 = 0, h = 0, length = 0.5, at=2), - env.new('xyshift', xt.XYShift, dx = 'dx', dy = 'dy', at = 3,), - env.new('xyshift2', 'xyshift', mode = 'replica', at = 4), - env.new('yrotation', xt.YRotation, angle = env['angle_x'], at = 5), - env.new('xrotation', xt.XRotation, angle = env['angle_y'], at = 6), - env.new('xyshift3', 'xyshift', mode = 'replica', at = 7), - # env.new('bend2', xt.Bend, k0 = 'k0', h = 'h', length = 0.5, at=8), - # env.new('xyshift3', 'xyshift', mode = 'replica', at = 9), - env.new('end', xt.Marker, at = 10)]) - -# line.cut_at_s(np.linspace(0, line.get_length(), 1001)) - -################################################################################ -# Horizontal Mult vs YRotation -################################################################################ - -######################################## -# Rotation -######################################## -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line(name = 'line', components=[ - env.new('yrotation', xt.YRotation, angle = 0, at = 1), - env.new('xyshift', xt.XYShift, dx = 1E-6, dy = 1E-6, at = 2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -# line.config.XTRACK_USE_EXACT_DRIFTS = True - -# NB Minus sign here as rotation works opposite direction -line['yrotation'].angle = -1E-3 - -sv, tw = madpoint_twiss_survey(line) -summary_plot(sv, tw, 'Rotation Y') - -######################################## -# Bend -######################################## -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line(name = 'line', components=[ - env.new('mult', xt.Multipole, hxl = 0, at = 1), - env.new('xyshift', xt.XYShift, dx = 1E-6, dy = 1E-6, at = 2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -# line.config.XTRACK_USE_EXACT_DRIFTS = True - -line['mult'].hxl = np.deg2rad(1E-3) - -sv, tw = madpoint_twiss_survey(line) -summary_plot(sv, tw, 'Horiztonal Mult') - -################################################################################ -# Vertical Bend vs XRotation -################################################################################ - -######################################## -# Rotation -######################################## -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line(name = 'line', components=[ - env.new('xrotation', xt.XRotation, angle = 0, at = 1), - env.new('xyshift', xt.XYShift, dx = 1E-6, dy = 1E-6, at = 2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -# line.config.XTRACK_USE_EXACT_DRIFTS = True - -line['xrotation'].angle = 1E-3 - -sv, tw = madpoint_twiss_survey(line) -summary_plot(sv, tw, 'Rotation X') - -######################################## -# Bend -######################################## -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line(name = 'line', components=[ - env.new('mult', xt.Multipole, hxl = 0, at = 1, rot_s_rad = np.pi / 2), - env.new('xyshift', xt.XYShift, dx = 1E-6, dy = 1E-6, at = 2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -# line.config.XTRACK_USE_EXACT_DRIFTS = True - -line['mult'].hxl = np.deg2rad(1E-3) - -sv, tw = madpoint_twiss_survey(line) -summary_plot(sv, tw, 'Vertical Mult') - -################################################################################ -# Vertical shift vs SRotation + Horizontal shift -################################################################################ - -######################################## -# Vertical Shift -######################################## -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line(name = 'line', components=[ - env.new('xyshift', xt.XYShift, dy = -1E-6, at = 2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -# line.config.XTRACK_USE_EXACT_DRIFTS = True - -sv, tw = madpoint_twiss_survey(line) -summary_plot(sv, tw, 'Vertical Shift') - -######################################## -# Bend -######################################## -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line(name = 'line', components=[ - env.new('srot', xt.SRotation, angle = 90, at=1), - env.new('xyshift', xt.XYShift, dx = 1E-6, at = 2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -# line.config.XTRACK_USE_EXACT_DRIFTS = True - -sv, tw = madpoint_twiss_survey(line) -summary_plot(sv, tw, 'SRot + Horizontal Shift') - -################################################################################ -# SRotation then Horizontal Bend -################################################################################ - -######################################## -# Vertical Shift -######################################## -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line(name = 'line', components=[ - env.new('bend', xt.Bend, k0 = 1E-3, h = 1E-3, length = 0.5, at=2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -# line.config.XTRACK_USE_EXACT_DRIFTS = True - -sv, tw = madpoint_twiss_survey(line) -summary_plot(sv, tw, 'Horizontal Bend') - -######################################## -# Bend -######################################## -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line(name = 'line', components=[ - env.new('srot', xt.SRotation, angle=90, at=1), - env.new('bend', xt.Bend, k0 = 1E-3, h = 1E-3, length = 0.5, at=2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -# line.config.XTRACK_USE_EXACT_DRIFTS = True - -sv, tw = madpoint_twiss_survey(line) -summary_plot(sv, tw, 'SRot + Horizontal Bend') - -################################################################################ -# Multiple -################################################################################ - -######################################## -# Rotation -######################################## -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) -line = env.new_line( - name = 'line', - components= [env.new('drift', xt.Drift, length = 1E-3)] +\ - [env.new('xrotation', xt.XRotation, angle = 1E-2), - env.new('drift', xt.Drift, length = 1E-3)]*1000 +\ - [env.new('xyshift', xt.XYShift, dy = 1E-2)] +\ - [env.new('xrotation', xt.XRotation, angle = 1E-2), - env.new('drift', xt.Drift, length = 1E-3)]*1000 +\ - [env.new('end', xt.Marker)]) -line.configure_bend_model(edge = 'suppressed') -sv, tw = madpoint_twiss_survey(line) -summary_plot(sv, tw, 'Check that XYShift is suitably rotated') - -# TODO: Replica -# TODO: Reverse - - -plt.show() \ No newline at end of file diff --git a/examples/survey/001_bend.py b/examples/survey/001_bend.py index 20ab17cff..3d90fda00 100644 --- a/examples/survey/001_bend.py +++ b/examples/survey/001_bend.py @@ -8,24 +8,33 @@ import matplotlib.pyplot as plt from matplotlib.lines import Line2D -from _helpers import madpoint_twiss_survey, summary_plot, zero_small_values +from _helpers import madpoint_twiss_survey, add_to_plot import xtrack as xt ################################################################################ # User variables ################################################################################ -PLOT_COMPARISONS = True TOLERANCE = 1E-12 -# Small angle needed to paraxial approximation tests +# Small angle needed for paraxial approximation tests BEND_ANGLE = 1E-3 BEND_LENGTH = 0.5 ################################################################################ -# Create line +# Plot setup ################################################################################ -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +fig_h = plt.figure(figsize = (16, 8)) +gs_h = fig_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_h = gs_h.subplots(sharex = 'row', sharey = True) + +fig_v = plt.figure(figsize = (16, 8)) +gs_v = fig_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_v = gs_v.subplots(sharex = 'row', sharey = True) +################################################################################ +# Create line +################################################################################ +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) line = env.new_line( name = 'line', components = [ @@ -33,51 +42,14 @@ env.new('end', xt.Marker, at = 2)]) ######################################## -# Configure Bend and Drift Model +# Configure Bend Model ######################################## line.configure_bend_model(edge = 'suppressed') -line.config.XTRACK_USE_EXACT_DRIFTS = True ################################################################################ # Horizontal Bend ################################################################################ -######################################## -# Plot setup -######################################## -fig = plt.figure(figsize=(16, 8)) -gs = fig.add_gridspec(5, 3, hspace = 0, wspace=0.5) -axs = gs.subplots(sharex=True) - -def add_to_plot(axes, survey, twiss, index, tol = 1E-12): - axes[index, 0].plot( - zero_small_values(survey.Z, tol), - zero_small_values(survey.X, tol), - c = 'k') - axes[index, 0].plot( - zero_small_values(survey.Z, tol), - zero_small_values(survey.Y, tol), - c = 'r') - - axes[index, 1].plot( - zero_small_values(twiss.s, tol), - zero_small_values(twiss.x, tol), - c = 'k') - axes[index, 1].plot( - zero_small_values(twiss.s, tol), - zero_small_values(twiss.y, tol), - c = 'r') - - axes[index, 2].plot( - zero_small_values(survey.s, tol), - zero_small_values(survey.xx, tol), - c = 'k') - axes[index, 2].plot( - zero_small_values(survey.s, tol), - zero_small_values(survey.yy, tol), - c = 'r') - - ######################################## # h = k0 = 0 ######################################## @@ -86,16 +58,15 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = 0 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Horizontal Bend: h = 0, k0 = 0') - add_to_plot(axs, sv, tw, 0) - -# Everything zero here -assert np.allclose(sv.X[-1], 0, atol = TOLERANCE) -assert np.allclose(sv.xx[-1], 0, atol = TOLERANCE) -assert np.allclose(tw.x[-1], 0, atol = TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) + +add_to_plot(axs_h, sv, tw, 0) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) ######################################## # h = 0, k0 != 0 @@ -105,11 +76,13 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = 0 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Horizontal Bend: h = 0, k0 != 0') - add_to_plot(axs, sv, tw, 1) -# No survey with h = 0 +add_to_plot(axs_h, sv, tw, 1) + +#################### +# Tests +#################### +# No bending with h = 0, so the survey must be zero assert np.allclose(sv.X[-1], 0, atol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) @@ -124,11 +97,13 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = 0 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Horizontal Bend: h != 0, k0 = 0') - add_to_plot(axs, sv, tw, 2) -# Survey negative of Twiss for h != 0, k0 = 0 +add_to_plot(axs_h, sv, tw, 2) + +#################### +# Tests +#################### +# With h!=0, k0=0, the survey must be the negative of the Twiss assert np.allclose(sv.X[-1], -tw.x[-1], rtol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) @@ -143,10 +118,12 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = 0 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Horizontal Bend: h = k0 / 2') - add_to_plot(axs, sv, tw, 3) +add_to_plot(axs_h, sv, tw, 3) + +#################### +# Tests +#################### # MadPoint sum of twiss and survey (paraxial approximation) assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) # All y related quantities are zero @@ -160,56 +137,19 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = 0 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Horizontal Bend: h = k0 != 0') - add_to_plot(axs, sv, tw, 4) -# No orbit with h = k0 +add_to_plot(axs_h, sv, tw, 4) + +#################### +# Tests +#################### +# With h=k0, there should be no residual orbit on the twiss assert np.allclose(tw.x[-1], 0, atol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) # All y related quantities are zero assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) -######################################## -# Plot -######################################## -# Titles -axs[0, 0].set_title('Survey') -axs[0, 1].set_title('Twiss') -axs[0, 2].set_title('MadPoint') - -# x labels -axs[4, 0].set_xlabel('Z [m]') -axs[4, 1].set_xlabel('s [m]') -axs[4, 2].set_xlabel('s [m]') - -# y labels -axs[2, 0].set_ylabel('X [m]') -axs[2, 1].set_ylabel('x [m]') -axs[2, 2].set_ylabel('x [m]') - -# Row labels -for j, label in enumerate([ - 'h = 0, k0 = 0', - 'h = 0, k0 != 0', - 'h != 0, k0 = 0', - 'h = k0 / 2 != 0', - 'h = k0 != 0']): - fig.text(0.05, 0.2 * j, label, va='center', ha='center', fontsize=12) - -legend_elements = [ - Line2D([0], [0], color='black', lw=2, label='x'), - Line2D([0], [0], color='red', lw=2, label='y')] - -# Place the custom legend below the x-axis labels -fig.legend(handles=legend_elements, loc='lower center', ncol=2, frameon=False, bbox_to_anchor=(0.5, -0.05)) - -fig.suptitle('Horizontal Bend') - -plt.tight_layout() -plt.show() - ################################################################################ # Vertical Bend ################################################################################ @@ -222,15 +162,15 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = np.pi / 2 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Vertical Bend: h = 0, k0 = 0') -# Everything zero here -assert np.allclose(sv.Y[-1], 0, atol = TOLERANCE) -assert np.allclose(sv.yy[-1], 0, atol = TOLERANCE) -assert np.allclose(tw.y[-1], 0, atol = TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) +add_to_plot(axs_v, sv, tw, 0) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) ######################################## # h = 0, k0 != 0 @@ -240,11 +180,16 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = np.pi / 2 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Vertical Bend: h = 0, k0 != 0') +add_to_plot(axs_v, sv, tw, 1) + +#################### +# Tests +#################### +# No bending with h = 0, so the survey must be zero assert np.allclose(sv.Y[-1], 0, atol = TOLERANCE) -assert np.allclose(sv.yy[-1], tw.y[-1], rtol = TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) # All x related quantities are zero assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) @@ -256,10 +201,13 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = np.pi / 2 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Vertical Bend: h != 0, k0 = 0') -# Survey negative of Twiss +add_to_plot(axs_v, sv, tw, 2) + +#################### +# Tests +#################### +# With h!=0, k0=0, the survey must be the negative of the Twiss assert np.allclose(sv.Y[-1], -tw.y[-1], rtol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) @@ -274,9 +222,12 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = np.pi / 2 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Vertical Bend: h = k0 / 2') +add_to_plot(axs_v, sv, tw, 3) + +#################### +# Tests +#################### # MadPoint sum of twiss and survey (paraxial approximation) assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) # All x related quantities are zero @@ -290,17 +241,87 @@ def add_to_plot(axes, survey, twiss, index, tol = 1E-12): line['bend'].rot_s_rad = np.pi / 2 sv, tw = madpoint_twiss_survey(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Vertical Bend: h = k0 != 0') +add_to_plot(axs_v, sv, tw, 4) + +#################### +# Tests +#################### +# With h=k0, there should be no residual orbit on the twiss assert np.allclose(tw.y[-1], 0, atol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) # All x related quantities are zero assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) +################################################################################ +# Show Plots +################################################################################ + +######################################## +# Horizontal Bends +######################################## +# y labels +axs_h[0, 0].set_ylabel('Survey x,y [m]') +axs_h[1, 0].set_ylabel('Twiss x,y [m]') +axs_h[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_h[0, 2].set_xlabel('Z [m]') +axs_h[1, 2].set_xlabel('s [m]') +axs_h[2, 2].set_xlabel('s [m]') + +# Titles +axs_h[0, 0].set_title('h = 0, k0 = 0') +axs_h[0, 1].set_title('h = 0, k0 != 0') +axs_h[0, 2].set_title('h != 0, k0 = 0') +axs_h[0, 3].set_title('h = k0 / 2 != 0') +axs_h[0, 4].set_title('h = k0 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_h.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_h.suptitle('Horizontal Bend') + ######################################## -# Show plots +# Vertical Bends ######################################## -if PLOT_COMPARISONS: - plt.show() +# y labels +axs_v[0, 0].set_ylabel('Survey x,y [m]') +axs_v[1, 0].set_ylabel('Twiss x,y [m]') +axs_v[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_v[0, 2].set_xlabel('Z [m]') +axs_v[1, 2].set_xlabel('s [m]') +axs_v[2, 2].set_xlabel('s [m]') + +# Titles +axs_v[0, 0].set_title('h = 0, k0 = 0') +axs_v[0, 1].set_title('h = 0, k0 != 0') +axs_v[0, 2].set_title('h != 0, k0 = 0') +axs_v[0, 3].set_title('h = k0 / 2 != 0') +axs_v[0, 4].set_title('h = k0 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_v.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_v.suptitle('Vertical Bend') + +######################################## +# Show +######################################## +plt.tight_layout() +plt.show() diff --git a/examples/survey/002_multipole.py b/examples/survey/002_multipole.py index f3144289d..6a2faae6d 100644 --- a/examples/survey/002_multipole.py +++ b/examples/survey/002_multipole.py @@ -6,103 +6,312 @@ ################################################################################ import numpy as np import matplotlib.pyplot as plt +from matplotlib.lines import Line2D -from _helpers import survey_test, summary_plot +from _helpers import madpoint_twiss_survey, add_to_plot import xtrack as xt ################################################################################ # User variables ################################################################################ -PLOT_COMPARISONS = True TOLERANCE = 1E-12 -REL_TOLERANCE = 1E-3 +HXL = 1E-3 -hxl = 1E-3 +################################################################################ +# Plot setup +################################################################################ +fig_h = plt.figure(figsize = (16, 8)) +gs_h = fig_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_h = gs_h.subplots(sharex = 'row', sharey = True) + +fig_v = plt.figure(figsize = (16, 8)) +gs_v = fig_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_v = gs_v.subplots(sharex = 'row', sharey = True) ################################################################################ # Create line ################################################################################ -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) line = env.new_line(name = 'line', components=[ env.new('mult', xt.Multipole, hxl = 0, length = 0.5, at = 1), env.new('end', xt.Marker, at = 2)]) -line.configure_bend_model(edge = 'suppressed') -line.config.XTRACK_USE_EXACT_DRIFTS = True ################################################################################ -# Horizontal +# Horizontal Bend ################################################################################ ######################################## -# hxl = 0 +# h = k0 = 0 ######################################## +line['mult'].knl[0] = 0 line['mult'].hxl = 0 line['mult'].rot_s_rad = 0 -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'hxl = 0') +sv, tw = madpoint_twiss_survey(line) -# All x related quantities are zero +add_to_plot(axs_h, sv, tw, 0) + +#################### +# Tests +#################### +# Element off must have no effect assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -# All y related quantities are zero assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) ######################################## -# hxl != 0 +# h = 0, k0 != 0 ######################################## -line['mult'].hxl = hxl +line['mult'].knl[0] = HXL +line['mult'].hxl = 0 line['mult'].rot_s_rad = 0 -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'hxl') +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_h, sv, tw, 1) -# Survey negative of Twiss for h != 0, k0 = 0 -assert np.allclose(sv.X[-1], -tw.x[-1], rtol=REL_TOLERANCE) +#################### +# Tests +#################### +# No bending with h = 0, so the survey must be zero +assert np.allclose(sv.X[-1], 0, atol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) # All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) + +######################################## +# h != 0, k0 = 0 +######################################## +line['mult'].knl[0] = 0 +line['mult'].hxl = HXL +line['mult'].rot_s_rad = 0 + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_h, sv, tw, 2) + +#################### +# Tests +#################### +# With h!=0, k0=0, the survey must be the negative of the Twiss +assert np.allclose(sv.X[-1], -tw.x[-1], rtol = TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) + +######################################## +# h != k0 !=0 +######################################## +line['mult'].knl[0] = HXL * 2 +line['mult'].hxl = HXL +line['mult'].rot_s_rad = 0 + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_h, sv, tw, 3) + +#################### +# Tests +#################### +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) + +######################################## +# h = k0 != 0 +######################################## +line['mult'].knl[0] = HXL +line['mult'].hxl = HXL +line['mult'].rot_s_rad = 0 + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_h, sv, tw, 4) + +#################### +# Tests +#################### +# With h=k0, there should be no residual orbit on the twiss +assert np.allclose(tw.x[-1], 0, atol = TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) +# All y related quantities are zero +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) ################################################################################ -# Vertical +# Vertical Bend ################################################################################ ######################################## -# hyl = 0 +# h = k0 = 0 ######################################## +line['mult'].knl[0] = 0 line['mult'].hxl = 0 line['mult'].rot_s_rad = np.pi / 2 -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'hyl = 0') +sv, tw = madpoint_twiss_survey(line) -# All x related quantities are zero +add_to_plot(axs_v, sv, tw, 0) + +#################### +# Tests +#################### +# Element off must have no effect assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -# All y related quantities are zero assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) ######################################## -# hyl != 0 +# h = 0, k0 != 0 ######################################## -line['mult'].hxl = hxl +line['mult'].knl[0] = HXL +line['mult'].hxl = 0 line['mult'].rot_s_rad = np.pi / 2 -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'hyl') +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_v, sv, tw, 1) -# Survey negative of Twiss for h != 0, k0 = 0 -assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=REL_TOLERANCE) +#################### +# Tests +#################### +# No bending with h = 0, so the survey must be zero +assert np.allclose(sv.Y[-1], 0, atol = TOLERANCE) # MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) + +######################################## +# h != 0, k0 = 0 +######################################## +line['mult'].knl[0] = 0 +line['mult'].hxl = HXL +line['mult'].rot_s_rad = np.pi / 2 + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_v, sv, tw, 2) + +#################### +# Tests +#################### +# With h!=0, k0=0, the survey must be the negative of the Twiss +assert np.allclose(sv.Y[-1], -tw.y[-1], rtol = TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) + +######################################## +# h != k0 !=0 +######################################## +line['mult'].knl[0] = HXL * 2 +line['mult'].hxl = HXL +line['mult'].rot_s_rad = np.pi / 2 + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_v, sv, tw, 3) + +#################### +# Tests +#################### +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) + +######################################## +# h = k0 != 0 +######################################## +line['mult'].knl[0] = HXL +line['mult'].hxl = HXL +line['mult'].rot_s_rad = np.pi / 2 + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_v, sv, tw, 4) + +#################### +# Tests +#################### +# With h=k0, there should be no residual orbit on the twiss +assert np.allclose(tw.y[-1], 0, atol = TOLERANCE) +# MadPoint sum of twiss and survey (paraxial approximation) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) +# All x related quantities are zero +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) + +################################################################################ +# Show Plots +################################################################################ + +######################################## +# Horizontal Bends +######################################## +# y labels +axs_h[0, 0].set_ylabel('Survey x,y [m]') +axs_h[1, 0].set_ylabel('Twiss x,y [m]') +axs_h[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_h[0, 2].set_xlabel('Z [m]') +axs_h[1, 2].set_xlabel('s [m]') +axs_h[2, 2].set_xlabel('s [m]') + +# Titles +axs_h[0, 0].set_title('hxl = 0, k0l = 0') +axs_h[0, 1].set_title('hxl = 0, k0l != 0') +axs_h[0, 2].set_title('hxl != 0, k0l = 0') +axs_h[0, 3].set_title('hxl = k0l / 2 != 0') +axs_h[0, 4].set_title('hxl = k0l != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_h.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_h.suptitle('Horizontal Multipole') + +######################################## +# Vertical Bends +######################################## +# y labels +axs_v[0, 0].set_ylabel('Survey x,y [m]') +axs_v[1, 0].set_ylabel('Twiss x,y [m]') +axs_v[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_v[0, 2].set_xlabel('Z [m]') +axs_v[1, 2].set_xlabel('s [m]') +axs_v[2, 2].set_xlabel('s [m]') + +# Titles +axs_v[0, 0].set_title('hyl = 0, k0sl = 0') +axs_v[0, 1].set_title('hyl = 0, k0sl != 0') +axs_v[0, 2].set_title('hyl != 0, k0sl = 0') +axs_v[0, 3].set_title('hyl = k0sl / 2 != 0') +axs_v[0, 4].set_title('hyl = k0sl != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_v.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_v.suptitle('Vertical Bend') ######################################## -# Show plots +# Show ######################################## -if PLOT_COMPARISONS: - plt.show() +plt.tight_layout() +plt.show() diff --git a/examples/survey/003_xyshift.py b/examples/survey/003_xyshift.py index a1e7a1e19..d811bae36 100644 --- a/examples/survey/003_xyshift.py +++ b/examples/survey/003_xyshift.py @@ -6,156 +6,147 @@ ################################################################################ import numpy as np import matplotlib.pyplot as plt +from matplotlib.lines import Line2D -from _helpers import survey_test, summary_plot +from _helpers import madpoint_twiss_survey, add_to_plot import xtrack as xt ################################################################################ # User variables ################################################################################ -PLOT_COMPARISONS = True TOLERANCE = 1E-12 -REL_TOLERANCE = 1E-3 + +################################################################################ +# Plot setup +################################################################################ +fig = plt.figure(figsize = (16, 8)) +gs = fig.add_gridspec(3, 4, hspace = 0.3, wspace = 0) +axs = gs.subplots(sharex = 'row', sharey = True) ################################################################################ # Create line ################################################################################ -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) line = env.new_line(name = 'line', components=[ - env.new('bend', xt.Bend, k0 = 0, h = 0, length = 0.5, at = 1), - env.new('xyshift', xt.XYShift, dx = 0, dy = 0, at = 2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -line.config.XTRACK_USE_EXACT_DRIFTS = True + env.new('xyshift', xt.XYShift, dx = 0, dy = 0, at = 1), + env.new('end', xt.Marker, at = 2)]) ################################################################################ -# X Shift +# No Shift ################################################################################ - -######################################## -# Shift only -######################################## -line['bend'].k0 = 0 -line['bend'].h = 0 -line['bend'].rot_s_rad = 0 - -line['xyshift'].dx = 0.1 +line['xyshift'].dx = 0 line['xyshift'].dy = 0 -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Shift X, bend off') +sv, tw = madpoint_twiss_survey(line) -# Survey negative of Twiss for h != 0, k0 = 0 -assert np.allclose(sv.X[-1], -tw.x[-1], rtol=REL_TOLERANCE) -# MadPoint sum of twiss and survey -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) +add_to_plot(axs, sv, tw, 0, ylims = (-0.15, 0.15)) -######################################## -# Shift + horizontal bend -######################################## -line['bend'].k0 = 1 -line['bend'].h = 1 -line['bend'].rot_s_rad = 0 +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) +################################################################################ +# X Shift +################################################################################ line['xyshift'].dx = 0.1 line['xyshift'].dy = 0 -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Shift X, bend off') +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs, sv, tw, 1, ylims = (-0.15, 0.15)) +#################### +# Tests +#################### +# No bending of particles, so must remain zero +assert np.allclose( sv.xx[-1], 0, atol=TOLERANCE) +# Survey negative of Twiss for h != 0, k0 = 0 +assert np.allclose(sv.X[-1], -tw.x[-1], rtol=TOLERANCE) # MadPoint sum of twiss and survey -# TODO: This is broken? -# assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=TOLERANCE) # All y related quantities are zero assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) -######################################## -# Shift + vertical bend -######################################## -line['bend'].k0 = 1 -line['bend'].h = 1 -line['bend'].rot_s_rad = np.pi / 2 - -line['xyshift'].dx = 0.1 -line['xyshift'].dy = 0 - -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Shift X, bend off') - -# MadPoint sum of twiss and survey -# TODO: This is broken? -# assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=REL_TOLERANCE) - ################################################################################ # Y Shift ################################################################################ - -######################################## -# Shift only -######################################## -line['bend'].k0 = 0 -line['bend'].h = 0 -line['bend'].rot_s_rad = 0 - line['xyshift'].dx = 0 line['xyshift'].dy = 0.1 -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Shift X, bend off') +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs, sv, tw, 2, ylims = (-0.15, 0.15)) +#################### +# Tests +#################### +# No bending of particles, so must remain zero +assert np.allclose(sv.yy[-1], 0, atol=TOLERANCE) # Survey negative of Twiss for h != 0, k0 = 0 -assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=TOLERANCE) # MadPoint sum of twiss and survey -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=TOLERANCE) # All x related quantities are zero assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -######################################## -# Shift + horizontal bend -######################################## -line['bend'].k0 = 1 -line['bend'].h = 1 -line['bend'].rot_s_rad = 0 - -line['xyshift'].dx = 0 +################################################################################ +# X and Y Shift +################################################################################ +line['xyshift'].dx = 0.1 line['xyshift'].dy = 0.1 -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Shift X, bend off') +sv, tw = madpoint_twiss_survey(line) -# MadPoint sum of twiss and survey -# TODO: This is broken? -# assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) - -######################################## -# Shift + vertical bend -######################################## -line['bend'].k0 = 1 -line['bend'].h = 1 -line['bend'].rot_s_rad = np.pi / 2 - -line['xyshift'].dx = 0 -line['xyshift'].dy = 0.1 - -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Shift X, bend off') +add_to_plot(axs, sv, tw, 3, ylims = (-0.15, 0.15)) +#################### +# Tests +#################### +# No bending of particles, so must remain zero +assert np.allclose( sv.xx[-1], 0, atol=TOLERANCE) +assert np.allclose(sv.yy[-1], 0, atol=TOLERANCE) +# Survey negative of Twiss for h != 0, k0 = 0 +assert np.allclose(sv.X[-1], -tw.x[-1], rtol=TOLERANCE) +assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=TOLERANCE) # MadPoint sum of twiss and survey -# TODO: This is broken? -# assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=REL_TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=TOLERANCE) +assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=TOLERANCE) + +################################################################################ +# Show Plots +################################################################################ +# y labels +axs[0, 0].set_ylabel('Survey x,y [m]') +axs[1, 0].set_ylabel('Twiss x,y [m]') +axs[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs[0, 2].set_xlabel('Z [m]') +axs[1, 2].set_xlabel('s [m]') +axs[2, 2].set_xlabel('s [m]') + +# Titles +axs[0, 0].set_title('DX = 0, DY = 0') +axs[0, 1].set_title('DX != 0, DY = 0') +axs[0, 2].set_title('DX = 0, DY != 0') +axs[0, 3].set_title('DX != 0, DY != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig.suptitle('XYShift') ######################################## -# Show plots +# Show ######################################## -if PLOT_COMPARISONS: - plt.show() +plt.tight_layout() +plt.show() diff --git a/examples/survey/004_xrotation.py b/examples/survey/004_xrotation.py index c9a1107a9..b79c3a296 100644 --- a/examples/survey/004_xrotation.py +++ b/examples/survey/004_xrotation.py @@ -6,72 +6,570 @@ ################################################################################ import numpy as np import matplotlib.pyplot as plt +from matplotlib.lines import Line2D -from _helpers import survey_test, summary_plot +from _helpers import madpoint_twiss_survey, add_to_plot import xtrack as xt ################################################################################ # User variables ################################################################################ -PLOT_COMPARISONS = True TOLERANCE = 1E-12 -REL_TOLERANCE = 1E-3 +ROT_DEG = 30 +HXL = 0.1 +DX = 0.5 +PLOT_X_LIMS = (-0.1, 4.1) +PLOT_Y_LIMS = (-1, 1) -hxl = 1E-3 +################################################################################ +# Plot setup +################################################################################ +fig_mult_h = plt.figure(figsize = (16, 8)) +gs_mult_h = fig_mult_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_mult_h = gs_mult_h.subplots(sharex = 'row', sharey = True) + +fig_mult_v = plt.figure(figsize = (16, 8)) +gs_mult_v = fig_mult_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_mult_v = gs_mult_v.subplots(sharex = 'row', sharey = True) +fig_shift_h = plt.figure(figsize = (16, 8)) +gs_shift_h = fig_shift_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_shift_h = gs_shift_h.subplots(sharex = 'row', sharey = True) + +fig_shift_v = plt.figure(figsize = (16, 8)) +gs_shift_v = fig_shift_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_shift_v = gs_shift_v.subplots(sharex = 'row', sharey = True) + +################################################################################ +# Horizontal Multipole Bends ################################################################################ + +######################################## # Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1), + env.new('xrotation', xt.XRotation, angle = 0, at = 2), + env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['xrotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = 0 +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = HXL +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = HXL +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Vertical Multipole Bends ################################################################################ -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) + +######################################## +# Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) line = env.new_line(name = 'line', components=[ - env.new('mult', xt.Multipole, hxl = 0, length = 0.5, at = 1), + env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1, rot_s_rad = np.pi / 2), env.new('xrotation', xt.XRotation, angle = 0, at = 2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -line.config.XTRACK_USE_EXACT_DRIFTS = True + env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3, rot_s_rad = np.pi / 2), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['xrotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = 0 +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = HXL +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = HXL +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO ################################################################################ +# Horizontal Shift +################################################################################ + +######################################## +# Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), + env.new('xrotation', xt.XRotation, angle = 0, at = 2), + env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['shift1'].dx = 0 +line['shift2'].dx = 0 +line['xrotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['shift1'].dx = 0 +line['shift2'].dx = 0 +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['shift1'].dx = DX +line['shift2'].dx = 0 +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['shift1'].dx = 0 +line['shift2'].dx = DX +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### # Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['shift1'].dx = DX +line['shift2'].dx = DX +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Vertical Shift ################################################################################ ######################################## -# Rotation only +# Create line ######################################## -line['mult'].hxl = 0 -line['mult'].rot_s_rad = 0 +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), + env.new('xrotation', xt.XRotation, angle = 0, at = 2), + env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), + env.new('end', xt.Marker, at = 4)]) -line['xrotation'].angle = np.rad2deg(hxl) +######################################## +# Off +######################################## +line['shift1'].dy = 0 +line['shift2'].dy = 0 +line['xrotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Rotation X, bend off') +add_to_plot(axs_shift_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) ######################################## -# Shift + horizontal bend +# Rotation Only ######################################## -line['mult'].hxl = hxl -line['mult'].rot_s_rad = 0 +line['shift1'].dy = 0 +line['shift2'].dy = 0 +line['xrotation'].angle = ROT_DEG -line['xrotation'].angle = np.rad2deg(hxl) +sv, tw = madpoint_twiss_survey(line) -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Rotation X, Horiztonal Mult') +add_to_plot(axs_shift_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['shift1'].dy = DX +line['shift2'].dy = 0 +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO ######################################## -# Shift + vertical bend +# Rot then mult ######################################## -line['mult'].hxl = hxl -line['mult'].rot_s_rad = np.pi / 2 +line['shift1'].dy = 0 +line['shift2'].dy = DX +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['shift1'].dy = DX +line['shift2'].dy = DX +line['xrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Show Plots +################################################################################ + +######################################## +# Horizontal Mult +######################################## +# y labels +axs_mult_h[0, 0].set_ylabel('Survey x,y [m]') +axs_mult_h[1, 0].set_ylabel('Twiss x,y [m]') +axs_mult_h[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_mult_h[0, 2].set_xlabel('Z [m]') +axs_mult_h[1, 2].set_xlabel('s [m]') +axs_mult_h[2, 2].set_xlabel('s [m]') + +# Titles +axs_mult_h[0, 0].set_title('Rot = 0,\n HXL1 = 0, HXL2 = 0') +axs_mult_h[0, 1].set_title('Rot != 0,\n HXL1 = 0, HXL2 = 0') +axs_mult_h[0, 2].set_title('Rot != 0,\n HXL1 != 0, HXL2 = 0') +axs_mult_h[0, 3].set_title('Rot != 0,\n HXL1 = 0, HXL2 != 0') +axs_mult_h[0, 4].set_title('Rot != 0,\n HXL1 != 0, HXL2 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_mult_h.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_mult_h.suptitle('Horizontal Multipole and XRotation') + +######################################## +# Vertical Mult +######################################## +# y labels +axs_mult_v[0, 0].set_ylabel('Survey x,y [m]') +axs_mult_v[1, 0].set_ylabel('Twiss x,y [m]') +axs_mult_v[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_mult_v[0, 2].set_xlabel('Z [m]') +axs_mult_v[1, 2].set_xlabel('s [m]') +axs_mult_v[2, 2].set_xlabel('s [m]') + +# Titles +axs_mult_v[0, 0].set_title('Rot = 0,\n HYL1 = 0, HYL2 = 0') +axs_mult_v[0, 1].set_title('Rot != 0,\n HYL1 = 0, HYL2 = 0') +axs_mult_v[0, 2].set_title('Rot != 0,\n HYL1 != 0, HYL2 = 0') +axs_mult_v[0, 3].set_title('Rot != 0,\n HYL1 = 0, HYL2 != 0') +axs_mult_v[0, 4].set_title('Rot != 0,\n HYL1 != 0, HYL2 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_mult_v.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_mult_v.suptitle('Vertical Multipole and XRotation') + +######################################## +# Horizontal XYShift +######################################## +# y labels +axs_shift_h[0, 0].set_ylabel('Survey x,y [m]') +axs_shift_h[1, 0].set_ylabel('Twiss x,y [m]') +axs_shift_h[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_shift_h[0, 2].set_xlabel('Z [m]') +axs_shift_h[1, 2].set_xlabel('s [m]') +axs_shift_h[2, 2].set_xlabel('s [m]') + +# Titles +axs_shift_h[0, 0].set_title('Rot = 0,\n DX1 = 0, DX2 = 0') +axs_shift_h[0, 1].set_title('Rot != 0,\n DX1 = 0, DX2 = 0') +axs_shift_h[0, 2].set_title('Rot != 0,\n DX1 != 0, DX2 = 0') +axs_shift_h[0, 3].set_title('Rot != 0,\n DX1 = 0, DX2 != 0') +axs_shift_h[0, 4].set_title('Rot != 0,\n DX1 != 0, DX2 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_shift_h.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_shift_h.suptitle('Horizontal XYShift and XRotation') + +######################################## +# Vertical XYShift +######################################## +# y labels +axs_shift_v[0, 0].set_ylabel('Survey x,y [m]') +axs_shift_v[1, 0].set_ylabel('Twiss x,y [m]') +axs_shift_v[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_shift_v[0, 2].set_xlabel('Z [m]') +axs_shift_v[1, 2].set_xlabel('s [m]') +axs_shift_v[2, 2].set_xlabel('s [m]') -line['xrotation'].angle = np.rad2deg(hxl) +# Titles +axs_shift_v[0, 0].set_title('Rot = 0,\n DY1 = 0, DY2 = 0') +axs_shift_v[0, 1].set_title('Rot != 0,\n DY1 = 0, DY2 = 0') +axs_shift_v[0, 2].set_title('Rot != 0,\n DY1 != 0, DY2 = 0') +axs_shift_v[0, 3].set_title('Rot != 0,\n DY1 = 0, DY2 != 0') +axs_shift_v[0, 4].set_title('Rot != 0,\n DY1 != 0, DY2 != 0') -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Rotation X, Vertical Mult') +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_shift_v.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_shift_v.suptitle('Vertical XYShift and XRotation') ######################################## -# Show plots +# Show ######################################## -if PLOT_COMPARISONS: - plt.show() +plt.tight_layout() +plt.show() diff --git a/examples/survey/005_yrotation.py b/examples/survey/005_yrotation.py index 4c1c8a2e6..12fb46c1d 100644 --- a/examples/survey/005_yrotation.py +++ b/examples/survey/005_yrotation.py @@ -6,73 +6,570 @@ ################################################################################ import numpy as np import matplotlib.pyplot as plt +from matplotlib.lines import Line2D -from _helpers import survey_test, summary_plot +from _helpers import madpoint_twiss_survey, add_to_plot import xtrack as xt ################################################################################ # User variables ################################################################################ -PLOT_COMPARISONS = True TOLERANCE = 1E-12 -REL_TOLERANCE = 1E-3 +ROT_DEG = 30 +HXL = 0.1 +DX = 0.5 +PLOT_X_LIMS = (-0.1, 4.1) +PLOT_Y_LIMS = (-1, 1) ################################################################################ +# Plot setup +################################################################################ +fig_mult_h = plt.figure(figsize = (16, 8)) +gs_mult_h = fig_mult_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_mult_h = gs_mult_h.subplots(sharex = 'row', sharey = True) + +fig_mult_v = plt.figure(figsize = (16, 8)) +gs_mult_v = fig_mult_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_mult_v = gs_mult_v.subplots(sharex = 'row', sharey = True) + +fig_shift_h = plt.figure(figsize = (16, 8)) +gs_shift_h = fig_shift_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_shift_h = gs_shift_h.subplots(sharex = 'row', sharey = True) + +fig_shift_v = plt.figure(figsize = (16, 8)) +gs_shift_v = fig_shift_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_shift_v = gs_shift_v.subplots(sharex = 'row', sharey = True) + +################################################################################ +# Horizontal Multipole Bends +################################################################################ + +######################################## +# Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1), + env.new('yrotation', xt.YRotation, angle = 0, at = 2), + env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['yrotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = 0 +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = HXL +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = HXL +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Vertical Multipole Bends +################################################################################ + +######################################## # Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1, rot_s_rad = np.pi / 2), + env.new('yrotation', xt.YRotation, angle = 0, at = 2), + env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3, rot_s_rad = np.pi / 2), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['yrotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = 0 +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = HXL +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = HXL +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + ################################################################################ -env = xt.Environment(particle_ref=xt.Particles(p0c=1e9)) +# Horizontal Shift +################################################################################ + +######################################## +# Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) line = env.new_line(name = 'line', components=[ - env.new('bend', xt.Bend, k0 = 0, h = 0, length = 0.5, at = 1), + env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), env.new('yrotation', xt.YRotation, angle = 0, at = 2), - env.new('end', xt.Marker, at = 3)]) -line.configure_bend_model(edge = 'suppressed') -line.config.XTRACK_USE_EXACT_DRIFTS = True + env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['shift1'].dx = 0 +line['shift2'].dx = 0 +line['yrotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['shift1'].dx = 0 +line['shift2'].dx = 0 +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['shift1'].dx = DX +line['shift2'].dx = 0 +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) +add_to_plot(axs_shift_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['shift1'].dx = 0 +line['shift2'].dx = DX +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['shift1'].dx = DX +line['shift2'].dx = DX +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Vertical Shift ################################################################################ + +######################################## +# Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), + env.new('yrotation', xt.YRotation, angle = 0, at = 2), + env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['shift1'].dy = 0 +line['shift2'].dy = 0 +line['yrotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['shift1'].dy = 0 +line['shift2'].dy = 0 +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['shift1'].dy = DX +line['shift2'].dy = 0 +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['shift1'].dy = 0 +line['shift2'].dy = DX +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['shift1'].dy = DX +line['shift2'].dy = DX +line['yrotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### # Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Show Plots ################################################################################ ######################################## -# Rotation only +# Horizontal Mult ######################################## -line['bend'].k0 = 0 -line['bend'].h = 0 -line['bend'].rot_s_rad = 0 +# y labels +axs_mult_h[0, 0].set_ylabel('Survey x,y [m]') +axs_mult_h[1, 0].set_ylabel('Twiss x,y [m]') +axs_mult_h[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_mult_h[0, 2].set_xlabel('Z [m]') +axs_mult_h[1, 2].set_xlabel('s [m]') +axs_mult_h[2, 2].set_xlabel('s [m]') -line['yrotation'].angle = 10 +# Titles +axs_mult_h[0, 0].set_title('Rot = 0,\n HXL1 = 0, HXL2 = 0') +axs_mult_h[0, 1].set_title('Rot != 0,\n HXL1 = 0, HXL2 = 0') +axs_mult_h[0, 2].set_title('Rot != 0,\n HXL1 != 0, HXL2 = 0') +axs_mult_h[0, 3].set_title('Rot != 0,\n HXL1 = 0, HXL2 != 0') +axs_mult_h[0, 4].set_title('Rot != 0,\n HXL1 != 0, HXL2 != 0') -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Rotation Y, bend off') +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_mult_h.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_mult_h.suptitle('Horizontal Multipole and YRotation') ######################################## -# Shift + horizontal bend +# Vertical Mult ######################################## -line['bend'].k0 = 1 -line['bend'].h = 1 -line['bend'].rot_s_rad = 0 +# y labels +axs_mult_v[0, 0].set_ylabel('Survey x,y [m]') +axs_mult_v[1, 0].set_ylabel('Twiss x,y [m]') +axs_mult_v[2, 0].set_ylabel('MadPoint xx,yy [m]') -line['yrotation'].angle = 10 +# x labels +axs_mult_v[0, 2].set_xlabel('Z [m]') +axs_mult_v[1, 2].set_xlabel('s [m]') +axs_mult_v[2, 2].set_xlabel('s [m]') -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Rotation Y, Horiztonal Bend') +# Titles +axs_mult_v[0, 0].set_title('Rot = 0,\n HYL1 = 0, HYL2 = 0') +axs_mult_v[0, 1].set_title('Rot != 0,\n HYL1 = 0, HYL2 = 0') +axs_mult_v[0, 2].set_title('Rot != 0,\n HYL1 != 0, HYL2 = 0') +axs_mult_v[0, 3].set_title('Rot != 0,\n HYL1 = 0, HYL2 != 0') +axs_mult_v[0, 4].set_title('Rot != 0,\n HYL1 != 0, HYL2 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_mult_v.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_mult_v.suptitle('Vertical Multipole and YRotation') ######################################## -# Shift + vertical bend +# Horizontal XYShift ######################################## -line['bend'].k0 = 1 -line['bend'].h = 1 -line['bend'].rot_s_rad = np.pi / 2 +# y labels +axs_shift_h[0, 0].set_ylabel('Survey x,y [m]') +axs_shift_h[1, 0].set_ylabel('Twiss x,y [m]') +axs_shift_h[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_shift_h[0, 2].set_xlabel('Z [m]') +axs_shift_h[1, 2].set_xlabel('s [m]') +axs_shift_h[2, 2].set_xlabel('s [m]') + +# Titles +axs_shift_h[0, 0].set_title('Rot = 0,\n DX1 = 0, DX2 = 0') +axs_shift_h[0, 1].set_title('Rot != 0,\n DX1 = 0, DX2 = 0') +axs_shift_h[0, 2].set_title('Rot != 0,\n DX1 != 0, DX2 = 0') +axs_shift_h[0, 3].set_title('Rot != 0,\n DX1 = 0, DX2 != 0') +axs_shift_h[0, 4].set_title('Rot != 0,\n DX1 != 0, DX2 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_shift_h.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_shift_h.suptitle('Horizontal XYShift and YRotation') + +######################################## +# Vertical XYShift +######################################## +# y labels +axs_shift_v[0, 0].set_ylabel('Survey x,y [m]') +axs_shift_v[1, 0].set_ylabel('Twiss x,y [m]') +axs_shift_v[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_shift_v[0, 2].set_xlabel('Z [m]') +axs_shift_v[1, 2].set_xlabel('s [m]') +axs_shift_v[2, 2].set_xlabel('s [m]') -line['yrotation'].angle = 10 +# Titles +axs_shift_v[0, 0].set_title('Rot = 0,\n DY1 = 0, DY2 = 0') +axs_shift_v[0, 1].set_title('Rot != 0,\n DY1 = 0, DY2 = 0') +axs_shift_v[0, 2].set_title('Rot != 0,\n DY1 != 0, DY2 = 0') +axs_shift_v[0, 3].set_title('Rot != 0,\n DY1 = 0, DY2 != 0') +axs_shift_v[0, 4].set_title('Rot != 0,\n DY1 != 0, DY2 != 0') -sv, tw = survey_test(line) -if PLOT_COMPARISONS: - summary_plot(sv, tw, 'Rotation Y, Vertical Bend') +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_shift_v.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_shift_v.suptitle('Vertical XYShift and YRotation') ######################################## -# Show plots +# Show ######################################## -if PLOT_COMPARISONS: - plt.show() +plt.tight_layout() +plt.show() diff --git a/examples/survey/006_srotation.py b/examples/survey/006_srotation.py new file mode 100644 index 000000000..a9c732913 --- /dev/null +++ b/examples/survey/006_srotation.py @@ -0,0 +1,575 @@ +""" +Test the survey on SRotation +""" +################################################################################ +# Packages +################################################################################ +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D + +from _helpers import madpoint_twiss_survey, add_to_plot +import xtrack as xt + +################################################################################ +# User variables +################################################################################ +TOLERANCE = 1E-12 +ROT_DEG = 30 +HXL = 0.1 +DX = 0.5 +PLOT_X_LIMS = (-0.1, 4.1) +PLOT_Y_LIMS = (-1, 1) + +################################################################################ +# Plot setup +################################################################################ +fig_mult_h = plt.figure(figsize = (16, 8)) +gs_mult_h = fig_mult_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_mult_h = gs_mult_h.subplots(sharex = 'row', sharey = True) + +fig_mult_v = plt.figure(figsize = (16, 8)) +gs_mult_v = fig_mult_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_mult_v = gs_mult_v.subplots(sharex = 'row', sharey = True) + +fig_shift_h = plt.figure(figsize = (16, 8)) +gs_shift_h = fig_shift_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_shift_h = gs_shift_h.subplots(sharex = 'row', sharey = True) + +fig_shift_v = plt.figure(figsize = (16, 8)) +gs_shift_v = fig_shift_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) +axs_shift_v = gs_shift_v.subplots(sharex = 'row', sharey = True) + +################################################################################ +# Horizontal Multipole Bends +################################################################################ + +######################################## +# Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1), + env.new('srotation', xt.SRotation, angle = 0, at = 2), + env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['srotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = 0 +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = HXL +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = HXL +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Vertical Multipole Bends +################################################################################ + +######################################## +# Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1, rot_s_rad = np.pi / 2), + env.new('srotation', xt.SRotation, angle = 0, at = 2), + env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3, rot_s_rad = np.pi / 2), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['srotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = 0 +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = 0 +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['mult1'].hxl = 0 +line['mult2'].hxl = HXL +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['mult1'].hxl = HXL +line['mult2'].hxl = HXL +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_mult_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Horizontal Shift +################################################################################ + +######################################## +# Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), + env.new('srotation', xt.SRotation, angle = 0, at = 2), + env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['shift1'].dx = 0 +line['shift2'].dx = 0 +line['srotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['shift1'].dx = 0 +line['shift2'].dx = 0 +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['shift1'].dx = DX +line['shift2'].dx = 0 +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['shift1'].dx = 0 +line['shift2'].dx = DX +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['shift1'].dx = DX +line['shift2'].dx = DX +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Vertical Shift +################################################################################ + +######################################## +# Create line +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) +line = env.new_line(name = 'line', components=[ + env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), + env.new('srotation', xt.SRotation, angle = 0, at = 2), + env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), + env.new('end', xt.Marker, at = 4)]) + +######################################## +# Off +######################################## +line['shift1'].dy = 0 +line['shift2'].dy = 0 +line['srotation'].angle = np.rad2deg(0) + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Element off must have no effect +assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) +assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) + +######################################## +# Rotation Only +######################################## +line['shift1'].dy = 0 +line['shift2'].dy = 0 +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot +######################################## +line['shift1'].dy = DX +line['shift2'].dy = 0 +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Rot then mult +######################################## +line['shift1'].dy = 0 +line['shift2'].dy = DX +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +######################################## +# Mult then rot then mult +######################################## +line['shift1'].dy = DX +line['shift2'].dy = DX +line['srotation'].angle = ROT_DEG + +sv, tw = madpoint_twiss_survey(line) + +add_to_plot(axs_shift_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) + +#################### +# Tests +#################### +# Only rotation must have no effect +# TODO + +################################################################################ +# Show Plots +################################################################################ + +######################################## +# Horizontal Mult +######################################## +# y labels +axs_mult_h[0, 0].set_ylabel('Survey x,y [m]') +axs_mult_h[1, 0].set_ylabel('Twiss x,y [m]') +axs_mult_h[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_mult_h[0, 2].set_xlabel('Z [m]') +axs_mult_h[1, 2].set_xlabel('s [m]') +axs_mult_h[2, 2].set_xlabel('s [m]') + +# Titles +axs_mult_h[0, 0].set_title('Rot = 0,\n HXL1 = 0, HXL2 = 0') +axs_mult_h[0, 1].set_title('Rot != 0,\n HXL1 = 0, HXL2 = 0') +axs_mult_h[0, 2].set_title('Rot != 0,\n HXL1 != 0, HXL2 = 0') +axs_mult_h[0, 3].set_title('Rot != 0,\n HXL1 = 0, HXL2 != 0') +axs_mult_h[0, 4].set_title('Rot != 0,\n HXL1 != 0, HXL2 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_mult_h.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_mult_h.suptitle('Horizontal Multipole and SRotation') + +######################################## +# Vertical Mult +######################################## +# y labels +axs_mult_v[0, 0].set_ylabel('Survey x,y [m]') +axs_mult_v[1, 0].set_ylabel('Twiss x,y [m]') +axs_mult_v[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_mult_v[0, 2].set_xlabel('Z [m]') +axs_mult_v[1, 2].set_xlabel('s [m]') +axs_mult_v[2, 2].set_xlabel('s [m]') + +# Titles +axs_mult_v[0, 0].set_title('Rot = 0,\n HYL1 = 0, HYL2 = 0') +axs_mult_v[0, 1].set_title('Rot != 0,\n HYL1 = 0, HYL2 = 0') +axs_mult_v[0, 2].set_title('Rot != 0,\n HYL1 != 0, HYL2 = 0') +axs_mult_v[0, 3].set_title('Rot != 0,\n HYL1 = 0, HYL2 != 0') +axs_mult_v[0, 4].set_title('Rot != 0,\n HYL1 != 0, HYL2 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_mult_v.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_mult_v.suptitle('Vertical Multipole and SRotation') + +######################################## +# Horizontal XYShift +######################################## +# y labels +axs_shift_h[0, 0].set_ylabel('Survey x,y [m]') +axs_shift_h[1, 0].set_ylabel('Twiss x,y [m]') +axs_shift_h[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_shift_h[0, 2].set_xlabel('Z [m]') +axs_shift_h[1, 2].set_xlabel('s [m]') +axs_shift_h[2, 2].set_xlabel('s [m]') + +# Titles +axs_shift_h[0, 0].set_title('Rot = 0,\n DX1 = 0, DX2 = 0') +axs_shift_h[0, 1].set_title('Rot != 0,\n DX1 = 0, DX2 = 0') +axs_shift_h[0, 2].set_title('Rot != 0,\n DX1 != 0, DX2 = 0') +axs_shift_h[0, 3].set_title('Rot != 0,\n DX1 = 0, DX2 != 0') +axs_shift_h[0, 4].set_title('Rot != 0,\n DX1 != 0, DX2 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_shift_h.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_shift_h.suptitle('Horizontal XYShift and SRotation') + +######################################## +# Vertical XYShift +######################################## +# y labels +axs_shift_v[0, 0].set_ylabel('Survey x,y [m]') +axs_shift_v[1, 0].set_ylabel('Twiss x,y [m]') +axs_shift_v[2, 0].set_ylabel('MadPoint xx,yy [m]') + +# x labels +axs_shift_v[0, 2].set_xlabel('Z [m]') +axs_shift_v[1, 2].set_xlabel('s [m]') +axs_shift_v[2, 2].set_xlabel('s [m]') + +# Titles +axs_shift_v[0, 0].set_title('Rot = 0,\n DY1 = 0, DY2 = 0') +axs_shift_v[0, 1].set_title('Rot != 0,\n DY1 = 0, DY2 = 0') +axs_shift_v[0, 2].set_title('Rot != 0,\n DY1 != 0, DY2 = 0') +axs_shift_v[0, 3].set_title('Rot != 0,\n DY1 = 0, DY2 != 0') +axs_shift_v[0, 4].set_title('Rot != 0,\n DY1 != 0, DY2 != 0') + +legend_elements = [ + Line2D([0], [0], color = 'black', lw = 2, label = 'x'), + Line2D([0], [0], color = 'red', lw = 2, label = 'y')] +fig_shift_v.legend( + handles = legend_elements, + loc = 'lower center', + ncol = 2, + frameon = False, + bbox_to_anchor = (0.5, -0.05)) +fig_shift_v.suptitle('Vertical XYShift and SRotation') + +######################################## +# Show +######################################## +plt.tight_layout() +plt.show() diff --git a/examples/survey/_helpers.py b/examples/survey/_helpers.py index d1a97d366..c12c98ad1 100644 --- a/examples/survey/_helpers.py +++ b/examples/survey/_helpers.py @@ -2,7 +2,6 @@ Helpers for testing the MadPoint class """ import numpy as np -import matplotlib.pyplot as plt from _madpoint import MadPoint def zero_small_values(arr, tol): @@ -35,100 +34,40 @@ def madpoint_twiss_survey(line): return survey, twiss -# def summary_plot(survey, twiss, title, zero_tol = 1E-12): -# """ -# Summary plot comparing survey, twiss and MadPoint -# """ -# def zero_small_values(arr, tol = zero_tol): -# return np.where(np.abs(arr) < tol, 0, arr) - -# fig, axs = plt.subplots(3, 2, figsize=(16, 8)) -# axs[0, 0].plot( -# zero_small_values(survey.Z), -# zero_small_values(survey.X)) -# axs[1, 0].plot( -# zero_small_values(twiss.s), -# zero_small_values(twiss.x)) -# axs[2, 0].plot( -# zero_small_values(survey.s), -# zero_small_values(survey.xx)) - -# axs[0, 1].plot( -# zero_small_values(survey.Z), -# zero_small_values(survey.Y)) -# axs[1, 1].plot( -# zero_small_values(twiss.s), -# zero_small_values(twiss.y)) -# axs[2, 1].plot( -# zero_small_values(survey.s), -# zero_small_values(survey.yy)) - -# axs[0, 0].set_xlabel('Z [m]') -# axs[0, 0].set_ylabel('X [m]') -# axs[1, 0].set_xlabel('s [m]') -# axs[1, 0].set_ylabel('x [m]') -# axs[2, 0].set_xlabel('s [m]') -# axs[2, 0].set_ylabel('xx [m]') - -# axs[0, 1].set_xlabel('Z [m]') -# axs[0, 1].set_ylabel('Y [m]') -# axs[1, 1].set_xlabel('s [m]') -# axs[1, 1].set_ylabel('y [m]') -# axs[2, 1].set_xlabel('s [m]') -# axs[2, 1].set_ylabel('yy [m]') - -# fig.suptitle(title) - -# plt.tight_layout() - -def summary_plot(survey, twiss, title, zero_tol = 1E-12): +def add_to_plot(axes, survey, twiss, index, tol = 1E-12, xlims = (-0.1, 2.1), ylims = (-2E-3, 2E-3)): """ - Summary plot comparing survey, twiss and MadPoint + Add a line to the overall plot to show a specific case """ - - fig, axs = plt.subplots(3, 1, figsize=(16, 8)) - axs[0].plot( - zero_small_values(survey.Z, zero_tol), - zero_small_values(survey.X, zero_tol), + axes[0, index].plot( + zero_small_values(survey.Z, tol), + zero_small_values(survey.X, tol), c = 'k') - axs[0].plot( - zero_small_values(survey.Z, zero_tol), - zero_small_values(survey.Y, zero_tol), + axes[0, index].plot( + zero_small_values(survey.Z, tol), + zero_small_values(survey.Y, tol), c = 'r') - - axs[1].plot( - zero_small_values(twiss.s, zero_tol), - zero_small_values(twiss.x, zero_tol), + + axes[1, index].plot( + zero_small_values(twiss.s, tol), + zero_small_values(twiss.x, tol), c = 'k') - axs[1].plot( - zero_small_values(twiss.s, zero_tol), - zero_small_values(twiss.y, zero_tol), + axes[1, index].plot( + zero_small_values(twiss.s, tol), + zero_small_values(twiss.y, tol), c = 'r') - - axs[2].plot( - zero_small_values(survey.s, zero_tol), - zero_small_values(survey.xx, zero_tol), + + axes[2, index].plot( + zero_small_values(survey.s, tol), + zero_small_values(survey.xx, tol), c = 'k') - axs[2].plot( - zero_small_values(survey.s, zero_tol), - zero_small_values(survey.yy, zero_tol), + axes[2, index].plot( + zero_small_values(survey.s, tol), + zero_small_values(survey.yy, tol), c = 'r') - axs[0].set_xlabel('Z [m]') - axs[0].set_ylabel('X [m]') - axs[0].set_title('Survey') - - axs[1].set_xlabel('s [m]') - axs[1].set_ylabel('x [m]') - axs[1].set_title('Twiss') - - axs[2].set_xlabel('s [m]') - axs[2].set_ylabel('xx [m]') - axs[2].set_title('MadPoint') - - axs[2].legend(loc='upper center', bbox_to_anchor=(0.5, -0.05), - fancybox=True, ncol=2) - - fig.suptitle(title) - - plt.tight_layout() + axes[0, index].set_xlim(xlims) + axes[0, index].set_ylim(ylims) + axes[1, index].set_xlim(xlims) + axes[1, index].set_ylim(ylims) + axes[2, index].set_xlim(xlims) + axes[2, index].set_ylim(ylims) From cec3552fcbfe326d75a74ee2f3594d29b2a03b3d Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Thu, 14 Nov 2024 18:14:23 +0100 Subject: [PATCH 13/21] Solenoid Multipole rotations and shifts + tests --- examples/solenoid/mult_rot_test.py | 129 ++++++++++++++++++ examples/solenoid/mult_shift_test.py | 131 +++++++++++++++++++ xtrack/beam_elements/elements.py | 3 + xtrack/beam_elements/elements_src/solenoid.h | 47 +++++-- 4 files changed, 299 insertions(+), 11 deletions(-) create mode 100644 examples/solenoid/mult_rot_test.py create mode 100644 examples/solenoid/mult_shift_test.py diff --git a/examples/solenoid/mult_rot_test.py b/examples/solenoid/mult_rot_test.py new file mode 100644 index 000000000..ec0c40e28 --- /dev/null +++ b/examples/solenoid/mult_rot_test.py @@ -0,0 +1,129 @@ +""" +Test implementation of the multipole shifts in solenoid slices +============================================= +Author(s): John P T Salvesen +Email: john.salvesen@cern.ch +Date: 14-11-2024 +""" +################################################################################ +# Required Modules +################################################################################ +import xtrack as xt +import numpy as np +import matplotlib.pyplot as plt + +################################################################################ +# User Parameters +################################################################################ +N_SLICES = int(1E3) + +BETX = 100E-3 +BETY = 1E-3 +PX0 = 0 + +KS = 0.00 +K0 = 1E-3 +K1 = 1E-3 +K2 = 1E-3 + +X_ROT = 1E-3 +Y_ROT = 1E-3 + +################################################################################ +# Build Test Elements +################################################################################ +drift0 = xt.Drift(length = 1) +drift1 = xt.Drift(length = 1) + +bend = xt.Bend(length = 1, k0 = K0) +quad = xt.Quadrupole(length = 1, k1 = K1) +sext = xt.Sextupole(length = 1, k2 = K2) + +bend_sol = xt.Solenoid(length = 1 / N_SLICES, ks = KS, + knl = [K0 * (1/N_SLICES), 0, 0], num_multipole_kicks = 1) +quad_sol = xt.Solenoid(length = 1 / N_SLICES, ks = KS, + knl = [0, K1 * (1/N_SLICES), 0], num_multipole_kicks = 1) +sext_sol = xt.Solenoid(length = 1 / N_SLICES, ks = KS, + knl = [0, 0, K2 * (1/N_SLICES)], num_multipole_kicks = 1) + +################################################################################ +# Comparisons +################################################################################ + +for test_element, test_sol, title in zip( + [bend, quad, sext], [bend_sol, quad_sol, sext_sol], ['Bend', 'Quadrupole', 'Sextupole']): + ######################################## + # Build Lines + ######################################## + line = xt.Line( + elements = [drift0] + [test_element] + [drift0], + particle_ref = xt.Particles(p0c = 1E9, mass0 = xt.ELECTRON_MASS_EV)) + line.configure_bend_model(edge = 'suppressed') + + sol_line = xt.Line( + elements = [drift1] + [test_sol] * N_SLICES + [drift1], + particle_ref = xt.Particles(p0c = 1E9, mass0 = xt.ELECTRON_MASS_EV)) + + # Slice test line + line.slice_thick_elements( + slicing_strategies=[ + xt.Strategy(slicing = xt.Uniform(N_SLICES, mode='thin'), element_type = xt.Bend), + xt.Strategy(slicing = xt.Uniform(N_SLICES, mode='thin'), element_type = xt.Quadrupole), + xt.Strategy(slicing = xt.Uniform(N_SLICES, mode='thin'), element_type = xt.Sextupole)]) + + ######################################## + # Initialise Plot + ######################################## + fig, axs = plt.subplots(2, 2, figsize = (10, 10)) + axs = axs.flatten() + fig.suptitle(title) + + ######################################## + # Test and plot with shifts + ######################################## + for i, (rot_x, rot_y) in enumerate( + zip([0, X_ROT, 0, X_ROT],[0, 0, Y_ROT, Y_ROT])): + # TODO: No way to compare this currently?!? + test_sol.mult_rot_x = rot_x + test_sol.mult_rot_y = rot_y + + tw = line.twiss( + _continue_if_lost = True, + start = xt.START, + end = xt.END, + betx = BETX, + bety = BETY, + px = PX0) + tw_sol = sol_line.twiss( + _continue_if_lost = True, + start = xt.START, + end = xt.END, + betx = BETX, + bety = BETY, + px = PX0) + + axs[i].plot(tw.s, tw.x, color = 'k', label = 'Element x') + axs[i].plot(tw.s, tw.y, color = 'r', label = 'Element y') + axs[i].plot(tw_sol.s, tw_sol.x, color = 'b', linestyle = ':', label = 'Sol x') + axs[i].plot(tw_sol.s, tw_sol.y, color = 'g', linestyle = ':', label = 'Sol y') + + axs[i].set_title(f'Rot x = {rot_x} [TBD], Rot y = {rot_y} [TBD]') + + ######################################## + # Assertions + ######################################## + assert np.isclose(tw.x[-1], tw_sol.x[-1], rtol = 1E-6) + assert np.isclose(tw.y[-1], tw_sol.y[-1], rtol = 1E-6) + + ######################################## + # Figure adjustments + ######################################## + fig.legend( + labels = ['Element x', 'Element y', 'Sol x', 'Sol y'], + loc = 'upper center', + ncol = 4) + +######################################## +# Show Plots +######################################## +plt.show() diff --git a/examples/solenoid/mult_shift_test.py b/examples/solenoid/mult_shift_test.py new file mode 100644 index 000000000..9f083fb82 --- /dev/null +++ b/examples/solenoid/mult_shift_test.py @@ -0,0 +1,131 @@ +""" +Test implementation of the multipole shifts in solenoid slices +============================================= +Author(s): John P T Salvesen +Email: john.salvesen@cern.ch +Date: 14-11-2024 +""" +################################################################################ +# Required Modules +################################################################################ +import xtrack as xt +import numpy as np +import matplotlib.pyplot as plt + +################################################################################ +# User Parameters +################################################################################ +N_SLICES = int(1E3) + +BETX = 100E-3 +BETY = 1E-3 +PX0 = 0 + +KS = 0.00 +K0 = 1E-3 +K1 = 1E-3 +K2 = 1E-3 + +X_SHIFT = 1E-3 +Y_SHIFT = 1E-3 + +################################################################################ +# Build Test Elements +################################################################################ +drift0 = xt.Drift(length = 1) +drift1 = xt.Drift(length = 1) + +bend = xt.Bend(length = 1, k0 = K0) +quad = xt.Quadrupole(length = 1, k1 = K1) +sext = xt.Sextupole(length = 1, k2 = K2) + +bend_sol = xt.Solenoid(length = 1 / N_SLICES, ks = KS, + knl = [K0 * (1/N_SLICES), 0, 0], num_multipole_kicks = 1) +quad_sol = xt.Solenoid(length = 1 / N_SLICES, ks = KS, + knl = [0, K1 * (1/N_SLICES), 0], num_multipole_kicks = 1) +sext_sol = xt.Solenoid(length = 1 / N_SLICES, ks = KS, + knl = [0, 0, K2 * (1/N_SLICES)], num_multipole_kicks = 1) + +################################################################################ +# Comparisons +################################################################################ + +for test_element, test_sol, title in zip( + [bend, quad, sext], [bend_sol, quad_sol, sext_sol], ['Bend', 'Quadrupole', 'Sextupole']): + ######################################## + # Build Lines + ######################################## + line = xt.Line( + elements = [drift0] + [test_element] + [drift0], + particle_ref = xt.Particles(p0c = 1E9, mass0 = xt.ELECTRON_MASS_EV)) + line.configure_bend_model(edge = 'suppressed') + + sol_line = xt.Line( + elements = [drift1] + [test_sol] * N_SLICES + [drift1], + particle_ref = xt.Particles(p0c = 1E9, mass0 = xt.ELECTRON_MASS_EV)) + + # Slice test line + line.slice_thick_elements( + slicing_strategies=[ + xt.Strategy(slicing = xt.Uniform(N_SLICES, mode='thin'), element_type = xt.Bend), + xt.Strategy(slicing = xt.Uniform(N_SLICES, mode='thin'), element_type = xt.Quadrupole), + xt.Strategy(slicing = xt.Uniform(N_SLICES, mode='thin'), element_type = xt.Sextupole)]) + + ######################################## + # Initialise Plot + ######################################## + fig, axs = plt.subplots(2, 2, figsize = (10, 10)) + axs = axs.flatten() + fig.suptitle(title) + + ######################################## + # Test and plot with shifts + ######################################## + for i, (shift_x, shift_y) in enumerate( + zip([0, X_SHIFT, 0, X_SHIFT],[0, 0, Y_SHIFT, Y_SHIFT])): + + test_element.shift_x = shift_x + test_element.shift_y = shift_y + test_sol.mult_shift_x = shift_x + test_sol.mult_shift_y = shift_y + + tw = line.twiss( + _continue_if_lost = True, + start = xt.START, + end = xt.END, + betx = BETX, + bety = BETY, + px = PX0) + tw_sol = sol_line.twiss( + _continue_if_lost = True, + start = xt.START, + end = xt.END, + betx = BETX, + bety = BETY, + px = PX0) + + axs[i].plot(tw.s, tw.x, color = 'k', label = 'Element x') + axs[i].plot(tw.s, tw.y, color = 'r', label = 'Element y') + axs[i].plot(tw_sol.s, tw_sol.x, color = 'b', linestyle = ':', label = 'Sol x') + axs[i].plot(tw_sol.s, tw_sol.y, color = 'g', linestyle = ':', label = 'Sol y') + + axs[i].set_title(f'Shift x = {shift_x * 1000} [mm], Shift y = {shift_y * 1000} [mm]') + + ######################################## + # Assertions + ######################################## + assert np.isclose(tw.x[-1], tw_sol.x[-1], rtol = 1E-6) + assert np.isclose(tw.y[-1], tw_sol.y[-1], rtol = 1E-6) + + ######################################## + # Figure adjustments + ######################################## + fig.legend( + labels = ['Element x', 'Element y', 'Sol x', 'Sol y'], + loc = 'upper center', + ncol = 4) + +######################################## +# Show Plots +######################################## +plt.show() diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index a0ca37179..1466d05a0 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -1293,8 +1293,11 @@ class Solenoid(BeamElement): 'inv_factorial_order': xo.Float64, 'knl': xo.Float64[ALLOCATED_MULTIPOLE_ORDER + 1], 'ksl': xo.Float64[ALLOCATED_MULTIPOLE_ORDER + 1], + 'mult_rot_x_rad': xo.Float64, 'mult_rot_y_rad': xo.Float64, 'mult_shift_x': xo.Float64, + 'mult_shift_y': xo.Float64, + 'mult_shift_s': xo.Float64, } _skip_in_to_dict = ['_order', 'inv_factorial_order'] # defined by knl, etc. diff --git a/xtrack/beam_elements/elements_src/solenoid.h b/xtrack/beam_elements/elements_src/solenoid.h index a5a0f73f2..f74306675 100644 --- a/xtrack/beam_elements/elements_src/solenoid.h +++ b/xtrack/beam_elements/elements_src/solenoid.h @@ -37,18 +37,33 @@ void Solenoid_track_local_particle(SolenoidData el, LocalParticle* part0) { const double slice_length = length / (num_multipole_kicks + 1); const double kick_weight = 1. / num_multipole_kicks; + double mult_rot_x_rad = SolenoidData_get_mult_rot_x_rad(el); double mult_rot_y_rad = SolenoidData_get_mult_rot_y_rad(el); double mult_shift_x = SolenoidData_get_mult_shift_x(el); - double sin_angle, cos_angle, tan_angle; + double mult_shift_y = SolenoidData_get_mult_shift_y(el); + double mult_shift_s = SolenoidData_get_mult_shift_s(el); + + double sin_x_rot, cos_x_rot, tan_x_rot; + double sin_y_rot, cos_y_rot, tan_y_rot; + if (mult_rot_x_rad != 0) { + sin_x_rot = sin(mult_rot_x_rad); + cos_x_rot = cos(mult_rot_x_rad); + tan_x_rot = sin_x_rot / cos_x_rot; + } + else { + sin_x_rot = 0; + cos_x_rot = 1; + tan_x_rot = 0; + } if (mult_rot_y_rad != 0) { - sin_angle = sin(mult_rot_y_rad); - cos_angle = cos(mult_rot_y_rad); - tan_angle = sin_angle / cos_angle; + sin_y_rot = sin(mult_rot_y_rad); + cos_y_rot = cos(mult_rot_y_rad); + tan_y_rot = sin_y_rot / cos_y_rot; } else { - sin_angle = 0; - cos_angle = 1; - tan_angle = 0; + sin_y_rot = 0; + cos_y_rot = 1; + tan_y_rot = 0; } @@ -65,17 +80,27 @@ void Solenoid_track_local_particle(SolenoidData el, LocalParticle* part0) { Solenoid_thick_track_single_particle(part, slice_length, ks, radiation_flag); LocalParticle_add_to_x(part, -mult_shift_x); - if (sin_angle != 0) { - YRotation_single_particle(part, sin_angle, cos_angle, tan_angle); + LocalParticle_add_to_y(part, -mult_shift_y); + LocalParticle_add_to_s(part, -mult_shift_s); + if (sin_x_rot != 0) { + XRotation_single_particle(part, sin_x_rot, cos_x_rot, tan_x_rot); + } + if (sin_y_rot != 0) { + YRotation_single_particle(part, sin_y_rot, cos_y_rot, tan_y_rot); } track_multipolar_kick_bend( part, order, inv_factorial_order, knl, ksl, factor_knl_ksl, kick_weight, 0, 0, 0, 0); - if (sin_angle != 0) { - YRotation_single_particle(part, -sin_angle, cos_angle, -tan_angle); + if (sin_y_rot != 0) { + YRotation_single_particle(part, -sin_y_rot, -cos_y_rot, -tan_y_rot); + } + if (sin_x_rot != 0) { + XRotation_single_particle(part, -sin_x_rot, -cos_x_rot, -tan_x_rot); } + LocalParticle_add_to_s(part, mult_shift_s); + LocalParticle_add_to_y(part, mult_shift_y); LocalParticle_add_to_x(part, mult_shift_x); } From 9802d0ecb34214ed7494ad24e6c18126ab6ff399 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Sun, 24 Nov 2024 17:49:38 +0100 Subject: [PATCH 14/21] Corrected sign on rotation --- xtrack/beam_elements/elements_src/solenoid.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xtrack/beam_elements/elements_src/solenoid.h b/xtrack/beam_elements/elements_src/solenoid.h index f74306675..4475ba7d2 100644 --- a/xtrack/beam_elements/elements_src/solenoid.h +++ b/xtrack/beam_elements/elements_src/solenoid.h @@ -94,10 +94,10 @@ void Solenoid_track_local_particle(SolenoidData el, LocalParticle* part0) { kick_weight, 0, 0, 0, 0); if (sin_y_rot != 0) { - YRotation_single_particle(part, -sin_y_rot, -cos_y_rot, -tan_y_rot); + YRotation_single_particle(part, -sin_y_rot, cos_y_rot, -tan_y_rot); } if (sin_x_rot != 0) { - XRotation_single_particle(part, -sin_x_rot, -cos_x_rot, -tan_x_rot); + XRotation_single_particle(part, -sin_x_rot, cos_x_rot, -tan_x_rot); } LocalParticle_add_to_s(part, mult_shift_s); LocalParticle_add_to_y(part, mult_shift_y); From d2164aa540f64fe8527e0b7ed550112f6b2b607a Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Sun, 24 Nov 2024 17:50:00 +0100 Subject: [PATCH 15/21] Check removed as forces the specific case of equal length rings --- xtrack/multiline/multiline.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/xtrack/multiline/multiline.py b/xtrack/multiline/multiline.py index b7c4df137..6ef663dfe 100644 --- a/xtrack/multiline/multiline.py +++ b/xtrack/multiline/multiline.py @@ -459,11 +459,11 @@ def install_beambeam_interactions(self, clockwise_line, anticlockwise_line, for nn, ll in self.lines.items(): ll.unfreeze() - if clockwise_line is not None and anticlockwise_line is not None: - circumference_cw = self.lines[clockwise_line].get_length() - circumference_acw = self.lines[anticlockwise_line].get_length() - assert np.isclose(circumference_cw, circumference_acw, - atol=1e-4, rtol=0) + # if clockwise_line is not None and anticlockwise_line is not None: + # circumference_cw = self.lines[clockwise_line].get_length() + # circumference_acw = self.lines[anticlockwise_line].get_length() + # assert np.isclose(circumference_cw, circumference_acw, + # atol=1e-4, rtol=0) import xfields as xf bb_df_cw, bb_df_acw = xf.install_beambeam_elements_in_lines( @@ -490,6 +490,8 @@ def install_beambeam_interactions(self, clockwise_line, anticlockwise_line, def configure_beambeam_interactions(self, num_particles, nemitt_x, nemitt_y, crab_strong_beam=True, + nemitt_x_cw=None, nemitt_y_cw=None, + nemitt_x_acw=None, nemitt_y_acw=None, use_antisymmetry=False, separation_bumps=None): @@ -548,6 +550,8 @@ def configure_beambeam_interactions(self, num_particles, num_particles=num_particles, nemitt_x=nemitt_x, nemitt_y=nemitt_y, crab_strong_beam=crab_strong_beam, + nemitt_x_cw=nemitt_x_cw, nemitt_y_cw=nemitt_y_cw, + nemitt_x_acw=nemitt_x_acw, nemitt_y_acw=nemitt_y_acw, ip_names=self._bb_config['ip_names'], use_antisymmetry=use_antisymmetry, separation_bumps=separation_bumps) From 2b90e7f6bbad779158eb80ec4ab6bed1e6998355 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Sun, 24 Nov 2024 17:50:31 +0100 Subject: [PATCH 16/21] Rotation test of solenoid Tests vertical and horizontal reference frame shifts --- examples/solenoid/mult_rot_test.py | 215 ++++++++++++++++------------- 1 file changed, 120 insertions(+), 95 deletions(-) diff --git a/examples/solenoid/mult_rot_test.py b/examples/solenoid/mult_rot_test.py index ec0c40e28..ba93859c4 100644 --- a/examples/solenoid/mult_rot_test.py +++ b/examples/solenoid/mult_rot_test.py @@ -3,7 +3,7 @@ ============================================= Author(s): John P T Salvesen Email: john.salvesen@cern.ch -Date: 14-11-2024 +Date: 18-11-2024 """ ################################################################################ # Required Modules @@ -15,113 +15,138 @@ ################################################################################ # User Parameters ################################################################################ -N_SLICES = int(1E3) +N_SLICES = int(1E2) +K0 = 1E-3 +L_SOL = 1 +XING_RAD = 1E-3 BETX = 100E-3 BETY = 1E-3 -PX0 = 0 - -KS = 0.00 -K0 = 1E-3 -K1 = 1E-3 -K2 = 1E-3 - -X_ROT = 1E-3 -Y_ROT = 1E-3 ################################################################################ -# Build Test Elements +# Build Test Lines ################################################################################ -drift0 = xt.Drift(length = 1) -drift1 = xt.Drift(length = 1) -bend = xt.Bend(length = 1, k0 = K0) -quad = xt.Quadrupole(length = 1, k1 = K1) -sext = xt.Sextupole(length = 1, k2 = K2) +######################################## +# Build Environment +######################################## +env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) + +######################################## +# Line (beamline frame) +######################################## +bl_components_in = [env.new('bl_drift0', xt.Drift, length = 1)] +bl_components_out = [env.new('bl_drift1', xt.Drift, length = 1)] + +bl_components_sol = [ + env.new(f'bl_sol.{i}', xt.Solenoid, + length = (L_SOL / N_SLICES), + ks = 0, + knl = [K0 * (L_SOL / N_SLICES), 0, 0], + num_multipole_kicks = 1) + for i in range(N_SLICES)] -bend_sol = xt.Solenoid(length = 1 / N_SLICES, ks = KS, - knl = [K0 * (1/N_SLICES), 0, 0], num_multipole_kicks = 1) -quad_sol = xt.Solenoid(length = 1 / N_SLICES, ks = KS, - knl = [0, K1 * (1/N_SLICES), 0], num_multipole_kicks = 1) -sext_sol = xt.Solenoid(length = 1 / N_SLICES, ks = KS, - knl = [0, 0, K2 * (1/N_SLICES)], num_multipole_kicks = 1) +bl_line = env.new_line( + components = bl_components_in + bl_components_sol + bl_components_out) + +######################################## +# Line (horizontal rotated frame) +######################################## +hrot_components_in = [ + env.new('hrot_drift0', xt.Drift, length = 1), + env.new('hshift_in', xt.XYShift, dx = np.sin(XING_RAD) * L_SOL / 2), + env.new('hrot_in', xt.YRotation, angle = -np.rad2deg(XING_RAD))] + +hrot_components_out = [ + env.new('hrot_out', xt.YRotation, angle = np.rad2deg(XING_RAD)), + env.new('hshift_out', xt.XYShift, dx = np.sin(XING_RAD) * L_SOL / 2), + env.new('hrot_drift1', xt.Drift, length = 1)] + +hrot_components_sol = [ + env.new(f'hrot_sol.{i}', xt.Solenoid, + length = (L_SOL / N_SLICES) * np.cos(XING_RAD), + ks = 0, + knl = [K0 * (L_SOL / N_SLICES), 0, 0], + num_multipole_kicks = 1, + mult_rot_y_rad = XING_RAD, + mult_shift_x = np.sin(XING_RAD) * L_SOL * (i/N_SLICES - 1/2)) + for i in range(N_SLICES)] + +hrot_line = env.new_line( + components = hrot_components_in + hrot_components_sol + hrot_components_out) + +######################################## +# Line (vertical rotated frame) +######################################## +vrot_components_in = [ + env.new('vrot_drift0', xt.Drift, length = 1), + env.new('vshift_in', xt.XYShift, dy = np.sin(XING_RAD) * L_SOL / 2), + env.new('vrot_in', xt.XRotation, angle = np.rad2deg(XING_RAD))] +# TODO: Minus sign difference here as still inconsistent definition with XRotation and YRotation +vrot_components_out = [ + env.new('vrot_out', xt.XRotation, angle = -np.rad2deg(XING_RAD)), + env.new('vshift_out', xt.XYShift, dy = np.sin(XING_RAD) * L_SOL / 2), + env.new('vrot_drift1', xt.Drift, length = 1)] + +vrot_components_sol = [ + env.new(f'vrot_sol.{i}', xt.Solenoid, + length = (L_SOL / N_SLICES) * np.cos(XING_RAD), + ks = 0, + knl = [K0 * (L_SOL / N_SLICES), 0, 0], + num_multipole_kicks = 1, + mult_rot_x_rad = XING_RAD, + mult_shift_y = np.sin(XING_RAD) * L_SOL * (i/N_SLICES - 1/2)) + for i in range(N_SLICES)] + +vrot_line = env.new_line( + components = vrot_components_in + vrot_components_sol + vrot_components_out) ################################################################################ # Comparisons ################################################################################ +bl_twiss = bl_line.twiss( + method = '4d', + start = xt.START, + end = xt.END, + betx = BETX, + bety = BETY) + +hrot_twiss = hrot_line.twiss( + method = '4d', + start = xt.START, + end = xt.END, + betx = BETX, + bety = BETY) + +vrot_twiss = vrot_line.twiss( + method = '4d', + start = xt.START, + end = xt.END, + betx = BETX, + bety = BETY) + +bl_twiss.plot() +plt.title('Beamline Frame') +bl_twiss.plot('x y') +plt.title('Beamline Frame') + +hrot_twiss.plot() +plt.title('Horizontal Rotated Frame') +hrot_twiss.plot('x y') +plt.title('Horizontal Rotated Frame') + +vrot_twiss.plot() +plt.title('Vertical Rotated Frame') +vrot_twiss.plot('x y') +plt.title('Vertical Rotated Frame') -for test_element, test_sol, title in zip( - [bend, quad, sext], [bend_sol, quad_sol, sext_sol], ['Bend', 'Quadrupole', 'Sextupole']): - ######################################## - # Build Lines - ######################################## - line = xt.Line( - elements = [drift0] + [test_element] + [drift0], - particle_ref = xt.Particles(p0c = 1E9, mass0 = xt.ELECTRON_MASS_EV)) - line.configure_bend_model(edge = 'suppressed') - - sol_line = xt.Line( - elements = [drift1] + [test_sol] * N_SLICES + [drift1], - particle_ref = xt.Particles(p0c = 1E9, mass0 = xt.ELECTRON_MASS_EV)) - - # Slice test line - line.slice_thick_elements( - slicing_strategies=[ - xt.Strategy(slicing = xt.Uniform(N_SLICES, mode='thin'), element_type = xt.Bend), - xt.Strategy(slicing = xt.Uniform(N_SLICES, mode='thin'), element_type = xt.Quadrupole), - xt.Strategy(slicing = xt.Uniform(N_SLICES, mode='thin'), element_type = xt.Sextupole)]) - - ######################################## - # Initialise Plot - ######################################## - fig, axs = plt.subplots(2, 2, figsize = (10, 10)) - axs = axs.flatten() - fig.suptitle(title) - - ######################################## - # Test and plot with shifts - ######################################## - for i, (rot_x, rot_y) in enumerate( - zip([0, X_ROT, 0, X_ROT],[0, 0, Y_ROT, Y_ROT])): - # TODO: No way to compare this currently?!? - test_sol.mult_rot_x = rot_x - test_sol.mult_rot_y = rot_y - - tw = line.twiss( - _continue_if_lost = True, - start = xt.START, - end = xt.END, - betx = BETX, - bety = BETY, - px = PX0) - tw_sol = sol_line.twiss( - _continue_if_lost = True, - start = xt.START, - end = xt.END, - betx = BETX, - bety = BETY, - px = PX0) - - axs[i].plot(tw.s, tw.x, color = 'k', label = 'Element x') - axs[i].plot(tw.s, tw.y, color = 'r', label = 'Element y') - axs[i].plot(tw_sol.s, tw_sol.x, color = 'b', linestyle = ':', label = 'Sol x') - axs[i].plot(tw_sol.s, tw_sol.y, color = 'g', linestyle = ':', label = 'Sol y') - - axs[i].set_title(f'Rot x = {rot_x} [TBD], Rot y = {rot_y} [TBD]') - - ######################################## - # Assertions - ######################################## - assert np.isclose(tw.x[-1], tw_sol.x[-1], rtol = 1E-6) - assert np.isclose(tw.y[-1], tw_sol.y[-1], rtol = 1E-6) - - ######################################## - # Figure adjustments - ######################################## - fig.legend( - labels = ['Element x', 'Element y', 'Sol x', 'Sol y'], - loc = 'upper center', - ncol = 4) +################################################################################ +# Test Assertions +################################################################################ +assert np.isclose(bl_twiss['x'][-1], hrot_twiss['x'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['y'][-1], hrot_twiss['y'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['x'][-1], vrot_twiss['x'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['y'][-1], vrot_twiss['y'][-1], rtol = 1E-6) ######################################## # Show Plots From 27a4519f19f168bc4569fe3e3b1718be061777e7 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Sun, 24 Nov 2024 18:18:11 +0100 Subject: [PATCH 17/21] Additional Assertions --- examples/solenoid/mult_rot_test.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/examples/solenoid/mult_rot_test.py b/examples/solenoid/mult_rot_test.py index ba93859c4..e19413816 100644 --- a/examples/solenoid/mult_rot_test.py +++ b/examples/solenoid/mult_rot_test.py @@ -143,10 +143,32 @@ ################################################################################ # Test Assertions ################################################################################ -assert np.isclose(bl_twiss['x'][-1], hrot_twiss['x'][-1], rtol = 1E-6) -assert np.isclose(bl_twiss['y'][-1], hrot_twiss['y'][-1], rtol = 1E-6) -assert np.isclose(bl_twiss['x'][-1], vrot_twiss['x'][-1], rtol = 1E-6) -assert np.isclose(bl_twiss['y'][-1], vrot_twiss['y'][-1], rtol = 1E-6) +# Tolerances lower for derivative quantities (alfx, alfy, dpx, dpy) +assert np.isclose(bl_twiss['x'][-1], hrot_twiss['x'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['y'][-1], hrot_twiss['y'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['betx'][-1], hrot_twiss['betx'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['bety'][-1], hrot_twiss['bety'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['alfx'][-1], hrot_twiss['alfx'][-1], rtol = 1E-4) +assert np.isclose(bl_twiss['alfy'][-1], hrot_twiss['alfy'][-1], rtol = 1E-4) +assert np.isclose(bl_twiss['dx'][-1], hrot_twiss['dx'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['dy'][-1], hrot_twiss['dy'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['dpx'][-1], hrot_twiss['dpx'][-1], rtol = 1E-4) +assert np.isclose(bl_twiss['dpy'][-1], hrot_twiss['dpy'][-1], rtol = 1E-4) +assert np.isclose(bl_twiss['mux'][-1], hrot_twiss['mux'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['muy'][-1], hrot_twiss['muy'][-1], rtol = 1E-6) + +assert np.isclose(bl_twiss['x'][-1], vrot_twiss['x'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['y'][-1], vrot_twiss['y'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['betx'][-1], vrot_twiss['betx'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['bety'][-1], vrot_twiss['bety'][-1], rtol = 1E-4) +assert np.isclose(bl_twiss['alfx'][-1], vrot_twiss['alfx'][-1], rtol = 1E-4) +assert np.isclose(bl_twiss['alfy'][-1], vrot_twiss['alfy'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['dx'][-1], vrot_twiss['dx'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['dy'][-1], vrot_twiss['dy'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['dpx'][-1], vrot_twiss['dpx'][-1], rtol = 1E-4) +assert np.isclose(bl_twiss['dpy'][-1], vrot_twiss['dpy'][-1], rtol = 1E-4) +assert np.isclose(bl_twiss['mux'][-1], vrot_twiss['mux'][-1], rtol = 1E-6) +assert np.isclose(bl_twiss['muy'][-1], vrot_twiss['muy'][-1], rtol = 1E-6) ######################################## # Show Plots From b610e03e221718fa4c46803221dd180507257cda Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Mon, 2 Dec 2024 13:28:38 +0900 Subject: [PATCH 18/21] Cleaning up branch to be solenoid specific --- ..._shift_test.py => 001_multipole_shifts.py} | 0 ...rot_test.py => 002_multipole_rotations.py} | 0 examples/survey/001_bend.py | 327 ---------- examples/survey/002_multipole.py | 317 ---------- examples/survey/003_xyshift.py | 152 ----- examples/survey/004_xrotation.py | 575 ------------------ examples/survey/005_yrotation.py | 575 ------------------ examples/survey/006_srotation.py | 575 ------------------ examples/survey/_helpers.py | 73 --- examples/survey/_madpoint.py | 137 ----- 10 files changed, 2731 deletions(-) rename examples/solenoid/{mult_shift_test.py => 001_multipole_shifts.py} (100%) rename examples/solenoid/{mult_rot_test.py => 002_multipole_rotations.py} (100%) delete mode 100644 examples/survey/001_bend.py delete mode 100644 examples/survey/002_multipole.py delete mode 100644 examples/survey/003_xyshift.py delete mode 100644 examples/survey/004_xrotation.py delete mode 100644 examples/survey/005_yrotation.py delete mode 100644 examples/survey/006_srotation.py delete mode 100644 examples/survey/_helpers.py delete mode 100644 examples/survey/_madpoint.py diff --git a/examples/solenoid/mult_shift_test.py b/examples/solenoid/001_multipole_shifts.py similarity index 100% rename from examples/solenoid/mult_shift_test.py rename to examples/solenoid/001_multipole_shifts.py diff --git a/examples/solenoid/mult_rot_test.py b/examples/solenoid/002_multipole_rotations.py similarity index 100% rename from examples/solenoid/mult_rot_test.py rename to examples/solenoid/002_multipole_rotations.py diff --git a/examples/survey/001_bend.py b/examples/survey/001_bend.py deleted file mode 100644 index 3d90fda00..000000000 --- a/examples/survey/001_bend.py +++ /dev/null @@ -1,327 +0,0 @@ -""" -Test the survey on Bends -""" -################################################################################ -# Packages -################################################################################ -import numpy as np -import matplotlib.pyplot as plt -from matplotlib.lines import Line2D - -from _helpers import madpoint_twiss_survey, add_to_plot -import xtrack as xt - -################################################################################ -# User variables -################################################################################ -TOLERANCE = 1E-12 - -# Small angle needed for paraxial approximation tests -BEND_ANGLE = 1E-3 -BEND_LENGTH = 0.5 - -################################################################################ -# Plot setup -################################################################################ -fig_h = plt.figure(figsize = (16, 8)) -gs_h = fig_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_h = gs_h.subplots(sharex = 'row', sharey = True) - -fig_v = plt.figure(figsize = (16, 8)) -gs_v = fig_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_v = gs_v.subplots(sharex = 'row', sharey = True) - -################################################################################ -# Create line -################################################################################ -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line( - name = 'line', - components = [ - env.new('bend', xt.Bend, k0 = 0, h = 0, length = BEND_LENGTH, at = 1), - env.new('end', xt.Marker, at = 2)]) - -######################################## -# Configure Bend Model -######################################## -line.configure_bend_model(edge = 'suppressed') - -################################################################################ -# Horizontal Bend -################################################################################ - -######################################## -# h = k0 = 0 -######################################## -line['bend'].k0 = 0 -line['bend'].h = 0 -line['bend'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 0) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# h = 0, k0 != 0 -######################################## -line['bend'].k0 = BEND_ANGLE / BEND_LENGTH -line['bend'].h = 0 -line['bend'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 1) - -#################### -# Tests -#################### -# No bending with h = 0, so the survey must be zero -assert np.allclose(sv.X[-1], 0, atol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) - -######################################## -# h != 0, k0 = 0 -######################################## -line['bend'].k0 = 0 -line['bend'].h = BEND_ANGLE / BEND_LENGTH -line['bend'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 2) - -#################### -# Tests -#################### -# With h!=0, k0=0, the survey must be the negative of the Twiss -assert np.allclose(sv.X[-1], -tw.x[-1], rtol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) - -######################################## -# h != k0 !=0 -######################################## -line['bend'].k0 = BEND_ANGLE / BEND_LENGTH * 2 -line['bend'].h = BEND_ANGLE / BEND_LENGTH -line['bend'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 3) - -#################### -# Tests -#################### -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) - -######################################## -# h = k0 != 0 -######################################## -line['bend'].k0 = BEND_ANGLE / BEND_LENGTH -line['bend'].h = BEND_ANGLE / BEND_LENGTH -line['bend'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 4) - -#################### -# Tests -#################### -# With h=k0, there should be no residual orbit on the twiss -assert np.allclose(tw.x[-1], 0, atol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) - -################################################################################ -# Vertical Bend -################################################################################ - -######################################## -# h = k0 = 0 -######################################## -line['bend'].k0 = 0 -line['bend'].h = 0 -line['bend'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 0) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# h = 0, k0 != 0 -######################################## -line['bend'].k0 = BEND_ANGLE / BEND_LENGTH -line['bend'].h = 0 -line['bend'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 1) - -#################### -# Tests -#################### -# No bending with h = 0, so the survey must be zero -assert np.allclose(sv.Y[-1], 0, atol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) - -######################################## -# h != 0, k0 = 0 -######################################## -line['bend'].k0 = 0 -line['bend'].h = BEND_ANGLE / BEND_LENGTH -line['bend'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 2) - -#################### -# Tests -#################### -# With h!=0, k0=0, the survey must be the negative of the Twiss -assert np.allclose(sv.Y[-1], -tw.y[-1], rtol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) - -######################################## -# h != k0 !=0 -######################################## -line['bend'].k0 = BEND_ANGLE / BEND_LENGTH * 2 -line['bend'].h = BEND_ANGLE / BEND_LENGTH -line['bend'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 3) - -#################### -# Tests -#################### -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) - -######################################## -# h = k0 != 0 -######################################## -line['bend'].k0 = BEND_ANGLE / BEND_LENGTH -line['bend'].h = BEND_ANGLE / BEND_LENGTH -line['bend'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 4) - -#################### -# Tests -#################### -# With h=k0, there should be no residual orbit on the twiss -assert np.allclose(tw.y[-1], 0, atol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) - -################################################################################ -# Show Plots -################################################################################ - -######################################## -# Horizontal Bends -######################################## -# y labels -axs_h[0, 0].set_ylabel('Survey x,y [m]') -axs_h[1, 0].set_ylabel('Twiss x,y [m]') -axs_h[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_h[0, 2].set_xlabel('Z [m]') -axs_h[1, 2].set_xlabel('s [m]') -axs_h[2, 2].set_xlabel('s [m]') - -# Titles -axs_h[0, 0].set_title('h = 0, k0 = 0') -axs_h[0, 1].set_title('h = 0, k0 != 0') -axs_h[0, 2].set_title('h != 0, k0 = 0') -axs_h[0, 3].set_title('h = k0 / 2 != 0') -axs_h[0, 4].set_title('h = k0 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_h.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_h.suptitle('Horizontal Bend') - -######################################## -# Vertical Bends -######################################## -# y labels -axs_v[0, 0].set_ylabel('Survey x,y [m]') -axs_v[1, 0].set_ylabel('Twiss x,y [m]') -axs_v[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_v[0, 2].set_xlabel('Z [m]') -axs_v[1, 2].set_xlabel('s [m]') -axs_v[2, 2].set_xlabel('s [m]') - -# Titles -axs_v[0, 0].set_title('h = 0, k0 = 0') -axs_v[0, 1].set_title('h = 0, k0 != 0') -axs_v[0, 2].set_title('h != 0, k0 = 0') -axs_v[0, 3].set_title('h = k0 / 2 != 0') -axs_v[0, 4].set_title('h = k0 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_v.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_v.suptitle('Vertical Bend') - -######################################## -# Show -######################################## -plt.tight_layout() -plt.show() diff --git a/examples/survey/002_multipole.py b/examples/survey/002_multipole.py deleted file mode 100644 index 6a2faae6d..000000000 --- a/examples/survey/002_multipole.py +++ /dev/null @@ -1,317 +0,0 @@ -""" -Test the survey on Bends -""" -################################################################################ -# Packages -################################################################################ -import numpy as np -import matplotlib.pyplot as plt -from matplotlib.lines import Line2D - -from _helpers import madpoint_twiss_survey, add_to_plot -import xtrack as xt - -################################################################################ -# User variables -################################################################################ -TOLERANCE = 1E-12 -HXL = 1E-3 - -################################################################################ -# Plot setup -################################################################################ -fig_h = plt.figure(figsize = (16, 8)) -gs_h = fig_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_h = gs_h.subplots(sharex = 'row', sharey = True) - -fig_v = plt.figure(figsize = (16, 8)) -gs_v = fig_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_v = gs_v.subplots(sharex = 'row', sharey = True) - -################################################################################ -# Create line -################################################################################ -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('mult', xt.Multipole, hxl = 0, length = 0.5, at = 1), - env.new('end', xt.Marker, at = 2)]) - -################################################################################ -# Horizontal Bend -################################################################################ - -######################################## -# h = k0 = 0 -######################################## -line['mult'].knl[0] = 0 -line['mult'].hxl = 0 -line['mult'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 0) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# h = 0, k0 != 0 -######################################## -line['mult'].knl[0] = HXL -line['mult'].hxl = 0 -line['mult'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 1) - -#################### -# Tests -#################### -# No bending with h = 0, so the survey must be zero -assert np.allclose(sv.X[-1], 0, atol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) - -######################################## -# h != 0, k0 = 0 -######################################## -line['mult'].knl[0] = 0 -line['mult'].hxl = HXL -line['mult'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 2) - -#################### -# Tests -#################### -# With h!=0, k0=0, the survey must be the negative of the Twiss -assert np.allclose(sv.X[-1], -tw.x[-1], rtol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) - -######################################## -# h != k0 !=0 -######################################## -line['mult'].knl[0] = HXL * 2 -line['mult'].hxl = HXL -line['mult'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 3) - -#################### -# Tests -#################### -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) - -######################################## -# h = k0 != 0 -######################################## -line['mult'].knl[0] = HXL -line['mult'].hxl = HXL -line['mult'].rot_s_rad = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_h, sv, tw, 4) - -#################### -# Tests -#################### -# With h=k0, there should be no residual orbit on the twiss -assert np.allclose(tw.x[-1], 0, atol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol = TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol = TOLERANCE) - -################################################################################ -# Vertical Bend -################################################################################ - -######################################## -# h = k0 = 0 -######################################## -line['mult'].knl[0] = 0 -line['mult'].hxl = 0 -line['mult'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 0) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# h = 0, k0 != 0 -######################################## -line['mult'].knl[0] = HXL -line['mult'].hxl = 0 -line['mult'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 1) - -#################### -# Tests -#################### -# No bending with h = 0, so the survey must be zero -assert np.allclose(sv.Y[-1], 0, atol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) - -######################################## -# h != 0, k0 = 0 -######################################## -line['mult'].knl[0] = 0 -line['mult'].hxl = HXL -line['mult'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 2) - -#################### -# Tests -#################### -# With h!=0, k0=0, the survey must be the negative of the Twiss -assert np.allclose(sv.Y[-1], -tw.y[-1], rtol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) - -######################################## -# h != k0 !=0 -######################################## -line['mult'].knl[0] = HXL * 2 -line['mult'].hxl = HXL -line['mult'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 3) - -#################### -# Tests -#################### -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) - -######################################## -# h = k0 != 0 -######################################## -line['mult'].knl[0] = HXL -line['mult'].hxl = HXL -line['mult'].rot_s_rad = np.pi / 2 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_v, sv, tw, 4) - -#################### -# Tests -#################### -# With h=k0, there should be no residual orbit on the twiss -assert np.allclose(tw.y[-1], 0, atol = TOLERANCE) -# MadPoint sum of twiss and survey (paraxial approximation) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol = TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol = TOLERANCE) - -################################################################################ -# Show Plots -################################################################################ - -######################################## -# Horizontal Bends -######################################## -# y labels -axs_h[0, 0].set_ylabel('Survey x,y [m]') -axs_h[1, 0].set_ylabel('Twiss x,y [m]') -axs_h[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_h[0, 2].set_xlabel('Z [m]') -axs_h[1, 2].set_xlabel('s [m]') -axs_h[2, 2].set_xlabel('s [m]') - -# Titles -axs_h[0, 0].set_title('hxl = 0, k0l = 0') -axs_h[0, 1].set_title('hxl = 0, k0l != 0') -axs_h[0, 2].set_title('hxl != 0, k0l = 0') -axs_h[0, 3].set_title('hxl = k0l / 2 != 0') -axs_h[0, 4].set_title('hxl = k0l != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_h.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_h.suptitle('Horizontal Multipole') - -######################################## -# Vertical Bends -######################################## -# y labels -axs_v[0, 0].set_ylabel('Survey x,y [m]') -axs_v[1, 0].set_ylabel('Twiss x,y [m]') -axs_v[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_v[0, 2].set_xlabel('Z [m]') -axs_v[1, 2].set_xlabel('s [m]') -axs_v[2, 2].set_xlabel('s [m]') - -# Titles -axs_v[0, 0].set_title('hyl = 0, k0sl = 0') -axs_v[0, 1].set_title('hyl = 0, k0sl != 0') -axs_v[0, 2].set_title('hyl != 0, k0sl = 0') -axs_v[0, 3].set_title('hyl = k0sl / 2 != 0') -axs_v[0, 4].set_title('hyl = k0sl != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_v.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_v.suptitle('Vertical Bend') - -######################################## -# Show -######################################## -plt.tight_layout() -plt.show() diff --git a/examples/survey/003_xyshift.py b/examples/survey/003_xyshift.py deleted file mode 100644 index d811bae36..000000000 --- a/examples/survey/003_xyshift.py +++ /dev/null @@ -1,152 +0,0 @@ -""" -Test the survey on XY Shifts -""" -################################################################################ -# Packages -################################################################################ -import numpy as np -import matplotlib.pyplot as plt -from matplotlib.lines import Line2D - -from _helpers import madpoint_twiss_survey, add_to_plot -import xtrack as xt - -################################################################################ -# User variables -################################################################################ -TOLERANCE = 1E-12 - -################################################################################ -# Plot setup -################################################################################ -fig = plt.figure(figsize = (16, 8)) -gs = fig.add_gridspec(3, 4, hspace = 0.3, wspace = 0) -axs = gs.subplots(sharex = 'row', sharey = True) - -################################################################################ -# Create line -################################################################################ -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('xyshift', xt.XYShift, dx = 0, dy = 0, at = 1), - env.new('end', xt.Marker, at = 2)]) - -################################################################################ -# No Shift -################################################################################ -line['xyshift'].dx = 0 -line['xyshift'].dy = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs, sv, tw, 0, ylims = (-0.15, 0.15)) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -################################################################################ -# X Shift -################################################################################ -line['xyshift'].dx = 0.1 -line['xyshift'].dy = 0 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs, sv, tw, 1, ylims = (-0.15, 0.15)) - -#################### -# Tests -#################### -# No bending of particles, so must remain zero -assert np.allclose( sv.xx[-1], 0, atol=TOLERANCE) -# Survey negative of Twiss for h != 0, k0 = 0 -assert np.allclose(sv.X[-1], -tw.x[-1], rtol=TOLERANCE) -# MadPoint sum of twiss and survey -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=TOLERANCE) -# All y related quantities are zero -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -################################################################################ -# Y Shift -################################################################################ -line['xyshift'].dx = 0 -line['xyshift'].dy = 0.1 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs, sv, tw, 2, ylims = (-0.15, 0.15)) - -#################### -# Tests -#################### -# No bending of particles, so must remain zero -assert np.allclose(sv.yy[-1], 0, atol=TOLERANCE) -# Survey negative of Twiss for h != 0, k0 = 0 -assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=TOLERANCE) -# MadPoint sum of twiss and survey -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=TOLERANCE) -# All x related quantities are zero -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) - -################################################################################ -# X and Y Shift -################################################################################ -line['xyshift'].dx = 0.1 -line['xyshift'].dy = 0.1 - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs, sv, tw, 3, ylims = (-0.15, 0.15)) - -#################### -# Tests -#################### -# No bending of particles, so must remain zero -assert np.allclose( sv.xx[-1], 0, atol=TOLERANCE) -assert np.allclose(sv.yy[-1], 0, atol=TOLERANCE) -# Survey negative of Twiss for h != 0, k0 = 0 -assert np.allclose(sv.X[-1], -tw.x[-1], rtol=TOLERANCE) -assert np.allclose(sv.Y[-1], -tw.y[-1], rtol=TOLERANCE) -# MadPoint sum of twiss and survey -assert np.allclose(sv.X[-1] + tw.x[-1], sv.xx[-1], rtol=TOLERANCE) -assert np.allclose(sv.Y[-1] + tw.y[-1], sv.yy[-1], rtol=TOLERANCE) - -################################################################################ -# Show Plots -################################################################################ -# y labels -axs[0, 0].set_ylabel('Survey x,y [m]') -axs[1, 0].set_ylabel('Twiss x,y [m]') -axs[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs[0, 2].set_xlabel('Z [m]') -axs[1, 2].set_xlabel('s [m]') -axs[2, 2].set_xlabel('s [m]') - -# Titles -axs[0, 0].set_title('DX = 0, DY = 0') -axs[0, 1].set_title('DX != 0, DY = 0') -axs[0, 2].set_title('DX = 0, DY != 0') -axs[0, 3].set_title('DX != 0, DY != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig.suptitle('XYShift') - -######################################## -# Show -######################################## -plt.tight_layout() -plt.show() diff --git a/examples/survey/004_xrotation.py b/examples/survey/004_xrotation.py deleted file mode 100644 index b79c3a296..000000000 --- a/examples/survey/004_xrotation.py +++ /dev/null @@ -1,575 +0,0 @@ -""" -Test the survey on XRotation -""" -################################################################################ -# Packages -################################################################################ -import numpy as np -import matplotlib.pyplot as plt -from matplotlib.lines import Line2D - -from _helpers import madpoint_twiss_survey, add_to_plot -import xtrack as xt - -################################################################################ -# User variables -################################################################################ -TOLERANCE = 1E-12 -ROT_DEG = 30 -HXL = 0.1 -DX = 0.5 -PLOT_X_LIMS = (-0.1, 4.1) -PLOT_Y_LIMS = (-1, 1) - -################################################################################ -# Plot setup -################################################################################ -fig_mult_h = plt.figure(figsize = (16, 8)) -gs_mult_h = fig_mult_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_mult_h = gs_mult_h.subplots(sharex = 'row', sharey = True) - -fig_mult_v = plt.figure(figsize = (16, 8)) -gs_mult_v = fig_mult_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_mult_v = gs_mult_v.subplots(sharex = 'row', sharey = True) - -fig_shift_h = plt.figure(figsize = (16, 8)) -gs_shift_h = fig_shift_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_shift_h = gs_shift_h.subplots(sharex = 'row', sharey = True) - -fig_shift_v = plt.figure(figsize = (16, 8)) -gs_shift_v = fig_shift_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_shift_v = gs_shift_v.subplots(sharex = 'row', sharey = True) - -################################################################################ -# Horizontal Multipole Bends -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1), - env.new('xrotation', xt.XRotation, angle = 0, at = 2), - env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['xrotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = 0 -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = HXL -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = HXL -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Vertical Multipole Bends -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1, rot_s_rad = np.pi / 2), - env.new('xrotation', xt.XRotation, angle = 0, at = 2), - env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3, rot_s_rad = np.pi / 2), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['xrotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = 0 -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = HXL -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = HXL -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Horizontal Shift -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), - env.new('xrotation', xt.XRotation, angle = 0, at = 2), - env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['shift1'].dx = 0 -line['shift2'].dx = 0 -line['xrotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['shift1'].dx = 0 -line['shift2'].dx = 0 -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['shift1'].dx = DX -line['shift2'].dx = 0 -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['shift1'].dx = 0 -line['shift2'].dx = DX -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['shift1'].dx = DX -line['shift2'].dx = DX -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Vertical Shift -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), - env.new('xrotation', xt.XRotation, angle = 0, at = 2), - env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['shift1'].dy = 0 -line['shift2'].dy = 0 -line['xrotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['shift1'].dy = 0 -line['shift2'].dy = 0 -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['shift1'].dy = DX -line['shift2'].dy = 0 -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['shift1'].dy = 0 -line['shift2'].dy = DX -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['shift1'].dy = DX -line['shift2'].dy = DX -line['xrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Show Plots -################################################################################ - -######################################## -# Horizontal Mult -######################################## -# y labels -axs_mult_h[0, 0].set_ylabel('Survey x,y [m]') -axs_mult_h[1, 0].set_ylabel('Twiss x,y [m]') -axs_mult_h[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_mult_h[0, 2].set_xlabel('Z [m]') -axs_mult_h[1, 2].set_xlabel('s [m]') -axs_mult_h[2, 2].set_xlabel('s [m]') - -# Titles -axs_mult_h[0, 0].set_title('Rot = 0,\n HXL1 = 0, HXL2 = 0') -axs_mult_h[0, 1].set_title('Rot != 0,\n HXL1 = 0, HXL2 = 0') -axs_mult_h[0, 2].set_title('Rot != 0,\n HXL1 != 0, HXL2 = 0') -axs_mult_h[0, 3].set_title('Rot != 0,\n HXL1 = 0, HXL2 != 0') -axs_mult_h[0, 4].set_title('Rot != 0,\n HXL1 != 0, HXL2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_mult_h.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_mult_h.suptitle('Horizontal Multipole and XRotation') - -######################################## -# Vertical Mult -######################################## -# y labels -axs_mult_v[0, 0].set_ylabel('Survey x,y [m]') -axs_mult_v[1, 0].set_ylabel('Twiss x,y [m]') -axs_mult_v[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_mult_v[0, 2].set_xlabel('Z [m]') -axs_mult_v[1, 2].set_xlabel('s [m]') -axs_mult_v[2, 2].set_xlabel('s [m]') - -# Titles -axs_mult_v[0, 0].set_title('Rot = 0,\n HYL1 = 0, HYL2 = 0') -axs_mult_v[0, 1].set_title('Rot != 0,\n HYL1 = 0, HYL2 = 0') -axs_mult_v[0, 2].set_title('Rot != 0,\n HYL1 != 0, HYL2 = 0') -axs_mult_v[0, 3].set_title('Rot != 0,\n HYL1 = 0, HYL2 != 0') -axs_mult_v[0, 4].set_title('Rot != 0,\n HYL1 != 0, HYL2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_mult_v.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_mult_v.suptitle('Vertical Multipole and XRotation') - -######################################## -# Horizontal XYShift -######################################## -# y labels -axs_shift_h[0, 0].set_ylabel('Survey x,y [m]') -axs_shift_h[1, 0].set_ylabel('Twiss x,y [m]') -axs_shift_h[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_shift_h[0, 2].set_xlabel('Z [m]') -axs_shift_h[1, 2].set_xlabel('s [m]') -axs_shift_h[2, 2].set_xlabel('s [m]') - -# Titles -axs_shift_h[0, 0].set_title('Rot = 0,\n DX1 = 0, DX2 = 0') -axs_shift_h[0, 1].set_title('Rot != 0,\n DX1 = 0, DX2 = 0') -axs_shift_h[0, 2].set_title('Rot != 0,\n DX1 != 0, DX2 = 0') -axs_shift_h[0, 3].set_title('Rot != 0,\n DX1 = 0, DX2 != 0') -axs_shift_h[0, 4].set_title('Rot != 0,\n DX1 != 0, DX2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_shift_h.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_shift_h.suptitle('Horizontal XYShift and XRotation') - -######################################## -# Vertical XYShift -######################################## -# y labels -axs_shift_v[0, 0].set_ylabel('Survey x,y [m]') -axs_shift_v[1, 0].set_ylabel('Twiss x,y [m]') -axs_shift_v[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_shift_v[0, 2].set_xlabel('Z [m]') -axs_shift_v[1, 2].set_xlabel('s [m]') -axs_shift_v[2, 2].set_xlabel('s [m]') - -# Titles -axs_shift_v[0, 0].set_title('Rot = 0,\n DY1 = 0, DY2 = 0') -axs_shift_v[0, 1].set_title('Rot != 0,\n DY1 = 0, DY2 = 0') -axs_shift_v[0, 2].set_title('Rot != 0,\n DY1 != 0, DY2 = 0') -axs_shift_v[0, 3].set_title('Rot != 0,\n DY1 = 0, DY2 != 0') -axs_shift_v[0, 4].set_title('Rot != 0,\n DY1 != 0, DY2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_shift_v.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_shift_v.suptitle('Vertical XYShift and XRotation') - -######################################## -# Show -######################################## -plt.tight_layout() -plt.show() diff --git a/examples/survey/005_yrotation.py b/examples/survey/005_yrotation.py deleted file mode 100644 index 12fb46c1d..000000000 --- a/examples/survey/005_yrotation.py +++ /dev/null @@ -1,575 +0,0 @@ -""" -Test the survey on YRotation -""" -################################################################################ -# Packages -################################################################################ -import numpy as np -import matplotlib.pyplot as plt -from matplotlib.lines import Line2D - -from _helpers import madpoint_twiss_survey, add_to_plot -import xtrack as xt - -################################################################################ -# User variables -################################################################################ -TOLERANCE = 1E-12 -ROT_DEG = 30 -HXL = 0.1 -DX = 0.5 -PLOT_X_LIMS = (-0.1, 4.1) -PLOT_Y_LIMS = (-1, 1) - -################################################################################ -# Plot setup -################################################################################ -fig_mult_h = plt.figure(figsize = (16, 8)) -gs_mult_h = fig_mult_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_mult_h = gs_mult_h.subplots(sharex = 'row', sharey = True) - -fig_mult_v = plt.figure(figsize = (16, 8)) -gs_mult_v = fig_mult_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_mult_v = gs_mult_v.subplots(sharex = 'row', sharey = True) - -fig_shift_h = plt.figure(figsize = (16, 8)) -gs_shift_h = fig_shift_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_shift_h = gs_shift_h.subplots(sharex = 'row', sharey = True) - -fig_shift_v = plt.figure(figsize = (16, 8)) -gs_shift_v = fig_shift_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_shift_v = gs_shift_v.subplots(sharex = 'row', sharey = True) - -################################################################################ -# Horizontal Multipole Bends -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1), - env.new('yrotation', xt.YRotation, angle = 0, at = 2), - env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['yrotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = 0 -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = HXL -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = HXL -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Vertical Multipole Bends -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1, rot_s_rad = np.pi / 2), - env.new('yrotation', xt.YRotation, angle = 0, at = 2), - env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3, rot_s_rad = np.pi / 2), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['yrotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = 0 -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = HXL -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = HXL -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Horizontal Shift -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), - env.new('yrotation', xt.YRotation, angle = 0, at = 2), - env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['shift1'].dx = 0 -line['shift2'].dx = 0 -line['yrotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['shift1'].dx = 0 -line['shift2'].dx = 0 -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['shift1'].dx = DX -line['shift2'].dx = 0 -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['shift1'].dx = 0 -line['shift2'].dx = DX -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['shift1'].dx = DX -line['shift2'].dx = DX -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Vertical Shift -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), - env.new('yrotation', xt.YRotation, angle = 0, at = 2), - env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['shift1'].dy = 0 -line['shift2'].dy = 0 -line['yrotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['shift1'].dy = 0 -line['shift2'].dy = 0 -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['shift1'].dy = DX -line['shift2'].dy = 0 -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['shift1'].dy = 0 -line['shift2'].dy = DX -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['shift1'].dy = DX -line['shift2'].dy = DX -line['yrotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Show Plots -################################################################################ - -######################################## -# Horizontal Mult -######################################## -# y labels -axs_mult_h[0, 0].set_ylabel('Survey x,y [m]') -axs_mult_h[1, 0].set_ylabel('Twiss x,y [m]') -axs_mult_h[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_mult_h[0, 2].set_xlabel('Z [m]') -axs_mult_h[1, 2].set_xlabel('s [m]') -axs_mult_h[2, 2].set_xlabel('s [m]') - -# Titles -axs_mult_h[0, 0].set_title('Rot = 0,\n HXL1 = 0, HXL2 = 0') -axs_mult_h[0, 1].set_title('Rot != 0,\n HXL1 = 0, HXL2 = 0') -axs_mult_h[0, 2].set_title('Rot != 0,\n HXL1 != 0, HXL2 = 0') -axs_mult_h[0, 3].set_title('Rot != 0,\n HXL1 = 0, HXL2 != 0') -axs_mult_h[0, 4].set_title('Rot != 0,\n HXL1 != 0, HXL2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_mult_h.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_mult_h.suptitle('Horizontal Multipole and YRotation') - -######################################## -# Vertical Mult -######################################## -# y labels -axs_mult_v[0, 0].set_ylabel('Survey x,y [m]') -axs_mult_v[1, 0].set_ylabel('Twiss x,y [m]') -axs_mult_v[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_mult_v[0, 2].set_xlabel('Z [m]') -axs_mult_v[1, 2].set_xlabel('s [m]') -axs_mult_v[2, 2].set_xlabel('s [m]') - -# Titles -axs_mult_v[0, 0].set_title('Rot = 0,\n HYL1 = 0, HYL2 = 0') -axs_mult_v[0, 1].set_title('Rot != 0,\n HYL1 = 0, HYL2 = 0') -axs_mult_v[0, 2].set_title('Rot != 0,\n HYL1 != 0, HYL2 = 0') -axs_mult_v[0, 3].set_title('Rot != 0,\n HYL1 = 0, HYL2 != 0') -axs_mult_v[0, 4].set_title('Rot != 0,\n HYL1 != 0, HYL2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_mult_v.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_mult_v.suptitle('Vertical Multipole and YRotation') - -######################################## -# Horizontal XYShift -######################################## -# y labels -axs_shift_h[0, 0].set_ylabel('Survey x,y [m]') -axs_shift_h[1, 0].set_ylabel('Twiss x,y [m]') -axs_shift_h[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_shift_h[0, 2].set_xlabel('Z [m]') -axs_shift_h[1, 2].set_xlabel('s [m]') -axs_shift_h[2, 2].set_xlabel('s [m]') - -# Titles -axs_shift_h[0, 0].set_title('Rot = 0,\n DX1 = 0, DX2 = 0') -axs_shift_h[0, 1].set_title('Rot != 0,\n DX1 = 0, DX2 = 0') -axs_shift_h[0, 2].set_title('Rot != 0,\n DX1 != 0, DX2 = 0') -axs_shift_h[0, 3].set_title('Rot != 0,\n DX1 = 0, DX2 != 0') -axs_shift_h[0, 4].set_title('Rot != 0,\n DX1 != 0, DX2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_shift_h.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_shift_h.suptitle('Horizontal XYShift and YRotation') - -######################################## -# Vertical XYShift -######################################## -# y labels -axs_shift_v[0, 0].set_ylabel('Survey x,y [m]') -axs_shift_v[1, 0].set_ylabel('Twiss x,y [m]') -axs_shift_v[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_shift_v[0, 2].set_xlabel('Z [m]') -axs_shift_v[1, 2].set_xlabel('s [m]') -axs_shift_v[2, 2].set_xlabel('s [m]') - -# Titles -axs_shift_v[0, 0].set_title('Rot = 0,\n DY1 = 0, DY2 = 0') -axs_shift_v[0, 1].set_title('Rot != 0,\n DY1 = 0, DY2 = 0') -axs_shift_v[0, 2].set_title('Rot != 0,\n DY1 != 0, DY2 = 0') -axs_shift_v[0, 3].set_title('Rot != 0,\n DY1 = 0, DY2 != 0') -axs_shift_v[0, 4].set_title('Rot != 0,\n DY1 != 0, DY2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_shift_v.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_shift_v.suptitle('Vertical XYShift and YRotation') - -######################################## -# Show -######################################## -plt.tight_layout() -plt.show() diff --git a/examples/survey/006_srotation.py b/examples/survey/006_srotation.py deleted file mode 100644 index a9c732913..000000000 --- a/examples/survey/006_srotation.py +++ /dev/null @@ -1,575 +0,0 @@ -""" -Test the survey on SRotation -""" -################################################################################ -# Packages -################################################################################ -import numpy as np -import matplotlib.pyplot as plt -from matplotlib.lines import Line2D - -from _helpers import madpoint_twiss_survey, add_to_plot -import xtrack as xt - -################################################################################ -# User variables -################################################################################ -TOLERANCE = 1E-12 -ROT_DEG = 30 -HXL = 0.1 -DX = 0.5 -PLOT_X_LIMS = (-0.1, 4.1) -PLOT_Y_LIMS = (-1, 1) - -################################################################################ -# Plot setup -################################################################################ -fig_mult_h = plt.figure(figsize = (16, 8)) -gs_mult_h = fig_mult_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_mult_h = gs_mult_h.subplots(sharex = 'row', sharey = True) - -fig_mult_v = plt.figure(figsize = (16, 8)) -gs_mult_v = fig_mult_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_mult_v = gs_mult_v.subplots(sharex = 'row', sharey = True) - -fig_shift_h = plt.figure(figsize = (16, 8)) -gs_shift_h = fig_shift_h.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_shift_h = gs_shift_h.subplots(sharex = 'row', sharey = True) - -fig_shift_v = plt.figure(figsize = (16, 8)) -gs_shift_v = fig_shift_v.add_gridspec(3, 5, hspace = 0.3, wspace = 0) -axs_shift_v = gs_shift_v.subplots(sharex = 'row', sharey = True) - -################################################################################ -# Horizontal Multipole Bends -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1), - env.new('srotation', xt.SRotation, angle = 0, at = 2), - env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['srotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = 0 -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = HXL -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = HXL -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Vertical Multipole Bends -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('mult1', xt.Multipole, hxl = 0, length = 0.5, at = 1, rot_s_rad = np.pi / 2), - env.new('srotation', xt.SRotation, angle = 0, at = 2), - env.new('mult2', xt.Multipole, hxl = 0, length = 0.5, at = 3, rot_s_rad = np.pi / 2), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['srotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = 0 -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = 0 -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['mult1'].hxl = 0 -line['mult2'].hxl = HXL -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['mult1'].hxl = HXL -line['mult2'].hxl = HXL -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_mult_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Horizontal Shift -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), - env.new('srotation', xt.SRotation, angle = 0, at = 2), - env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['shift1'].dx = 0 -line['shift2'].dx = 0 -line['srotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['shift1'].dx = 0 -line['shift2'].dx = 0 -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['shift1'].dx = DX -line['shift2'].dx = 0 -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['shift1'].dx = 0 -line['shift2'].dx = DX -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['shift1'].dx = DX -line['shift2'].dx = DX -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_h, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Vertical Shift -################################################################################ - -######################################## -# Create line -######################################## -env = xt.Environment(particle_ref = xt.Particles(p0c = 1E9)) -line = env.new_line(name = 'line', components=[ - env.new('shift1', xt.XYShift, dx = 0, dy = 0, at = 1), - env.new('srotation', xt.SRotation, angle = 0, at = 2), - env.new('shift2', xt.XYShift, dx = 0, dy = 0, at = 3), - env.new('end', xt.Marker, at = 4)]) - -######################################## -# Off -######################################## -line['shift1'].dy = 0 -line['shift2'].dy = 0 -line['srotation'].angle = np.rad2deg(0) - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 0, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Element off must have no effect -assert np.allclose(np.array([sv.X[-1], sv.xx[-1], tw.x[-1]]), 0, atol=TOLERANCE) -assert np.allclose(np.array([sv.Y[-1], sv.yy[-1], tw.y[-1]]), 0, atol=TOLERANCE) - -######################################## -# Rotation Only -######################################## -line['shift1'].dy = 0 -line['shift2'].dy = 0 -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 1, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot -######################################## -line['shift1'].dy = DX -line['shift2'].dy = 0 -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 2, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Rot then mult -######################################## -line['shift1'].dy = 0 -line['shift2'].dy = DX -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 3, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -######################################## -# Mult then rot then mult -######################################## -line['shift1'].dy = DX -line['shift2'].dy = DX -line['srotation'].angle = ROT_DEG - -sv, tw = madpoint_twiss_survey(line) - -add_to_plot(axs_shift_v, sv, tw, 4, xlims = PLOT_X_LIMS, ylims = PLOT_Y_LIMS) - -#################### -# Tests -#################### -# Only rotation must have no effect -# TODO - -################################################################################ -# Show Plots -################################################################################ - -######################################## -# Horizontal Mult -######################################## -# y labels -axs_mult_h[0, 0].set_ylabel('Survey x,y [m]') -axs_mult_h[1, 0].set_ylabel('Twiss x,y [m]') -axs_mult_h[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_mult_h[0, 2].set_xlabel('Z [m]') -axs_mult_h[1, 2].set_xlabel('s [m]') -axs_mult_h[2, 2].set_xlabel('s [m]') - -# Titles -axs_mult_h[0, 0].set_title('Rot = 0,\n HXL1 = 0, HXL2 = 0') -axs_mult_h[0, 1].set_title('Rot != 0,\n HXL1 = 0, HXL2 = 0') -axs_mult_h[0, 2].set_title('Rot != 0,\n HXL1 != 0, HXL2 = 0') -axs_mult_h[0, 3].set_title('Rot != 0,\n HXL1 = 0, HXL2 != 0') -axs_mult_h[0, 4].set_title('Rot != 0,\n HXL1 != 0, HXL2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_mult_h.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_mult_h.suptitle('Horizontal Multipole and SRotation') - -######################################## -# Vertical Mult -######################################## -# y labels -axs_mult_v[0, 0].set_ylabel('Survey x,y [m]') -axs_mult_v[1, 0].set_ylabel('Twiss x,y [m]') -axs_mult_v[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_mult_v[0, 2].set_xlabel('Z [m]') -axs_mult_v[1, 2].set_xlabel('s [m]') -axs_mult_v[2, 2].set_xlabel('s [m]') - -# Titles -axs_mult_v[0, 0].set_title('Rot = 0,\n HYL1 = 0, HYL2 = 0') -axs_mult_v[0, 1].set_title('Rot != 0,\n HYL1 = 0, HYL2 = 0') -axs_mult_v[0, 2].set_title('Rot != 0,\n HYL1 != 0, HYL2 = 0') -axs_mult_v[0, 3].set_title('Rot != 0,\n HYL1 = 0, HYL2 != 0') -axs_mult_v[0, 4].set_title('Rot != 0,\n HYL1 != 0, HYL2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_mult_v.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_mult_v.suptitle('Vertical Multipole and SRotation') - -######################################## -# Horizontal XYShift -######################################## -# y labels -axs_shift_h[0, 0].set_ylabel('Survey x,y [m]') -axs_shift_h[1, 0].set_ylabel('Twiss x,y [m]') -axs_shift_h[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_shift_h[0, 2].set_xlabel('Z [m]') -axs_shift_h[1, 2].set_xlabel('s [m]') -axs_shift_h[2, 2].set_xlabel('s [m]') - -# Titles -axs_shift_h[0, 0].set_title('Rot = 0,\n DX1 = 0, DX2 = 0') -axs_shift_h[0, 1].set_title('Rot != 0,\n DX1 = 0, DX2 = 0') -axs_shift_h[0, 2].set_title('Rot != 0,\n DX1 != 0, DX2 = 0') -axs_shift_h[0, 3].set_title('Rot != 0,\n DX1 = 0, DX2 != 0') -axs_shift_h[0, 4].set_title('Rot != 0,\n DX1 != 0, DX2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_shift_h.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_shift_h.suptitle('Horizontal XYShift and SRotation') - -######################################## -# Vertical XYShift -######################################## -# y labels -axs_shift_v[0, 0].set_ylabel('Survey x,y [m]') -axs_shift_v[1, 0].set_ylabel('Twiss x,y [m]') -axs_shift_v[2, 0].set_ylabel('MadPoint xx,yy [m]') - -# x labels -axs_shift_v[0, 2].set_xlabel('Z [m]') -axs_shift_v[1, 2].set_xlabel('s [m]') -axs_shift_v[2, 2].set_xlabel('s [m]') - -# Titles -axs_shift_v[0, 0].set_title('Rot = 0,\n DY1 = 0, DY2 = 0') -axs_shift_v[0, 1].set_title('Rot != 0,\n DY1 = 0, DY2 = 0') -axs_shift_v[0, 2].set_title('Rot != 0,\n DY1 != 0, DY2 = 0') -axs_shift_v[0, 3].set_title('Rot != 0,\n DY1 = 0, DY2 != 0') -axs_shift_v[0, 4].set_title('Rot != 0,\n DY1 != 0, DY2 != 0') - -legend_elements = [ - Line2D([0], [0], color = 'black', lw = 2, label = 'x'), - Line2D([0], [0], color = 'red', lw = 2, label = 'y')] -fig_shift_v.legend( - handles = legend_elements, - loc = 'lower center', - ncol = 2, - frameon = False, - bbox_to_anchor = (0.5, -0.05)) -fig_shift_v.suptitle('Vertical XYShift and SRotation') - -######################################## -# Show -######################################## -plt.tight_layout() -plt.show() diff --git a/examples/survey/_helpers.py b/examples/survey/_helpers.py deleted file mode 100644 index c12c98ad1..000000000 --- a/examples/survey/_helpers.py +++ /dev/null @@ -1,73 +0,0 @@ -""" -Helpers for testing the MadPoint class -""" -import numpy as np -from _madpoint import MadPoint - -def zero_small_values(arr, tol): - """ - Set values within a tolerance of 0, to 0 - """ - return np.where(np.abs(arr) < tol, 0, arr) - -def madpoint_twiss_survey(line): - """ - Produce twiss and survey for a line, including MadPoint - """ - survey = line.survey() - twiss = line.twiss4d(_continue_if_lost = True, betx = 1, bety = 1) - - madpoints = [] - xx = [] - yy = [] - zz = [] - for nn in twiss.name: - madpoints.append( - MadPoint(name = nn, xsuite_twiss = twiss, xsuite_survey = survey)) - xx.append(madpoints[-1].p[0]) - yy.append(madpoints[-1].p[1]) - zz.append(madpoints[-1].p[2]) - - survey['xx'] = np.array(xx) - survey['yy'] = np.array(yy) - survey['zz'] = np.array(zz) - - return survey, twiss - -def add_to_plot(axes, survey, twiss, index, tol = 1E-12, xlims = (-0.1, 2.1), ylims = (-2E-3, 2E-3)): - """ - Add a line to the overall plot to show a specific case - """ - axes[0, index].plot( - zero_small_values(survey.Z, tol), - zero_small_values(survey.X, tol), - c = 'k') - axes[0, index].plot( - zero_small_values(survey.Z, tol), - zero_small_values(survey.Y, tol), - c = 'r') - - axes[1, index].plot( - zero_small_values(twiss.s, tol), - zero_small_values(twiss.x, tol), - c = 'k') - axes[1, index].plot( - zero_small_values(twiss.s, tol), - zero_small_values(twiss.y, tol), - c = 'r') - - axes[2, index].plot( - zero_small_values(survey.s, tol), - zero_small_values(survey.xx, tol), - c = 'k') - axes[2, index].plot( - zero_small_values(survey.s, tol), - zero_small_values(survey.yy, tol), - c = 'r') - - axes[0, index].set_xlim(xlims) - axes[0, index].set_ylim(ylims) - axes[1, index].set_xlim(xlims) - axes[1, index].set_ylim(ylims) - axes[2, index].set_xlim(xlims) - axes[2, index].set_ylim(ylims) diff --git a/examples/survey/_madpoint.py b/examples/survey/_madpoint.py deleted file mode 100644 index 56437c86a..000000000 --- a/examples/survey/_madpoint.py +++ /dev/null @@ -1,137 +0,0 @@ -import numpy as np - -class MadPoint(object): - @classmethod - def from_survey(cls, name, mad=None, xsuite_survey=None): - return cls(name, mad=mad, use_twiss=False, use_survey=True, - xsuite_survey=xsuite_survey) - - @classmethod - def from_twiss(cls, name, mad): - return cls(name, mad, use_twiss=True, use_survey=False) - - def __init__(self, name, mad=None, use_twiss=True, use_survey=True, - xsuite_twiss=None, xsuite_survey=None): - - self.use_twiss = use_twiss - self.use_survey = use_survey - - if not (use_survey) and not (use_twiss): - raise ValueError( - "use_survey and use_twiss cannot be False at the same time" - ) - - self.tx = None - self.ty = None - self.tpx = None - self.tpy = None - - self.sx = None - self.sy = None - self.sz = None - self.sp = None - theta = 0.0 - phi = 0.0 - psi = 0.0 - - if mad is not None: - - self.name = name - if use_twiss: - assert xsuite_survey is None - twiss = mad.table.twiss - names = twiss.name - if use_survey: - assert xsuite_twiss is None - survey = mad.table.survey - names = survey.name - # patch for this issue https://github.com/hibtc/cpymad/issues/91 - for ii, nn in enumerate(names): - if not nn.endswith(':1'): - names[ii] = nn+':1' - - idx = np.where(names == name)[0][0] - - if use_twiss: - self.tx = twiss.x[idx] - self.ty = twiss.y[idx] - self.tpx = twiss.px[idx] - self.tpy = twiss.py[idx] - - if use_survey: - self.sx = survey.x[idx] - self.sy = survey.y[idx] - self.sz = survey.z[idx] - self.sp = np.array([self.sx, self.sy, self.sz]) - theta = survey.theta[idx] - phi = survey.phi[idx] - psi = survey.psi[idx] - else: - - - if use_twiss: - assert xsuite_twiss is not None - idx = np.where(np.array(xsuite_twiss['name']) == name)[0][0] - self.tx = xsuite_twiss.x[idx] - self.ty = xsuite_twiss.y[idx] - self.tpx = xsuite_twiss.px[idx] - self.tpy = xsuite_twiss.py[idx] - - if use_survey: - assert xsuite_survey is not None - idx = np.where(np.array(xsuite_survey['name']) == name)[0][0] - self.sx = xsuite_survey.X[idx] - self.sy = xsuite_survey.Y[idx] - self.sz = xsuite_survey.Z[idx] - self.sp = np.array([self.sx, self.sy, self.sz]) - theta = xsuite_survey.theta[idx] - phi = xsuite_survey.phi[idx] - psi = xsuite_survey.psi[idx] - - thetam = np.array( - [ - [np.cos(theta), 0, np.sin(theta)], - [0, 1, 0], - [-np.sin(theta), 0, np.cos(theta)], - ] - ) - phim = np.array( - [ - [1, 0, 0], - [0, np.cos(phi), np.sin(phi)], - [0, -np.sin(phi), np.cos(phi)], - ] - ) - psim = np.array( - [ - [np.cos(psi), -np.sin(psi), 0], - [np.sin(psi), np.cos(psi), 0], - [0, 0, 1], - ] - ) - wm = np.dot(thetam, np.dot(phim, psim)) - self.ex = np.dot(wm, np.array([1, 0, 0])) - self.ey = np.dot(wm, np.array([0, 1, 0])) - self.ez = np.dot(wm, np.array([0, 0, 1])) - - self.p = np.array([0.0, 0.0, 0.0]) - - if use_twiss: - self.p += self.ex * self.tx + self.ey * self.ty - - if use_survey: - self.p += self.sp - - def shift_survey(self, delta): - self.sx -= delta[0] - self.sy -= delta[1] - self.sz -= delta[2] - self.sp -= delta - self.p -= delta - - def dist(self, other): - return np.sqrt(np.sum((self.p - other.p) ** 2)) - - def distxy(self, other): - dd = self.p - other.p - return np.dot(dd, self.ex), np.dot(dd, self.ey) From b1c312a68823609a109373cf9ade1072f22360b3 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Mon, 2 Dec 2024 13:32:51 +0900 Subject: [PATCH 19/21] Restored on this branch --- xtrack/multiline/multiline.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/xtrack/multiline/multiline.py b/xtrack/multiline/multiline.py index 6ef663dfe..b7c4df137 100644 --- a/xtrack/multiline/multiline.py +++ b/xtrack/multiline/multiline.py @@ -459,11 +459,11 @@ def install_beambeam_interactions(self, clockwise_line, anticlockwise_line, for nn, ll in self.lines.items(): ll.unfreeze() - # if clockwise_line is not None and anticlockwise_line is not None: - # circumference_cw = self.lines[clockwise_line].get_length() - # circumference_acw = self.lines[anticlockwise_line].get_length() - # assert np.isclose(circumference_cw, circumference_acw, - # atol=1e-4, rtol=0) + if clockwise_line is not None and anticlockwise_line is not None: + circumference_cw = self.lines[clockwise_line].get_length() + circumference_acw = self.lines[anticlockwise_line].get_length() + assert np.isclose(circumference_cw, circumference_acw, + atol=1e-4, rtol=0) import xfields as xf bb_df_cw, bb_df_acw = xf.install_beambeam_elements_in_lines( @@ -490,8 +490,6 @@ def install_beambeam_interactions(self, clockwise_line, anticlockwise_line, def configure_beambeam_interactions(self, num_particles, nemitt_x, nemitt_y, crab_strong_beam=True, - nemitt_x_cw=None, nemitt_y_cw=None, - nemitt_x_acw=None, nemitt_y_acw=None, use_antisymmetry=False, separation_bumps=None): @@ -550,8 +548,6 @@ def configure_beambeam_interactions(self, num_particles, num_particles=num_particles, nemitt_x=nemitt_x, nemitt_y=nemitt_y, crab_strong_beam=crab_strong_beam, - nemitt_x_cw=nemitt_x_cw, nemitt_y_cw=nemitt_y_cw, - nemitt_x_acw=nemitt_x_acw, nemitt_y_acw=nemitt_y_acw, ip_names=self._bb_config['ip_names'], use_antisymmetry=use_antisymmetry, separation_bumps=separation_bumps) From 9e8d60c88e4ee083be04e954ff942b5910ee6cec Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Mon, 2 Dec 2024 13:34:56 +0900 Subject: [PATCH 20/21] Restored survey specific changes --- xtrack/line.py | 27 --------------------------- 1 file changed, 27 deletions(-) diff --git a/xtrack/line.py b/xtrack/line.py index 6ea8dad46..1e30c0a33 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -4153,16 +4153,6 @@ def _get_attr_cache(self): '_own_k4sl': ('ksl', 4), '_own_k5sl': ('ksl', 5), - '_own_dx': 'dx', - '_own_dy': 'dy', - - '_own_sin_angle': 'sin_angle', - '_own_cos_angle': 'cos_angle', - - # TODO: This should be corrected at the element level - '_own_sin_z': 'sin_z', - '_own_cos_z': 'cos_z', - '_parent_length': (('_parent', 'length'), None), '_parent_sin_rot_s': (('_parent', '_sin_rot_s'), None), '_parent_cos_rot_s': (('_parent', '_cos_rot_s'), None), @@ -4201,16 +4191,6 @@ def _get_attr_cache(self): '_parent_k4sl': (('_parent', 'ksl'), 4), '_parent_k5sl': (('_parent', 'ksl'), 5), - '_parent_dx': (('_parent', 'dx'), None), - '_parent_dy': (('_parent', 'dy'), None), - - '_parent_sin_angle': (('_parent', 'sin_angle'), None), - '_parent_cos_angle': (('_parent', 'cos_angle'), None), - - # TODO: This should be corrected at the element level - '_parent_sin_z': (('_parent', 'sin_z'), None), - '_parent_cos_z': (('_parent', 'cos_z'), None), - }, derived_fields={ 'length': lambda attr: @@ -4288,13 +4268,6 @@ def _get_attr_cache(self): + attr['_parent_k5s'] * attr['_parent_length'] * attr['weight'] * attr._inherit_strengths), 'hkick': lambda attr: attr["angle_rad"] - attr["k0l"], 'vkick': lambda attr: attr["k0sl"], - 'dx': lambda attr: attr['_own_dx'] + attr['_parent_dx'], - 'dy': lambda attr: attr['_own_dy'] + attr['_parent_dy'], - 'transform_angle_rad': lambda attr: np.arctan2( - attr['_own_sin_angle'] + attr['_parent_sin_angle'] +\ - attr['_own_sin_z'] + attr['_parent_sin_z'], - attr['_own_cos_angle'] + attr['_parent_cos_angle'] +\ - attr['_own_cos_z'] + attr['_parent_cos_z']), } ) return cache From 06d7555281a2660d29cbd5a98527d793d7f9f0a5 Mon Sep 17 00:00:00 2001 From: John Salvesen Date: Mon, 2 Dec 2024 13:35:50 +0900 Subject: [PATCH 21/21] Restored survey --- xtrack/survey.py | 417 +++++++++++------------------------------------ 1 file changed, 99 insertions(+), 318 deletions(-) diff --git a/xtrack/survey.py b/xtrack/survey.py index 3ae893b88..7d9e0d06c 100644 --- a/xtrack/survey.py +++ b/xtrack/survey.py @@ -14,7 +14,7 @@ # Required functions # ================================================== -def get_w_from_angles(theta, phi, psi, reverse_xs = False): +def get_w_from_angles(theta, phi, psi, reverse_xs=False): """W matrix, see MAD-X manual""" costhe = np.cos(theta) cosphi = np.cos(phi) @@ -40,7 +40,7 @@ def get_w_from_angles(theta, phi, psi, reverse_xs = False): return w -def get_angles_from_w(w, reverse_xs = False): +def get_angles_from_w(w, reverse_xs=False): """Inverse function of get_w_from_angles()""" # w[0, 2]/w[2, 2] = (sinthe * cosphi)/(costhe * cosphi) # w[1, 0]/w[1, 1] = (cosphi * sinpsi)/(cosphi * cospsi) @@ -55,8 +55,7 @@ def get_angles_from_w(w, reverse_xs = False): psi = np.arctan2(w[1, 0], w[1, 1]) phi = np.arctan2(w[1, 2], w[1, 1] / np.cos(psi)) - # TODO: arctan2 returns angle between [-pi,pi]. - # Hence theta ends up not at 2pi after a full survey + # TODO: arctan2 returns angle between [-pi,pi]. Hence theta ends up not at 2pi after a full survey return theta, phi, psi @@ -66,80 +65,19 @@ def advance_bend(v, w, R, S): return np.dot(w, R) + v, np.dot(w, S) -def advance_rotation(v, w, S): - """Advancing through rotation element: - Rotate w matrix according to transformation matrix S""" - return v, np.dot(w, S) - - def advance_drift(v, w, R): """Advancing through drift element, see MAD-X manual: v2 = w1*R + v1 | w2 = w1*S -> S is unity""" return np.dot(w, R) + v, w -def advance_element( - v, w, length = 0, angle = 0, tilt = 0, - dx = 0, dy = 0, transf_x_rad = 0, transf_y_rad = 0, transf_s_rad = 0): - """Computing the advance element-by-element. - See MAD-X manual for generation of R and S""" - # XYShift Handling - if dx != 0 or dy != 0: - assert angle == 0, "dx and dy are only supported for angle = 0" - assert tilt == 0, "dx and dy are only supported for tilt = 0" - assert length == 0, "dx and dy are only supported for length = 0" - - R = np.array([dx, dy, 0]) - # XYShift tarnsforms as a drift - return advance_drift(v, w, R) - - # XRotation Handling - if transf_x_rad != 0: - assert angle == 0, "rot_x_rad is only supported for angle = 0" - assert tilt == 0, "rot_x_rad is only supported for tilt = 0" - assert length == 0, "rot_x_rad is only supported for length = 0" - - # Rotation sine/cosine - cr = np.cos(-transf_x_rad) - sr = np.sin(-transf_x_rad) - # ------ - S = np.array([[1, 0, 0], [0, cr, sr], [0, -sr, cr]]) # x rotation matrix - return advance_rotation(v, w, S) - - # YRotation Handling - if transf_y_rad != 0: - assert angle == 0, "rot_y_rad is only supported for angle = 0" - assert tilt == 0, "rot_y_rad is only supported for tilt = 0" - assert length == 0, "rot_y_rad is only supported for length = 0" - - # Rotation sine/cosine - cr = np.cos(transf_y_rad) - sr = np.sin(transf_y_rad) - # ------ - S = np.array([[cr, 0, -sr], [0, 1, 0], [sr, 0, cr]]) # y rotation matrix - return advance_rotation(v, w, S) - - # SRotation Handling - if transf_s_rad != 0: - assert angle == 0, "rot_s_rad is only supported for angle = 0" - assert tilt == 0, "rot_s_rad is only supported for tilt = 0" - assert length == 0, "rot_s_rad is only supported for length = 0" - - # Rotation sine/cosine - cr = np.cos(transf_s_rad) - sr = np.sin(transf_s_rad) - # ------ - S = np.array([[cr, sr, 0], [-sr, cr, 0], [0, 0, 1]]) # z rotation matrix - return advance_rotation(v, w, S) - - # Non bending elements +def advance_element(v, w, length=0, angle=0, tilt=0): + """Computing the advance element-by-element. See MAD-X manual for generation of R and S""" if angle == 0: R = np.array([0, 0, length]) return advance_drift(v, w, R) - - # Horizontal bending elements elif tilt == 0: - # Angle sine/cosine + # Relevant sine/cosine ca = np.cos(angle) sa = np.sin(angle) # ------ @@ -148,12 +86,10 @@ def advance_element( S = np.array([[ca, 0, -sa], [0, 1, 0], [sa, 0, ca]]) return advance_bend(v, w, R, S) - # Tilted bending elements else: - # Angle sine/cosine + # Relevant sine/cosine ca = np.cos(angle) sa = np.sin(angle) - # Tilt sine/cosine ct = np.cos(tilt) st = np.sin(tilt) # ------ @@ -162,26 +98,18 @@ def advance_element( S = np.array([[ca, 0, -sa], [0, 1, 0], [sa, 0, ca]]) # Orthogonal rotation matrix for tilt - T = np.array([[ct, -st, 0], [st, ct, 0], [0, 0, 1]]) - Tinv = np.array([[ct, st, 0], [-st, ct, 0], [0, 0, 1]]) + T = np.array([[ct, -st, 0], [st, ct, 0], [0, 0, 1]]) + Tinv = np.array([[ct, st, 0], [-st, ct, 0], [0, 0, 1]]) return advance_bend(v, w, np.dot(T, R), np.dot(T, np.dot(S, Tinv))) class SurveyTable(Table): - """ - Table for survey data. - """ _error_on_row_not_found = True - def reverse( - self, - X0 = None, Y0 = None, Z0 = None, - theta0 = None, phi0 = None, psi0 = None, element0 = None): - """ - Reverse the survey. - """ + def reverse(self, X0=None, Y0=None, Z0=None, theta0=None, + phi0=None, psi0=None, element0=None): if element0 is None: element0 = len(self.name) - self.element0 - 1 @@ -202,75 +130,44 @@ def reverse( psi0 = self.psi[self.element0] # We cut away the last marker (added by survey) and reverse the order - out_drift_length = list(self.drift_length[:-1][::-1]) - out_angle = list(-self.angle[:-1][::-1]) - out_tilt = list(-self.tilt[:-1][::-1]) - out_dx = list(-self.dx[:-1][::-1]) - out_dy = list(-self.dy[:-1][::-1]) - out_transf_x_rad = list(-self.transf_x_rad[:-1][::-1]) - out_transf_y_rad = list(-self.transf_y_rad[:-1][::-1]) - out_transf_s_rad = list(-self.transf_s_rad[:-1][::-1]) - out_name = list(self.name[:-1][::-1]) - - if isinstance(element0, str): + out_drift_length = list(self.drift_length[:-1][::-1]) + out_angle = list(-self.angle[:-1][::-1]) + out_tilt = list(-self.tilt[:-1][::-1]) + out_name = list(self.name[:-1][::-1]) + + if type(element0) is str: element0 = out_name.index(element0) X, Y, Z, theta, phi, psi = compute_survey( - X0 = X0, - Y0 = Y0, - Z0 = Z0, - theta0 = theta0, - phi0 = phi0, - psi0 = psi0, - drift_length = out_drift_length, - angle = out_angle, - tilt = out_tilt, - dx = out_dx, - dy = out_dy, - transf_x_rad = out_transf_x_rad, - transf_y_rad = out_transf_y_rad, - transf_s_rad = out_transf_s_rad, - element0 = element0, - reverse_xs = False) + X0, Y0, Z0, theta0, phi0, psi0, + out_drift_length, out_angle, out_tilt, + element0=element0) # Initializing dictionary out_columns = {} - out_scalars = {} + out_columns["X"] = np.array(X) + out_columns["Y"] = np.array(Y) + out_columns["Z"] = np.array(Z) + out_columns["theta"] = np.unwrap(theta) + out_columns["phi"] = np.unwrap(phi) + out_columns["psi"] = np.unwrap(psi) + + out_columns["name"] = np.array(list(out_name) + ["_end_point"]) + out_columns["s"] = self.s[-1] - self.s[::-1] - # Fill survey data - out_columns["X"] = np.array(X) - out_columns["Y"] = np.array(Y) - out_columns["Z"] = np.array(Z) - out_columns["theta"] = np.unwrap(theta) - out_columns["phi"] = np.unwrap(phi) - out_columns["psi"] = np.unwrap(psi) - out_columns["name"] = np.array(list(out_name) + ["_end_point"]) - out_columns["s"] = self.s[-1] - self.s[::-1] out_columns['drift_length'] = np.array(out_drift_length + [0.]) - out_columns['angle'] = np.array(out_angle + [0.]) - out_columns['tilt'] = np.array(out_tilt + [0.]) - out_columns["dx"] = np.array(out_dx + [0.]) - out_columns["dy"] = np.array(out_dy + [0.]) - out_columns["transf_x_rad"] = np.array(out_transf_x_rad + [0.]) - out_columns["transf_y_rad"] = np.array(out_transf_y_rad + [0.]) - out_columns["transf_s_rad"] = np.array(out_transf_s_rad + [0.]) + out_columns['angle'] = np.array(out_angle + [0.]) + out_columns['tilt'] = np.array(out_tilt + [0.]) + out_scalars = {} out_scalars["element0"] = element0 - out = SurveyTable( - data=(out_columns | out_scalars), - col_names=out_columns.keys()) + out = SurveyTable(data=(out_columns | out_scalars), + col_names=out_columns.keys()) return out - def plot(self, element_width = None, legend = True, **kwargs): - """ - Plot the survey using xplt.FloorPlot - """ - # Import the xplt module here - # (Not at the top as not default installation with xsuite) - import xplt - + def plot(self, element_width=None, legend=True, **kwargs): # Shallow copy of self out_sv_table = SurveyTable.__new__(SurveyTable) out_sv_table.__dict__.update(self.__dict__) @@ -279,19 +176,13 @@ def plot(self, element_width = None, legend = True, **kwargs): # Removing the count for repeated elements out_sv_table.name = np.array([nn.split('::')[0] for nn in out_sv_table.name]) - # Setting element width for plotting if element_width is None: x_range = max(self.X) - min(self.X) y_range = max(self.Y) - min(self.Y) z_range = max(self.Z) - min(self.Z) - element_width = max([x_range, y_range, z_range]) * 0.03 - - xplt.FloorPlot( - survey = out_sv_table, - line = self.line, - element_width = element_width, - **kwargs) - + element_width = max([x_range, y_range, z_range]) * 0.03 + import xplt + xplt.FloorPlot(out_sv_table, self.line, element_width=element_width, **kwargs) if legend: import matplotlib.pyplot as plt plt.legend() @@ -301,10 +192,10 @@ def plot(self, element_width = None, legend = True, **kwargs): # Main function # ================================================== -def survey_from_line( - line, - X0 = 0, Y0 = 0, Z0 = 0, theta0 = 0, phi0 = 0, psi0 = 0, - element0 = 0, values_at_element_exit = False, reverse = True): +def survey_from_line(line, X0=0, Y0=0, Z0=0, theta0=0, phi0=0, psi0=0, + element0=0, + values_at_element_exit=False, + reverse=True): """Execute SURVEY command. Based on MADX equivalent. Attributes, must be given in this order in the dictionary: X0 (float) Initial X position in meters. @@ -320,187 +211,87 @@ def survey_from_line( assert not values_at_element_exit, "Not implemented yet" - # Get line table to extract attributes - tt = line.get_table(attr = True) - # Extract angle and tilt from elements - angle = tt.angle_rad - tilt = tt.rot_s_rad - - # Extract drift lengths + tt = line.get_table(attr = True) + angle = tt.angle_rad + tilt = tt.rot_s_rad drift_length = tt.length drift_length[~tt.isthick] = 0 - # Extract xy shifts from elements - dx = tt.dx - dy = tt.dy - - # Handling of XYSRotation elements - transf_angle_rad = tt.transform_angle_rad - transf_x_rad = transf_angle_rad * np.array(tt.element_type == 'XRotation') - transf_y_rad = transf_angle_rad * np.array(tt.element_type == 'YRotation') - transf_s_rad = transf_angle_rad * np.array(tt.element_type == 'SRotation') - - if isinstance(element0, str): + if type(element0) == str: element0 = line.element_names.index(element0) X, Y, Z, theta, phi, psi = compute_survey( - X0 = X0, - Y0 = Y0, - Z0 = Z0, - theta0 = theta0, - phi0 = phi0, - psi0 = psi0, - drift_length = drift_length[:-1], - angle = angle[:-1], - tilt = tilt[:-1], - dx = dx[:-1], - dy = dy[:-1], - transf_x_rad = transf_x_rad[:-1], - transf_y_rad = transf_y_rad[:-1], - transf_s_rad = transf_s_rad[:-1], - element0 = element0, - reverse_xs = False) + X0, Y0, Z0, theta0, phi0, psi0, drift_length[:-1], angle[:-1], tilt[:-1], + element0=element0) # Initializing dictionary out_columns = {} out_scalars = {} - - # Fill survey data - out_columns["X"] = np.array(X) - out_columns["Y"] = np.array(Y) - out_columns["Z"] = np.array(Z) - out_columns["theta"] = np.unwrap(theta) - out_columns["phi"] = np.unwrap(phi) - out_columns["psi"] = np.unwrap(psi) - out_columns["name"] = tt.name - out_columns["s"] = tt.s + out_columns["X"] = np.array(X) + out_columns["Y"] = np.array(Y) + out_columns["Z"] = np.array(Z) + out_columns["theta"] = np.unwrap(theta) + out_columns["phi"] = np.unwrap(phi) + out_columns["psi"] = np.unwrap(psi) + + out_columns["name"] = tt.name + out_columns["s"] = tt.s out_columns['drift_length'] = drift_length - out_columns['angle'] = angle - out_columns['tilt'] = tilt - out_columns['dx'] = dx - out_columns['dy'] = dy - out_columns['transf_x_rad'] = transf_x_rad - out_columns['transf_y_rad'] = transf_y_rad - out_columns['transf_s_rad'] = transf_s_rad - - out_scalars['element0'] = element0 - - out = SurveyTable( - data = {**out_columns, **out_scalars}, # this is a merge - col_names = out_columns.keys()) + out_columns['angle'] = angle + out_columns['tilt'] = tilt + + out_scalars['element0'] = element0 + + out = SurveyTable(data={**out_columns, **out_scalars}, # this is a merge + col_names=out_columns.keys()) out._data['line'] = line return out -def compute_survey( - X0, Y0, Z0, theta0, phi0, psi0, - drift_length, angle, tilt, - dx, dy, transf_x_rad, transf_y_rad, transf_s_rad, - element0 = 0, reverse_xs = False): - """ - Compute survey from initial position and orientation. - """ +def compute_survey(X0, Y0, Z0, theta0, phi0, psi0, drift_length, angle, tilt, + element0=0, reverse_xs=False): - # If element0 is not the first element, split the survey if element0 != 0: - # Assert that reverse_xs is not implemented yet assert not(reverse_xs), "Not implemented yet" - - # Forward section of survey - drift_forward = drift_length[element0:] - angle_forward = angle[element0:] - tilt_forward = tilt[element0:] - dx_forward = dx[element0:] - dy_forward = dy[element0:] - transf_x_rad_forward = transf_x_rad[element0:] - transf_y_rad_forward = transf_y_rad[element0:] - transf_s_rad_forward = transf_s_rad[element0:] - - # Evaluate forward survey + drift_forward = drift_length[element0:] + angle_forward = angle[element0:] + tilt_forward = tilt[element0:] (X_forward, Y_forward, Z_forward, theta_forward, phi_forward, - psi_forward) = compute_survey( - X0 = X0, - Y0 = Y0, - Z0 = Z0, - theta0 = theta0, - phi0 = phi0, - psi0 = psi0, - drift_length = drift_forward, - angle = angle_forward, - tilt = tilt_forward, - dx = dx_forward, - dy = dy_forward, - transf_x_rad = transf_x_rad_forward, - transf_y_rad = transf_y_rad_forward, - transf_s_rad = transf_s_rad_forward, - element0 = 0, - reverse_xs = False) - - # Backward section of survey - drift_backward = drift_length[:element0][::-1] - angle_backward = -np.array(angle[:element0][::-1]) - tilt_backward = -np.array(tilt[:element0][::-1]) - dx_backward = -np.array(dx[:element0][::-1]) - dy_backward = -np.array(dy[:element0][::-1]) - transf_x_rad_backward = -np.array(transf_x_rad[:element0][::-1]) - transf_y_rad_backward = -np.array(transf_y_rad[:element0][::-1]) - transf_s_rad_backward = -np.array(transf_s_rad[:element0][::-1]) - - # Evaluate backward survey + psi_forward) = compute_survey(X0, Y0, Z0, theta0, phi0, psi0, + drift_forward, angle_forward, tilt_forward) + + drift_backward = drift_length[:element0][::-1] + angle_backward = -np.array(angle[:element0][::-1]) + tilt_backward = -np.array(tilt[:element0][::-1]) (X_backward, Y_backward, Z_backward, theta_backward, phi_backward, - psi_backward) = compute_survey( - X0 = X0, - Y0 = Y0, - Z0 = Z0, - theta0 = theta0, - phi0 = phi0, - psi0 = psi0, - drift_length = drift_backward, - angle = angle_backward, - tilt = tilt_backward, - dx = dx_backward, - dy = dy_backward, - transf_x_rad = transf_x_rad_backward, - transf_y_rad = transf_y_rad_backward, - transf_s_rad = transf_s_rad_backward, - element0 = 0, - reverse_xs = True) - - # Concatenate forward and backward - X = np.array(X_backward[::-1][:-1] + X_forward) - Y = np.array(Y_backward[::-1][:-1] + Y_forward) - Z = np.array(Z_backward[::-1][:-1] + Z_forward) - theta = np.array(theta_backward[::-1][:-1] + theta_forward) - phi = np.array(phi_backward[::-1][:-1] + phi_forward) - psi = np.array(psi_backward[::-1][:-1]+ psi_forward) + psi_backward) = compute_survey(X0, Y0, Z0, theta0, phi0, psi0, + drift_backward, angle_backward, tilt_backward, + reverse_xs=True) + + X = np.array(X_backward[::-1][:-1] + X_forward) + Y = np.array(Y_backward[::-1][:-1] + Y_forward) + Z = np.array(Z_backward[::-1][:-1] + Z_forward) + theta = np.array(theta_backward[::-1][:-1] + theta_forward) + phi = np.array(phi_backward[::-1][:-1] + phi_forward) + psi = np.array(psi_backward[::-1][:-1]+ psi_forward) return X, Y, Z, theta, phi, psi - # Initialise lists for storing the survey - X = [] - Y = [] - Z = [] - theta = [] - phi = [] - psi = [] - - # Initial position and orientation - v = np.array([X0, Y0, Z0]) - w = get_w_from_angles( - theta = theta0, - phi = phi0, - psi = psi0, - reverse_xs = reverse_xs) - + X = [] + Y = [] + Z = [] + theta = [] + phi = [] + psi = [] + v = np.array([X0, Y0, Z0]) + w = get_w_from_angles(theta=theta0, phi=phi0, psi=psi0, + reverse_xs=reverse_xs) # Advancing element by element - for ll, aa, tt, xx, yy, tx, ty, ts, in zip( - drift_length, angle, tilt, dx, dy, transf_x_rad, transf_y_rad, transf_s_rad): + for ll, aa, tt in zip(drift_length, angle, tilt): - # Get angles from w matrix after previous element - th, ph, ps = get_angles_from_w(w, reverse_xs = reverse_xs) + th, ph, ps = get_angles_from_w(w, reverse_xs=reverse_xs) - # Store position and orientation at element entrance X.append(v[0]) Y.append(v[1]) Z.append(v[2]) @@ -509,20 +300,10 @@ def compute_survey( psi.append(ps) # Advancing - v, w = advance_element( - v = v, - w = w, - length = ll, - angle = aa, - tilt = tt, - dx = xx, - dy = yy, - transf_x_rad = tx, - transf_y_rad = ty, - transf_s_rad = ts) + v, w = advance_element(v, w, length=ll, angle=aa, tilt=tt) # Last marker - th, ph, ps = get_angles_from_w(w, reverse_xs = reverse_xs) + th, ph, ps = get_angles_from_w(w, reverse_xs=reverse_xs) X.append(v[0]) Y.append(v[1]) Z.append(v[2]) @@ -530,5 +311,5 @@ def compute_survey( phi.append(ph) psi.append(ps) - # Return data for SurveyTable object + # Returns as SurveyTable object return X, Y, Z, theta, phi, psi