xtrack 0.72.0__tar.gz → 0.72.2__tar.gz

{xtrack-0.72.0/xtrack.egg-info → xtrack-0.72.2}/PKG-INFO

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

xtrack-0.72.2/tests/test_acceleration.py

@@ -0,0 +1,366 @@
+# copyright ############################### #
+# This file is part of the Xtrack Package.  #
+# Copyright (c) CERN, 2021.                 #
+# ######################################### #
+
+import json
+import pathlib
+
+import numpy as np
+import pandas as pd
+from cpymad.madx import Madx
+from scipy.constants import c as clight
+
+import xobjects as xo
+import xpart as xp
+import xtrack as xt
+from xobjects.test_helpers import for_all_test_contexts
+
+test_data_folder = pathlib.Path(
+    __file__).parent.joinpath('../test_data').absolute()
+
+
+@for_all_test_contexts
+def test_acceleration(test_context):
+    Delta_p0c = 450e9/10*23e-6  # ramp rate 450GeV/10s
+
+    fname_line = test_data_folder.joinpath(
+        'sps_w_spacecharge/line_no_spacecharge_and_particle.json')
+
+    with open(fname_line, 'r') as fid:
+        input_data = json.load(fid)
+    line = xt.Line.from_dict(input_data['line'])
+
+    energy_increase = xt.ReferenceEnergyIncrease(Delta_p0c=Delta_p0c)
+    line.append_element(energy_increase, 'energy_increase')
+
+    line.build_tracker(_context=test_context)
+
+    # Assume only first cavity is active
+    frequency = line.get_elements_of_type(xt.Cavity)[0][0].frequency
+    voltage = line.get_elements_of_type(xt.Cavity)[0][0].voltage
+    # Assuming proton and beta=1
+    stable_z = np.arcsin(Delta_p0c/voltage)/frequency/2/np.pi*clight
+
+    p_co = line.find_closed_orbit(particle_ref=xp.Particles.from_dict(
+        input_data['particle']))
+
+    xo.assert_allclose(p_co._xobject.zeta[0], stable_z, atol=0, rtol=1e-2)
+
+
+@for_all_test_contexts(excluding=('ContextPyopencl',))
+def test_energy_program(test_context):
+
+    df = pd.read_csv(test_data_folder / 'psb_chicane/Ramp_and_RF_functions.dat',
+                     sep='\t', skiprows=2,
+                     names=['t_s', 'E_kin_GeV', 'V1_MV', 'phi1_rad', 'V2_MV', 'phi2_rad'])
+    E_kin_GeV = df.E_kin_GeV.values
+    t_s = df.t_s.values
+
+    # Shift to enhance change in revolution frequency
+    E_min = np.min(E_kin_GeV)
+    E_max = np.max(E_kin_GeV)
+    E_kin_GeV = E_min/100 + (E_kin_GeV - E_min)
+    # Shift the time scale for testing purposes
+    t_s = t_s
+
+    # Load mad model and apply element shifts
+    mad = Madx(stdout=False)
+    mad.call(str(test_data_folder / 'psb_chicane/psb.seq'))
+    mad.call(str(test_data_folder / 'psb_chicane/psb_fb_lhc.str'))
+    mad.input('''
+        beam, particle=PROTON, pc=0.5708301551893517;
+        use, sequence=psb1;
+        twiss;
+    ''')
+
+    line = xt.Line.from_madx_sequence(mad.sequence.psb1, allow_thick=True,
+                                      deferred_expressions=True)
+    line.particle_ref = xt.Particles(mass0=xt.PROTON_MASS_EV,
+                                     gamma0=mad.sequence.psb1.beam.gamma)
+
+    line.build_tracker(_context=test_context)
+
+    # Attach energy program
+    ep = xt.EnergyProgram(t_s=t_s, kinetic_energy0=E_kin_GeV*1e9)
+    line.energy_program = ep
+
+    line['br1.acwf7l1.1'].voltage = 2e6
+    line['br1.acwf7l1.1'].frequency = 1e3
+
+    # test to_dict and from_dict
+    lbefore = line
+    line_dict = line.to_dict()
+    line = xt.Line.from_dict(line_dict)
+    assert np.all(line.vars.get_table().name == lbefore.vars.get_table().name)
+
+    line = line.copy()
+    assert np.all(line.vars.get_table().name == lbefore.vars.get_table().name)
+
+    line.build_tracker(_context=test_context)
+
+    p_test = line.build_particles()
+    line.enable_time_dependent_vars = True
+    n_turn_test = 5000
+    monitor = xt.ParticlesMonitor(num_particles=len(p_test.zeta), start_at_turn=0,
+                                  stop_at_turn=n_turn_test,
+                                  _context=test_context)
+    for ii in range(n_turn_test):
+        line.track(p_test, turn_by_turn_monitor=monitor)
+
+    beta_at_turn = monitor.beta0[0, :]
+    gamma_at_turn = 1 / np.sqrt(1 - beta_at_turn**2)
+
+    t_turn_ref = np.cumsum(line.get_length()/clight/beta_at_turn)
+    t_turn_ref = t_turn_ref - t_turn_ref[0]
+    E_kin_turn = line.particle_ref.mass0 * (monitor.gamma0[0, :] - 1)
+
+    t_check = np.linspace(0, 20e-3, 1000)
+    E_check = np.interp(t_check, t_turn_ref, E_kin_turn)
+    E_check_ref = np.interp(t_check, t_s, E_kin_GeV*1e9)
+    xo.assert_allclose(E_check, E_check_ref, atol=0, rtol=2e-3)
+
+    t_turn_check = line.energy_program.get_t_s_at_turn(np.arange(n_turn_test))
+    xo.assert_allclose(t_turn_check, t_turn_ref, atol=0, rtol=6e-4)
+
+    p0c_check = line.energy_program.get_p0c_at_t_s(t_check)
+    p0c_ref = np.interp(t_check,
+                        t_turn_check,
+                        line.particle_ref.mass0 * gamma_at_turn * beta_at_turn)
+    xo.assert_allclose(p0c_check, p0c_ref, atol=0, rtol=1e-3)
+
+    kinetic_energy0_check = line.energy_program.get_kinetic_energy0_at_t_s(t_check)
+    kinetic_energy0_ref = np.interp(t_check,
+                        t_turn_check,
+                        line.particle_ref.mass0 * (gamma_at_turn - 1))
+    xo.assert_allclose(kinetic_energy0_check, kinetic_energy0_ref, atol=0, rtol=2e-3)
+
+    beta0_check = line.energy_program.get_beta0_at_t_s(t_check)
+    beta0_ref = np.interp(t_check, t_turn_check, beta_at_turn)
+    xo.assert_allclose(beta0_check, beta0_ref, atol=0, rtol=1e-3)
+
+    frev_check = line.energy_program.get_frev_at_t_s(t_check)
+    frev_ref = np.interp(t_check, t_turn_check[:-1], 1/np.diff(t_turn_ref))
+    xo.assert_allclose(frev_check, frev_ref, atol=0, rtol=4e-5)
+
+    p0c_increse_per_turn_check = line.energy_program.get_p0c_increse_per_turn_at_t_s(
+        t_check)
+    p0c_increse_per_turn_ref = np.interp(
+        t_check, t_turn_check[:-1], np.diff(monitor.p0c[0, :]))
+    xo.assert_allclose(p0c_increse_per_turn_check - p0c_increse_per_turn_ref, 0,
+                       atol=5e-5 * p0c_ref[0], rtol=0)
+
+    line.enable_time_dependent_vars = False
+    line.vars['t_turn_s'] = 20e-3
+
+    E_kin_expected = np.interp(line.vv['t_turn_s'], t_s, E_kin_GeV*1e9)
+    E_tot_expected = E_kin_expected + line.particle_ref.mass0
+    xo.assert_allclose(
+        E_tot_expected, line.particle_ref.energy0[0], rtol=1e-4, atol=0)
+    xo.assert_allclose(
+        E_kin_expected, line.particle_ref.kinetic_energy0[0], rtol=1e-4, atol=0)
+
+    tw = line.twiss(method='6d')
+    # To check that it does not change
+    xo.assert_allclose(tw.zeta[0], -13.48, rtol=0, atol=1e-4)
+    xo.assert_allclose(line.particle_ref.mass0 * tw.gamma0, E_tot_expected,
+                      atol=0, rtol=1e-12)
+
+    line.vars['t_turn_s'] = 0
+    line.vars['on_chicane_k0'] = 0
+    tw = line.twiss(method='6d')
+    xo.assert_allclose(tw.zeta[0], 0, rtol=0, atol=1e-12)
+    xo.assert_allclose(line.particle_ref.mass0 * tw.gamma0, line.particle_ref.mass0 + E_kin_turn[0],
+                       rtol=1e-10, atol=0)
+
+@for_all_test_contexts(excluding=('ContextPyopencl',))
+def test_acceleration_transverse_shrink(test_context):
+
+    mad = Madx(stdout=False)
+
+    # Load mad model and apply element shifts
+    mad.input(f'''
+    call, file = '{str(test_data_folder)}/psb_chicane/psb.seq';
+    call, file = '{str(test_data_folder)}/psb_chicane/psb_fb_lhc.str';
+    beam;
+    use, sequence=psb1;
+    ''')
+
+    line = xt.Line.from_madx_sequence(mad.sequence.psb1,
+                                        deferred_expressions=True)
+    e_kin_start_eV = 160e6
+    line.particle_ref = xt.Particles(mass0=xt.PROTON_MASS_EV, q0=1.,
+                                    energy0=xt.PROTON_MASS_EV + e_kin_start_eV)
+
+    # Slice to gain some tracking speed
+    line.slice_thick_elements(
+        slicing_strategies=[
+            xt.Strategy(slicing=xt.Teapot(1)),
+        ])
+    line.build_tracker(_context=test_context)
+
+    # User-defined energy ramp
+    t_s = np.array([0., 0.0006, 0.0008, 0.001 , 0.0012, 0.0014, 0.0016, 0.0018,
+                    0.002 , 0.0022, 0.0024, 0.0026, 0.0028, 0.003, 0.01, 0.1])
+    E_kin_GeV = np.array([0.16000000,0.16000000,
+        0.16000437, 0.16001673, 0.16003748, 0.16006596, 0.16010243, 0.16014637,
+        0.16019791, 0.16025666, 0.16032262, 0.16039552, 0.16047524, 0.16056165,
+        0.163586, 0.20247050000000014])
+
+    # Enhance energy swing to better see the effect of energy on beam size
+    E_kin_GeV -= 0.140
+
+    # Go away from half integer
+    opt = line.match(
+        #verbose=True,
+        method='4d',
+        solve=False,
+        vary=[
+            xt.Vary('kbrqfcorr', step=1e-4),
+            xt.Vary('kbrqdcorr', step=1e-4),
+        ],
+        targets = [
+            xt.Target('qx', value=4.15, tol=1e-5, scale=1),
+            xt.Target('qy', value=4.18, tol=1e-5, scale=1),
+        ]
+    )
+    opt.solve()
+
+
+    # Attach energy program to the line
+    line.energy_program = xt.EnergyProgram(t_s=t_s, kinetic_energy0=E_kin_GeV*1e9)
+
+    # Plot energy and revolution frequency vs time
+    t_plot = np.linspace(0, 10e-3, 20)
+    E_kin_plot = line.energy_program.get_kinetic_energy0_at_t_s(t_plot)
+    f_rev_plot = line.energy_program.get_frev_at_t_s(t_plot)
+
+    # import matplotlib.pyplot as plt
+    # plt.close('all')
+    # plt.figure(1, figsize=(6.4 * 1.5, 4.8))
+    # ax1 = plt.subplot(2,2,1)
+    # plt.plot(t_plot * 1e3, E_kin_plot * 1e-6)
+    # plt.ylabel(r'$E_{kin}$ [MeV]')
+    # ax2 = plt.subplot(2,2,3, sharex=ax1)
+    # plt.plot(t_plot * 1e3, f_rev_plot * 1e-3)
+    # plt.ylabel(r'$f_{rev}$ [kHz]')
+    # plt.xlabel('t [ms]')
+
+    # Setup frequency of the RF cavity  to stay on the second harmonic of the
+    # revolution frequency during the acceleration
+
+    t_rf = np.linspace(0, 3e-3, 100) # time samples for the frequency program
+    f_rev = line.energy_program.get_frev_at_t_s(t_rf)
+    h_rf = 2 # harmonic number
+    f_rf = h_rf * f_rev # frequency program
+
+    # Build a function with these samples and link it to the cavity
+    line.functions['fun_f_rf'] = xt.FunctionPieceWiseLinear(x=t_rf, y=f_rf)
+    line.element_refs['br1.acwf5l1.1'].frequency = line.functions['fun_f_rf'](
+                                                            line.vars['t_turn_s'])
+
+    # Setup voltage and lag
+    line.element_refs['br1.acwf5l1.1'].voltage = 3000 # V
+    line.element_refs['br1.acwf5l1.1'].lag = 0 # degrees (below transition energy)
+
+    # When setting line.vars['t_turn_s'] the reference energy and the rf frequency
+    # are updated automatically
+    line.vars['t_turn_s'] = 0
+    line.particle_ref.kinetic_energy0 # is 160.00000 MeV
+    line['br1.acwf5l1.1'].frequency # is 1983931.935 Hz
+
+    line.vars['t_turn_s'] = 3e-3
+    line.particle_ref.kinetic_energy0 # is 160.56165 MeV
+    line['br1.acwf5l1.1'].frequency # is 1986669.0559674294
+
+    # Back to zero for tracking!
+    line.vars['t_turn_s'] = 0
+
+    # Track a few particles to visualize the longitudinal phase space
+    p_test = line.build_particles(x_norm=0, zeta=np.linspace(0, line.get_length(), 101))
+
+    # Enable time-dependent variables (t_turn_s and all the variables that depend on
+    # it are automatically updated at each turn)
+    line.enable_time_dependent_vars = True
+
+    # Track
+    line.track(p_test, num_turns=9000, turn_by_turn_monitor=True, with_progress=True)
+    mon = line.record_last_track
+
+    # Plot
+    # plt.subplot2grid((2,2), (0,1), rowspan=2)
+    # plt.plot(mon.zeta[:, -2000:].T, mon.delta[:, -2000:].T, color='C0')
+    # plt.xlabel(r'$\zeta$ [m]')
+    # plt.ylabel('$\delta$')
+    # plt.xlim(-40, 30)
+    # plt.ylim(-0.0025, 0.0025)
+    # plt.title('Last 2000 turns')
+    # plt.subplots_adjust(left=0.08, right=0.95, wspace=0.26)
+
+
+    # Check transverse beam size reduction
+    line['t_turn_s'] = 0
+    line.enable_time_dependent_vars = False
+
+    n_part_test = 500
+    # Generate Gaussian distribution with fixed rng seed
+    rng = np.random.default_rng(seed=123)
+    x_norm = rng.normal(loc=0, scale=1, size=n_part_test)
+    px_norm = rng.normal(loc=0, scale=1, size=n_part_test)
+    y_norm = rng.normal(loc=0, scale=1, size=n_part_test)
+    py_norm = rng.normal(loc=0, scale=1, size=n_part_test)
+
+    # rescale to have exact std dev.
+    x_norm = x_norm / np.std(x_norm)
+    px_norm = px_norm / np.std(px_norm)
+    y_norm = y_norm / np.std(y_norm)
+    py_norm = py_norm / np.std(py_norm)
+
+    p_test2 = line.build_particles(x_norm=x_norm, px_norm=px_norm,
+                                y_norm=x_norm, py_norm=px_norm,
+                                nemitt_x=1e-6, nemitt_y=1e-6,
+                                delta=0)
+
+    line.enable_time_dependent_vars = True
+    line.track(p_test2, num_turns=10_000, turn_by_turn_monitor=True, with_progress=True)
+    mon2 = line.record_last_track
+
+    std_y = np.std(mon2.y, axis=0)
+    std_x = np.std(mon2.x, axis=0)
+
+    # Apply moving average filter
+    from scipy.signal import savgol_filter
+    std_y_smooth = savgol_filter(std_y, 10000, 2)
+    std_x_smooth = savgol_filter(std_x, 10000, 2)
+
+    i_turn_match = 1000
+    std_y_expected = std_y_smooth[i_turn_match] * np.sqrt(
+        mon2.gamma0[0, i_turn_match]* mon2.beta0[0, i_turn_match]
+        / mon2.gamma0[0, :] / mon2.beta0[0, :])
+    std_x_expected = std_x_smooth[i_turn_match] * np.sqrt(
+        mon2.gamma0[0, i_turn_match]* mon2.beta0[0, i_turn_match]
+        / mon2.gamma0[0, :] / mon2.beta0[0, :])
+
+    d_sigma_x = std_x_expected[0] - std_x_expected[-1]
+    d_sigma_y = std_y_expected[0] - std_y_expected[-1]
+
+    import xobjects as xo
+    xo.assert_allclose(std_y_expected[8000:9000].mean(),
+                    std_y_smooth[8000:9000].mean(),
+                    rtol=0, atol=0.03 * d_sigma_y)
+    xo.assert_allclose(std_x_expected[8000:9000].mean(),
+                    std_x_smooth[8000:9000].mean(),
+                    rtol=0, atol=0.03 * d_sigma_x)
+
+    # plt.figure(2)
+    # ax1 = plt.subplot(2,1,1)
+    # plt.plot(std_x, label='raw')
+    # plt.plot(std_x_smooth, label='smooth')
+    # plt.plot(std_x_expected, label='expected')
+    # plt.legend()
+
+    # ax2 = plt.subplot(2,1,2, sharex=ax1)
+    # plt.plot(std_y, label='raw')
+    # plt.plot(std_y_smooth, label='smooth')
+    # plt.plot(std_y_expected, label='expected')
+    # plt.show()

