xtrack 0.90.2__tar.gz → 0.93.0__tar.gz

{xtrack-0.90.2/xtrack.egg-info → xtrack-0.93.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.90.2
+Version: 0.93.0
 Summary: Tracking library for particle accelerators
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xtrack
@@ -25,6 +25,7 @@ Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-mock; extra == "tests"
 Requires-Dist: pymadng; extra == "tests"
 Requires-Dist: requests-mock; extra == "tests"
+Requires-Dist: tfs-pandas; extra == "tests"
 Provides-Extra: notebooks
 Requires-Dist: jupyter; extra == "notebooks"
 Requires-Dist: ipympl; extra == "notebooks"

{xtrack-0.90.2 → xtrack-0.93.0}/setup.py

@@ -50,7 +50,7 @@ setup(
     ],
     extras_require={
         'tests': ['cpymad', 'nafflib', 'PyHEADTAIL', 'pytest', 'pytest-mock',
-                  'pymadng', 'requests-mock'],
+                  'pymadng', 'requests-mock', 'tfs-pandas'],
         'notebooks': ['jupyter', 'ipympl', 'xplt'],
     },
 )

xtrack-0.93.0/tests/test_boris_spatial.py

@@ -0,0 +1,117 @@
+import numpy as np
+from scipy.constants import c as clight
+from scipy.constants import e as qe
+
+import xobjects as xo
+import xtrack as xt
+from xtrack._temp.boris_and_solenoid_map.solenoid_field import SolenoidField
+
+
+def test_boris_spatial():
+    delta=np.array([0, 4])
+    p0 = xt.Particles(mass0=xt.ELECTRON_MASS_EV, q0=1,
+                    energy0=45.6e9/1000,
+                    x=[-1e-3, -1e-3],
+                    px=-1e-3*(1+delta),
+                    y=1e-3,
+                    delta=delta)
+
+    p = p0.copy()
+
+    sf = SolenoidField(L=4, a=0.3, B0=1.5, z0=20)
+
+    integrator = xt.BorisSpatialIntegrator(fieldmap_callable=sf.get_field,
+                                            s_start=0,
+                                            s_end=30,
+                                            n_steps=15000)
+    p_boris = p.copy()
+    integrator.track(p_boris)
+
+
+    x_log = np.array(integrator.x_log)
+    y_log = np.array(integrator.y_log)
+    z_log = np.array(integrator.z_log)
+
+    b_seen = sf.get_field(x_log, y_log, z_log)
+
+
+    z_axis = np.linspace(0, 30, 1001)
+    Bz_axis = sf.get_field(0 * z_axis, 0 * z_axis, z_axis)[2]
+
+    P0_J = p.p0c[0] * qe / clight
+    brho = P0_J / qe / p.q0
+
+    # ks = 0.5 * (Bz_axis[:-1] + Bz_axis[1:]) / brho
+    ks = Bz_axis / brho
+    ks_entry = ks[:-1]
+    ks_exit = ks_entry*0
+    ks_exit = ks[1:]
+
+    dz = z_axis[1]-z_axis[0]
+
+    line = xt.Line(elements=[xt.VariableSolenoid(length=dz,
+                                        ks_profile=[ks_entry[ii], ks_exit[ii]])
+                                for ii in range(len(z_axis)-1)])
+    line.build_tracker()
+
+    p_xt = p0.copy()
+    line.track(p_xt, turn_by_turn_monitor='ONE_TURN_EBE')
+    mon = line.record_last_track
+
+    p_xt = p0.copy()
+    line.configure_radiation(model=None)
+    line.track(p_xt, turn_by_turn_monitor='ONE_TURN_EBE')
+    mon_no_rad = line.record_last_track
+
+    Bz_mid = 0.5 * (Bz_axis[:-1] + Bz_axis[1:])
+    Bz_mon = 0 * Bz_axis
+    Bz_mon[1:] = Bz_mid
+
+    # Wolsky Eq. 3.114
+    Ax = -0.5 * Bz_mon * mon.y
+    Ay =  0.5 * Bz_mon * mon.x
+
+    # Wolsky Eq. 2.74
+    ax_ref = Ax * p0.q0 * qe / P0_J
+    ay_ref = Ay * p0.q0 * qe / P0_J
+
+    px_mech = mon.px - ax_ref
+    py_mech = mon.py - ay_ref
+    pz_mech = np.sqrt((1 + mon.delta)**2 - px_mech**2 - py_mech**2)
+
+    xp = px_mech / pz_mech
+    yp = py_mech / pz_mech
+
+    dx_ds = np.diff(mon.x, axis=1) / np.diff(mon.s, axis=1)
+    dy_ds = np.diff(mon.y, axis=1) / np.diff(mon.s, axis=1)
+
+    z_check = sf.z0 + sf.L * np.linspace(-2, 2, 1001)
+
+    for i_part in range(z_log.shape[1]):
+
+        this_s_boris = 0.5 * (z_log[:-1, i_part] + z_log[1:, i_part])
+        dx_ds_boris = np.diff(x_log[:, i_part]) / np.diff(z_log[:, i_part])
+        dy_ds_boris = np.diff(y_log[:, i_part]) / np.diff(z_log[:, i_part])
+
+        s_xsuite = 0.5 * (mon.s[i_part, :-1] + mon.s[i_part, 1:])
+        dx_ds_xsuite = np.diff(mon.x[i_part, :]) / np.diff(mon.s[i_part, :])
+        dy_ds_xsuite = np.diff(mon.y[i_part, :]) / np.diff(mon.s[i_part, :])
+
+        dx_ds_xsuite_check = np.interp(z_check, s_xsuite, dx_ds_xsuite)
+        dy_ds_xsuite_check = np.interp(z_check, s_xsuite, dy_ds_xsuite)
+
+        dx_ds_boris_check = np.interp(z_check, this_s_boris, dx_ds_boris)
+        dy_ds_boris_check = np.interp(z_check, this_s_boris, dy_ds_boris)
+
+        this_dx_ds = dx_ds[i_part, :]
+        this_dy_ds = dy_ds[i_part, :]
+
+        xo.assert_allclose(dx_ds_xsuite_check, dx_ds_boris_check, rtol=0,
+                atol=2.8e-2 * (np.max(dx_ds_boris_check) - np.min(dx_ds_boris_check)))
+        xo.assert_allclose(dy_ds_xsuite_check, dy_ds_boris_check, rtol=0,
+                atol=2.8e-2 * (np.max(dy_ds_boris_check) - np.min(dy_ds_boris_check)))
+
+        xo.assert_allclose(ax_ref[i_part, :], mon.ax[i_part, :],
+                        rtol=0, atol=np.max(np.abs(ax_ref)*3e-2))
+        xo.assert_allclose(ay_ref[i_part, :], mon.ay[i_part, :],
+                        rtol=0, atol=np.max(np.abs(ay_ref)*3e-2))

