xtrack 0.84.0__tar.gz → 0.84.3__tar.gz

{xtrack-0.84.0/xtrack.egg-info → xtrack-0.84.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.84.0
+Version: 0.84.3
 Summary: Tracking library for particle accelerators
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xtrack

xtrack-0.84.3/tests/test_radiation_integrals.py

@@ -0,0 +1,399 @@
+import xtrack as xt
+import xobjects as xo
+
+import numpy as np
+import pathlib
+import pytest
+
+test_data_folder = pathlib.Path(
+        __file__).parent.joinpath('../test_data').absolute()
+
+def test_radiation_wiggler():
+    env = xt.load_madx_lattice(test_data_folder / 'sps_thick/sps.seq')
+    env.vars.load_madx(test_data_folder / 'sps_thick/lhc_q20.str')
+    line = env.sps
+
+    line['actcse.31632'].voltage = 4.2e+08
+    line['actcse.31632'].frequency = 3e6
+    line['actcse.31632'].lag = 180.
+
+    line.particle_ref = xt.Particles(energy0=20e9, mass0=xt.ELECTRON_MASS_EV)
+    env.particle_ref = line.particle_ref
+
+    # Wiggler parameters
+    k0_wig = 5e-3
+    tilt_rad = np.pi/2
+
+    lenwig = 25
+    numperiods = 20
+    lambdawig = lenwig / numperiods
+
+    wig = Wiggler(period=lambdawig, amplitude=k0_wig, num_periods=numperiods,
+                    angle_rad=tilt_rad, scheme='121a')
+
+    tt = line.get_table()
+    wig_elems = []
+    for name, element in wig.wiggler_dict.items():
+        env.elements[name] = element['element']
+        wig_elems.append(name)
+
+    wig_line = env.new_line(components=[
+                            env.new('s.wig', 'Marker'),
+                            wig_elems,
+                            env.new('e.wig', 'Marker'),
+    ])
+
+    line.insert(wig_line, anchor='start', at=1, from_='qd.31710@end')
+
+    env['sps_thick'] = env.sps.copy(shallow=True)
+
+    line.discard_tracker()
+    slicing_strategies = [
+        xt.Strategy(slicing=xt.Teapot(1)),  # Default
+        xt.Strategy(slicing=xt.Teapot(2), element_type=xt.Bend),
+        xt.Strategy(slicing=xt.Teapot(2), element_type=xt.RBend),
+        xt.Strategy(slicing=xt.Teapot(8), element_type=xt.Quadrupole),
+        xt.Strategy(slicing=xt.Teapot(20), name='mwp.*'),
+    ]
+    line.slice_thick_elements(slicing_strategies)
+
+    tw4d = line.twiss4d(radiation_integrals=True)
+    tw6d = line.twiss()
+
+    line.configure_radiation(model='mean')
+
+    tw_rad = line.twiss(eneloss_and_damping=True, strengths=True)
+
+    print('ex rad int:', tw4d.rad_int_eq_gemitt_x)
+    print('ex Chao:   ', tw_rad.eq_gemitt_x)
+    print('ey rad int:', tw4d.rad_int_eq_gemitt_y)
+    print('ey Chao:   ', tw_rad.eq_gemitt_y)
+
+    print('damping rate x [s^-1] rad int:   ', tw4d.rad_int_damping_constant_x_s)
+    print('damping rate x [s^-1] eigenval:  ', tw_rad.damping_constants_s[0])
+    print('damping rate y [s^-1] rad int:   ', tw4d.rad_int_damping_constant_y_s)
+    print('damping rate y [s^-1] eigenval:  ', tw_rad.damping_constants_s[1])
+    print('damping rate z [s^-1] rad int:   ', tw4d.rad_int_damping_constant_zeta_s)
+    print('damping rate z [s^-1] eigenval:  ', tw_rad.damping_constants_s[2])
+
+    xo.assert_allclose(
+        tw4d.rad_int_eq_gemitt_x, tw_rad.eq_gemitt_x, rtol=1e-3, atol=0)
+    xo.assert_allclose(
+        tw4d.rad_int_eq_gemitt_y, tw_rad.eq_gemitt_y, rtol=5e-3, atol=0)
+    xo.assert_allclose(
+        tw4d.rad_int_damping_constant_x_s, tw_rad.damping_constants_s[0],
+        rtol=1e-3, atol=0)
+    xo.assert_allclose(
+        tw4d.rad_int_damping_constant_y_s, tw_rad.damping_constants_s[1],
+        rtol=5e-3, atol=0)
+    xo.assert_allclose(
+        tw4d.rad_int_damping_constant_zeta_s, tw_rad.damping_constants_s[2],
+        rtol=1e-3, atol=0)
+
+def test_radiation_integrals_sls_combined_function_magnets():
+
+    env = xt.load_madx_lattice(test_data_folder / 'sls_2.0/b075_2024.09.25.madx')
+    line = env.ring
+    line.particle_ref = xt.Particles(energy0=2.7e9, mass0=xt.ELECTRON_MASS_EV)
+    line.configure_bend_model(num_multipole_kicks=20)
+
+    line['vrf'] = 1.8e6
+    line['frf'] = 499.6e6
+    line['lagrf'] = 180.
+
+    line.insert(
+        env.new('cav', 'Cavity', voltage='vrf', frequency='frf', lag='lagrf', at=0))
+
+    tt = line.get_table()
+    tw4d_thick = line.twiss4d()
+    tw6d_thick = line.twiss()
+
+    env['ring_thick'] = env.ring.copy(shallow=True)
+
+    line.discard_tracker()
+    slicing_strategies = [
+        xt.Strategy(slicing=None),  # Default
+        xt.Strategy(slicing=xt.Teapot(20), element_type=xt.Bend),
+        xt.Strategy(slicing=xt.Teapot(8), element_type=xt.Quadrupole),
+    ]
+    line.slice_thick_elements(slicing_strategies)
+
+    tw4d = line.twiss4d()
+    tw6d = line.twiss()
+
+    line.configure_radiation(model='mean')
+
+    tw_rad = line.twiss(eneloss_and_damping=True, strengths=True)
+
+    tw_integrals = line.twiss(radiation_integrals=True)
+
+    print('ex rad int:', tw_integrals.rad_int_eq_gemitt_x)
+    print('ex Chao:   ', tw_rad.eq_gemitt_x)
+    print('ey rad int:', tw_integrals.rad_int_eq_gemitt_y)
+    print('ey Chao:   ', tw_rad.eq_gemitt_y)
+
+    print('damping rate x [s^-1] rad int:   ', tw_integrals.rad_int_damping_constant_x_s)
+    print('damping rate x [s^-1] eigenval:  ', tw_rad.damping_constants_s[0])
+    print('damping rate y [s^-1] rad int:   ', tw_integrals.rad_int_damping_constant_y_s)
+    print('damping rate y [s^-1] eigenval:  ', tw_rad.damping_constants_s[1])
+    print('damping rate z [s^-1] rad int:   ', tw_integrals.rad_int_damping_constant_zeta_s)
+    print('damping rate z [s^-1] eigenval:  ', tw_rad.damping_constants_s[2])
+
+    xo.assert_allclose(
+        tw_integrals.rad_int_eq_gemitt_x, tw_rad.eq_gemitt_x, rtol=1e-3, atol=0)
+    xo.assert_allclose(
+        tw_integrals.rad_int_eq_gemitt_y, tw_rad.eq_gemitt_y, rtol=5e-3, atol=1e-20)
+    xo.assert_allclose(
+        tw_integrals.rad_int_damping_constant_x_s, tw_rad.damping_constants_s[0],
+        rtol=1e-3, atol=0)
+    xo.assert_allclose(
+        tw_integrals.rad_int_damping_constant_y_s, tw_rad.damping_constants_s[1],
+        rtol=5e-3, atol=0)
+    xo.assert_allclose(
+        tw_integrals.rad_int_damping_constant_zeta_s, tw_rad.damping_constants_s[2],
+        rtol=1e-3, atol=0)
+
+@pytest.mark.parametrize('tilt', [True, False], ids=['tilt', 'no_tilt'])
+def test_radiation_integrals_sps_vs_df(tilt):
+    env = xt.load_madx_lattice(test_data_folder / 'sps_thick/sps.seq')
+    env.vars.load_madx(test_data_folder / 'sps_thick/lhc_q20.str')
+    line = env.sps
+
+    line.particle_ref = xt.Particles(mass0=xt.ELECTRON_MASS_EV, energy0=10e9)
+
+    line.insert('zeta_shift', obj=xt.ZetaShift(), at=0)
+
+    # RF set tp stay in the linear region
+    env['actcse.31632'].voltage = 2500e6
+    env['actcse.31632'].frequency = 3e6
+    env['actcse.31632'].lag = 180.
+
+    if tilt:
+
+        tt = line.get_table()
+        tt_mb = tt.rows['mb.*']
+        tt_lsf = tt.rows['lsf.*']
+        tt_lsd = tt.rows['lsd.*']
+
+        for nn in tt_mb.name:
+            line[nn].rot_s_rad = np.deg2rad(90)
+
+        # I need skew sextupoles to correct the chromaticity when the dispersion is vertical
+        for nn in list(tt_lsf.name) + list(tt_lsd.name):
+            line[nn].rot_s_rad = np.deg2rad(-30)
+
+        opt_q = line.match(
+            solve=False,
+            vary=xt.VaryList(['kqf', 'kqd'], step=1e-4),
+            targets=xt.TargetSet(qx=20.18, qy=20.13, tol=1e-4))
+        opt_q.solve()
+
+        opt_chrom = line.match(
+            solve=False,
+            vary=xt.VaryList(['klsfb', 'klsfa', 'klsdb', 'klsda'], step=1e-4),
+            targets=xt.TargetSet(dqx=1., dqy=1, tol=1e-4))
+        opt_chrom.solve()
+
+
+    tw4d = line.twiss4d()
+    tw6d = line.twiss()
+
+    line.configure_radiation(model='mean')
+    line.compensate_radiation_energy_loss()
+
+    tw_rad = line.twiss(eneloss_and_damping=True)
+
+    # Prepare trim
+    env['frev0'] = 1. / tw4d.T_rev0
+    env['circum'] = tw4d.circumference
+    env['frev_trim'] = 0.
+
+    env['zeta_shift'].dzeta = 'circum * frev_trim / frev0'
+
+    dfrev = np.linspace(-0.7, 0.7, 21)
+    part_x = []
+    part_y = []
+    part_zeta = []
+    damp_cons_x_s = []
+    demp_const_y_s = []
+    damp_const_zeta_s = []
+    eq_gemitt_x = []
+    eq_gemitt_y = []
+    eq_gemitt_zeta = []
+    rad_int_dconst_x_s =[]
+    rad_int_dconst_y_s = []
+    rad_int_dconst_zeta_s = []
+    rad_int_ex = []
+    rad_int_ey = []
+    delta_ave = []
+    for dff in dfrev:
+        print(f'dfrev: {dff}')
+        env['frev_trim'] = dff
+        tw = line.twiss(eneloss_and_damping=True,
+                        radiation_integrals=True)
+        part_x.append(tw.partition_numbers[0])
+        part_y.append(tw.partition_numbers[1])
+        part_zeta.append(tw.partition_numbers[2])
+        eq_gemitt_x.append(tw.eq_gemitt_x)
+        eq_gemitt_y.append(tw.eq_gemitt_y)
+        eq_gemitt_zeta.append(tw.eq_gemitt_zeta)
+        delta_ave.append(tw.delta.mean())
+
+        damp_cons_x_s.append(tw.damping_constants_s[0])
+        damp_const_zeta_s.append(tw.damping_constants_s[2])
+        demp_const_y_s.append(tw.damping_constants_s[1])
+
+        rad_int_dconst_x_s.append(tw.rad_int_damping_constant_x_s)
+        rad_int_dconst_y_s.append(tw.rad_int_damping_constant_y_s)
+        rad_int_dconst_zeta_s.append(tw.rad_int_damping_constant_zeta_s)
+
+        rad_int_ex.append(tw.rad_int_eq_gemitt_x)
+        rad_int_ey.append(tw.rad_int_eq_gemitt_y)
+
+    # Cast to numpy arrays
+    part_x = np.array(part_x)
+    part_y = np.array(part_y)
+    part_zeta = np.array(part_zeta)
+    eq_gemitt_x = np.array(eq_gemitt_x)
+    eq_gemitt_y = np.array(eq_gemitt_y)
+    eq_gemitt_zeta = np.array(eq_gemitt_zeta)
+    delta_ave = np.array(delta_ave)
+    damp_cons_x_s = np.array(damp_cons_x_s)
+    demp_const_y_s = np.array(demp_const_y_s)
+    damp_const_zeta_s = np.array(damp_const_zeta_s)
+    rad_int_dconst_x_s = np.array(rad_int_dconst_x_s)
+    rad_int_dconst_y_s = np.array(rad_int_dconst_y_s)
+    rad_int_dconst_zeta_s = np.array(rad_int_dconst_zeta_s)
+    rad_int_ex = np.array(rad_int_ex)
+    rad_int_ey = np.array(rad_int_ey)
+
+    xo.assert_allclose(
+        rad_int_dconst_x_s, damp_cons_x_s, rtol=0.03, atol=0.05)
+    xo.assert_allclose(
+        rad_int_dconst_y_s, demp_const_y_s, rtol=0.03, atol=0.05)
+    xo.assert_allclose(
+        rad_int_dconst_zeta_s, damp_const_zeta_s, rtol=0.03, atol=0.05)
+
+    mask = np.abs(rad_int_dconst_x_s) > 0.25
+    xo.assert_allclose(
+        rad_int_ex[mask], eq_gemitt_x[mask], rtol=0.065, atol=1e-14)
+    mask = np.abs(rad_int_dconst_y_s) > 0.25
+    xo.assert_allclose(
+        rad_int_ey[mask], eq_gemitt_y[mask], rtol=0.065, atol=1e-14)
+
+    if tilt:
+        xo.assert_allclose(
+            rad_int_ex, 0, rtol=0, atol=1e-14)
+    else:
+        xo.assert_allclose(
+            rad_int_ey, 0, rtol=1e-14, atol=1e-14)
+
+class Wiggler:
+    def __init__(self, period, amplitude, num_periods, angle_rad=0,
+                scheme='121s'):
+        # The scheme_library is a list of all the possible schemes that can be
+        # used. The scheme determines the order of the dipoles in the wiggler.
+        # The 's' and 'a' stand for a symmetric/antisymmetric configuration
+        # respectively.
+        self.scheme_library = ['121s', '121a']
+
+        self.wiggler_period = period
+        self.wiggler_amplitude = amplitude
+        self.wiggler_num_periods = num_periods
+        self.angle_rad = angle_rad
+        self.scheme = scheme
+        self.spacing = 0
+        self.wiggler = self._build_wiggler_()
+        self.wiggler_dict = self._build_dict_()
+
+    def _build_wiggler_(self):
+        wiggler = []
+
+        if self.scheme == '121s':
+            for i in range(self.wiggler_num_periods + 1):
+                if i != 0 and i != self.wiggler_num_periods:
+                    wiggler += [
+                        xt.Bend(length=self.wiggler_period / 4,
+                                k0=-self.wiggler_amplitude, h=0,
+                                rot_s_rad=self.angle_rad),
+                        xt.Bend(length=self.wiggler_period / 4,
+                                k0=-self.wiggler_amplitude, h=0,
+                                rot_s_rad=self.angle_rad),
+                        xt.Bend(length=self.wiggler_period / 4,
+                                k0=self.wiggler_amplitude, h=0,
+                                rot_s_rad=self.angle_rad),
+                        xt.Bend(length=self.wiggler_period / 4,
+                                k0=self.wiggler_amplitude, h=0,
+                                rot_s_rad=self.angle_rad)
+                    ]
+
+                elif i == 0:
+                    wiggler += [
+                        xt.Bend(length=self.wiggler_period / 4,
+                                k0=self.wiggler_amplitude, h=0,
+                                rot_s_rad=self.angle_rad)
+                    ]
+
+                else:
+                    wiggler += [
+                        xt.Bend(length=self.wiggler_period / 4,
+                                k0=-self.wiggler_amplitude, h=0,
+                                rot_s_rad=self.angle_rad),
+                        xt.Bend(length=self.wiggler_period / 4,
+                                k0=-self.wiggler_amplitude, h=0,
+                                rot_s_rad=self.angle_rad),
+                        xt.Bend(length=self.wiggler_period / 4,
+                                k0=self.wiggler_amplitude, h=0,
+                                rot_s_rad=self.angle_rad)
+                    ]
+
+        if self.scheme == '121a':
+            for i in range(self.wiggler_num_periods):
+                sign = 1 if i % 2 == 0 else -1
+                wiggler += [
+                    xt.Bend(length=self.wiggler_period / 4,
+                            k0=-sign * self.wiggler_amplitude, h=0,
+                            rot_s_rad=self.angle_rad),
+                    xt.Bend(length=self.wiggler_period / 4,
+                            k0=sign * self.wiggler_amplitude, h=0,
+                            rot_s_rad=self.angle_rad),
+                    xt.Bend(length=self.wiggler_period / 4,
+                            k0=sign * self.wiggler_amplitude, h=0,
+                            rot_s_rad=self.angle_rad),
+                    xt.Bend(length=self.wiggler_period / 4,
+                            k0=-sign * self.wiggler_amplitude, h=0,
+                            rot_s_rad=self.angle_rad)
+                ]
+
+        print(f'wiggler.shape = {len(wiggler)}')
+
+        return wiggler
+
+    def _get_wiggler_names_(self, wiggler, wiggler_number='1'):
+        wiggler_names = []
+        for i in range(len(wiggler)):
+            wiggler_names += ['mwp' + str(i + 1) + '.' + wiggler_number]
+
+        print(f'wiggler_names.shape = {len(wiggler_names)}')
+
+        return wiggler_names
+
+    def _get_element_positions_(self, wiggler):
+        ele_pos = np.zeros(len(wiggler))
+        for i in range(1, len(wiggler)):
+            ele_pos[i] = ele_pos[i - 1] + wiggler[i - 1].length + self.spacing
+
+        print(f'ele_pos.shape = {ele_pos.shape}')
+
+        return ele_pos
+
+    def _build_dict_(self):
+        wiggler = self._build_wiggler_()
+        wiggler_names = self._get_wiggler_names_(wiggler)
+        ele_pos = self._get_element_positions_(wiggler)
+        wiggler_dict = {
+            name: {'element': obj, 'position': pos}
+            for name, obj, pos in zip(wiggler_names, wiggler, ele_pos)
+        }
+
+        return wiggler_dict

{xtrack-0.84.0 → xtrack-0.84.3}/tests/test_rf_track.py

@@ -26,7 +26,8 @@ def test_rf_track_lattice():
 
     # Define the RF-Track element
     vol = RFT.Volume()
-    vol.dt_mm = 0.1
+    vol.dt_mm = 1
+    vol.odeint_epsabs = 1e-10
     vol.odeint_algorithm = 'rk2'
     vol.set_static_Bfield(0.0, By, 0.0)
     vol.set_s0(rho, 0.0, 0.0, 0.0, 0.0, 0.0)
@@ -41,29 +42,29 @@ def test_rf_track_lattice():
     pi = np.pi
     elements = {
         'd1.1':  xt.Drift(length=1),
-        'mb1.1': xt.RFT_Element(element=vol),
+        'mb1.1': xt.RFT_Element(element=vol,update_ref=False),
         'd2.1':  xt.Drift(length=1),
 
         'mqd.1': xt.Quadrupole(length=0.3, k1=-0.7),
         'd3.1':  xt.Drift(length=1),
-        'mb2.1': xt.RFT_Element(element=vol),
+        'mb2.1': xt.RFT_Element(element=vol,update_ref=False),
         'd4.1':  xt.Drift(length=1),
 
         'd1.2':  xt.Drift(length=1),
-        'mb1.2': xt.RFT_Element(element=vol),
+        'mb1.2': xt.RFT_Element(element=vol,update_ref=False),
         'd2.2':  xt.Drift(length=1),
 
         'mqd.2': xt.Quadrupole(length=0.3, k1=-0.7),
         'd3.2':  xt.Drift(length=1),
-        'mb2.2': xt.RFT_Element(element=vol),
+        'mb2.2': xt.RFT_Element(element=vol,update_ref=False),
         'd4.2':  xt.Drift(length=1),
-    }
+        }
 
     # Build the ring
     line = xt.Line(elements=elements, element_names=list(elements.keys()))
     line.particle_ref = xt.Particles(p0c=p0c, mass0=xt.PROTON_MASS_EV)
     line.configure_bend_model(core='full', edge=None)
