xtrack 0.54.1__tar.gz → 0.55.0__tar.gz

{xtrack-0.54.1 → xtrack-0.55.0}/PKG-INFO

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

xtrack-0.55.0/tests/test_cavity_absolute_time.py

@@ -0,0 +1,47 @@
+import numpy as np
+import xtrack as xt
+from xobjects.test_helpers import for_all_test_contexts
+
+import pathlib
+
+test_data_folder = pathlib.Path(
+    __file__).parent.joinpath('../test_data').absolute()
+
+@for_all_test_contexts
+def test_cavity_absolute_time(test_context):
+    line = xt.Line.from_json(test_data_folder /
+                            'hllhc15_thick/lhc_thick_with_knobs.json')
+    line.build_tracker(_context=test_context)
+    line.vv['vrf400'] = 16
+
+    for vv in line.vars.get_table().rows[
+        'on_x.*|on_sep.*|on_crab.*|on_alice|on_lhcb|corr_.*'].name:
+        line.vars[vv] = 0
+
+    tw = line.twiss()
+
+    df_hz = -50
+    h_rf = 35640
+    f_rev = 1/tw.T_rev0
+    df_rev = df_hz / h_rf
+    eta = tw.slip_factor
+
+    f_rf0 = 1/tw.T_rev0 * h_rf
+
+    f_rf = f_rf0 + df_hz
+    line.vars['f_rf'] = f_rf
+    tt = line.get_table()
+    for nn in tt.rows[tt.element_type=='Cavity'].name:
+        line.element_refs[nn].absolute_time = 1 # Need property
+        line.element_refs[nn].frequency = line.vars['f_rf']
+
+    tw1 = line.twiss(search_for_t_rev=True)
+
+    f_rev_expected = f_rf / h_rf
+
+    assert np.isclose(f_rev_expected, 1/tw1.T_rev, atol=1e-5, rtol=0)
+    assert np.allclose(tw1.delta, tw1.delta[0], atol=1e-5, rtol=0) # Check that it is flat
+    delta_expected = -df_rev / f_rev / eta
+    assert np.allclose(tw1.delta, delta_expected, atol=2e-6, rtol=0)
+    tw_off_mom = line.twiss(method='4d', delta0=tw1.delta[0])
+    assert np.allclose(tw1.x, tw_off_mom.x, atol=1e-5, rtol=0)

xtrack-0.55.0/tests/test_coasting.py