xtrack-0.93.0/tests/test_json.py

@@ -0,0 +1,94 @@
+import io
+
+import numpy as np
+
+from xtrack import json as xjson
+
+
+def test_complex_scalar_round_trip():
+    data = {"value": 1 + 2j}
+
+    buffer = io.StringIO()
+    xjson.dump(data, buffer)
+    raw = buffer.getvalue()
+
+    buffer.seek(0)
+    loaded_from_file = xjson.load(file=buffer)
+    assert loaded_from_file == data
+    assert isinstance(loaded_from_file["value"], complex)
+
+    loaded_from_string = xjson.load(string=raw)
+    assert loaded_from_string == data
+    assert isinstance(loaded_from_string["value"], complex)
+
+
+def test_complex_array_round_trip_preserves_dtype_and_shape(tmp_path):
+    array = np.array([[1 + 2j, 3 + 4j], [5 + 6j, 7 + 8j]], dtype=np.complex64)
+
+    path = tmp_path / "complex.json"
+    xjson.dump({"array": array}, path)
+
+    loaded = xjson.load(path)
+    round_trip = loaded["array"]
+
+    assert isinstance(round_trip, np.ndarray)
+    assert round_trip.dtype == array.dtype
+    np.testing.assert_allclose(round_trip, array)
+
+
+def test_complex_array_round_trip_with_gzip(tmp_path):
+    array = np.array([1 + 2j, 3 + 4j], dtype=np.complex128)
+
+    path = tmp_path / "complex.json.gz"
+    xjson.dump({"array": array}, path)
+
+    loaded = xjson.load(path)
+    round_trip = loaded["array"]
+
+    assert isinstance(round_trip, np.ndarray)
+    assert round_trip.dtype == array.dtype
+    np.testing.assert_allclose(round_trip, array)
+
+
+def test_python_primitive_types_round_trip():
+    data = {
+        "int": 42,
+        "float": 3.14,
+        "bool": True,
+        "str": "xtrack",
+        "list": [1, 2, 3],
+        "dict": {"nested": "value"},
+        "none": None,
+    }
+
+    buffer = io.StringIO()
+    xjson.dump(data, buffer)
+
+    buffer.seek(0)
+    loaded = xjson.load(file=buffer)
+    assert loaded == data
+
+
+def test_numpy_scalar_round_trip():
+    data = {"value": np.int64(5)}
+
+    buffer = io.StringIO()
+    xjson.dump(data, buffer)
+
+    buffer.seek(0)
+    loaded = xjson.load(file=buffer)
+
+    assert loaded["value"] == 5
+    assert isinstance(loaded["value"], int)
+
+
+def test_numpy_real_array_round_trip():
+    array = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+
+    buffer = io.StringIO()
+    xjson.dump({"array": array}, buffer)
+
+    buffer.seek(0)
+    loaded = xjson.load(file=buffer)
+
+    assert loaded["array"] == array.tolist()