-
+    line.reset_s_at_end_turn = False
 
     ## Transfer lattice on context and compile tracking code=
     line.build_tracker()
@@ -74,13 +75,13 @@ def test_rf_track_lattice():
     rng = np.random.default_rng(2021)
 
     particles0 = xp.Particles(p0c=p0c, #eV
-                            q0=1, mass0=xp.PROTON_MASS_EV,
-                            x=rng.uniform(-1e-3, 1e-3, n_part),
-                            px=rng.uniform(-1e-5, 1e-5, n_part),
-                            y=rng.uniform(-2e-3, 2e-3, n_part),
-                            py=rng.uniform(-3e-5, 3e-5, n_part),
-                            zeta=rng.uniform(-1e-2, 1e-2, n_part),
-                            delta=rng.uniform(-1e-2, 1e-2, n_part))
+                                  q0=1, mass0=xp.PROTON_MASS_EV,
+                                  x=rng.uniform(-1e-3, 1e-3, n_part),
+                                  px=rng.uniform(-1e-5, 1e-5, n_part),
+                                  y=rng.uniform(-2e-3, 2e-3, n_part),
+                                  py=rng.uniform(-3e-5, 3e-5, n_part),
+                                  zeta=rng.uniform(-1e-2, 1e-2, n_part),
+                                  delta=rng.uniform(-1e-2, 1e-2, n_part))
 
     particles_rft = particles0.copy()
 
