xtrack 0.94.0__tar.gz → 0.94.2__tar.gz

{xtrack-0.94.0/xtrack.egg-info → xtrack-0.94.2}/PKG-INFO

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

{xtrack-0.94.0 → xtrack-0.94.2}/tests/test_native_madloader_lhc.py

@@ -109,6 +109,10 @@ def test_native_loader_lhc(line_mode, data_mode, tmpdir, lines_ref):
     elif data_mode == 'copy':
         env = env.copy()
     elif data_mode == 'py':
+        # Force k0_from_h to False (they are all provided)
+        for nn in list(env.elements.keys()):
+            if hasattr(env.elements[nn], 'k0_from_h'):
+                env.elements[nn].k0_from_h = False
         lpg.write_py_lattice_file(env,
                                   output_fname=tmpdir / f'lhc_{line_mode}.py')
         env = xt.Environment()
@@ -381,18 +385,18 @@ def test_native_loader_lhc(line_mode, data_mode, tmpdir, lines_ref):
             if nn == '_end_point':
                 continue
             nn_straight = nn[:-len(f'/lhcb{beam}')] if nn.endswith(f'/lhcb{beam}') else nn
-            e2 = lref[nn_straight]
-            e4 = ltest[nn]
-            d2 = e2.to_dict()
-            d4 = e4.to_dict()
-            is_rbend = isinstance(e4, xt.RBend)
+            eref = lref[nn_straight]
+            etest = ltest[nn]
+            dref = eref.to_dict()
+            dtest = etest.to_dict()
+            is_rbend = isinstance(etest, xt.RBend)
 
-            for kk in d2.keys():
+            for kk in dref.keys():
                 if kk in ('__class__', 'model', 'side'):
-                    assert d2[kk] == d4[kk]
+                    assert dref[kk] == dtest[kk]
                     continue
 
-                if kk == '_isthick' and e2.length == 0:
+                if kk == '_isthick' and eref.length == 0:
                     continue  # Skip the check for zero-length elements
 
                 if kk in {
@@ -403,21 +407,21 @@ def test_native_loader_lhc(line_mode, data_mode, tmpdir, lines_ref):
                     continue
 
                 if kk in {'knl', 'ksl'}:
-                    maxlen = max(len(d2[kk]), len(d4[kk]))
-                    lhs = np.pad(d2[kk], (0, maxlen - len(d2[kk])), mode='constant')
-                    rhs = np.pad(d4[kk], (0, maxlen - len(d4[kk])), mode='constant')
+                    maxlen = max(len(dref[kk]), len(dtest[kk]))
+                    lhs = np.pad(dref[kk], (0, maxlen - len(dref[kk])), mode='constant')
+                    rhs = np.pad(dtest[kk], (0, maxlen - len(dtest[kk])), mode='constant')
                     xo.assert_allclose(lhs, rhs, rtol=1e-10, atol=1e-16)
                     continue
 
                 if is_rbend and kk in ('length', 'length_straight'):
-                    xo.assert_allclose(d2[kk], d4[kk], rtol=1e-7, atol=1e-6)
+                    xo.assert_allclose(dref[kk], dtest[kk], rtol=1e-7, atol=1e-6)
                     continue
 
                 if is_rbend and kk in ('h', 'k0'):
-                    xo.assert_allclose(d2[kk], d4[kk], rtol=1e-7, atol=5e-10)
+                    xo.assert_allclose(dref[kk], dtest[kk], rtol=1e-7, atol=5e-10)
                     continue
 
-                xo.assert_allclose(d2[kk], d4[kk], rtol=1e-10, atol=1e-16)
+                xo.assert_allclose(dref[kk], dtest[kk], rtol=1e-10, atol=1e-16)
 
         twref = lref.twiss4d()
         twtest = ltest.twiss4d()

{xtrack-0.94.0 → xtrack-0.94.2}/tests/test_particles_basics.py

@@ -697,3 +697,14 @@ def test_LocalParticle_angles(test_context):
     particles.move(_context=xo.ContextCpu())
     xo.assert_allclose(particles.px, [23.99302e-3, 18.01805e-3], atol=1e-14, rtol=5e-7)
     xo.assert_allclose(particles.py, [-1.81976e-3, -4.73529e-3], atol=1e-14, rtol=5e-7)
+
+
+@for_all_test_contexts
+def test_update_rigidity0(test_context):
+    p = xt.Particles("proton",p0c=7000e9)
+    pb82 = xt.Particles("Pb208",q0=82,rigidity0=p.rigidity0)
+    assert np.allclose(pb82.p0c/82,p.p0c)
+    pb82.rigidity0=p.rigidity0
+    assert np.allclose(pb82.p0c/82,p.p0c)
+
+

{xtrack-0.94.0 → xtrack-0.94.2}/xtrack/__init__.py

@@ -6,7 +6,7 @@
 from .general import _pkg_root, _print, START, END
 
 from .particles import (Particles, PROTON_MASS_EV, ELECTRON_MASS_EV,
-                        enable_pyheadtail_interface, disable_pyheadtail_interface)
+                        enable_pyheadtail_interface, disable_pyheadtail_interface, masses)
 
 from .base_element import BeamElement, Replica
 from .beam_elements import *

