Merge pull request #704 from rdemaria/expose-masses

expose masses

Author: giadarol

tests/test_particles_basics.py

@@ -697,3 +697,14 @@ class KickAngExact(xt.BeamElement):
     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/__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/environment.py

@@ -1861,8 +1861,7 @@ def __setattr__(self, key, value):
 
     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/particles/particles.py

@@ -179,6 +179,10 @@ def __init__(
             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 @@ def __init__(
 
         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 @@ def __init__(
         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 @@ def __init__(
         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 @@ def __init__(
         # 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 @@ def p0c(self):
     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 @@ def gamma0(self):
     def gamma0(self, value):
         self.gamma0[:] = value
 
-    
 
     def update_beta0(self, new_beta0):
 
@@ -1528,13 +1542,13 @@ def gen_local_particle_api(cls, mode='no_local_copy'):
         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,15 +1598,15 @@ def gen_local_particle_api(cls, mode='no_local_copy'):
         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('}')
 
         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,13 +1624,13 @@ def gen_local_particle_api(cls, mode='no_local_copy'):
                 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(f'    return p{xx}*rpp;')
+                    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('    return p{xx}*rpp;')
                 src_angles_lines.append('}')
                 src_angles_lines.append('')
 
@@ -1625,15 +1639,15 @@ def gen_local_particle_api(cls, mode='no_local_copy'):
                 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 @@ def gen_local_particle_api(cls, mode='no_local_copy'):
                 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 @@ def gen_local_particle_api(cls, mode='no_local_copy'):
                 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 @@ def _setattr_if_consistent(self, varname, given_value, computed_value,
         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 @@ def _update_refs(self, p0c=None, energy0=None, gamma0=None, beta0=None,
             _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.')