@@ -96,6 +97,7 @@ def test_rf_track_lattice():
     ele_xt['mb2.2'] = b_xt
 
     line_ref = xt.Line(elements=ele_xt, element_names=line.element_names)
+    line_ref.reset_s_at_end_turn = False
     line_ref.build_tracker()
     particles_ref = particles0.copy()
 
@@ -105,18 +107,19 @@ def test_rf_track_lattice():
 
     assert_allclose = np.testing.assert_allclose
     assert np.all(particles_rft.particle_id == particles_ref.particle_id)
-    assert_allclose(particles_rft.x,  particles_ref.x,  atol=6e-5, rtol=0)
-    assert_allclose(particles_rft.px, particles_ref.px, atol=2e-5, rtol=0)
-    assert_allclose(particles_rft.y,  particles_ref.y,  atol=6e-7, rtol=0)
-    assert_allclose(particles_rft.py, particles_ref.py, atol=2e-6, rtol=0)
-    assert_allclose(particles_rft.rpp, particles_ref.rpp, atol=3e-5, rtol=0)
-    assert_allclose(particles_rft.rvv, particles_ref.rvv, atol=1e-5, rtol=0)
-    assert_allclose(particles_rft.ptau, particles_ref.ptau, atol=2e-5, rtol=0)
-    assert_allclose(particles_rft.delta, particles_ref.delta, atol=4e-5, rtol=0)
+    assert_allclose(particles_rft.x,  particles_ref.x,  atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.px, particles_ref.px, atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.y,  particles_ref.y,  atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.py, particles_ref.py, atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.rpp, particles_ref.rpp, atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.rvv, particles_ref.rvv, atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.ptau, particles_ref.ptau, atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.delta, particles_ref.delta, atol=1e-10, rtol=0)
     assert_allclose(particles_rft.chi, particles_ref.chi, atol=1e-10, rtol=0)