{xtrack-0.94.0 → xtrack-0.94.2}/xtrack/_temp/boris_and_solenoid_map/solenoid_field.py

@@ -1,7 +1,8 @@
 import numpy as np
 import scipy
-
-# https://par.nsf.gov/servlets/purl/10220882
+# Hampton et al., Closed-form expressions for the magnetic fields of rectangular
+# and circular finite-length solenoids and current loops
+# https://pubs.aip.org/aip/adv/article/10/6/065320/997382/
 
 def ellipp(n, m):
     """

{xtrack-0.94.0 → xtrack-0.94.2}/xtrack/_temp/python_lattice_writer/lattice_py_generation.py

@@ -6,7 +6,7 @@ from xtrack.mad_parser.loader import CONSTANTS
 
 formatter = xd.refs.CompactFormatter(scope=None)
 SKIP_PARAMS = ['order', 'model', '_edge_entry_model', '_edge_exit_model',
-               'k0_from_h', 'h']
+               'h']
 
 
 def _repr_arr_ref(arr_ref, formatter):
@@ -36,15 +36,17 @@ def _elem_to_tokens(env, nn, formatter):
     if hasattr(ee, 'ksl'):
         fields += ['ksl']
 
-    if isinstance(ee, xt.RBend):
-        assert 'length_straight' in fields
-        fields = [ff for ff in fields if ff != 'length']
-
     tt = env[nn].get_table()
     for kk in tt.name:
+        if kk in fields:
+            continue
         if tt['expr', kk] is not None and tt['expr', kk] != 'None':
             fields.append(kk)
 
+    if isinstance(ee, xt.RBend):
+        assert 'length_straight' in fields
+        fields = [ff for ff in fields if ff != 'length']
+
     params = []
     for kk in fields:
         if kk == '__class__':
@@ -55,6 +57,8 @@ def _elem_to_tokens(env, nn, formatter):
             continue
         if kk == 'prototype':
             continue
+        if kk == 'k0_from_h' and not ee.k0_from_h:
+            continue
         if kk == 'knl' or kk == 'ksl':
             arr_ref = getattr(ee_ref, kk)
             vv = _repr_arr_ref(arr_ref, formatter)
@@ -84,6 +88,13 @@ def _elem_to_tokens(env, nn, formatter):
 
 def write_py_lattice_file(env, output_fname):
 
+    file_content = gen_py_lattice(env)
+
+    with open(output_fname, 'w') as ff:
+        ff.write(file_content)
+
+def gen_py_lattice(env):
+
     ###################
     # Handle elements #
     ###################
@@ -301,5 +312,6 @@ def write_py_lattice_file(env, output_fname):
     builder_part,
     postamble])
 
-    with open(output_fname, 'w') as ff:
-        ff.write(file_content)
+    return file_content
+
+

xtrack-0.94.2/xtrack/_version.py

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

{xtrack-0.94.0 → xtrack-0.94.2}/xtrack/environment.py

@@ -1861,8 +1861,7 @@ class EnvParticleRef:
 
     def copy(self, **kwargs):
         return self._resolved.copy(**kwargs)
-
-
+   
 class EnvVars:
 
     def __init__(self, env):
@@ -2200,3 +2199,4 @@ def _disable_name_clash_checks(env):
         yield
     finally:
         env._enable_name_clash_check = old_value
+

{xtrack-0.94.0 → xtrack-0.94.2}/xtrack/mad_loader.py

@@ -810,8 +810,8 @@ class MadLoader:
 
         if self.enable_field_errors:
             kwargs = _prepare_field_errors_thick_elem(mad_el)
-            knl = kwargs['knl']
-            ksl = kwargs['ksl']
+            knl = kwargs.get('knl', [0, 0, 0])
+            ksl = kwargs.get('ksl', [0, 0, 0])
             num_multipole_kicks = 1
         else:
             knl = [0] * 3
@@ -988,6 +988,7 @@ class MadLoader:
     convert_vmonitor = convert_drift_like
     convert_collimator = convert_drift_like
     convert_rcollimator = convert_drift_like
+    convert_ecollimator = convert_drift_like
     convert_elseparator = convert_drift_like
     convert_instrument = convert_drift_like
 

{xtrack-0.94.0 → xtrack-0.94.2}/xtrack/mad_parser/loader.py

@@ -413,8 +413,8 @@ class MadxLoader:
                 if 'length' in params:
                     params['length_straight'] = params.pop('length')
 
-            # Default MAD-X behaviour is to take k0 from h only if k0 is not
-            # given. We need to replicate this behaviour. Ideally we should
+            # Default MAD-X behaviour is to take k0 from h only if k0 is zero.
+            # We need to replicate this behaviour. Ideally we should
             # evaluate expressions here, but that's tricky.
             if self._parameter_cache[name].get('k0', 0) == 0:
                 params['k0_from_h'] = True

{xtrack-0.94.0 → xtrack-0.94.2}/xtrack/mad_writer.py

@@ -565,6 +565,8 @@ def element_to_mad_str(
             tokens = xsuite_to_mad_converters[el._parent.__class__](eref, mad_type=mad_type, substituted_vars=substituted_vars)
             if isinstance(el, xt.beam_elements.slice_elements_edge._ThinSliceEdgeBase):
                 tokens.append(mad_assignment('kill_body', True, mad_type, substituted_vars=substituted_vars))
+        else:
+            raise NotImplementedError(f"Element of type {el.__class__} not supported yet in MAD writer")
     else:
         tokens = xsuite_to_mad_converters[el.__class__](eref, mad_type=mad_type, substituted_vars=substituted_vars)
 

{xtrack-0.94.0 → xtrack-0.94.2}/xtrack/particles/particles.py

@@ -179,6 +179,10 @@ class Particles(xo.HybridClass):
             Reference relativistic gamma
         beta0 : array_like of float, optional
             Reference relativistic beta
+        rigidity0 : array_like of float, optional
+            Reference magnetic rigidity [T.m]
+        kinetic_energy0 : array_like of float, optional
+            Reference kinetic energy [eV]
         mass_ratio : array_like of float, optional
             mass/mass0 (this is used to track particles of
             different species. Note that mass is the rest mass
@@ -224,7 +228,7 @@ class Particles(xo.HybridClass):
 
         accepted_args = set(self._xofields.keys()) | {
             'energy0', 'tau', 'pzeta', 'mass_ratio', 'mass', 'kinetic_energy0',
-            '_context', '_buffer', '_offset', 'p0', 'name',
+            '_context', '_buffer', '_offset', 'name', 'rigidity0',
         }
         if set(kwargs.keys()) - accepted_args:
             raise NameError(f'Invalid argument(s) provided: '
@@ -236,6 +240,8 @@ class Particles(xo.HybridClass):
         per_part_input_vars = (
             self.per_particle_vars +
             ((xo.Float64, 'energy0'),
+             (xo.Float64, 'kinetic_energy0'),
+             (xo.Float64, 'rigidity0'),
              (xo.Float64, 'tau'),
              (xo.Float64, 'pzeta'),
              (xo.Float64, 'mass_ratio'))
@@ -322,11 +328,6 @@ class Particles(xo.HybridClass):
         self.start_tracking_at_element = kwargs.get(
                             'start_tracking_at_element', -1)
 
-        # Init refs
-        if 'kinetic_energy0' in kwargs.keys():
-            assert kwargs.get('energy0') is None
-            kwargs['energy0'] = kwargs.pop('kinetic_energy0') + self.mass0
-
         # Ensure that all per particle inputs are numpy arrays of the same
         # length, and move them to the target context
         for xotype, field in per_part_input_vars:
@@ -351,20 +352,22 @@ class Particles(xo.HybridClass):
         # Init independent per particle vars
         self.init_independent_per_part_vars(kwargs)
 
+        # Init chi and charge ratio
+        self._update_chi_charge_ratio(
+            chi=kwargs.get('chi'),
+            charge_ratio=kwargs.get('charge_ratio'),
+            mass_ratio=kwargs.get('mass_ratio'),
+            mask=input_mask,
+        )
 
+        # Init reference momentum and related vars
         self._update_refs(
             p0c=kwargs.get('p0c'),
             energy0=kwargs.get('energy0'),
             gamma0=kwargs.get('gamma0'),
             beta0=kwargs.get('beta0'),
-            mask=input_mask,
-        )
-
-        # Init chi and charge ratio
-        self._update_chi_charge_ratio(
-            chi=kwargs.get('chi'),
-            charge_ratio=kwargs.get('charge_ratio'),
-            mass_ratio=kwargs.get('mass_ratio'),
+            kinetic_energy0=kwargs.get('kinetic_energy0'),
+            rigidity0=kwargs.get('rigidity0'),
             mask=input_mask,
         )
 
@@ -1235,13 +1238,25 @@ class Particles(xo.HybridClass):
     def p0c(self, value):
         self.p0c[:] = value
 
+    def _rigidity0_setitem(self, indx, val):
+        ctx = self._buffer.context
+        temp_rigidity0 = ctx.zeros(shape=self._p0c.shape, dtype=np.float64)
+        temp_rigidity0[:] = np.nan
+        temp_rigidity0[indx] = val
+        self.update_p0c(temp_rigidity0 * clight * self.q0)
+
     @property
     def rigidity0(self):
         rigidity0 = self.p0c / clight / self.q0
         return self._buffer.context.linked_array_type.from_array(
             rigidity0,
-            mode='readonly',
-            container=self)
+            mode='setitem_from_container',
+            container=self,
+            container_setitem_name='_rigidity0_setitem')
+
+    @rigidity0.setter
+    def rigidity0(self, value):
+        self.rigidity0[:] = value
 
     def update_gamma0(self, new_gamma0):
 
@@ -1277,7 +1292,6 @@ class Particles(xo.HybridClass):
     def gamma0(self, value):
         self.gamma0[:] = value
 
-    
 
     def update_beta0(self, new_beta0):
 
@@ -1528,13 +1542,13 @@ class Particles(xo.HybridClass):
         src_lines.append('if (set_scalar){')
         for _, vv in cls.size_vars + cls.scalar_vars:
             src_lines.append(
-                f'  ParticlesData_set_' + vv + '(dest,'
+                '  ParticlesData_set_' + vv + '(dest,'
                                                f'      LocalParticle_get_{vv}(source));')
         src_lines.append('}')
 
         for _, vv in cls.per_particle_vars:
             src_lines.append(
-                f'  ParticlesData_set_' + vv + '(dest, id, '
+                '  ParticlesData_set_' + vv + '(dest, id, '
                                                f'      LocalParticle_get_{vv}(source));')
         src_lines.append('}')
         src_local_to_particles = '\n'.join(src_lines)
@@ -1584,7 +1598,7 @@ class Particles(xo.HybridClass):
         for tt, vv in cls.size_vars + cls.scalar_vars:
             src_lines.append('/*gpufun*/')
             src_lines.append(f'{tt._c_type} LocalParticle_get_' + vv
-                             + f'(LocalParticle* part)'
+                             + '(LocalParticle* part)'
                              + '{')
             src_lines.append(f'  return part->{vv};')
             src_lines.append('}')
@@ -1592,7 +1606,7 @@ class Particles(xo.HybridClass):
         for tt, vv in cls.per_particle_vars:
             src_lines.append('/*gpufun*/')
             src_lines.append(f'{tt._c_type} LocalParticle_get_' + vv
-                             + f'(LocalParticle* part)'
+                             + '(LocalParticle* part)'
                              + '{')
             src_lines.append(f'  return part->{vv}[part->ipart];')
             src_lines.append('}')
@@ -1610,12 +1624,12 @@ class Particles(xo.HybridClass):
                 src_angles_lines.append(f'    double const p{xx} = LocalParticle_get_p{xx}(part);')
                 if exact == 'exact_':
                     src_angles_lines.append(f'    double const p{yy} = LocalParticle_get_p{yy}(part);')
-                    src_angles_lines.append(f'    double const one_plus_delta = 1. + LocalParticle_get_delta(part);')
+                    src_angles_lines.append('    double const one_plus_delta = 1. + LocalParticle_get_delta(part);')
                     src_angles_lines.append(
-                        f'    double const rpp = 1./sqrt(one_plus_delta*one_plus_delta - px*px - py*py);')
+                        '    double const rpp = 1./sqrt(one_plus_delta*one_plus_delta - px*px - py*py);')
                 else:
-                    src_angles_lines.append(f'    double const rpp = LocalParticle_get_rpp(part);')
-                src_angles_lines.append(f'    // INFO: this is not the angle, but sin(angle)')
+                    src_angles_lines.append('    double const rpp = LocalParticle_get_rpp(part);')
+                src_angles_lines.append('    // INFO: this is not the angle, but sin(angle)')
                 src_angles_lines.append(f'    return p{xx}*rpp;')
                 src_angles_lines.append('}')
                 src_angles_lines.append('')
@@ -1625,15 +1639,15 @@ class Particles(xo.HybridClass):
                 src_angles_lines.append('/*gpufun*/')
                 src_angles_lines.append(f'void LocalParticle_set_{exact}{xx}p(LocalParticle* part, double {xx}p){{')
                 src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}')
-                src_angles_lines.append(f'    double rpp = LocalParticle_get_rpp(part);')
+                src_angles_lines.append('    double rpp = LocalParticle_get_rpp(part);')
                 if exact == 'exact_':
                     src_angles_lines.append(
                         f'    // Careful! If {yy}p also changes, use LocalParticle_set_{exact}xp_yp!')
                     src_angles_lines.append(f'    double const {yy}p = LocalParticle_get_{exact}{yy}p(part);')
-                    src_angles_lines.append(f'    rpp *= sqrt(1 + xp*xp + yp*yp);')
+                    src_angles_lines.append('    rpp *= sqrt(1 + xp*xp + yp*yp);')
                 src_angles_lines.append(f'    // INFO: {xx}p is not the angle, but sin(angle)')
                 src_angles_lines.append(f'    LocalParticle_set_p{xx}(part, {xx}p/rpp);')
-                src_angles_lines.append(f'#endif')
+                src_angles_lines.append('#endif')
                 src_angles_lines.append('}')
                 src_angles_lines.append('')
 
@@ -1644,7 +1658,7 @@ class Particles(xo.HybridClass):
                 src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}')
                 src_angles_lines.append(f'    LocalParticle_set_{exact}{xx}p(part, '
                                         + f'LocalParticle_get_{exact}{xx}p(part) + {xx}p);')
-                src_angles_lines.append(f'#endif')
+                src_angles_lines.append('#endif')
                 src_angles_lines.append('}')
                 src_angles_lines.append('')
                 # Scaler
@@ -1653,19 +1667,19 @@ class Particles(xo.HybridClass):
                 src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}')
                 src_angles_lines.append(f'    LocalParticle_set_{exact}{xx}p(part, '
                                         + f'LocalParticle_get_{exact}{xx}p(part) * value);')
-                src_angles_lines.append(f'#endif')
+                src_angles_lines.append('#endif')
                 src_angles_lines.append('}')
                 src_angles_lines.append('')
             # Double setter, adder, scaler
             src_angles_lines.append('/*gpufun*/')
             src_angles_lines.append(f'void LocalParticle_set_{exact}xp_yp(LocalParticle* part, double xp, double yp){{')
-            src_angles_lines.append(f'    double rpp = LocalParticle_get_rpp(part);')
+            src_angles_lines.append('    double rpp = LocalParticle_get_rpp(part);')
             if exact == 'exact_':
-                src_angles_lines.append(f'    rpp *= sqrt(1 + xp*xp + yp*yp);')
+                src_angles_lines.append('    rpp *= sqrt(1 + xp*xp + yp*yp);')
             for xx in ['x', 'y']:
                 src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}')
                 src_angles_lines.append(f'    LocalParticle_set_p{xx}(part, {xx}p/rpp);')
-                src_angles_lines.append(f'#endif')
+                src_angles_lines.append('#endif')
             src_angles_lines.append('}')
             src_angles_lines.append('')
             src_angles_lines.append('/*gpufun*/')
@@ -2099,8 +2113,10 @@ class Particles(xo.HybridClass):
         getattr(self, varname)[mask] = target_val[mask]
 
     def _update_refs(self, p0c=None, energy0=None, gamma0=None, beta0=None,
+                     kinetic_energy0=None, rigidity0=None,
                      mask=None):
-        if not any(ff is not None for ff in (p0c, energy0, gamma0, beta0)):
+        if not any(ff is not None for ff in (p0c, energy0, gamma0, beta0, 
+                                            kinetic_energy0, rigidity0)):
             self._p0c = 1e9
             p0c = self._p0c
 
@@ -2125,6 +2141,16 @@ class Particles(xo.HybridClass):
             _energy0 = self.mass0 * _gamma0
             _p0c = _energy0 * beta0
             _beta0 = beta0
+        elif kinetic_energy0 is not None:
+            _energy0 = kinetic_energy0 + self.mass0
+            _p0c = _sqrt(_energy0 ** 2 - self.mass0 ** 2)
+            _beta0 = _p0c / _energy0
+            _gamma0 = _energy0 / self.mass0
+        elif rigidity0 is not None:
+            _p0c = rigidity0 * abs(self.q0) * clight
+            _energy0 = _sqrt(_p0c ** 2 + self.mass0 ** 2)
+            _beta0 = _p0c / _energy0
+            _gamma0 = _energy0 / self.mass0
         else:
             raise RuntimeError('This statement is unreachable.')
 

{xtrack-0.94.0 → xtrack-0.94.2/xtrack.egg-info}/PKG-INFO

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

xtrack-0.94.0/xtrack/_version.py

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