xtrack 0.82.0__tar.gz → 0.82.2__tar.gz

{xtrack-0.82.0/xtrack.egg-info → xtrack-0.82.2}/PKG-INFO

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

{xtrack-0.82.0 → xtrack-0.82.2}/tests/test_particles_basics.py

@@ -32,23 +32,22 @@ def _check_consistency_energy_variables(particles):
 
     beta = particles.beta0 * particles.rvv
     gamma = 1/np.sqrt(1 - beta**2)
-    pc = particles.mass0 * gamma * beta
+    pc = particles.mass * gamma * beta
 
     # Check consistency of delta with rvv
-    xo.assert_allclose(particles.delta, (pc-particles.p0c)/(particles.p0c),
+    xo.assert_allclose(particles.delta, (pc/particles.mass_ratio - particles.p0c)/(particles.p0c),
                        rtol=1e-14, atol=1e-14)
 
     # Check consistency of ptau with rvv
-    energy = particles.mass0 * gamma
-    xo.assert_allclose(particles.ptau, (energy - particles.energy0)/particles.p0c,
+    energy = particles.mass * gamma
+    xo.assert_allclose(particles.ptau, (energy/particles.mass_ratio - particles.energy0)/particles.p0c,
                        rtol=1e-14, atol=1e-14)
 
     # Check consistency of pzeta
-    energy = particles.mass0 * gamma
-    xo.assert_allclose(particles.pzeta, (energy - particles.energy0)/(particles.beta0 * particles.p0c),
+    energy = particles.mass * gamma
+    xo.assert_allclose(particles.pzeta, (energy/particles.mass_ratio  - particles.energy0)/(particles.beta0 * particles.p0c),
                        rtol=1e-14, atol=1e-14)
 
-
     # Check energy property
     xo.assert_allclose(particles.energy, energy, rtol=1e-14, atol=1e-14)
 
@@ -294,6 +293,25 @@ def test_python_add_to_energy(test_context):
 
     assert np.all(particles.zeta == zeta_before)
 
+    particles = xp.Particles(_context=test_context, mass0=xp.PROTON_MASS_EV,
+                             q0=1, p0c=1.4e9, x=[1e-3, 0], px=[1e-6, -1e-6],
+                             y=[0, 1e-3], py=[2e-6, 0], zeta=[1e-2, 2e-2],
+                             delta=[0, 1e-4], mass_ratio=[0.2, 1.7], charge_ratio=[0.5, 2])
+
+    energy_before = particles.copy(_context=xo.ContextCpu()).energy
+    zeta_before = particles.copy(_context=xo.ContextCpu()).zeta
+
+    particles.add_to_energy(3e6)
+
+    expected_energy = energy_before + 3e6
+    particles.move(_context=xo.ContextCpu())
+    xo.assert_allclose(particles.energy, expected_energy,
+                       atol=1e-14, rtol=1e-14)
+
+    _check_consistency_energy_variables(particles)
+
+    assert np.all(particles.zeta == zeta_before)
+
 
 @for_all_test_contexts
 def test_python_delta_setter(test_context):

{xtrack-0.82.0 → xtrack-0.82.2}/tests/test_twiss_vs_madx_psb.py

@@ -31,8 +31,8 @@ def test_twiss_psb(test_context):
     # With the approximation beta ~= beta0 we have delta ~= pzeta ~= 1/beta0 ptau
     # ==> ptau ~= beta0 delta ==> dptau / ddelta~= beta0
 
-    xo.assert_allclose(tw.dqx, twmad.summary.dq1 * beta0, rtol=0, atol=1e-5)
-    xo.assert_allclose(tw.dqy, twmad.summary.dq2 * beta0, rtol=0, atol=1e-5)
+    xo.assert_allclose(tw.dqx, twmad.summary.dq1 * beta0, rtol=0, atol=1e-3)
+    xo.assert_allclose(tw.dqy, twmad.summary.dq2 * beta0, rtol=0, atol=1e-3)
 
     dx_ref = np.interp(tw.s, twmad.s, twmad.dx * beta0)
     betx_ref = np.interp(tw.s, twmad.s, twmad.betx)

{xtrack-0.82.0 → xtrack-0.82.2}/tests/test_vs_madx.py

@@ -620,7 +620,7 @@ def test_low_beta_twiss(test_context):
     xo.assert_allclose(mad.table.summ['q1'][0], tw['qx'], rtol=0, atol=1e-6)
     xo.assert_allclose(mad.table.summ['q2'][0], tw['qy'], rtol=0, atol=1e-6)
     xo.assert_allclose(mad.table.summ['dq1'][0]*beta0, tw['dqx'], rtol=0,
-                        atol=1e-6)
+                        atol=1e-3)
     xo.assert_allclose(mad.table.summ['dq2'][0]*beta0, tw['dqy'], rtol=0,
-                        atol=1e-6)
+                        atol=1e-3)
     xo.assert_allclose(tw.qs, emitdf.qs.iloc[0], rtol=0, atol=1e-8)

xtrack-0.82.2/xtrack/_version.py

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

{xtrack-0.82.0 → xtrack-0.82.2}/xtrack/particles/particles.py

@@ -349,6 +349,14 @@ class Particles(xo.HybridClass):
             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'),
+            mask=input_mask,
+        )
+
         # Init energy deviations
         self._update_energy_deviations(
             delta=kwargs.get('delta'),
@@ -366,14 +374,6 @@ class Particles(xo.HybridClass):
             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'),
-            mask=input_mask,
-        )
-
         self.unhide_first_n_particles()
         if isinstance(self._context, xo.ContextCpu) and not _no_reorganize:
             self.reorganize()