-    assert_allclose(particles_rft.p0c, particles_ref.p0c, atol=4e-5, rtol=0)
-    assert_allclose(particles_rft.energy0, particles_ref.energy0, atol=4e-5, rtol=0)
-    assert_allclose(particles_rft.zeta, particles_ref.zeta, atol=4e-5, rtol=0)
+    assert_allclose(particles_rft.p0c, particles_ref.p0c, atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.energy0, particles_ref.energy0, atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.zeta, particles_ref.zeta, atol=1e-10, rtol=0)
     assert_allclose(particles_rft.beta0, particles_ref.beta0, atol=1e-10, rtol=0)
     assert_allclose(particles_rft.mass0, particles_ref.mass0, atol=1e-10, rtol=0)
     assert_allclose(particles_rft.gamma0, particles_ref.gamma0, atol=1e-10, rtol=0)
+    assert_allclose(particles_rft.s, particles_ref.s, atol=1e-10, rtol=0)

xtrack-0.84.3/xtrack/_version.py

@@ -0,0 +1 @@
+__version__ = '0.84.3'

{xtrack-0.84.0 → xtrack-0.84.3}/xtrack/beam_elements/apertures_src/limitpolygon.h

@@ -56,30 +56,30 @@ void LimitPolygon_track_local_particle(LimitPolygonData el,
 GPUKERN
 void LimitPolygon_impact_point_and_normal(
 		             LimitPolygonData el,
-                /*gpuglmem*/ const double* x_in,
-                /*gpuglmem*/ const double* y_in, 
-                /*gpuglmem*/ const double* z_in,
-                /*gpuglmem*/ const double* x_out,
-                /*gpuglmem*/ const double* y_out, 
-                /*gpuglmem*/ const double* z_out,
-                             const int64_t n_impacts,
-                /*gpuglmem*/       double* x_inters,
-                /*gpuglmem*/       double* y_inters, 
-                /*gpuglmem*/       double* z_inters,
-                /*gpuglmem*/       double* Nx_inters,
-                /*gpuglmem*/       double* Ny_inters,
-                /*gpuglmem*/       int64_t* i_found){ 
-			           
+                GPUGLMEM const double* x_in,
+                GPUGLMEM const double* y_in,
+                GPUGLMEM const double* z_in,
+                GPUGLMEM const double* x_out,
+                GPUGLMEM const double* y_out,
+                GPUGLMEM const double* z_out,
+                         const int64_t n_impacts,
+                GPUGLMEM double* x_inters,
+                GPUGLMEM double* y_inters,
+                GPUGLMEM double* z_inters,
+                GPUGLMEM double* Nx_inters,
+                GPUGLMEM double* Ny_inters,
+                GPUGLMEM int64_t* i_found){
+
     double const tol = 1e-13;
 
     int64_t N_edg = LimitPolygonData_len_x_vertices(el);
-    /*gpuglmem*/ const double* Vx = LimitPolygonData_getp1_x_vertices(el, 0);
-    /*gpuglmem*/ const double* Vy = LimitPolygonData_getp1_y_vertices(el, 0);
-    /*gpuglmem*/ const double* Nx = LimitPolygonData_getp1_x_normal(el, 0);
-    /*gpuglmem*/ const double* Ny = LimitPolygonData_getp1_y_normal(el, 0);
+    GPUGLMEM const double* Vx = LimitPolygonData_getp1_x_vertices(el, 0);
+    GPUGLMEM const double* Vy = LimitPolygonData_getp1_y_vertices(el, 0);
+    GPUGLMEM const double* Nx = LimitPolygonData_getp1_x_normal(el, 0);
+    GPUGLMEM const double* Ny = LimitPolygonData_getp1_y_normal(el, 0);
     double resc_fac = LimitPolygonData_get_resc_fac(el);
 
-    for (int64_t i_imp=0; i_imp<n_impacts; i_imp++){ //vectorize_over i_imp n_impacts
+    VECTORIZE_OVER(i_imp, n_impacts);
 
         double t_min_curr = 1.;
         int64_t i_found_curr = -1;
@@ -89,9 +89,8 @@ void LimitPolygon_impact_point_and_normal(
         double y_out_curr = y_out[i_imp];
 
         for (int64_t ii=0; ii<N_edg; ii++){
-
-	    double t_border;
-	    double t_ii;
+	        double t_border;
+	        double t_ii;
             double const den = ((y_out_curr-y_in_curr)*(Vx[ii+1]-Vx[ii])
 			    +(x_in_curr-x_out_curr)*(Vy[ii+1]-Vy[ii]));
             if (den == 0.){
@@ -99,23 +98,23 @@ void LimitPolygon_impact_point_and_normal(
                 // the case case overlapping the edge is not possible (this would not allow Pin inside and Pout outside - a point on the edge is condidered outside)
                 // the only case left is segment parallel to tue edge => no intersection
                 t_border = -2.;
-	    }
+	        }
             else{
                 t_border=((y_out_curr-y_in_curr)*(x_in_curr-Vx[ii])
 		         +(x_in_curr-x_out_curr)*(y_in_curr-Vy[ii]))/den;
-	    }
+	        }
 
             if (t_border>=0.-tol && t_border<=1.+tol){
                 t_ii = (Nx[ii]*(Vx[ii]-x_in_curr)
-		       +Ny[ii]*(Vy[ii]-y_in_curr)) 
-		       /(Nx[ii]*(x_out_curr-x_in_curr)
-	                 +Ny[ii]*(y_out_curr-y_in_curr));
+		            +Ny[ii]*(Vy[ii]-y_in_curr))
+		            /(Nx[ii]*(x_out_curr-x_in_curr)
+	                +Ny[ii]*(y_out_curr-y_in_curr));
                 if (t_ii>=0.-tol && t_ii<t_min_curr+tol){
-                    t_min_curr=t_ii;
+                    t_min_curr = t_ii;
                     i_found_curr = ii;
-		}
+		        }
             }
-	}
+	    }
 
         t_min_curr=resc_fac*t_min_curr;
         x_inters[i_imp]=t_min_curr*x_out_curr+(1.-t_min_curr)*x_in_curr;