{xtrack-0.72.0 → xtrack-0.72.2}/tests/test_line.py

@@ -357,7 +357,6 @@ def test_to_pandas():
 def test_check_aperture():
 
     class ThickElement:
-
         length = 2.
         isthick = True
 
@@ -380,11 +379,16 @@ def test_check_aperture():
             'th3_ap_front': xt.LimitEllipse(a=1e-2, b=1e-2),
             'th3': ThickElement(),
             'dr6': xt.Drift(length=1),
+            'th4_ap_entry': xt.LimitEllipse(a=1e-2, b=1e-2),
+            'th4': ThickElement(),
+            'th4_ap_exit': xt.LimitEllipse(a=1e-2, b=1e-2),
         },
         element_names=['dr1', 'm1_ap', 'dum', 'm1', 'dr2', 'm2', 'dr3',
                        'th1_ap_front', 'dum', 'th1', 'dum', 'th1_ap_back',
                        'dr4', 'th2', 'th2_ap_back',
-                       'dr5', 'th3_ap_front', 'th3'])
+                       'dr5', 'th3_ap_front', 'th3', 'dr6',
+                       'th4_ap_entry', 'th4', 'th4_ap_exit'])
+    
     df = line.check_aperture()
 
     expected_miss_upstream = [nn in ('m2', 'th2') for nn in df['name'].values]
