More complete Multipole Rotations and Shifts within solenoid elements

examples/solenoid/001_multipole_shifts.py

@@ -0,0 +1,131 @@
+"""
+Test implementation of the multipole shifts in solenoid slices
+=============================================
+Author(s): John P T Salvesen
+Email:  [email protected]
+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()

examples/solenoid/002_multipole_rotations.py

@@ -0,0 +1,176 @@
+"""
+Test implementation of the multipole shifts in solenoid slices
+=============================================
+Author(s): John P T Salvesen
+Email:  [email protected]
+Date:   18-11-2024
+"""
+################################################################################
+# Required Modules
+################################################################################
+import xtrack as xt
+import numpy as np
+import matplotlib.pyplot as plt
+
+################################################################################
+# User Parameters
+################################################################################
+N_SLICES    = int(1E2)
+K0          = 1E-3
+L_SOL       = 1
+XING_RAD    = 1E-3
+
+BETX        = 100E-3
+BETY        = 1E-3
+
+################################################################################
+# Build Test Lines
+################################################################################
+
+########################################
+# 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)]
+
+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')
+
+################################################################################
+# Test Assertions
+################################################################################
+# 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
+########################################
+plt.show()

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.

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);
     }