{xtrack-0.90.2 → xtrack-0.93.0}/tests/test_mad_writer.py

@@ -29,7 +29,6 @@ def test_mad_writer(case):
     mad_seq = line.to_madx_sequence(sequence_name='myseq')
 
     mad = Madx(stdout=True)
-    mad.options.rbarc = False
     mad.input(mad_seq)
     mad.beam(particle='proton', energy=7000e9)
     mad.use('myseq')

{xtrack-0.90.2 → xtrack-0.93.0}/tests/test_rbend_straight_body.py

@@ -1,6 +1,7 @@
 import xtrack as xt
 import numpy as np
 import xobjects as xo
+from cpymad.madx import Madx
 from xobjects.test_helpers import for_all_test_contexts
 
 import pathlib
@@ -1786,4 +1787,61 @@ def test_rbend_straight_body_v_shift():
         xo.assert_allclose(sv_test.Z, sv_curved.Z, atol=5e-14)
         xo.assert_allclose(sv_test.phi, sv_curved.phi, atol=5e-14)
         xo.assert_allclose(sv_test.psi, sv_curved.psi, atol=5e-14)
-        xo.assert_allclose(sv_test.theta, sv_curved.theta, atol=5e-14)
+        xo.assert_allclose(sv_test.theta, sv_curved.theta, atol=5e-14)
+
+def test_rbend_loading_writing_in_interfaces():
+
+    madx = Madx()
+
+    l_straight = 1.0
+    angle = 0.5
+    l_arc = l_straight / np.sinc(angle / np.pi / 2) # np.sinc is sin(pi*x)/(pi*x)
+
+    madx.input('''
+    b: rbend, l=1.0, angle=0.5;
+
+    ss: sequence, l=5.0, refer=centre;
+        b, at=2.5;
+    endsequence;
+    beam;
+    use, sequence=ss;
+    twiss, betx=1, bety=1;
+    survey;
+    ''')
+    sv_madx = xt.Table(madx.table.survey, _copy_cols=True)
+    xo.assert_allclose(madx.elements['b'].l, l_straight, rtol=1e-12)
+    xo.assert_allclose(sv_madx['s', 'b:1'], 2.5 + l_arc / 2, rtol=1e-12) # b si at exit
+
+    # Check cpymad loader
+    line = xt.Line.from_madx_sequence(madx.sequence.ss)
+    line.remove('ss$start')
+    line.remove('ss$end')
+    line.particle_ref = xt.Particles(p0c=1e9)
+    sv_xs = line.survey()
+    xo.assert_allclose(line['b'].length_straight, l_straight, rtol=1e-12)
+    xo.assert_allclose(line['b'].length, l_arc, rtol=1e-12)
+    xo.assert_allclose(sv_xs['s', 'b'], 2.5 - l_arc / 2, rtol=1e-12) # b si at entrance
+
+    # Check to madx
+    mad_str = line.to_madx_sequence(sequence_name='ggg')
+    madx2 = Madx()
+    madx2.input(mad_str)
+    madx2.input('''
+    beam;
+    use, sequence=ggg;
+    twiss, betx=1, bety=1;
+    survey;
+    ''')
+    sv_madx2 = xt.Table(madx2.table.survey, _copy_cols=True)
+    xo.assert_allclose(madx2.elements['b'].l, l_straight, rtol=1e-12)
+
+    # Check native madx loader
+    env2 = xt.load(string=mad_str, format='madx')
+    sv_xs2 = env2.ggg.survey()
+    xo.assert_allclose(env2['b'].length_straight, l_straight, rtol=1e-12)
+    xo.assert_allclose(env2['b'].length, l_arc, rtol=1e-12)
+    xo.assert_allclose(sv_xs2['s', 'b'], 2.5 - l_arc / 2, rtol=1e-12) # b si at entrance
+
+    # Check madng
+    sv_ng = line.madng_survey()
+    xo.assert_allclose(sv_ng['s', 'b'], 2.5 - l_arc / 2, rtol=1e-12) # b si at entrance