@@ -1092,4 +1096,14 @@ def test_insert_repeated_names():
     line.insert_element("m2",xt.Marker(),at="d")
     assert line.element_names[0]=="m2"
 
+def test_line_table_unique_names():
+    line = xt.Line(
+        elements = {"obm": xt.Bend(length=0.5)},
+        element_names= ["obm","obm"]
+    )
+    table = line.get_table()
+    assert np.all(np.unique_counts(table.name).counts == 1), "Not all elements are unique"
+    for name, env_name in zip(table.name, table.env_name):
+        if name == '_end_point': continue
+        assert line[name] == line[env_name]
 

{xtrack-0.72.0 → xtrack-0.72.2}/tests/test_madnginterface.py

@@ -7,7 +7,6 @@ import numpy as np
 test_data_folder = pathlib.Path(
     __file__).parent.joinpath('../test_data').absolute()
 
-# @pytest.mark.xfail(strict=False) # sometimes madng fails to start the process (to be investigated)
 def test_madng_twiss():
     rdts = ["f4000", "f3100", "f2020", "f1120"]
 
@@ -46,7 +45,6 @@ def test_madng_twiss():
     assert np.abs(tw_rdt.f2020).max() > 0
     assert np.abs(tw_rdt.f1120).max() > 0
 
-# @pytest.mark.xfail(strict=False) # sometimes madng fails to start the process (to be investigated)
 def test_madng_interface_with_multipole_errors_and_misalignments():
     line = xt.Line.from_json(test_data_folder /
                             'hllhc15_thick/lhc_thick_with_knobs.json')