@@ -0,0 +1,155 @@
+import numpy as np
+import xtrack as xt
+import pathlib
+
+from scipy.constants import c as clight
+
+test_data_folder = pathlib.Path(
+    __file__).parent.joinpath('../test_data').absolute()
+
+
+def test_coasting():
+    delta0 = 1e-2
+
+    line = xt.Line.from_json(test_data_folder /
+                             'psb_injection/line_and_particle.json')
+
+    # RF off!
+    tt = line.get_table()
+    ttcav = tt.rows[tt.element_type == 'Cavity']
+    for nn in ttcav.name:
+        line.element_refs[nn].voltage=0
+
+    line.configure_bend_model(core='bend-kick-bend', edge='full')
+    line.twiss_default['method'] = '4d'
+
+    tw = line.twiss()
+    twom = line.twiss(delta0=delta0)
+    line.discard_tracker()
+
+    # Install dummy collective elements
+    s_sync = np.linspace(0, tw.circumference, 10)
+    line.cut_at_s(s_sync)
+    for ii, ss in enumerate(s_sync):
+        nn = f'sync_here_{ii}'
+        line.insert_element(element=xt.Marker(), name=nn, at_s=ss)
+        line[nn].iscollective = True
+
+    import xtrack.synctime as st
+    st.install_sync_time_at_collective_elements(line)
+
+    import xobjects as xo
+    line.build_tracker(_context=xo.ContextCpu(omp_num_threads='auto'))
+
+    beta1 = tw.beta0 / 0.9
+
+    circumference = tw.circumference
+    zeta_min0 = -circumference/2*tw.beta0/beta1
+    zeta_max0 = circumference/2*tw.beta0/beta1
+
+    num_particles = 50000
+    p = line.build_particles(
+        delta=delta0 + 0 * np.random.uniform(-1, 1, num_particles),
+        x_norm=0, y_norm=0
+    )
+
+    # Need to take beta of actual particles to convert the distribution along the
+    # circumference to a distribution in time
+    p.zeta = (np.random.uniform(0, circumference, num_particles) / p.rvv
+            + (zeta_max0 - circumference) / p.rvv)
+
+    st.prepare_particles_for_sync_time(p, line)
+
+    p.y[(p.zeta > 1) & (p.zeta < 2)] = 1e-3  # kick
+    p.weight[(p.zeta > 5) & (p.zeta < 10)] *= 1.3
+
+    p0 = p.copy()
+
+    def particles(_, p):
+        return p.copy()
+
+    def intensity(line, particles):
+        return (np.sum(particles.weight[particles.state > 0])
+                    / ((zeta_max0 - zeta_min0)/tw.beta0/clight))
+
+    def z_range(line, particles):
+        mask_alive = particles.state > 0
+        return particles.zeta[mask_alive].min(), particles.zeta[mask_alive].max()
+
+    def long_density(line, particles):
+        mask_alive = particles.state > 0
+        if not(np.any(particles.at_turn[mask_alive] == 0)): # don't check at the first turn
+            assert np.all(particles.zeta[mask_alive] > zeta_min0)
+            assert np.all(particles.zeta[mask_alive] < zeta_max0)
+        return np.histogram(particles.zeta[mask_alive], bins=200,
+                            range=(zeta_min0, zeta_max0),
+                            weights=particles.weight[mask_alive])
+
+    def y_mean_hist(line, particles):
+
+        mask_alive = particles.state > 0
+        if not(np.any(particles.at_turn[mask_alive] == 0)): # don't check at the first turn
+            assert np.all(particles.zeta[mask_alive] > zeta_min0)
+            assert np.all(particles.zeta[mask_alive] < zeta_max0)
+        return np.histogram(particles.zeta[mask_alive], bins=200,
+                            range=(zeta_min0, zeta_max0), weights=particles.y[mask_alive])
+
+
+    line.enable_time_dependent_vars = True
+    num_turns=200
+    line.track(p, num_turns=num_turns, log=xt.Log(intensity=intensity,
+                                            long_density=long_density,
+                                            y_mean_hist=y_mean_hist,
+                                            z_range=z_range,
+                                            particles=particles
+                                            ), with_progress=10)
+
+    inten = line.log_last_track['intensity']
+
+    f_rev_ave = 1 / tw.T_rev0 * (1 - tw.slip_factor * p.delta.mean())
+    t_rev_ave = 1 / f_rev_ave
+
+    inten_exp =  np.sum(p0.weight) / t_rev_ave
+
+    z_axis = line.log_last_track['long_density'][0][1]
+    hist_mat = np.array([rr[0] for rr in line.log_last_track['long_density']])
+    hist_y = np.array([rr[0] for rr in line.log_last_track['y_mean_hist']])
+
+    dz = z_axis[1] - z_axis[0]
+    y_vs_t = np.fliplr(hist_y).flatten() # need to flip because of the minus in z = -beta0 c t
+    intensity_vs_t = np.fliplr(hist_mat).flatten()
+    z_unwrapped = np.arange(0, len(y_vs_t)) * dz
+    t_unwrapped = z_unwrapped / (tw.beta0 * clight)
+
+    z_range_size = z_axis[-1] - z_axis[0]
+    t_range_size = z_range_size / (tw.beta0 * clight)
+
+    import nafflib
+    f_harmons = nafflib.get_tunes(intensity_vs_t, N=50)[0] / (t_unwrapped[1] - t_unwrapped[0])
+    f_nominal = 1 / tw.T_rev0
+    dt_expected = -(twom.zeta[-1] - twom.zeta[0]) / tw.beta0 / clight
+    f_expected = 1 / (tw.T_rev0 + dt_expected)
+
+    f_measured = f_harmons[np.argmin(np.abs(f_harmons - f_nominal))]
+
+    print('f_nominal:  ', f_nominal, ' Hz')
+    print('f_expected: ', f_expected, ' Hz')
+    print('f_measured: ', f_measured, ' Hz')
+    print('Error:      ', f_measured - f_expected, 'Hz')
+
+    assert np.isclose(f_expected, f_measured, rtol=0, atol=2) # 2 Hz tolerance (to account for random fluctuations)
+    assert np.isclose(np.mean(inten), inten_exp, rtol=1e-2, atol=0)
+    assert np.allclose(p.at_turn, num_turns*0.9, rtol=3e-2, atol=0) #beta1 defaults to 0.1
+
+    tt = line.get_table()
+    tt_synch = tt.rows[tt.element_type=='SyncTime']
+    assert len(tt_synch) == 12
+    assert tt_synch.name[0] == 'synctime_start'
+    assert tt_synch.name[-1] == 'synctime_end'
+    assert np.all(tt_synch.name[5] == 'synctime_4')
+    assert line['synctime_start'].at_start
+    assert not line['synctime_end'].at_start
+    assert not line['synctime_4'].at_start
+    assert line['synctime_end'].at_end
+    assert not line['synctime_start'].at_end
+    assert not line['synctime_4'].at_end

{xtrack-0.54.1 → xtrack-0.55.0}/tests/test_elements_thick.py