{xtrack-0.90.2 → xtrack-0.93.0}/tests/test_twiss.py

@@ -2058,4 +2058,4 @@ def test_twiss_disable_apertures():
     tw = line.twiss(betx=1, bety=1, x=0.2, y=0)
     xo.assert_allclose(tw.x[-1], 0.2)
     with pytest.raises(AssertionError):
-        line.twiss(betx=1, bety=1, x=0.2, y=0, disable_apertures=False)
+        line.twiss(betx=1, bety=1, x=0.2, y=0, disable_apertures=False)

xtrack-0.93.0/tests/test_twiss_to_file.py

@@ -0,0 +1,88 @@
+import xtrack as xt
+import xobjects as xo
+import numpy as np
+import pytest
+import pathlib
+
+test_data_folder = pathlib.Path(
+    __file__).parent.joinpath('../test_data').absolute()
+
+@pytest.mark.parametrize("check_type", ['csv', 'hdf5', 'tfs', 'json'])
+def test_twiss_table_file(check_type, tmp_path):
+
+    env = xt.load(test_data_folder / 'sps_thick/sps.seq')
+    env.vars.load(test_data_folder / 'sps_thick/lhc_q20.str')
+    line = env.sps
+    line.set_particle_ref('p', energy0=26e9)
+    tw = line.twiss4d()
+
+    if check_type == 'csv':
+        tw.to_csv(tmp_path / 'twiss_test.csv')
+        tw_test = xt.load(tmp_path / 'twiss_test.csv')
+    elif check_type == 'hdf5':
+        tw.to_hdf5(tmp_path / 'twiss_test.h5')
+        tw_test = xt.load(tmp_path / 'twiss_test.h5')
+    elif check_type == 'tfs':
+        tw.to_tfs(tmp_path / 'twiss_test.tfs')
+        tw_test = xt.load(tmp_path / 'twiss_test.tfs')
+    elif check_type == 'json':
+        tw.to_json(tmp_path / 'twiss_test.json')
+        tw_test = xt.load(tmp_path / 'twiss_test.json')
+    else:
+        raise ValueError(f'check_type {check_type} not supported')
+
+    assert isinstance(tw_test, xt.TwissTable)
+
+    assert np.all(np.array(tw_test._col_names) == np.array(tw._col_names))
+    if check_type == 'tfs':
+        assert set(tw_test.keys()) - set(tw.keys()) == {
+            '__class__', 'attrs_serialization', 'column_serialization', 'xtrack_version'
+        }
+    elif check_type == 'json':
+        assert set(tw_test.keys()) - set(tw.keys()) == {'xtrack_version'}
+    else:
+        assert set(tw_test.keys()) - set(tw.keys()) == {'__class__', 'xtrack_version'}
+    assert set(tw.keys()) - set(tw_test.keys()) == {'_action'}
+
+    for kk in tw._data:
+        if kk == '_action':
+            continue
+        if kk in ['particle_on_co', 'steps_r_matrix', 'line_config']:
+            continue # To be checked separately
+        if tw[kk] is None:
+            assert tw_test[kk] is None
+            continue
+        if isinstance(tw[kk], np.ndarray) and isinstance(tw[kk][0], str):
+            assert np.all(tw[kk] == tw_test[kk])
+            continue
+        elif isinstance(tw[kk], str):
+            assert tw[kk] == tw_test[kk]
+            continue
+        xo.assert_allclose(tw[kk], tw_test[kk], rtol=1e-7, atol=1e-15)
+
+    # Check particle_on_co
+    assert isinstance(tw.particle_on_co, xt.Particles)
+    assert isinstance(tw_test.particle_on_co, xt.Particles)
+    dct_ref = line.particle_ref.to_dict()
+    dct_test = tw_test.particle_on_co.to_dict()
+    for kk in dct_ref:
+        if isinstance(dct_ref[kk], np.ndarray):
+            xo.assert_allclose(dct_ref[kk], dct_test[kk], rtol=1e-10, atol=1e-15)
+        else:
+            assert dct_ref[kk] == dct_test[kk]
+
+    # Check steps_r_matrix
+    assert isinstance(tw.steps_r_matrix, dict)
+    assert isinstance(tw_test.steps_r_matrix, dict)
+    assert set(tw.steps_r_matrix.keys()) == set(tw_test.steps_r_matrix.keys())
+    for kk in tw.steps_r_matrix:
+        rmat_ref = tw.steps_r_matrix[kk]
+        rmat_test = tw_test.steps_r_matrix[kk]
+        xo.assert_allclose(rmat_ref, rmat_test, rtol=1e-10, atol=1e-15)
+
+    # Check line_config
+    assert isinstance(tw.line_config, dict)
+    assert isinstance(tw_test.line_config, dict)
+    assert set(tw.line_config.keys()) == set(tw_test.line_config.keys())
+    for kk in tw.line_config:
+        assert tw.line_config[kk] == tw_test.line_config[kk]