{xtrack-0.72.0 → xtrack-0.72.2}/tests/test_radiation_equilibrium_emittances.py

@@ -145,7 +145,7 @@ def test_eq_emitt(conf):
         xo.assert_allclose(ez, 3.5766e-6,  atol=0,     rtol=1e-4)
         checked = True
     elif tilt_machine_by_90_degrees and vertical_orbit_distortion and not wiggler_on:
-        xo.assert_allclose(ex, 2.5039e-12, atol=0,     rtol=4e-3)
+        xo.assert_allclose(ex, 2.5039e-12, atol=0,     rtol=5e-3)
         xo.assert_allclose(ey, 7.0576e-10, atol=0,     rtol=1e-4)
         xo.assert_allclose(ez, 3.5763e-6,  atol=0,     rtol=1e-4)
         checked = True

xtrack-0.72.2/xtrack/_version.py

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

{xtrack-0.72.0 → xtrack-0.72.2}/xtrack/beam_elements/elements.py

@@ -51,6 +51,7 @@ class Marker(BeamElement):
     allow_loss_refinement = True
     has_backtrack = True
     allow_rot_and_shift = False
+    _skip_in_repr = ['_dummy']
 
     _extra_c_sources = [
         "/*gpufun*/\n"
@@ -83,6 +84,11 @@ class Drift(BeamElement):
         _pkg_root.joinpath('beam_elements/elements_src/drift_elem.h'),
         ]
 
