xtrack 0.98.13__tar.gz → 0.99.0__tar.gz

{xtrack-0.98.13/xtrack.egg-info → xtrack-0.99.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xtrack
-Version: 0.98.13
+Version: 0.99.0
 Summary: Tracking library for particle accelerators
 Author: G. Iadarola et al.
 License-Expression: Apache-2.0

{xtrack-0.98.13 → xtrack-0.99.0}/tests/test_coupling_edwards_teng.py

@@ -81,10 +81,12 @@ def test_coupling_edwards_teng():
     xo.assert_allclose(tw.bety_edw_teng, bety_mad_at_s, atol=0, rtol=1e-6)
     xo.assert_allclose(tw.alfy_edw_teng, alfy_mad_at_s, atol=1e-4, rtol=1e-6)
 
-    xo.assert_allclose(tw.f1001, rdt_mad_at_s['f1001'], rtol=1e-3,
+    # Note different convention:
+    xo.assert_allclose(tw.f1001, -np.conj(rdt_mad_at_s['f1001']), rtol=1e-3,
                     atol=5e-4 * np.max(np.abs(rdt_mad_at_s['f1001'])))
-    xo.assert_allclose(tw.f1010, rdt_mad_at_s['f1010'], rtol=1e-3,
+    xo.assert_allclose(tw.f1010, -np.conj(rdt_mad_at_s['f1010']), rtol=1e-3,
                     atol=5e-3 * np.max(np.abs(rdt_mad_at_s['f1010'])))
+    xo.assert_allclose(tw.f0110, np.conj(tw.f1001), rtol=0, atol=1e-12)
 
 
 

{xtrack-0.98.13 → xtrack-0.99.0}/tests/test_monitor.py

@@ -492,3 +492,128 @@ def test_beam_position_monitor(test_context):
     expected_x_centroid[18:20] = np.nan
     assert_allclose(monitor.x_cen, expected_x_centroid, err_msg="Monitor x centroid does not match expected values")
     assert_allclose(monitor.y_cen, -expected_x_centroid, err_msg="Monitor y centroid does not match expected values")
+
+@pytest.mark.parametrize("with_progress_bar", [False, True])
+def test_multi_element_monitor(with_progress_bar):
+
+    line = xt.load(test_data_folder / 'hllhc15_thick/lhc_thick_with_knobs.json')
+
+    # Get names of all BPMs in the line
+    tt = line.get_table()
+    tt_obs = tt.rows.match(name='bpm.*')
+    tt_obs.name # is ['bpmw.4r7.b1', 'bpmwe.4r7.b1', 'bpmw.5r7.b1', ...,
+
+    # Generate some particle
+    p0 = xt.Particles(p0c=7e12, x=1e-6*np.arange(20),
+                            y=2e-6*np.arange(20),
+                            delta=1e-5*np.arange(20)
+    )
+
+    # Track particles for 10 turns, monitoring at all BPMs
+    num_turns = 10
+    p = p0.copy()
+    line.track(p, num_turns=num_turns, multi_element_monitor_at=tt_obs.name,
+               with_progress=(1 if with_progress_bar else None))
+
+    # Get the recorded data
+    mon = line.record_multi_element_last_track
+
+    line_with_monitors = line.copy(shallow=True)
+
+    nn_check_1 = 'bpm.17r3.b1'
+    nn_check_2 = 'bpm.23l5.b1'
+
+    mon_1 = xt.ParticlesMonitor(
+        start_at_turn=0,
+        stop_at_turn=num_turns,
+        num_particles=len(p0.x))
+    mon_2 = xt.ParticlesMonitor(
+        start_at_turn=0,
+        stop_at_turn=num_turns,
+        num_particles=len(p0.x))
+
+    line_with_monitors.insert('mon_1', mon_1, at=nn_check_1+'@start')
+    line_with_monitors.insert('mon_2', mon_2, at=nn_check_2+'@start')
+    p_ref = p0.copy()
+    line_with_monitors.track(p_ref, num_turns=num_turns)
+    xo.assert_allclose(mon.get('x', obs_name=nn_check_1), mon_1.x.T, atol=1e-14)
+    xo.assert_allclose(mon.get('px', obs_name=nn_check_1), mon_1.px.T, atol=1e-14)
+    xo.assert_allclose(mon.get('y', obs_name=nn_check_1), mon_1.y.T, atol=1e-14)
+    xo.assert_allclose(mon.get('py', obs_name=nn_check_1), mon_1.py.T, atol=1e-14)
+    xo.assert_allclose(mon.get('zeta', obs_name=nn_check_1), mon_1.zeta.T, atol=1e-14)
+    xo.assert_allclose(mon.get('delta', obs_name=nn_check_1), mon_1.delta.T, atol=1e-14)
+
+    xo.assert_allclose(mon.get('x', obs_name=nn_check_2), mon_2.x.T, atol=1e-14)
+    xo.assert_allclose(mon.get('px', obs_name=nn_check_2), mon_2.px.T, atol=1e-14)
+    xo.assert_allclose(mon.get('y', obs_name=nn_check_2), mon_2.y.T, atol=1e-14)
+    xo.assert_allclose(mon.get('py', obs_name=nn_check_2), mon_2.py.T, atol=1e-14)
+    xo.assert_allclose(mon.get('zeta', obs_name=nn_check_2), mon_2.zeta.T, atol=1e-14)
+    xo.assert_allclose(mon.get('delta', obs_name=nn_check_2), mon_2.delta.T, atol=1e-14)
+
+    # Access all data for a given coordinate
+    assert mon.data.shape == (num_turns, len(p0.x), 6, len(tt_obs))
+    xo.assert_allclose(mon.get('x'), mon.data[:,:,0,:], atol=1e-14)
+    xo.assert_allclose(mon.get('px'), mon.data[:,:,1,:], atol=1e-14)
+    xo.assert_allclose(mon.get('y'), mon.data[:,:,2,:], atol=1e-14)
+    xo.assert_allclose(mon.get('py'), mon.data[:,:,3,:], atol=1e-14)
+    xo.assert_allclose(mon.get('zeta'), mon.data[:,:,4,:], atol=1e-14)
+    xo.assert_allclose(mon.get('delta'), mon.data[:,:,5,:], atol=1e-14)
+
+    idx = tt_obs.rows.indices['bpm.17r3.b1'][0]
+
+    # # Access data for a given coordinate and observation point
+    # mon.get('x', obs_name='bpmw.4r7.b1') # shape ddddis (nom_turns, n_particles)
+    xo.assert_allclose(mon.get('x', obs_name=nn_check_1), mon.data[:,:,0,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('px', obs_name=nn_check_1), mon.data[:,:,1,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('y', obs_name=nn_check_1), mon.data[:,:,2,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('py', obs_name=nn_check_1), mon.data[:,:,3,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('delta', obs_name=nn_check_1), mon.data[:,:,5,idx], atol=1e-14)
+
+
+    # # Access data for a given coordinate, observation point and particle_id
+    # mon.get('x', obs_name='bpmw.4r7.b1', particle_id=10) # shape is (nom_turns,)
+    particle_id = 10
+    xo.assert_allclose(mon.get('x', obs_name=nn_check_1, particle_id=particle_id),
+                    mon.data[:,particle_id - mon.part_id_start,0,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('px', obs_name=nn_check_1, particle_id=particle_id),
+                    mon.data[:,particle_id - mon.part_id_start,1,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('y', obs_name=nn_check_1, particle_id=particle_id),
+                    mon.data[:,particle_id - mon.part_id_start,2,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('py', obs_name=nn_check_1, particle_id=particle_id),
+                    mon.data[:,particle_id - mon.part_id_start,3,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('zeta', obs_name=nn_check_1, particle_id=particle_id),
+                    mon.data[:,particle_id - mon.part_id_start,4,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('delta', obs_name=nn_check_1, particle_id=particle_id),
+                    mon.data[:,particle_id - mon.part_id_start,5,idx], atol=1e-14)
+
+    # # Access data for a given coordinate, observation point, particle_id and turn
+    # mon.get('x', obs_name='bpmw.4r7.b1', particle_id=10, turn=5) # is a scalar
+    turn = 5
+    xo.assert_allclose(mon.get('x', obs_name=nn_check_1, particle_id=particle_id, turn=turn),
+                        mon.data[turn,particle_id - mon.part_id_start,0,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('px', obs_name=nn_check_1, particle_id=particle_id, turn=turn),
+                        mon.data[turn,particle_id - mon.part_id_start,1,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('y', obs_name=nn_check_1, particle_id=particle_id, turn=turn),
+                        mon.data[turn,particle_id - mon.part_id_start,2,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('py', obs_name=nn_check_1, particle_id=particle_id, turn=turn),
+                        mon.data[turn,particle_id - mon.part_id_start,3,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('zeta', obs_name=nn_check_1, particle_id=particle_id, turn=turn),
+                    mon.data[turn,particle_id - mon.part_id_start,4,idx], atol=1e-14)
+    xo.assert_allclose(mon.get('delta', obs_name=nn_check_1, particle_id=particle_id, turn=turn),
+                        mon.data[turn,particle_id - mon.part_id_start,5,idx], atol=1e-14)
+
+    # # Get all BPMs for a given particles and turn
+    # mon.get('x', particle_id=10, turn=5) # shape is (n_obs_points,)
+    xo.assert_allclose(mon.get('x', particle_id=particle_id, turn=turn),
+                    mon.data[turn,particle_id - mon.part_id_start,0,:], atol=1e-14)
+    xo.assert_allclose(mon.get('px', particle_id=particle_id, turn=turn),
+                    mon.data[turn,particle_id - mon.part_id_start,1,:], atol=1e-14)
+    xo.assert_allclose(mon.get('y', particle_id=particle_id, turn=turn),
+                    mon.data[turn,particle_id - mon.part_id_start,2,:], atol=1e-14)
+    xo.assert_allclose(mon.get('py', particle_id=particle_id, turn=turn),
+                    mon.data[turn,particle_id - mon.part_id_start,3,:], atol=1e-14)
+    xo.assert_allclose(mon.get('zeta', particle_id=particle_id, turn=turn),
+                    mon.data[turn,particle_id - mon.part_id_start,4,:], atol=1e-14)
+    xo.assert_allclose(mon.get('delta', particle_id=particle_id, turn=turn),
+                    mon.data[turn,particle_id - mon.part_id_start,5,:], atol=1e-14)
+

xtrack-0.99.0/tests/test_rdt_first_order_perturb.py

@@ -0,0 +1,416 @@
+import xtrack as xt
+import xtrack.feed_down as fd
+from xtrack.rdt import tracking_from_rdt
+import numpy as np
+import xobjects as xo
+import pytest
+import pathlib
+
+test_data_folder = pathlib.Path(
+        __file__).parent.joinpath('../test_data').absolute()
+
+def test_feed_down_rotation_dipole():
+    kn = np.array([1, 0])
+    ks = np.array([0, 0])
+    psi = np.pi / 2 # 90 degrees tilt
+
+    kn_eff, kskew_eff = fd.feed_down(
+        kn=kn[None, :],
+        kskew=ks[None, :],
+        shift_x=0.,
+        shift_y=0.,
+        psi=psi,
+        x0=0.,
+        y0=0.,
+        max_output_order=None,
+    )
+
+    assert kn_eff.shape == (1, 2)
+    assert kskew_eff.shape == (1, 2)
+
+    xo.assert_allclose(kn_eff, np.array([[0, 0]]), atol=1e-15)
+    xo.assert_allclose(kskew_eff, np.array([[-1, 0]]), atol=1e-15)
+
+    m = xt.Multipole(knl=kn, ksl=ks, rot_s_rad=psi)
+    m_fdown = xt.Multipole(knl=kn_eff[0, :], ksl=kskew_eff[0, :], rot_s_rad=0.0)
+
+    p = xt.Particles()
+    p_fdown = p.copy()
+    m.track(p)
+    m_fdown.track(p_fdown)
+
+    xo.assert_allclose(p.px, p_fdown.px, atol=1e-15)
+    xo.assert_allclose(p.py, p_fdown.py, atol=1e-15)
+
+def test_feed_down_rotation_quadrupole():
+    kn = np.array([0, 1])
+    ks = np.array([0, 0])
+    psi = np.pi / 4 # 45 degrees tilt
+
+    kn_eff, kskew_eff = fd.feed_down(
+        kn=kn[None, :],
+        kskew=ks[None, :],
+        shift_x=0.,
+        shift_y=0.,
+        psi=psi,
+        x0=0.,
+        y0=0.,
+        max_output_order=None,
+    )
+
+    m = xt.Multipole(knl=kn, ksl=ks, rot_s_rad=psi)
+    m_fdown = xt.Multipole(knl=kn_eff[0, :], ksl=kskew_eff[0, :], rot_s_rad=0.0)
+
+    p = xt.Particles()
+    p_fdown = p.copy()
+    m.track(p)
+    m_fdown.track(p_fdown)
+
+    xo.assert_allclose(p.px, p_fdown.px, atol=1e-15)
+    xo.assert_allclose(p.py, p_fdown.py, atol=1e-15)
+
+
+def test_feed_down_rotation_higher_order():
+    kn = np.array([0, 1, 2, 10])
+    ks = np.array([0, 3, 2, 20])
+    psi = np.pi / 4 # 45 degrees tilt
+    shift_x = 0.001
+    shift_y = -0.002
+
+    kn_eff, kskew_eff = fd.feed_down(
+        kn=kn[None, :],
+        kskew=ks[None, :],
+        shift_x=shift_x,
+        shift_y=shift_y,
+        psi=psi,
+        x0=0.,
+        y0=0.,
+        max_output_order=None,
+    )
+
+    m = xt.Multipole(knl=kn, ksl=ks, rot_s_rad=psi, shift_x=shift_x, shift_y=shift_y)
+    m_fdown = xt.Multipole(knl=kn_eff[0, :], ksl=kskew_eff[0, :])
+
+    p = xt.Particles(x=4e-3, y=5e-3)
+    p_fdown = p.copy()
+    m.track(p)
+    m_fdown.track(p_fdown)
+
+    xo.assert_allclose(p.px, p_fdown.px, atol=1e-15)
+    xo.assert_allclose(p.py, p_fdown.py, atol=1e-15)
+
+PIMMS_RTD_CONFIGURATIONS = [
+    'sextupole',
+    'skew_sextupole',
+    'skew_quadrupole',
+    'tilted_quadrupole',
+    'shifted_octupole_x',
+    'shifted_octupole_y',
+    'orbit_in_octupole_x',
+    'orbit_in_octupole_y',
+    'orbit_in_skew_octupole_x',
+    'orbit_in_skew_octupole_y',
+]
+
+@pytest.mark.parametrize("configuration", PIMMS_RTD_CONFIGURATIONS)
+def test_rdt_first_order_perturbation_against_tracking(configuration):
+
+    env = xt.load(test_data_folder / 'pimms/PIMM.seq')
+    line = env.pimms
+    line.set_particle_ref('proton', kinetic_energy0=100e6)
+    line.replace_all_repeated_elements()
+
+    line.env.new('skew_quad', xt.Quadrupole, length=0.2)
+    line.env.new('octup', xt.Octupole, length=0.2)
+
+    line.insert('skew_quad', anchor='start', at='xrra@end')
+    line.insert('octup', anchor='start', at='skew_quad@end')
+
+    env['qf1k1'] =  3.15396e-01
+    env['qd1k1'] = -5.24626e-01
+    env['qf2k1'] =  5.22717e-01
+
+    # Build knobs to control the orbit in the straight section with the non-linear
+    # elements
+    line.vars.default_to_zero = True
+    env['qd.3'].knl[0] = 'kx1'
+    env['qf1.3'].knl[0] = 'kx2'
+    env['qf1.4'].knl[0] = 'kx3'
+    env['qd.4'].knl[0] = 'kx4'
+    env['qd.3'].ksl[0] = 'ky1'
+    env['qf1.3'].ksl[0] = 'ky2'
+    env['qf1.4'].ksl[0] = 'ky3'
+    env['qd.4'].ksl[0] = 'ky4'
+    line.vars.default_to_zero = False
+
+    opt = line.match_knob(
+        method='4d',
+        betx=1, bety=1,
+        knob_name='x_bump_mm',
+        vary=xt.VaryList(['kx1', 'kx2', 'kx3', 'kx4'], step=1e-6),
+        targets=[
+            xt.TargetSet(x=1e-3, px=0, at='skew_quad'),
+            xt.TargetSet(x=0.0, px=0, at='qd.5')
+        ],
+        run=False)
+    opt.solve()
+    opt.generate_knob()
+
+    opt = line.match_knob(
+        method='4d',
+        betx=1, bety=1,
+        knob_name='y_bump_mm',
+        vary=xt.VaryList(['ky1', 'ky2', 'ky3', 'ky4'], step=1e-6),
+        targets=[
+            xt.TargetSet(y=-1e-3, py=0, at='skew_quad'),
+            xt.TargetSet(y=0.0, py=0, at='qd.5')
+        ],
+        run=False)
+    opt.solve()
+    opt.generate_knob()
+
+    tw0 = line.twiss4d()
+
+    rtols = {}
+    if configuration == 'sextupole':
+        env['xrra'].k2 = 0.8
+        rdts = ['f3000', 'f1200', 'f1020', 'f0111',
+                # 'f0120' # this one in quite bad
+                ]
+    elif configuration == 'skew_sextupole':
+        env['xrrb'].k2s = 0.8
+        rdts = ['f0030', 'f0012', 'f2010', 'f0210', 'f1110']
+    elif configuration == 'skew_quadrupole':
+        env['skew_quad'].k1s = 0.02
+        rdts = ['f1001', 'f1010', 'f0110']
+    elif configuration == 'tilted_quadrupole':
+        env['qd.4'].rot_s_rad = 0.002
+        rdts = ['f1001', 'f1010', 'f0110']
+    elif configuration == 'shifted_octupole_x':
+        env['octup'].k3 = 80.
+        env['octup'].shift_x = 0.005
+        rdts = ['f3000', 'f1200', 'f1020', 'f0111',
+                # 'f0120' # this one in quite bad
+                ]
+        rtols = {'f3000': 0.5, 'f1200': 0.2, 'f0111': 0.2}
+    elif configuration == 'shifted_octupole_y':
+        env['octup'].k3 = 80.
+        env['octup'].shift_y = -0.005
+        rdts = ['f0030', 'f0012', 'f2010', 'f0210', 'f1110']
+    elif configuration == 'orbit_in_octupole_x':
+        env['octup'].k3 = 80.
+        line['x_bump_mm'] = -5.0 # mm
+        rdts = ['f3000', 'f1200', 'f1020', 'f0111',
+                # 'f0120' # this one in quite bad
+                ]
+        rtols = {'f1200': 0.2, 'f0111': 0.2}
+    elif configuration == 'orbit_in_octupole_y':
+        env['octup'].k3 = 80.
+        line['y_bump_mm'] = 5.0 # mm
+        rdts = ['f0030', 'f0012', 'f2010', 'f0210', 'f1110']
+    elif configuration == 'orbit_in_skew_octupole_x':
+        env['octup'].k3s = 80.
+        line['x_bump_mm'] = 5.0 # mm
+        rdts = ['f0030', 'f0012', 'f2010', 'f0210', 'f1110']
+    elif configuration == 'orbit_in_skew_octupole_y':
+        env['octup'].k3s = 80.
+        line['y_bump_mm'] = -5.0 # mm
+        rdts = ['f3000', 'f1200', 'f1020', 'f0111',
+                # 'f0120' # this one in quite bad
+                ]
+        rtols = {'f1200': 0.2, 'f0111': 0.2}
+    else:
+        raise ValueError(f'Unknown configuration {configuration}')
+
+
+    tw1 = line.twiss4d()
+
+    # Compute strengths
+    strengths = line.get_table(attr=True)
+
+    # Generate 20 particles on the x axis
+    # particles = line.build_particles(x=3e-3, px=5e-4, y=0, py=0, zeta=0, delta=0)
+    particles = line.build_particles(x=3e-3, px=5e-4, y=2e-3, py=3e-5, zeta=0, delta=0)
+
+    # Inspect the particles
+    particles.get_table()
+
+    # Track 1000 turns logging turn-by-turn data
+    num_turns = 20_000
+    line.track(particles, num_turns=num_turns, turn_by_turn_monitor=True,
+            with_progress=100)
+    rec = line.record_last_track
+
+    nc  = tw0.get_normalized_coordinates(rec)
+
+    # Compute RDTs via first-order perturbation theory
+    rdt_vals = xt.rdt_first_order_perturbation(rdt=rdts, twiss=tw0,
+                                            orbit=tw1,
+                                            strengths=strengths)
+
+    i_part_analyze = 0
+    x_norm = nc.x_norm[i_part_analyze, :]
+    px_norm = nc.px_norm[i_part_analyze, :]
+    y_norm = nc.y_norm[i_part_analyze, :]
+    py_norm = nc.py_norm[i_part_analyze, :]
+
+    zx_norm = x_norm - 1j * px_norm
+    zy_norm = y_norm - 1j * py_norm
+    # tracking from RDTs
+    Ix = 0.5 * (zx_norm[0].real**2 + zx_norm[0].imag**2)
+    Iy = 0.5 * (zy_norm[0].real**2 + zy_norm[0].imag**2)
+    psi_x0 = np.angle(zx_norm[0].real + 1j * zx_norm[0].imag)
+    psi_y0 = np.angle(zy_norm[0].real + 1j * zy_norm[0].imag)
+
+    def initial_conditions(Ix, Iy, psi_x0, psi_y0):
+
+        hx_minus, hy_minus = tracking_from_rdt(
+            rdts={rr: (rdt_vals[rr][0]) for rr in rdts},
+            Ix=Ix,
+            Iy=Iy,
+            Qx=tw0.qx,
+            Qy=tw0.qy,
+            psi_x0=psi_x0,
+            psi_y0=psi_y0,
+            num_turns=1
+        )
+
+        return np.array([hx_minus[0].real,
+                        hx_minus[0].imag,
+                        hy_minus[0].real,
+                        hy_minus[0].imag])
+
+    opt = xt.match.opt_from_callable(
+            lambda xx: initial_conditions(xx[0], xx[1], xx[2], xx[3]),
+            x0=np.array([Ix, Iy, psi_x0, psi_y0]),
+            steps=np.array([Ix*1e-4, Ix*1e-4, 1e-4, 1e-4]),
+            tar=np.array([zx_norm[0].real, zx_norm[0].imag,
+                            zy_norm[0].real, zy_norm[0].imag]),
+            tols=[1e-10, 1e-10, 1e-10, 1e-10],
+        )
+    opt.step()
+    res = opt.get_knob_values()
+    Ix = res[0]
+    Iy = res[1]
+    psi_x0 = res[2]
+    psi_y0 = res[3]
+
+
+    hx_minus, hy_minus = tracking_from_rdt(
+        rdts={rr: (rdt_vals[rr][0]) for rr in rdts},
+        Ix=Ix,
+        Iy=Iy,
+        Qx=tw0.qx,
+        Qy=tw0.qy,
+        psi_x0=psi_x0,
+        psi_y0=psi_y0,
+        num_turns=num_turns
+    )
+
+    zx_spectrum = np.fft.fft(zx_norm)
+    zy_spectrum = np.fft.fft(zy_norm)
+
+    hx_spectrum = np.fft.fft(hx_minus)
+    hy_spectrum = np.fft.fft(hy_minus)
+
+    freqs = np.fft.fftfreq(num_turns)
+    freqs[freqs < 0] += 1.0
+
+    import nafflib
+    f_hx, s_hx = nafflib.get_tunes_all(hx_minus, N=100)
+    f_hy, s_hy = nafflib.get_tunes_all(hy_minus, N=100)
+    f_x, s_x = nafflib.get_tunes_all(zx_norm, N=100)
+    f_y, s_y = nafflib.get_tunes_all(zy_norm, N=100)
+
+    f_x[f_x < 0] += 1.0
+    f_hx[f_hx < 0] += 1.0
+    f_y[f_y < 0] += 1.0
+    f_hy[f_hy < 0] += 1.0
+
+    f_h = {'x': f_hx, 'y': f_hy}
+    s_h = {'x': s_hx, 'y': s_hy}
+    f_z = {'x': f_x, 'y': f_y}
+    s_z = {'x': s_x, 'y': s_y}
+
+    dq_search = 0.001
+    print()
+    for rr in rdts:
+        print(f'Spectral lines excited by {rr}:')
+        for pp in ['x', 'y']:
+            if rdt_vals[rr + f'_ampl_{pp}_expr'] == '0':
+                print(f'  Expected {pp} freq: not excited')
+            elif rdt_vals[rr + f"_freq_{pp}"] == 0:
+                print(f'  Expected {pp} freq: zero frequency')
+            else:
+                print(f'  Expected {pp} freq: {rdt_vals[rr + f"_freq_{pp}_expr"]} = {rdt_vals[rr + f"_freq_{pp}"]:.6f}')
+                mask_search_z = (np.abs(f_z[pp] - rdt_vals[rr + f'_freq_{pp}']) < dq_search)
+                i_max_z = np.argmax(np.abs(s_z[pp][mask_search_z]))
+                f_z_max = f_z[pp][mask_search_z][i_max_z]
+                s_z_max = s_z[pp][mask_search_z][i_max_z]
+                mask_search_h = (np.abs(f_h[pp] - rdt_vals[rr + f'_freq_{pp}']) < dq_search)
+                i_max_h = np.argmax(np.abs(s_h[pp][mask_search_h]))
+                f_h_max = f_h[pp][mask_search_h][i_max_h]
+                s_h_max = s_h[pp][mask_search_h][i_max_h]
+                print(f'    From tracking: freq={f_z_max:.6f}, amp={np.abs(s_z_max):.6e}, phase={np.angle(s_z_max, deg=True):.2f} deg')
+                print(f'    From RDTs:     freq={f_h_max:.6f}, amp={np.abs(s_h_max):.6e}, phase={np.angle(s_h_max, deg=True):.2f} deg')
+                xo.assert_allclose(f_z_max, f_h_max, atol=5e-4)
+                xo.assert_allclose(s_z_max, s_h_max, rtol=rtols.get(rr, 0.1))
+        print('')
+
+def test_rdt_lhc_vs_madng():
+
+    env = xt.load(test_data_folder / 'hllhc15_thick/hllhc15_collider_thick.json')
+
+    line = env.lhcb1
+    tw0 = line.twiss(method='4d')
+
+    # Coupling RDTs
+    line['kqs.a45b1'] = 1e-4 # Introduce a skew quadrupole
+    rdts = xt.rdt_first_order_perturbation(
+        rdt=['f1001', 'f1010', 'f0110'],
+        twiss=tw0,
+        strengths=line.get_table(attr=True)
+    )
+    tw_edw_teng = line.twiss4d(coupling_edw_teng=True)
+    line['kqs.a45b1'] = 0.0 # Remove skew quadrupole
+
+    # outliers came from the fact that at the sources the RDT jumps at entry
+    # instead of exit of the element or viceversa
+    xo.assert_allclose(rdts.f1001, tw_edw_teng.f1001, atol=0.07*np.max(np.abs(tw_edw_teng.f1001)), max_outliers=5)
+    xo.assert_allclose(rdts.f1010, tw_edw_teng.f1010, atol=0.07*np.max(np.abs(tw_edw_teng.f1010)), max_outliers=5)
+    xo.assert_allclose(rdts.f0110, tw_edw_teng.f0110, atol=0.07*np.max(np.abs(tw_edw_teng.f0110)), max_outliers=5)
+    xo.assert_allclose(rdts.f1001, np.conj(rdts.f0110), atol=1e-12)
+
+    # RDTs from chromaticity sextupoles
+    rdt_names = ['f3000', 'f1200', 'f1020', 'f0120', 'f0111']
+    rdts = xt.rdt_first_order_perturbation(
+        rdt=rdt_names,
+        twiss=tw0,
+        strengths=line.get_table(attr=True)
+    )
+    tw_ng = line.madng_twiss(rdts=rdt_names)
+    for nn in rdt_names:
+        # outliers came from the fact that at the sources the RDT jumps at entry
+        # instead of exit of the element or viceversa
+        xo.assert_allclose(rdts[nn], tw_ng[nn],
+                        atol=0.1*np.max(np.abs(tw_ng[nn])),
+                        max_outliers=0.01*len(tw_ng))
+
+    line['on_x1'] = 0
+    line['on_x5'] = 250
+    line['on_disp'] = 2
+    tw1 = line.twiss(method='4d', coupling_edw_teng=True)
+    rdt_names = ['f1001', 'f1010', 'f0110']
+    rdts = xt.rdt_first_order_perturbation(
+        rdt=rdt_names,
+        twiss=tw0,
+        orbit=tw1,
+        strengths=line.get_table(attr=True)
+    )
+    for nn in rdt_names:
+        # outliers came from the fact that at the sources the RDT jumps at entry
+        # instead of exit of the element or viceversa
+        xo.assert_allclose(rdts[nn], tw1[nn],
+                        atol=0.06*np.max(np.abs(tw1[nn])),
+                        max_outliers=0.005*len(tw1))
+

{xtrack-0.98.13 → xtrack-0.99.0}/xtrack/__init__.py

@@ -22,6 +22,7 @@ from .match import (Vary, Target, TargetList, VaryList, TargetInequality, Action
 from .targets import (TargetLuminosity, TargetSeparationOrthogonalToCrossing,
                       TargetSeparation)
 from .twiss import TwissInit, TwissTable
+from .rdt import rdt_first_order_perturbation
 from .loss_location_refinement import LossLocationRefinement
 from .internal_record import (RecordIdentifier, RecordIndex, new_io_buffer,
                              start_internal_logging, stop_internal_logging)

xtrack-0.99.0/xtrack/_version.py

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

{xtrack-0.98.13 → xtrack-0.99.0}/xtrack/environment.py

@@ -37,7 +37,7 @@ DEFAULT_REF_STRENGTH_NAME = {
 class Environment:
 
     def __init__(self, element_dict=None, particle_ref=None, lines=None,
-                 _var_management_dct=None):
+                 _var_management_dct=None, particles=None):
 
         '''
         Create an environment.
@@ -91,10 +91,7 @@ class Environment:
         self._particles = {}
         self.particle_ref = particle_ref
 
-        self._var_management = _make_var_management(
-            element_dict=self._element_dict,
-            particles=self._particles,
-            dct=_var_management_dct)
+        self._init_var_management(dct=_var_management_dct)
         self._line_vars = EnvVars(self)
 
         self.lines = EnvLines(self)
@@ -136,8 +133,17 @@ class Environment:
                     line_name=nn, rename_elements=rename_elements)
                 self.lines[nn]._renamed_elements = rename_elements
 
+        if particles is not None:
+            self._particles.update(particles)
+
         self.metadata = {}
 
+    def _init_var_management(self, dct=None):
+        self._var_management = _make_var_management(
+            element_dict=self._element_dict,
+            particles=self._particles,
+            dct=dct)
+
     def __repr__(self):
         line_names = list(self.lines.keys())
         n_lines = len(line_names)
@@ -792,6 +798,10 @@ class Environment:
 
         elements = _deserialize_elements(dct=dct, classes=classes,
                                          _buffer=_buffer, _context=_context)
+        particles = {}
+        if 'particles' in dct:
+            for nn, ppd in dct['particles'].items():
+               particles[nn] = xt.Particles.from_dict(ppd)
 
         particle_ref = None
         if 'particle_ref' in dct.keys():
@@ -806,7 +816,7 @@ class Environment:
             _var_management_dct = None
 
         out = cls(element_dict=elements, particle_ref=particle_ref,
-                _var_management_dct=_var_management_dct)
+                _var_management_dct=_var_management_dct, particles=particles)
 
         dct_lines = dct.copy()
         dct_lines.pop('elements', None)
@@ -833,10 +843,6 @@ class Environment:
         if "metadata" in dct:
             out.metadata = dct["metadata"]
 
-        if 'particles' in dct:
-            for nn, ppd in dct['particles'].items():
-               out._particles[nn] = xt.Particles.from_dict(ppd)
-
         return out
 
     @classmethod
@@ -1113,6 +1119,13 @@ class Environment:
         try:
             eva_obj = self._xdeps_eval_obj
         except AttributeError:
+            eva_obj = None
+
+        # in case the var manager has been replaced, invalidate the cached eval obj
+        if eva_obj is not None and eva_obj.variables is not self._xdeps_vref:
+            eva_obj = None
+
+        if eva_obj is None:
             eva_obj = xd.madxutils.MadxEval(variables=self._xdeps_vref,
                                             functions=self._xdeps_fref,
                                             elements=self._element_dict,
@@ -1824,7 +1837,7 @@ def _reverse_element(env, name):
     _reverse_field('rot_s_rad')
 
     if ee.__class__.__name__ == 'CrabCavity':
-        ee_ref.voltage = -ee_ref.voltage._expr or ee_ref.voltage._value
+        ee_ref.crab_voltage = -(ee_ref.crab_voltage._expr or ee_ref.crab_voltage._value)
 
     if hasattr(ee, 'lag'):
         ee_ref.lag = 180 - (ee_ref.lag._expr or ee_ref.lag._value)
@@ -2160,6 +2173,7 @@ class EnvVars:
         filename : str or list of str
             Path to the MAD-X file(s) to load.
         '''
+        old_default_to_zero = self.default_to_zero
         loader = xt.mad_parser.MadxLoader(env=self.env)
         if filename is not None:
             assert string is None, 'Cannot specify both filename and string'
@@ -2167,6 +2181,7 @@ class EnvVars:
         elif string is not None:
             assert filename is None, 'Cannot specify both filename and string'
             loader.load_string(string)
+        self.default_to_zero = old_default_to_zero # restore (in case changed by loader)
 
     def load_madx_optics_file(self, filename=None, string=None):
         self.set_from_madx_file(filename, string)