{xtrack-0.90.2 → xtrack-0.93.0}/xtrack/__init__.py

@@ -47,11 +47,12 @@ from .trajectory_correction import TrajectoryCorrection
 from .mad_parser.loader import load_madx_lattice
 from . import json
 from .load import load
+from .boris import BorisSpatialIntegrator
 
 from .multiline import Multiline
 
-from xdeps import Table, FunctionPieceWiseLinear
+from xdeps import FunctionPieceWiseLinear
+from .table import Table
 
 from ._version import __version__
 
-

xtrack-0.93.0/xtrack/_version.py

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

xtrack-0.93.0/xtrack/boris.py

@@ -0,0 +1,161 @@
+import numpy as np
+
+from scipy.constants import c as clight
+from scipy.constants import e as qe
+
+class BorisSpatialIntegrator:
+
+    isthick = True
+
+    def __init__(self, fieldmap_callable, s_start, s_end, n_steps, verbose=False):
+        self.fieldmap_callable = fieldmap_callable
+        self.s_start = s_start
+        self.s_end = s_end
+        self.ds = (s_end - s_start) / n_steps
+        self.n_steps = n_steps
+        self.verbose = verbose
+        self.length = s_end - s_start
+
+    def track(self, p):
+
+        x_log = []
+        y_log = []
+        z_log = []
+        s_in = p.s.copy()
+        p.s=self.s_start
+
+        for ii in range(self.n_steps):
+
+            if self.verbose:
+                print(f's_in = {s_in[0]:.3f} s_in_map = {p.s[0]:.3f}', end='\r', flush=True)
+
+            q0 = p.q0
+            x = p.x.copy()
+            y = p.y.copy()
+            z = p.s.copy()
+            px = p.px.copy()
+            py = p.py.copy()
+            delta = p.delta.copy()
+            p0c = p.p0c
+
+            charge0_coulomb = q0 * qe
+
+            P0 = p0c * qe / clight # in kg m/s
+            P = P0 * (1 + delta)
+
+            Px = px * P0
+            Py = py * P0
+
+            w = np.zeros((Px.shape[0], 3), dtype=Px.dtype)
+            w[:, 0] = Px
+            w[:, 1] = Py
+            w[:, 2] = P
+
+            x_new, y_new, z_new, w_new = step_spatial_boris_B(x, y, z, w,
+                charge0_coulomb, self.ds,
+                field_fn=self.fieldmap_callable)
+            p.x = x_new.copy()
+            p.y = y_new.copy()
+            p.s = z_new.copy()
+            p.px = w_new[:, 0] / P0
+            p.py = w_new[:, 1] / P0
+
+            x_log.append(p.x.copy())
+            y_log.append(p.y.copy())
+            z_log.append(p.s.copy())
+        p.s = s_in + self.length
+        self.x_log = np.array(x_log)
+        self.y_log = np.array(y_log)
+        self.z_log = np.array(z_log)
+
+
+import numpy as np
+c = 299_792_458.0  # m/s
+
+def step_spatial_boris_B(x, y, z, w, q, dz, field_fn):
+    """
+    Spatial Boris step for magnetic fields only (E = 0),
+    using constant total momentum magnitude P = w[:,2].
+
+    Parameters
+    ----------
+    x, y, z : (N,) arrays
+        Particle positions [m].
+    w : (N, 3) array
+        Momentum-like state: (px, py, P), where
+        P = |p| = sqrt(px^2 + py^2 + pz^2) is constant.
+    q : float
+        Particle charge [C].
+    dz : float
+        Spatial step along z [m].
+    field_fn : callable(x, y, z) -> (Bx, By, Bz)
+        Function returning magnetic-field components [T].
+
+    Returns
+    -------
+    x1, y1, z1 : (N,) arrays
+        Updated positions [m].
+    w1 : (N, 3) array
+        Updated momentum state (px, py, P).
+    """
+
+    # --- Unpack and ensure arrays
+    x = np.asarray(x)
+    y = np.asarray(y)
+    z = np.asarray(z)
+    w = np.asarray(w)
+    px, py, P = w[:, 0], w[:, 1], w[:, 2]
+
+    # --- Compute longitudinal momentum from constant magnitude
+    pz = np.sqrt(np.maximum(P**2 - px**2 - py**2, 0.0))
+
+    # --- Half drift in (x, y)
+    xh = x + (px / pz) * (dz * 0.5)
+    yh = y + (py / pz) * (dz * 0.5)
+    zh = z + dz * 0.5
+
+    # --- Evaluate magnetic field at mid-step
+    Bx, By, Bz = field_fn(xh, yh, zh)
+
+    # ============================================
+    # (1) FIRST HALF-KICK from (Bx, By)
+    # ============================================
+    pxm = px - 0.5 * q * dz * By
+    pym = py + 0.5 * q * dz * Bx
+
+    # --- Recompute pz to maintain |p| = P
+    pz = np.sqrt(np.maximum(P**2 - pxm**2 - pym**2, 0.0))
+
+    # ============================================
+    # (2) ROTATION due to Bz
+    # ============================================
+    t = 0.5 * q * Bz * dz / pz
+    t2 = t * t
+    s = 2.0 * t / (1.0 + t2)
+    c0 = (1.0 - t2) / (1.0 + t2)
+
+    # pxp = c0 * pxm - s * pym
+    # pyp = s * pxm + c0 * pym
+    pxp = c0 * pxm + s * pym
+    pyp = -s * pxm + c0 * pym
+
+    # ============================================
+    # (3) SECOND HALF-KICK from (Bx, By)
+    # ============================================
+    px1 = pxp - 0.5 * q * dz * By
+    py1 = pyp + 0.5 * q * dz * Bx
+
+    # --- Recompute pz after full step
+    pz1 = np.sqrt(np.maximum(P**2 - px1**2 - py1**2, 0.0))
+
+    # ============================================
+    # (4) SECOND HALF-DRIFT (x, y)
+    # ============================================
+    x1 = xh + (px1 / pz1) * (dz * 0.5)
+    y1 = yh + (py1 / pz1) * (dz * 0.5)
+    z1 = z + dz
+
+    # --- Pack updated (px, py, P)
+    w1 = np.stack([px1, py1, P], axis=1)
+
+    return x1, y1, z1, w1