@@ -1394,7 +1394,7 @@ class Particles(xo.HybridClass):
         Add `delta_energy` to the `energy` of the particles object. `delta`,
         'ptau', `rvv` and `rpp` are updated accordingly.
         """
-        self.ptau += delta_energy / self.p0c * self.mass_ratio
+        self.ptau += delta_energy / self.p0c / self.mass_ratio
 
     def set_particle(self, index, set_scalar_vars=False, **kwargs):
         raise NotImplementedError('This functionality has been removed')
@@ -1873,9 +1873,9 @@ class Particles(xo.HybridClass):
                         double const p0c = LocalParticle_get_p0c(part);
                         double const charge_ratio = LocalParticle_get_charge_ratio(part);
                         double const chi = LocalParticle_get_chi(part);
-                        double const mass_ratio = chi / charge_ratio;
-                        
-                        ptau += delta_energy/p0c * mass_ratio;
+                        double const mass_ratio = charge_ratio / chi;
+
+                        ptau += delta_energy/p0c / mass_ratio;
 
                         double const old_rpp = LocalParticle_get_rpp(part);
 
@@ -2083,6 +2083,9 @@ class Particles(xo.HybridClass):
             if _rpp is not None or _rvv is not None:
                 raise ValueError('Setting `delta` and `ptau` by only giving '
                                  '`_rpp` and `_rvv` is not supported.')
+            if any(self.mass_ratio != 1.0):
+                raise ValueError('Need to provide `delta` or `ptau` with '
+                                 'non-default mass ratios.')
             self._delta = 0.0
             delta = self._delta  # Cupy complains if we later assign LinkedArray
 

{xtrack-0.82.0 → xtrack-0.82.2}/xtrack/twiss.py

@@ -680,6 +680,43 @@ def twiss_line(line, particle_ref=None, method=None,
         _initial_particles=_initial_particles,
         _ebe_monitor=_ebe_monitor)
 
+
+
+    if (not only_orbit and (
+        (compute_chromatic_properties is True)
+        or (compute_chromatic_properties is None and periodic))):
+
+        with xt.line._preserve_config(line):
+            line.freeze_energy(force=True)
+            cols_chrom, scalars_chrom = _compute_chromatic_functions(
+                line=line,
+                init=init,
+                delta_chrom=delta_chrom,
+                steps_r_matrix=steps_r_matrix,
+                matrix_responsiveness_tol=matrix_responsiveness_tol,
+                matrix_stability_tol=matrix_stability_tol,
+                symplectify=symplectify,
+                method=method,
+                use_full_inverse=use_full_inverse,
+                nemitt_x=nemitt_x,
+                nemitt_y=nemitt_y,
+                on_momentum_twiss_res=twiss_res,
+                r_sigma=r_sigma,
+                delta_disp=delta_disp,
+                zeta_disp=zeta_disp,
+                start=start,
+                end=end,
+                num_turns=num_turns,
+                hide_thin_groups=hide_thin_groups,
+                only_markers=only_markers,
+                periodic=periodic,
+                periodic_mode=periodic_mode,
+                include_collective=include_collective,
+            )
+        twiss_res._data.update(cols_chrom)
+        twiss_res._data.update(scalars_chrom)
+        twiss_res._col_names += list(cols_chrom.keys())
+
     if not skip_global_quantities and not only_orbit:
         twiss_res._data['R_matrix'] = R_matrix
         twiss_res._data['steps_r_matrix'] = steps_r_matrix
@@ -691,39 +728,6 @@ def twiss_line(line, particle_ref=None, method=None,
         twiss_res._data['eigenvalues'] = eigenvalues.copy()
         twiss_res._data['rotation_matrix'] = Rot.copy()
 
-    if (not only_orbit and (
-        (compute_chromatic_properties is True)
-        or (compute_chromatic_properties is None and periodic))):
-
-        cols_chrom, scalars_chrom = _compute_chromatic_functions(
-            line=line,
-            init=init,
-            delta_chrom=delta_chrom,
-            steps_r_matrix=steps_r_matrix,
-            matrix_responsiveness_tol=matrix_responsiveness_tol,
-            matrix_stability_tol=matrix_stability_tol,
-            symplectify=symplectify,
-            method=method,
-            use_full_inverse=use_full_inverse,
-            nemitt_x=nemitt_x,
-            nemitt_y=nemitt_y,
-            on_momentum_twiss_res=twiss_res,
-            r_sigma=r_sigma,
-            delta_disp=delta_disp,
-            zeta_disp=zeta_disp,
-            start=start,
-            end=end,
-            num_turns=num_turns,
-            hide_thin_groups=hide_thin_groups,
-            only_markers=only_markers,
-            periodic=periodic,
-            periodic_mode=periodic_mode,
-            include_collective=include_collective,
-        )
-        twiss_res._data.update(cols_chrom)
-        twiss_res._data.update(scalars_chrom)
-        twiss_res._col_names += list(cols_chrom.keys())
-
     if eneloss_and_damping and not only_orbit:
         assert 'R_matrix' in twiss_res._data
         if radiation_method != 'full' or twiss_res._data['R_matrix_ebe'] is None:
@@ -814,14 +818,16 @@ def twiss_line(line, particle_ref=None, method=None,
         if ((twiss_res.orientation == 'forward' and not reverse)
                 or (twiss_res.orientation == 'backward' and reverse)):
             twiss_res.muzeta += init.muzeta - twiss_res.muzeta[0]
-            twiss_res.dzeta += init.dzeta - twiss_res.dzeta[0]
+            if 'dzeta' in twiss_res._data:
+                twiss_res.dzeta += init.dzeta - twiss_res.dzeta[0]
             if 'mux' in twiss_res._data:
                 twiss_res.mux += init.mux - twiss_res.mux[0]
                 twiss_res.muy += init.muy - twiss_res.muy[0]
         elif ((twiss_res.orientation == 'forward' and reverse)
             or (twiss_res.orientation == 'backward' and not reverse)):
             twiss_res.muzeta += init.muzeta - twiss_res.muzeta[-1]
-            twiss_res.dzeta += init.dzeta - twiss_res.dzeta[-1]
+            if 'dzeta' in twiss_res._data:
+                twiss_res.dzeta += init.dzeta - twiss_res.dzeta[-1]
             if 'mux' in twiss_res._data:
                 twiss_res.mux += init.mux - twiss_res.mux[-1]
                 twiss_res.muy += init.muy - twiss_res.muy[-1]
@@ -1013,14 +1019,6 @@ def _twiss_open(line, init,
     Ws[:, 4, :] -= 0.5 * (line.record_last_track.zeta[7:13, i_start:i_stop+1] - zeta_co).T / scale_eigen
     Ws[:, 5, :] -= 0.5 * (line.record_last_track.ptau[7:13, i_start:i_stop+1] - ptau_co).T / particle_on_co._xobject.beta0[0] / scale_eigen
 
-    dzeta = (((line.record_last_track.zeta[6, i_start:i_stop+1] - zeta_co).T
-            - (line.record_last_track.zeta[12, i_start:i_stop+1] - zeta_co).T )
-            / ((line.record_last_track.delta[6, i_start:i_stop+1] - delta_co).T
-            - (line.record_last_track.delta[12, i_start:i_stop+1] - delta_co).T))
-
-    dzeta -= dzeta[0]
-    dzeta = np.array(dzeta)
-
     name_co = np.array(line._element_names_unique[i_start:i_stop] + ('_end_point',))
     name_co_env = np.array(line.element_names[i_start:i_stop] + ('_end_point',))
 
@@ -1052,8 +1050,6 @@ def _twiss_open(line, init,
         lattice_functions, i_replace = _compute_lattice_functions(Ws, use_full_inverse, s_co)
         twiss_res_element_by_element.update(lattice_functions)
 
-    twiss_res_element_by_element['dzeta'] = dzeta
-
     extra_data = {}
     extra_data['only_markers'] = only_markers
     if _keep_tracking_data:
@@ -1067,7 +1063,7 @@ def _twiss_open(line, init,
         'x', 'px', 'y', 'py', 'zeta', 'delta', 'ptau',
         'betx', 'bety', 'alfx', 'alfy', 'gamx', 'gamy',
         'betx1', 'bety1', 'betx2', 'bety2',
-        'dx', 'dpx', 'dy', 'dzeta', 'dpy',
+        'dx', 'dpx', 'dy', 'dpy',
         ]
 
         for key in _vars_hide_changes:
@@ -1228,20 +1224,17 @@ def _compute_global_quantities(line, twiss_res):
         part_on_co = twiss_res['particle_on_co']
         W_matrix = twiss_res['W_matrix']
 
-        dzeta = twiss_res['dzeta']
-        eta = -dzeta[-1]/circumference
-        alpha = eta + 1/part_on_co._xobject.gamma0[0]**2
-
         beta0 = part_on_co._xobject.beta0[0]
         T_rev0 = circumference/clight/beta0
         bets0 = W_matrix[0, 4, 4]**2 + W_matrix[0, 4, 5]**2
-        if eta < 0: # below transition
-            bets0 = -bets0
-        ptau_co = twiss_res['ptau']
-
+        if 'slip_factor' in twiss_res:
+            if twiss_res['slip_factor'] < 0: # below transition
+                bets0 = -bets0
+        else:
+            bets0 = 0
 
         twiss_res._data.update({
-            'slip_factor': eta, 'momentum_compaction_factor': alpha, 'bets0': bets0,
+            'bets0': bets0,
             'circumference': circumference, 'T_rev0': T_rev0,
             'particle_on_co':part_on_co.copy(_context=xo.context_default),
             'gamma0': part_on_co._xobject.gamma0[0],
@@ -1439,12 +1432,22 @@ def _compute_chromatic_functions(line, init, delta_chrom, steps_r_matrix,
     dqx = dmux[-1]
     dqy = dmuy[-1]
 
-    cols_chrom = {'dmux': dmux, 'dmuy': dmuy,
+    dzeta = (tw_chrom_res[1].zeta - tw_chrom_res[0].zeta)/(2*delta_chrom)
+    dzeta -= dzeta[0]
+    dzeta = np.array(dzeta)
+
+    slip_factor = -dzeta[-1] / tw_chrom_res[0].circumference
+    momentum_compaction_factor = (slip_factor
+                        + 1/tw_chrom_res[0].particle_on_co._xobject.gamma0[0]**2)
+
+    cols_chrom = {'dmux': dmux, 'dmuy': dmuy, 'dzeta': dzeta,
                   'bx_chrom': bx_chrom, 'by_chrom': by_chrom,
                   'ax_chrom': ax_chrom, 'ay_chrom': ay_chrom,
                   'wx_chrom': wx_chrom, 'wy_chrom': wy_chrom,
                   }
-    scalars_chrom = {'dqx': dqx, 'dqy': dqy}
+    scalars_chrom = {'dqx': dqx, 'dqy': dqy, 'slip_factor': slip_factor,
+                     'momentum_compaction_factor': momentum_compaction_factor,
+                     }
 
     if on_momentum_twiss_res is not None:
         mux = on_momentum_twiss_res.mux
@@ -2049,6 +2052,8 @@ def _handle_loop_around(kwargs):
         tw_res.mux -= tw_res['mux', ele_name_init] - init.mux
         tw_res.muy -= tw_res['muy', ele_name_init] - init.muy
         tw_res.muzeta -= tw_res['muzeta', ele_name_init] - init.muzeta
+
+    if 'dzeta' in tw_res.keys():
         tw_res.dzeta -= tw_res['dzeta', ele_name_init] - init.dzeta
 
     # Not yet supported
@@ -2098,7 +2103,9 @@ def _handle_init_inside_range(kwargs):
     tw_res.mux -= tw_res['mux', ele_name_init] - init.mux
     tw_res.muy -= tw_res['muy', ele_name_init] - init.muy
     tw_res.muzeta -= tw_res['muzeta', ele_name_init] - init.muzeta
-    tw_res.dzeta -= tw_res['dzeta', ele_name_init] - init.dzeta
+
+    if 'dzeta' in tw_res:
+        tw_res.dzeta -= tw_res['dzeta', ele_name_init] - init.dzeta
 
     # Not correctly handled yet
     if 'dmux' in tw_res.keys():
@@ -3009,11 +3016,14 @@ class TwissTable(Table):
             mux = self.mux[at_element]
             muy = self.muy[at_element]
             muzeta = self.muzeta[at_element]
-            dzeta = self.dzeta[at_element]
         else:
             mux = 0
             muy = 0
             muzeta = 0
+
+        if 'dzeta' in self.keys():
+            dzeta = self.dzeta[at_element]
+        else:
             dzeta = 0
 
         return TwissInit(particle_on_co=part, W_matrix=W,
@@ -3342,7 +3352,8 @@ class TwissTable(Table):
             out.dpx = out.dpx
             out.dy = out.dy
             out.dpy = -out.dpy
-            out.dzeta = -out.dzeta
+            if 'dzeta' in out:
+                out.dzeta = -out.dzeta
 
             if 'dx_zeta' in out._col_names:
                 out.dx_zeta = out.dx_zeta
@@ -3365,10 +3376,12 @@ class TwissTable(Table):
             out.mux = out.mux[0] - out.mux
             out.muy = out.muy[0] - out.muy
             out.muzeta = out.muzeta[0] - out.muzeta
-            out.dzeta = out.dzeta[0] - out.dzeta
             out.phix = -out.phix
             out.phiy = -out.phiy
 
+            if 'dzeta' in out:
+                out.dzeta = out.dzeta[0] - out.dzeta
+
         if 'ax_chrom' in out._col_names:
             out.ax_chrom = -out.ax_chrom
             out.ay_chrom = -out.ay_chrom
@@ -3471,7 +3484,8 @@ class TwissTable(Table):
 
     def zero_at(self, name):
         for kk in CYCLICAL_QUANTITIES:
-            self[kk] -= self[kk, name]
+            if kk in self:
+                self[kk] -= self[kk, name]
 
     def target(self, tars=None, value=None, at=None, **kwargs):
         if value is None:

{xtrack-0.82.0 → xtrack-0.82.2/xtrack.egg-info}/PKG-INFO

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

xtrack-0.82.0/xtrack/_version.py

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