@@ -1036,7 +1036,9 @@ def test_solenoid_against_madx(test_context, ks, ksi, length):
 @for_all_test_contexts
 def test_solenoid_thick_drift_like(test_context):
     solenoid = xt.Solenoid(ks=1.001e-9, length=1, _context=test_context)
-    drift = xt.Drift(length=1, _context=test_context)
+    l_drift = xt.Line(elements=[xt.Drift(length=1)])
+    l_drift.config.XTRACK_USE_EXACT_DRIFTS = True
+    l_drift.build_tracker(_context=test_context)
 
     p0 = xp.Particles(
         x=0.1, px=0.2, y=0.3, py=0.4, zeta=0.5, delta=0.6,
@@ -1047,7 +1049,7 @@ def test_solenoid_thick_drift_like(test_context):
     solenoid.track(p_sol)
 
     p_drift = p0.copy()
-    drift.track(p_drift)
+    l_drift.track(p_drift)
 
     p_sol.move(_context=xo.context_default)
     p_drift.move(_context=xo.context_default)

xtrack-0.55.0/tests/test_fcc_ee_solenoid_correction.py

@@ -0,0 +1,404 @@
+import numpy as np
+import xtrack as xt
+from scipy.constants import c as clight
+from scipy.constants import e as qe
+import pathlib
+
+from cpymad.madx import Madx
+
+
+test_data_folder = pathlib.Path(
+    __file__).parent.joinpath('../test_data').absolute()
+
+
+def test_fcc_ee_solenoid_correction():
+    fname = 'fccee_t'; pc_gev = 182.5
+
+    mad = Madx()
+    mad.call(str(test_data_folder) + '/fcc_ee/' + fname + '.seq')
+    mad.beam(particle='positron', pc=pc_gev)
+    mad.use('fccee_p_ring')
+
+    line = xt.Line.from_madx_sequence(mad.sequence.fccee_p_ring, allow_thick=True,
+                                    deferred_expressions=True)
+    line.particle_ref = xt.Particles(mass0=xt.ELECTRON_MASS_EV,
+                                    gamma0=mad.sequence.fccee_p_ring.beam.gamma)
+    line.cycle('ip.4', inplace=True)
+    line.append_element(element=xt.Marker(), name='ip.4.l')
+
+    tt = line.get_table()
+    bz_data_file = test_data_folder / 'fcc_ee/Bz_closed_before_quads.dat'
+
+    line.vars['voltca1_ref'] = line.vv['voltca1']
+    if 'voltca2' in line.vars.keys():
+        line.vars['voltca2_ref'] = line.vv['voltca2']
+    else:
+        line.vars['voltca2_ref'] = 0
+
+    line.vars['voltca1'] = 0
+    line.vars['voltca2'] = 0
+
+    import pandas as pd
+    bz_df = pd.read_csv(bz_data_file, sep=r'\s+', skiprows=1, names=['z', 'Bz'])
+
+    l_solenoid = 4.4
+    ds_sol_start = -l_solenoid / 2 * np.cos(15e-3)
+    ds_sol_end = +l_solenoid / 2 * np.cos(15e-3)
+    ip_sol = 'ip.1'
+
+    theta_tilt = 15e-3 # rad
+
+    s_sol_slices = np.linspace(ds_sol_start, ds_sol_end, 1001)
+    bz_sol_slices = np.interp(s_sol_slices, bz_df.z, bz_df.Bz)
+    bz_sol_slices[0] = 0
+    bz_sol_slices[-1] = 0
+
+    P0_J = line.particle_ref.p0c[0] * qe / clight
+    brho = P0_J / qe / line.particle_ref.q0
+    ks = 0.5 * (bz_sol_slices[:-1] + bz_sol_slices[1:]) / brho
+    l_sol_slices = np.diff(s_sol_slices)
+    s_sol_slices_entry = s_sol_slices[:-1]
+
+    sol_slices = []
+    for ii in range(len(s_sol_slices_entry)):
+        sol_slices.append(xt.Solenoid(length=l_sol_slices[ii], ks=0)) # Off for now
+
+    s_ip = tt['s', ip_sol]
+
+    line.discard_tracker()
+    line.insert_element(name='sol_start_'+ip_sol, element=xt.Marker(),
+                        at_s=s_ip + ds_sol_start)
+    line.insert_element(name='sol_end_'+ip_sol, element=xt.Marker(),
+                        at_s=s_ip + ds_sol_end)
+
+    sol_start_tilt = xt.YRotation(angle=-theta_tilt * 180 / np.pi)
+    sol_end_tilt = xt.YRotation(angle=+theta_tilt * 180 / np.pi)
+    sol_start_shift = xt.XYShift(dx=l_solenoid/2 * np.sin(theta_tilt))
+    sol_end_shift = xt.XYShift(dx=l_solenoid/2 * np.sin(theta_tilt))
+
+    line.element_dict['sol_start_tilt_'+ip_sol] = sol_start_tilt
+    line.element_dict['sol_end_tilt_'+ip_sol] = sol_end_tilt
+    line.element_dict['sol_start_shift_'+ip_sol] = sol_start_shift
+    line.element_dict['sol_end_shift_'+ip_sol] = sol_end_shift
+
+    line.element_dict['sol_entry_'+ip_sol] = xt.Solenoid(length=0, ks=0)
+    line.element_dict['sol_exit_'+ip_sol] = xt.Solenoid(length=0, ks=0)
+
+
+    sol_slice_names = []
+    sol_slice_names.append('sol_entry_'+ip_sol)
+    for ii in range(len(s_sol_slices_entry)):
+        nn = f'sol_slice_{ii}_{ip_sol}'
+        line.element_dict[nn] = sol_slices[ii]
+        sol_slice_names.append(nn)
+    sol_slice_names.append('sol_exit_'+ip_sol)
+
+    tt = line.get_table()
+    names_upstream = list(tt.rows[:'sol_start_'+ip_sol].name)
+    names_downstream = list(tt.rows['sol_end_'+ip_sol:].name[:-1]) # -1 to exclude '_end_point' added by the table
+
+    element_names = (names_upstream
+                    + ['sol_start_tilt_'+ip_sol, 'sol_start_shift_'+ip_sol]
+                    + sol_slice_names
+                    + ['sol_end_shift_'+ip_sol, 'sol_end_tilt_'+ip_sol]
+                    + names_downstream)
+
+    line.element_names = element_names
+
+    # re-insert the ip
+    line.element_dict.pop(ip_sol)
+    tt = line.get_table()
+    line.insert_element(name=ip_sol, element=xt.Marker(),
+            at_s = 0.5 * (tt['s', 'sol_start_'+ip_sol] + tt['s', 'sol_end_'+ip_sol]))
+
+    line.vars['on_corr_ip.1'] = 0
+
+    line.build_tracker()
+
+    # Set strength
+    line.vars['on_sol_'+ip_sol] = 0
+    for ii in range(len(s_sol_slices_entry)):
+        nn = f'sol_slice_{ii}_{ip_sol}'
+        line.element_refs[nn].ks = ks[ii] * line.vars['on_sol_'+ip_sol]
+
+    tt = line.get_table()
+
+    tt.rows['sol_start_ip.1':'sol_end_ip.1'].show()
+
+    line.vars['on_corr_ip.1'] = 1
+    line.vars['ks0.r1'] = 0
+    line.vars['ks1.r1'] = 0
+    line.vars['ks2.r1'] = 0
+    line.vars['ks3.r1'] = 0
+    line.vars['ks4.r1'] = 0
+    line.vars['ks0.l1'] = 0
+    line.vars['ks1.l1'] = 0
+    line.vars['ks2.l1'] = 0
+    line.vars['ks3.l1'] = 0
+    line.vars['ks4.l1'] = 0
+
+    line.element_refs['qc1r1.1'].k1s = line.vars['on_corr_ip.1'] * line.vars['ks0.r1']
+    line.element_refs['qc2r1.1'].k1s = line.vars['on_corr_ip.1'] * line.vars['ks1.r1']
+    line.element_refs['qc2r2.1'].k1s = line.vars['on_corr_ip.1'] * line.vars['ks2.r1']
+    line.element_refs['qc1r2.1'].k1s = line.vars['on_corr_ip.1'] * line.vars['ks3.r1']
+    line.element_refs['qc1l1.4'].k1s = line.vars['on_corr_ip.1'] * line.vars['ks0.l1']
+    line.element_refs['qc2l1.4'].k1s = line.vars['on_corr_ip.1'] * line.vars['ks1.l1']
+    line.element_refs['qc2l2.4'].k1s = line.vars['on_corr_ip.1'] * line.vars['ks2.l1']
+    line.element_refs['qc1l2.4'].k1s = line.vars['on_corr_ip.1'] * line.vars['ks3.l1']
+
+    line.vars['corr_k0.r1'] = 0
+    line.vars['corr_k1.r1'] = 0
+    line.vars['corr_k2.r1'] = 0
+    line.vars['corr_k3.r1'] = 0
+    line.vars['corr_k4.r1'] = 0
+    line.vars['corr_k0.l1'] = 0
+    line.vars['corr_k1.l1'] = 0
+    line.vars['corr_k2.l1'] = 0
+    line.vars['corr_k3.l1'] = 0
+    line.vars['corr_k4.l1'] = 0
+
+    line.element_refs['qc1r1.1'].k1 += line.vars['on_corr_ip.1'] * line.vars['corr_k0.r1']
+    line.element_refs['qc2r1.1'].k1 += line.vars['on_corr_ip.1'] * line.vars['corr_k1.r1']
+    line.element_refs['qc2r2.1'].k1 += line.vars['on_corr_ip.1'] * line.vars['corr_k2.r1']
+    line.element_refs['qc1r2.1'].k1 += line.vars['on_corr_ip.1'] * line.vars['corr_k3.r1']
+    line.element_refs['qc1l1.4'].k1 += line.vars['on_corr_ip.1'] * line.vars['corr_k0.l1']
+    line.element_refs['qc2l1.4'].k1 += line.vars['on_corr_ip.1'] * line.vars['corr_k1.l1']
+    line.element_refs['qc2l2.4'].k1 += line.vars['on_corr_ip.1'] * line.vars['corr_k2.l1']
+    line.element_refs['qc1l2.4'].k1 += line.vars['on_corr_ip.1'] * line.vars['corr_k3.l1']
+
+
+    Strategy = xt.Strategy
+    Teapot = xt.Teapot
+    slicing_strategies = [
+        Strategy(slicing=None),  # Default catch-all as in MAD-X
+        Strategy(slicing=Teapot(3), element_type=xt.Bend),
+        Strategy(slicing=Teapot(3), element_type=xt.CombinedFunctionMagnet),
+        # Strategy(slicing=Teapot(50), element_type=xt.Quadrupole), # Starting point
+        Strategy(slicing=Teapot(5), name=r'^qf.*'),
+        Strategy(slicing=Teapot(5), name=r'^qd.*'),
+        Strategy(slicing=Teapot(5), name=r'^qfg.*'),
+        Strategy(slicing=Teapot(5), name=r'^qdg.*'),
+        Strategy(slicing=Teapot(5), name=r'^ql.*'),
+        Strategy(slicing=Teapot(5), name=r'^qs.*'),
+        Strategy(slicing=Teapot(10), name=r'^qb.*'),
+        Strategy(slicing=Teapot(10), name=r'^qg.*'),
+        Strategy(slicing=Teapot(10), name=r'^qh.*'),
+        Strategy(slicing=Teapot(10), name=r'^qi.*'),
+        Strategy(slicing=Teapot(10), name=r'^qr.*'),
+        Strategy(slicing=Teapot(10), name=r'^qu.*'),
+        Strategy(slicing=Teapot(10), name=r'^qy.*'),
+        Strategy(slicing=Teapot(50), name=r'^qa.*'),
+        Strategy(slicing=Teapot(50), name=r'^qc.*'),
+        Strategy(slicing=Teapot(20), name=r'^sy\..*'),
+        Strategy(slicing=Teapot(30), name=r'^mwi\..*'),
+    ]
+    line.discard_tracker()
+    line.slice_thick_elements(slicing_strategies=slicing_strategies)
+
+    # Add dipole correctors
+    line.insert_element(name='mcb1.r1', element=xt.Multipole(knl=[0]),
+                        at='qc1r1.1_exit')
+    line.insert_element(name='mcb2.r1', element=xt.Multipole(knl=[0]),
+                        at='qc1r2.1_exit')
+    line.insert_element(name='mcb1.l1', element=xt.Multipole(knl=[0]),
+                        at='qc1l1.4_entry')
+    line.insert_element(name='mcb2.l1', element=xt.Multipole(knl=[0]),
+                        at='qc1l2.4_entry')
+
+    line.vars['acb1h.r1'] = 0
+    line.vars['acb1v.r1'] = 0
+    line.vars['acb2h.r1'] = 0
+    line.vars['acb2v.r1'] = 0
+    line.vars['acb1h.l1'] = 0
+    line.vars['acb1v.l1'] = 0
+    line.vars['acb2h.l1'] = 0
+    line.vars['acb2v.l1'] = 0
+
+    line.element_refs['mcb1.r1'].knl[0] = line.vars['on_corr_ip.1']*line.vars['acb1h.r1']
+    line.element_refs['mcb2.r1'].knl[0] = line.vars['on_corr_ip.1']*line.vars['acb2h.r1']
+    line.element_refs['mcb1.r1'].ksl[0] = line.vars['on_corr_ip.1']*line.vars['acb1v.r1']
+    line.element_refs['mcb2.r1'].ksl[0] = line.vars['on_corr_ip.1']*line.vars['acb2v.r1']
+    line.element_refs['mcb1.l1'].knl[0] = line.vars['on_corr_ip.1']*line.vars['acb1h.l1']
+    line.element_refs['mcb2.l1'].knl[0] = line.vars['on_corr_ip.1']*line.vars['acb2h.l1']
+    line.element_refs['mcb1.l1'].ksl[0] = line.vars['on_corr_ip.1']*line.vars['acb1v.l1']
+    line.element_refs['mcb2.l1'].ksl[0] = line.vars['on_corr_ip.1']*line.vars['acb2v.l1']
+
+    tw_thick_no_rad = line.twiss(method='4d')
+
+    assert line.element_names[-1] == 'ip.4.l'
+    assert line.element_names[0] == 'ip.4'
+
+    opt = line.match(
+        only_markers=True,
+        method='4d',
+        start='ip.4', end='ip.4.l',
+        init=tw_thick_no_rad,
+        vary=xt.VaryList(['k1qf4', 'k1qf2', 'k1qd3', 'k1qd1',], step=1e-8,
+        ),
+        targets=[
+            xt.TargetSet(at=xt.END, mux=tw_thick_no_rad.qx, muy=tw_thick_no_rad.qy, tol=1e-5),
+        ]
+    )
+    opt.solve()
+
+    line.vars['on_sol_ip.1'] = 0
+    tw_sol_off = line.twiss(method='4d')
+    line.vars['on_sol_ip.1'] = 1
+    tw_sol_on = line.twiss(method='4d')
+
+    opt_l = line.match(
+        solve=False,
+        method='4d', n_steps_max=30,
+        start='pqc2le.4', end='ip.1', init=tw_sol_off, init_at=xt.START,
+        vary=[
+            xt.VaryList(['acb1h.l1', 'acb2h.l1','acb1v.l1', 'acb2v.l1'], step=1e-8, tag='corr_l'),
+            xt.VaryList(['ks1.l1', 'ks2.l1', 'ks3.l1', 'ks0.l1'], step=1e-7, tag='skew_l'),
+            xt.VaryList(['corr_k1.l1', 'corr_k2.l1', 'corr_k3.l1', 'corr_k0.l1'], step=1e-6, tag='normal_l'),
+        ],
+        targets=[
+            xt.TargetSet(['x', 'y'], value=tw_sol_off, tol=1e-7, at='ip.1', tag='orbit'),
+            xt.TargetSet(['px', 'py'], value=tw_sol_off, tol=1e-10, at='ip.1', tag='orbit'),
+            xt.TargetRmatrix(
+                        r13=0, r14=0, r23=0, r24=0, # Y-X block
+                        r31=0, r32=0, r41=0, r42=0, # X-Y block,
+                        start='pqc2le.4', end='ip.1', tol=1e-6, tag='coupl'),
+            xt.Target('mux', value=tw_sol_off, at='ip.1', tag='mu_ip', weight=0.1, tol=1e-6),
+            xt.Target('muy', value=tw_sol_off, at='ip.1', tag='mu_ip', weight=0.1, tol=1e-6),
+            xt.Target('betx', value=tw_sol_off, at='ip.1', tag='bet_ip', weight=1, tol=1e-5),
+            xt.Target('bety', value=tw_sol_off, at='ip.1', tag='bet_ip', weight=10, tol=1e-6),
+            xt.Target('alfx', value=tw_sol_off, at='ip.1', tag='alf_ip', weight=0.1, tol=1e-4),
+            xt.Target('alfy', value=tw_sol_off, at='ip.1', tag='alf_ip', weight=0.1, tol=1e-4),
+
+        ]
+    )
+
+
+    for iter in range(2):
+        # Orbit alone
+        opt_l.disable_all_targets(); opt_l.disable_all_vary()
+        opt_l.enable_targets(tag='orbit'); opt_l.enable_vary(tag='corr_l'); opt_l.solve()
+
+        # Coupling alone
+        opt_l.disable_all_targets(); opt_l.disable_all_vary()
+        opt_l.enable_targets(tag='coupl'); opt_l.enable_vary(tag='skew_l'); opt_l.solve()
+
+        # phase, beta and alpha alone
+        opt_l.disable_all_targets(); opt_l.disable_all_vary()
+        opt_l.enable_vary(tag='normal_l')
+        opt_l.enable_targets(tag='mu_ip'); opt_l.solve()
+        opt_l.enable_targets(tag='bet_ip'); opt_l.solve()
+        opt_l.enable_targets(tag='alf_ip'); opt_l.solve()
+
+    # All together
+    opt_l.enable_all_targets()
+    opt_l.enable_all_vary()
+    opt_l.solve()
+
+
+    opt_r = line.match(
+        solve=False,
+        method='4d', n_steps_max=30,
+        start='ip.1', end='pqc2re.1', init=tw_sol_off, init_at=xt.END,
+        vary=[
+            xt.VaryList(['acb1h.r1', 'acb2h.r1','acb1v.r1', 'acb2v.r1'], step=1e-8, tag='corr_r'),
+            xt.VaryList(['ks1.r1', 'ks2.r1', 'ks3.r1', 'ks0.r1'], step=1e-7, tag='skew_r'),
+            xt.VaryList(['corr_k1.r1', 'corr_k2.r1', 'corr_k3.r1', 'corr_k0.r1'], step=1e-6, tag='normal_r'),
+        ],
+        targets=[
+            xt.TargetSet(['x', 'y'], value=tw_sol_off, tol=1e-7, at='ip.1', tag='orbit'),
+            xt.TargetSet(['px', 'py'], value=tw_sol_off, tol=1e-10, at='ip.1', tag='orbit'),
+            xt.TargetRmatrix(r13=0, r14=0, r23=0, r24=0, # Y-X block
+                            r31=0, r32=0, r41=0, r42=0, # X-Y block,
+                            start='ip.1', end='pqc2re.1', tol=1e-6, tag='coupl'),
+            xt.Target('mux', value=tw_sol_off, at='ip.1', tag='mu_ip', weight=0.1, tol=1e-6),
+            xt.Target('muy', value=tw_sol_off, at='ip.1', tag='mu_ip', weight=0.1, tol=1e-6),
+            xt.Target('betx', value=tw_sol_off, at='ip.1', tag='bet_ip', weight=1, tol=1e-5),
+            xt.Target('bety', value=tw_sol_off, at='ip.1', tag='bet_ip', weight=10, tol=1e-6),
+            xt.Target('alfx', value=tw_sol_off, at='ip.1', tag='alf_ip', weight=0.1, tol=1e-4),
+            xt.Target('alfy', value=tw_sol_off, at='ip.1', tag='alf_ip', weight=0.1, tol=1e-4),
+
+        ]
+    )
+
+    for iter in range(2):
+        # Orbit alone
+        opt_r.disable_all_targets(); opt_r.disable_all_vary()
+        opt_r.enable_targets(tag='orbit'); opt_r.enable_vary(tag='corr_r'); opt_r.solve()
+
+        # Coupling alone
+        opt_r.disable_all_targets(); opt_r.disable_all_vary()
+        opt_r.enable_targets(tag='coupl'); opt_r.enable_vary(tag='skew_r'); opt_r.solve()
+
+        # phase, beta and alpha alone
+        opt_r.disable_all_targets(); opt_r.disable_all_vary()
+        opt_r.enable_vary(tag='normal_r')
+        opt_r.enable_targets(tag='mu_ip'); opt_r.solve()
+        opt_r.enable_targets(tag='bet_ip'); opt_r.solve()
+        opt_r.enable_targets(tag='alf_ip'); opt_r.solve()
+
+    # All together
+    opt_r.enable_all_targets()
+    opt_r.enable_all_vary()
+    opt_r.solve()
+
+    line.to_json(fname + '_with_sol_corrected.json')
+
+    tw_sol_on_corrected = line.twiss(method='4d')
+
+    assert_allclose = np.testing.assert_allclose
+
+    # Check that tilt is present
+    assert_allclose(tw_sol_off['x_prime', 'ip.1'], np.tan(0.015), atol=1e-14, rtol=0)
+
+    # Check that solenoid introduces coupling
+    assert tw_sol_on.c_minus > 1e-4
+
+    # Check correction
+    tw_chk = tw_sol_on_corrected
+
+    assert_allclose(tw_chk['x', 'ip.1'], 0, atol=1e-8, rtol=0)
+    assert_allclose(tw_chk['y', 'ip.1'], 0, atol=1e-10, rtol=0)
+    assert_allclose(tw_chk['x_prime', 'ip.1'], tw_sol_off['x_prime', 'ip.1'],  atol=1e-9, rtol=0)
+    assert_allclose(tw_chk['y_prime', 'ip.1'], 0,  atol=1e-8, rtol=0)
+    assert_allclose(tw_chk['x', 'pqc2re.1'], 0, atol=5e-8, rtol=0)
+    assert_allclose(tw_chk['y', 'pqc2re.1'], 0, atol=5e-8, rtol=0)
+    assert_allclose(tw_chk['x_prime', 'pqc2re.1'], 0, atol=1e-8, rtol=0)
+    assert_allclose(tw_chk['y_prime', 'pqc2re.1'], 0, atol=1e-8, rtol=0)
+    assert_allclose(tw_chk['x', 'pqc2le.4'], 0, atol=5e-8, rtol=0)
+    assert_allclose(tw_chk['y', 'pqc2le.4'], 0, atol=5e-8, rtol=0)
+    assert_allclose(tw_chk['x_prime', 'pqc2le.4'], 0, atol=1e-8, rtol=0)
+    assert_allclose(tw_chk['y_prime', 'pqc2le.4'], 0, atol=1e-8, rtol=0)
+
+    assert tw_chk['ax', 'pqc2le.4'] == 0
+    assert tw_chk['ay', 'pqc2le.4'] == 0
+    assert tw_chk['ax', 'pqc2re.1'] == 0
+    assert tw_chk['ay', 'pqc2re.1'] == 0
+
+    assert_allclose(tw_chk['betx', 'ip.1'], tw_sol_off['betx', 'ip.1'], atol=0, rtol=5e-5)
+    assert_allclose(tw_chk['bety', 'ip.1'], tw_sol_off['bety', 'ip.1'], atol=0, rtol=5e-5)
+    assert_allclose(tw_chk['alfx', 'ip.1'], tw_sol_off['alfx', 'ip.1'], atol=1e-5, rtol=0)
+    assert_allclose(tw_chk['alfy', 'ip.1'], tw_sol_off['alfy', 'ip.1'], atol=1e-5, rtol=0)
+    assert_allclose(tw_chk['mux', 'ip.1'], tw_sol_off['mux', 'ip.1'], atol=2e-6, rtol=0)
+    assert_allclose(tw_chk['muy', 'ip.1'], tw_sol_off['muy', 'ip.1'], atol=2e-6, rtol=0)
+
+    assert_allclose(tw_chk['betx', 'pqc2re.1'], tw_sol_off['betx', 'pqc2re.1'], atol=0, rtol=5e-5)
+    assert_allclose(tw_chk['bety', 'pqc2re.1'], tw_sol_off['bety', 'pqc2re.1'], atol=0, rtol=5e-5)
+    assert_allclose(tw_chk['alfx', 'pqc2re.1'], tw_sol_off['alfx', 'pqc2re.1'], atol=1e-5, rtol=5e-5)
+    assert_allclose(tw_chk['alfy', 'pqc2re.1'], tw_sol_off['alfy', 'pqc2re.1'], atol=1e-5, rtol=5e-5)
+    assert_allclose(tw_chk['mux', 'pqc2re.1'], tw_sol_off['mux', 'pqc2re.1'], atol=2e-6, rtol=5e-5)
+    assert_allclose(tw_chk['muy', 'pqc2re.1'], tw_sol_off['muy', 'pqc2re.1'], atol=2e-6, rtol=5e-5)
+
+    assert_allclose(tw_chk['betx', 'pqc2le.4'], tw_sol_off['betx', 'pqc2le.4'], atol=0, rtol=5e-5)
+    assert_allclose(tw_chk['bety', 'pqc2le.4'], tw_sol_off['bety', 'pqc2le.4'], atol=0, rtol=5e-5)
+    assert_allclose(tw_chk['alfx', 'pqc2le.4'], tw_sol_off['alfx', 'pqc2le.4'], atol=1e-5, rtol=5e-5)
+    assert_allclose(tw_chk['alfy', 'pqc2le.4'], tw_sol_off['alfy', 'pqc2le.4'], atol=1e-5, rtol=5e-5)
+    assert_allclose(tw_chk['mux', 'pqc2le.4'], tw_sol_off['mux', 'pqc2le.4'], atol=2e-6, rtol=5e-5)
+    assert_allclose(tw_chk['muy', 'pqc2le.4'], tw_sol_off['muy', 'pqc2le.4'], atol=2e-6, rtol=5e-5)
+
+    assert tw_chk.c_minus < 1e-6
+    assert_allclose(tw_chk['betx2', 'ip.1'] / tw_chk['betx', 'ip.1'], 0, atol=1e-11)
+    assert_allclose(tw_chk['bety1', 'ip.1'] / tw_chk['bety', 'ip.1'], 0, atol=1e-11)
+    assert_allclose(tw_chk['betx2', 'pqc2re.1'] / tw_chk['betx', 'pqc2re.1'], 0, atol=1e-11)
+    assert_allclose(tw_chk['bety1', 'pqc2re.1'] / tw_chk['bety', 'pqc2re.1'], 0, atol=1e-11)
+    assert_allclose(tw_chk['betx2', 'pqc2le.4'] / tw_chk['betx', 'pqc2le.4'], 0, atol=1e-11)
+    assert_allclose(tw_chk['bety1', 'pqc2le.4'] / tw_chk['bety', 'pqc2le.4'], 0, atol=1e-11)

{xtrack-0.54.1 → xtrack-0.55.0}/tests/test_line.py

@@ -497,7 +497,7 @@ def test_to_pandas():
     df = line.to_pandas()
 
     assert tuple(df.columns) == (
-        's', 'element_type', 'name', 'isthick', 'compound_name', 'element')
+        's', 'element_type', 'name', 'isthick', 'iscollective', 'compound_name', 'element')
     assert len(df) == 4
 
 def test_check_aperture():

{xtrack-0.54.1 → xtrack-0.55.0}/tests/test_mad_writer.py

@@ -28,7 +28,7 @@ def test_mad_writer(case):
 
     mad = Madx()
     mad.input(mad_seq)
-    mad.beam(particle='proton', energy=7000e9) #!!!!!
+    mad.beam(particle='proton', energy=7000e9)
     mad.use('myseq')
 
     line2 = xt.Line.from_madx_sequence(mad.sequence.myseq, deferred_expressions=True)