xtrack 0.94.0__tar.gz → 0.94.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.94.0
+Version: 0.94.1
 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.1}/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.1}/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.1}/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.1/xtrack/_version.py

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

{xtrack-0.94.0 → xtrack-0.94.1}/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.1}/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.1}/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.1}/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.1/xtrack.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.94.0
+Version: 0.94.1
 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'