@@ -125,10 +124,10 @@ void LimitPolygon_impact_point_and_normal(
         if (i_found_curr>=0){
             Nx_inters[i_imp] = Nx[i_found_curr];
             Ny_inters[i_imp] = Ny[i_found_curr];
-	}
+	    }
         i_found[i_imp] = i_found_curr;
-    } //end_vectorize
-    
+
+    END_VECTORIZE;
 }
 
 #endif

{xtrack-0.84.0 → xtrack-0.84.3}/xtrack/beam_elements/rft_element.py

@@ -23,7 +23,7 @@ class RFT_Element:
     iscollective = True # <-- To state that the element uses a python track method
     isthick = True
 
-    def __init__(self, element):
+    def __init__(self, element, update_ref=False):
 
         import RF_Track as rft
         self.length = element.get_length()
@@ -33,6 +33,7 @@ class RFT_Element:
         self.bunch_out = rft.Bunch6d()
         self.arr_for_rft = np.empty(0)
         self.arr_for_xt = np.empty(0)
+        self.update_ref = update_ref
 
     def track(self, particles, increment_at_element=False):
 
@@ -67,10 +68,13 @@ class RFT_Element:
 
         # Update particles
         self.arr_for_xt = self.arr_for_xt[self.arr_for_xt[:,7].argsort()] # sort by particle id
+        if self.update_ref:
+            p.p0c = pref1[0].Pc * 1e6
+        p.s += self.length
         p.x  = self.arr_for_xt[:,0] / 1e3 # m
         p.px = self.arr_for_xt[:,1] * 1e6 / p.p0c # rad
         p.y  = self.arr_for_xt[:,2] / 1e3 # m
         p.py = self.arr_for_xt[:,3] * 1e6 / p.p0c # rad
-        p.zeta += self.length - self.arr_for_xt[:,4] * p.beta0 / 1e3 # m
+        p.zeta = self.length - self.arr_for_xt[:,4] * p.beta0 / 1e3 # m
         p.delta = (self.arr_for_xt[:,5] * 1e6 - p.p0c) / p.p0c #
         p.state[np.isnan(self.arr_for_xt[:,6])!=np.isnan(self.arr_for_rft[:,9])] = -400 # lost in RF-Track