+    def __init__(self, length=None, **kwargs):
+        if length:  # otherwise length cannot be set as a positional argument
+            kwargs['length'] = length
+        super().__init__(**kwargs)
+
     @property
     def _thin_slice_class(self):
         return None
@@ -2441,6 +2447,10 @@ class SecondOrderTaylorMap(BeamElement):
             A `SecondOrderTaylorMap` object.
 
         '''
+        if start == end:
+            # start == end will lead to compute_one_turn_matrix_finite_differences() computing a
+            # full one-turn response matrix (but here we would rather expect identity)
+            raise NotImplementedError('end element must be after start element')
 
         if twiss_table is None:
             tw = line.twiss(reverse=False)

{xtrack-0.72.0 → xtrack-0.72.2}/xtrack/beam_elements/exciter.py

@@ -99,12 +99,21 @@ class Exciter(BeamElement):
     _extra_c_sources = [_pkg_root.joinpath('beam_elements/elements_src/exciter.h')]
 
 
-    def __init__(self, *, samples=None, nsamples=None, sampling_frequency=0, frev=0, knl=[1], ksl=[], start_turn=0, duration=None, _xobject=None, **kwargs):
+    def __init__(self, *, samples=None, nsamples=None, sampling_frequency=0, frev=0, knl=None, ksl=None, start_turn=0, duration=None, _xobject=None, **kwargs):
 
         if _xobject is not None:
             super().__init__(_xobject=_xobject)
 
         else:
+            # The default is needed, as we wish to be able to instantiate
+            # elements without properties (empty elements can be templates).
+            if knl is None and ksl is None:
+                knl = [0]
+                ksl = []
+            elif knl is None:
+                knl = []
+            elif ksl is None:
+                ksl = []
 
             # sanitize knl and ksl array length
             n = max(len(knl), len(ksl))

{xtrack-0.72.0 → xtrack-0.72.2}/xtrack/line.py

@@ -3213,13 +3213,13 @@ class Line:
         i_next_aperture = 0
 
         for iee in progress(range(i_prev_aperture, num_elements), desc='Checking aperture'):
-            if dont_need_aperture[elements_df.loc[iee, 'name']]:
-                continue
-
             if elements_df.loc[iee, 'is_aperture']:
                 i_prev_aperture = iee
                 continue
 
+            if dont_need_aperture[elements_df.loc[iee, 'name']]:
+                continue
+
             if i_next_aperture < iee:
                 for ii in range(iee, num_elements):
                     if elements_df.loc[ii, 'is_aperture']:
@@ -4060,6 +4060,8 @@ class Line:
             return self.vv[key]
         elif hasattr(self, 'lines') and key in self.lines: # Want to reuse the method for the env
             return self.lines[key]
+        elif "::" in key and (env_name := key.split("::")[0]) in self.element_dict:
+            return self[env_name]
         else:
             raise KeyError(f'Name {key} not found')
 
@@ -5332,7 +5334,8 @@ class EnergyProgram:
     @t_turn_s_line.setter
     def t_turn_s_line(self, value):
         p0c = self.get_p0c_at_t_s(value)
-        self.line.particle_ref.update_p0c_and_energy_deviations(p0c=p0c)
+        self.line.particle_ref.update_p0c_and_energy_deviations(
+                                                    p0c=p0c, update_pxpy=True)
 
     def to_dict(self):
         assert not self.needs_complete, 'EnergyProgram not completed'

{xtrack-0.72.0 → xtrack-0.72.2}/xtrack/particles/particles.py

@@ -212,6 +212,14 @@ class Particles(xo.HybridClass):
             raise NameError('`psigma` is not supported anymore.'
                             'Please use `pzeta` instead.')
 
+        accepted_args = set(self._xofields.keys()) | {
+            'energy0', 'tau', 'pzeta', 'mass_ratio', 'mass', 'kinetic_energy0',
+            '_context', '_buffer', '_offset', 'p0',
+        }
+        if set(kwargs.keys()) - accepted_args:
+            raise NameError(f'Invalid argument(s) provided: '
+                            f'{set(kwargs.keys()) - accepted_args}')
+
         per_part_input_vars = (
             self.per_particle_vars +
             ((xo.Float64, 'energy0'),
@@ -2175,7 +2183,7 @@ class Particles(xo.HybridClass):
                                     computed_value=_charge_ratio,
                                     mask=mask)
 
-    def update_p0c_and_energy_deviations(self, p0c):
+    def update_p0c_and_energy_deviations(self, p0c, update_pxpy=False):
 
         assert np.isscalar(p0c), 'p0c must be a scalar'
 
@@ -2195,6 +2203,13 @@ class Particles(xo.HybridClass):
         self._update_energy_deviations(mask=mask, delta=new_delta)
         self._update_zeta(mask=mask, zeta=self.zeta * self.beta0 / old_beta0)
 
+        if update_pxpy:
+            if isinstance(self._context, xo.ContextPyopencl):
+                raise NotImplementedError # Issue wiht masking
+            scale_pxpy = old_p0c[mask] / new_p0c[mask]
+            self.px[mask] *= scale_pxpy
+            self.py[mask] *= scale_pxpy
+
 
 def _mask_to_where(mask, ctx):
     if hasattr(mask, 'get'):

{xtrack-0.72.0 → xtrack-0.72.2}/xtrack/tracker.py

@@ -1470,7 +1470,7 @@ class Tracker:
 
             if self.line.energy_program is not None:
                 p0c = self.line.particle_ref._xobject.p0c[0]
-                particles.update_p0c_and_energy_deviations(p0c)
+                particles.update_p0c_and_energy_deviations(p0c, update_pxpy=True)
 
     def _handle_log(self, _session_to_resume, particles, log):
         if _session_to_resume is not None:

{xtrack-0.72.0 → xtrack-0.72.2/xtrack.egg-info}/PKG-INFO

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