{xtrack-0.90.2 → xtrack-0.93.0}/xtrack/environment.py

@@ -658,11 +658,10 @@ class Environment:
 
     def _get_a_drift_name(self):
         self._drift_counter += 1
-        nn = f'drift_{self._drift_counter}'
-        if nn not in self.element_dict:
-            return nn
-        else:
-            return self._get_a_drift_name()
+        while nn := f'drift_{self._drift_counter}':
+            if nn not in self.element_dict:
+                return nn
+            self._drift_counter += 1
 
     def __setitem__(self, key, value):
 
@@ -679,6 +678,7 @@ class Environment:
     def to_dict(self, include_var_management=True, include_version=True):
 
         out = {}
+        out['__class__'] = self.__class__.__name__
 
         if include_version:
             out["xtrack_version"] = xt.__version__

{xtrack-0.90.2 → xtrack-0.93.0}/xtrack/general.py

@@ -23,7 +23,7 @@ class _LOC:
 START = _LOC('START')
 END = _LOC('END')
 
-def read_url(url, timeout=0.1):
+def read_url(url, timeout=0.1, binary=False):
     """
     Read content from a URL.
     """
@@ -31,9 +31,16 @@ def read_url(url, timeout=0.1):
         response = requests.get(url, timeout=timeout)
         response.raise_for_status()  # Raise an error for bad responses
         if url.endswith('.gz'):
-            return gzip.decompress(response.content).decode("utf-8")
+            out = gzip.decompress(response.content)
+            if binary:
+                return out
+            else:
+                return out.decode("utf-8")
         else:
-            return response.text
+            if binary:
+                return response.content
+            else:
+                return response.text
     except requests.exceptions.RequestException as e:
         raise RuntimeError(f"Failed to read from URL {url}: {e}")