xtrack 0.89.0__tar.gz → 0.89.1__tar.gz

{xtrack-0.89.0/xtrack.egg-info → xtrack-0.89.1}/PKG-INFO

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

{xtrack-0.89.0 → xtrack-0.89.1}/tests/test_madnginterface.py

@@ -126,4 +126,41 @@ def test_madng_survey():
     xo.assert_allclose(survey.angle, xsurvey.angle, atol=1e-5, rtol=0)
     xo.assert_allclose(survey.tilt, xsurvey.rot_s_rad, atol=1e-5, rtol=0)
     # Length doesn't work because of multipoles.
-    #xo.assert_allclose(survey.length, xsurvey.length, atol=1e-5, rtol=0)
+    #xo.assert_allclose(survey.length, xsurvey.length, atol=1e-5, rtol=0)
+
+def test_madng_conversion_drift_slice():
+
+    env = xt.Environment()
+    env.particle_ref = xt.Particles(p0c=1e9)
+
+    line = env.new_line(length=10, components=[
+        env.new('q1', xt.Quadrupole, length=1, k1=0.3, at=4),
+        env.new('q2', xt.Quadrupole, length=1, k1=-0.3, at=6)
+    ])
+
+    line.insert('m', xt.Marker(), at=2)
+
+    tt = line.get_table()
+    # is:
+    # Table: 8 rows, 11 cols
+    # name                   s element_type isthick isreplica parent_name ...
+    # drift_1..0             0 DriftSlice      True     False drift_1
+    # m                      2 Marker         False     False None
+    # drift_1..1             2 DriftSlice      True     False drift_1
+    # q1                   3.5 Quadrupole      True     False None
+    # drift_2              4.5 Drift           True     False None
+    # q2                   5.5 Quadrupole      True     False None
+    # drift_3              6.5 Drift           True     False None
+    # _end_point            10                False     False None
+
+    assert np.all(tt.name == [
+        'drift_1..0', 'm', 'drift_1..1', 'q1', 'drift_2', 'q2', 'drift_3', '_end_point'])
+    xo.assert_allclose(tt.s, [0, 2, 2, 3.5, 4.5, 5.5, 6.5, 10], atol=1e-10)
+    assert np.all(tt.element_type == [
+        'DriftSlice', 'Marker', 'DriftSlice', 'Quadrupole', 'Drift', 'Quadrupole', 'Drift', ''])
+
+    tw = line.twiss4d()
+    tw_ng = line.madng_twiss(normal_form=False)
+
+    xo.assert_allclose(tw_ng.beta11_ng, tw.betx, rtol=1e-8)
+    xo.assert_allclose(tw_ng.beta22_ng, tw.bety, rtol=1e-8)

xtrack-0.89.1/xtrack/_version.py

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

{xtrack-0.89.0 → xtrack-0.89.1}/xtrack/beam_elements/elements_src/linesegmentmap.h

@@ -204,6 +204,8 @@ void longitudinal_motion(LocalParticle *part0,
         double const bucket_length = LineSegmentMapData_get_bucket_length(el)*LocalParticle_get_beta0(part0)*C_LIGHT;
         double const sin_s = sin(2 * PI * qs);
         double const cos_s = cos(2 * PI * qs);
+        uint8_t kill_energy_kick = LocalParticle_check_track_flag(
+                    part0, XS_FLAG_KILL_CAVITY_KICK);
         START_PER_PARTICLE_BLOCK(part0, part);
             // We set cos_s = 999 if long map is to be skipped
             double shift = 0.0;
@@ -215,7 +217,9 @@ void longitudinal_motion(LocalParticle *part0,
             double const new_pzeta = sin_s * LocalParticle_get_zeta(part) / bets + cos_s * LocalParticle_get_pzeta(part);
 
             LocalParticle_set_zeta(part, new_zeta+shift);
-            LocalParticle_update_pzeta(part, new_pzeta);
+            if (!kill_energy_kick){
+                LocalParticle_update_pzeta(part, new_pzeta);
+            }
         END_PER_PARTICLE_BLOCK;
     }
     else if (mode_flag==2){ // non-linear motion
@@ -325,6 +329,8 @@ GPUFUN
 void correlated_radiation_damping(LocalParticle *part0,
             LineSegmentMapData el){
 
+    uint8_t kill_energy_kick = LocalParticle_check_track_flag(
+                    part0, XS_FLAG_KILL_CAVITY_KICK);
     START_PER_PARTICLE_BLOCK(part0, part);
         double in[6];
         in[0] = LocalParticle_get_x(part);
@@ -345,7 +351,8 @@ void correlated_radiation_damping(LocalParticle *part0,
         LocalParticle_set_y(part, out[2]);
         LocalParticle_set_py(part, out[3]);
         LocalParticle_set_zeta(part, out[4]);
-        LocalParticle_update_pzeta(part,out[5]);
+        if (!kill_energy_kick){
+            LocalParticle_update_pzeta(part,out[5]);}
     END_PER_PARTICLE_BLOCK;
 }
 

{xtrack-0.89.0 → xtrack-0.89.1}/xtrack/beam_elements/elements_src/track_magnet.template.h

@@ -348,7 +348,7 @@ void track_magnet_particles(
         if (rbend_model == 2){
             // straight body --> curvature in the edges
             START_PER_PARTICLE_BLOCK(part0, part);
-                YRotation_single_particle(part, sin_rbha, cos_rbha, sin_rbha/cos_rbha);
+                YRotation_single_particle(part, -sin_rbha, cos_rbha, -sin_rbha/cos_rbha);
                 LocalParticle_add_to_x(part, -dx_rb);
             END_PER_PARTICLE_BLOCK;
         }
@@ -500,7 +500,7 @@ void track_magnet_particles(
             // straight body --> curvature in the edges
             START_PER_PARTICLE_BLOCK(part0, part);
                 LocalParticle_add_to_x(part, dx_rb); // shift by half sagitta
-                YRotation_single_particle(part, sin_rbha, cos_rbha, sin_rbha/cos_rbha);
+                YRotation_single_particle(part, -sin_rbha, cos_rbha, -sin_rbha/cos_rbha);
             END_PER_PARTICLE_BLOCK;
         }
     }

{xtrack-0.89.0 → xtrack-0.89.1}/xtrack/beam_elements/elements_src/track_rf.h

@@ -28,7 +28,8 @@ void track_rf_kick_single_particle(
     GPUGLMEM const double* knl,
     GPUGLMEM const double* ksl,
     GPUGLMEM const double* pn,
-    GPUGLMEM const double* ps
+    GPUGLMEM const double* ps,
+    const uint8_t kill_energy_kick
 ){
 
     double phase0 = 0;
@@ -140,11 +141,13 @@ void track_rf_kick_single_particle(
     }
 
 
-    #ifdef XTRACK_CAVITY_PRESERVE_ANGLE
-    LocalParticle_add_to_energy(part, energy_kick + rfmultipole_energy_kick, 0);
-    #else
-    LocalParticle_add_to_energy(part, energy_kick + rfmultipole_energy_kick, 1);
-    #endif
+    if (!kill_energy_kick) {
+        #ifdef XTRACK_CAVITY_PRESERVE_ANGLE
+        LocalParticle_add_to_energy(part, energy_kick + rfmultipole_energy_kick, 0);
+        #else
+        LocalParticle_add_to_energy(part, energy_kick + rfmultipole_energy_kick, 1);
+        #endif
+    }
 
 }
 
@@ -167,7 +170,8 @@ void track_rf_body_single_particle(
     GPUGLMEM const double* ps,
     const int64_t num_kicks,
     const int8_t drift_model,
-    const int8_t integrator
+    const int8_t integrator,
+    const uint8_t kill_energy_kick
 ) {
 
     #define RF_KICK(part, kick_weight) \
@@ -175,7 +179,8 @@ void track_rf_body_single_particle(
             part, voltage * (kick_weight), frequency, lag,\
             transverse_voltage * (kick_weight), transverse_lag,\
             absolute_time, order, \
-            factor_knl_ksl * (kick_weight), knl, ksl, pn, ps\
+            factor_knl_ksl * (kick_weight), knl, ksl, pn, ps,\
+            kill_energy_kick\
         )
 
     #define RF_DRIFT(part, dlength) \
@@ -318,6 +323,8 @@ void track_rf_particles(
 ) {
 
     double factor_knl_ksl = 1.0;
+    uint8_t kill_energy_kick = LocalParticle_check_track_flag(
+                    part0, XS_FLAG_KILL_CAVITY_KICK);
 
     // Backtracking
     #ifdef XSUITE_BACKTRACK
@@ -397,7 +404,8 @@ void track_rf_particles(
                 ps,
                 num_kicks,
                 drift_model,
-                integrator
+                integrator,
+                kill_energy_kick
             );
         END_PER_PARTICLE_BLOCK;
 

{xtrack-0.89.0 → xtrack-0.89.1}/xtrack/beam_elements/elements_src/track_rf.template.h

@@ -28,7 +28,8 @@ void track_rf_kick_single_particle(
     GPUGLMEM const double* knl,
     GPUGLMEM const double* ksl,
     GPUGLMEM const double* pn,
-    GPUGLMEM const double* ps
+    GPUGLMEM const double* ps,
+    const uint8_t kill_energy_kick
 ){
 
     double phase0 = 0;
@@ -140,11 +141,13 @@ void track_rf_kick_single_particle(
     }
 
 
-    #ifdef XTRACK_CAVITY_PRESERVE_ANGLE
-    LocalParticle_add_to_energy(part, energy_kick + rfmultipole_energy_kick, 0);
-    #else
-    LocalParticle_add_to_energy(part, energy_kick + rfmultipole_energy_kick, 1);
-    #endif
+    if (!kill_energy_kick) {
+        #ifdef XTRACK_CAVITY_PRESERVE_ANGLE
+        LocalParticle_add_to_energy(part, energy_kick + rfmultipole_energy_kick, 0);
+        #else
+        LocalParticle_add_to_energy(part, energy_kick + rfmultipole_energy_kick, 1);
+        #endif
+    }
 
 }
 
@@ -167,7 +170,8 @@ void track_rf_body_single_particle(
     GPUGLMEM const double* ps,
     const int64_t num_kicks,
     const int8_t drift_model,
-    const int8_t integrator
+    const int8_t integrator,
+    const uint8_t kill_energy_kick
 ) {
 
     #define RF_KICK(part, kick_weight) \
@@ -175,7 +179,8 @@ void track_rf_body_single_particle(
             part, voltage * (kick_weight), frequency, lag,\
             transverse_voltage * (kick_weight), transverse_lag,\
             absolute_time, order, \
-            factor_knl_ksl * (kick_weight), knl, ksl, pn, ps\
+            factor_knl_ksl * (kick_weight), knl, ksl, pn, ps,\
+            kill_energy_kick\
         )
 
     #define RF_DRIFT(part, dlength) \
@@ -237,6 +242,8 @@ void track_rf_particles(
 ) {
 
     double factor_knl_ksl = 1.0;
+    uint8_t kill_energy_kick = LocalParticle_check_track_flag(
+                    part0, XS_FLAG_KILL_CAVITY_KICK);
 
     // Backtracking
     #ifdef XSUITE_BACKTRACK
@@ -316,7 +323,8 @@ void track_rf_particles(
                 ps,
                 num_kicks,
                 drift_model,
-                integrator
+                integrator,
+                kill_energy_kick
             );
         END_PER_PARTICLE_BLOCK;
 

{xtrack-0.89.0 → xtrack-0.89.1}/xtrack/line.py

@@ -5513,6 +5513,15 @@ def _preserve_config(ln_or_trk):
         ln_or_trk.config.clear()
         ln_or_trk.config.update(config)
 
+@contextmanager
+def _preserve_track_flags(line):
+    old_flags = line.tracker.track_flags.flags.copy()
+    try:
+        yield
+    finally:
+        line.tracker.track_flags.flags.clear()
+        line.tracker.track_flags.flags.update(old_flags)
+
 
 @contextmanager
 def freeze_longitudinal(ln_or_trk):

{xtrack-0.89.0 → xtrack-0.89.1}/xtrack/mad_writer.py

@@ -232,6 +232,34 @@ def drift_to_mad_str(name, line, mad_type=MadType.MADX, substituted_vars=None):
         tokens = _handle_tokens_madng(tokens, substituted_vars)
     return ', '.join(tokens)
 
+def drift_slice_to_mad_str(name, line, mad_type=MadType.MADX, substituted_vars=None):
+    """
+    Convert a drift element to a MADX/MAD-NG string representation.
+
+    Parameters:
+    - name: Name of the drift element.
+    - line: The line containing the element.
+    - mad_type: Type of MAD (MADX or MADNG).
+    - substituted_vars: List of substituted variables for MADNG.
+
+    Returns:
+    - A string representation of the drift in MADX/MAD-NG format.
+    """
+
+    drift_slice = _get_eref(line, name)
+    tokens = []
+    tokens.append('drift')
+    if mad_type == MadType.MADNG:
+        tokens.append(f"'{name.replace(':', '__')}'")  # replace ':' with '__' for MADNG
+
+    tokens.append(mad_assignment('l', (float(drift_slice._parent.length._value)
+                                     * float(drift_slice.weight._value)),
+                                 mad_type, substituted_vars=substituted_vars))
+
+    if mad_type == MadType.MADNG:
+        tokens = _handle_tokens_madng(tokens, substituted_vars)
+    return ', '.join(tokens)
+
 def multipole_to_mad_str(name, line, mad_type=MadType.MADX, substituted_vars=None):
     """ Convert a multipole element to a MADX/MAD-NG string representation.
 
@@ -562,7 +590,8 @@ xsuite_to_mad_converters = {
     xt.UniformSolenoid: solenoid_to_mad_str,
     xt.SRotation: srotation_to_mad_str,
     xt.RFMultipole: rfmultipole_to_mad_str,
-    xt.CrabCavity: crabcavity_to_mad_str
+    xt.CrabCavity: crabcavity_to_mad_str,
+    xt.DriftSlice: drift_slice_to_mad_str
 }
 
 def to_madx_sequence(line, name='seq', mode='sequence'):

{xtrack-0.89.0 → xtrack-0.89.1}/xtrack/particles/particles.py

@@ -1429,6 +1429,7 @@ class Particles(xo.HybridClass):
 
         src_lines.append('                 int64_t ipart;')
         src_lines.append('                 int64_t endpart;')
+        src_lines.append('                 uint64_t track_flags;')
         src_lines.append('    /*gpuglmem*/ int8_t* io_buffer;')
         src_lines.append('} LocalParticle;')
         src_typedef = '\n'.join(src_lines)
@@ -1443,6 +1444,14 @@ class Particles(xo.HybridClass):
 
             ''')
 
+        # Get track flag
+        src_lines.append('''
+            /*gpufun*/
+            uint64_t LocalParticle_check_track_flag(LocalParticle* part, uint8_t index){
+                return (part->track_flags >> index) & 1;
+            }
+        ''')
+
         # Particles_to_LocalParticle
         src_lines.append('''
             /*gpufun*/

xtrack-0.89.1/xtrack/track_flags.py

@@ -0,0 +1,68 @@
+import numpy as np
+
+flag_mapping = {
+    'XS_FLAG_BACKTRACK': 0,
+    'XS_FLAG_TAPER': 1,
+    'XS_FLAG_KILL_CAVITY_KICK': 2,
+    'XS_FLAG_IGNORE_GLOBAL_APERTURE': 3,
+    'XS_FLAG_IGNORE_LOCAL_APERTURE': 4
+}
+
+flag_defaults = {
+    'XS_FLAG_BACKTRACK': False,
+    'XS_FLAG_TAPER': False,
+    'XS_FLAG_KILL_CAVITY_KICK': False,
+    'XS_FLAG_IGNORE_GLOBAL_APERTURE': False,
+    'XS_FLAG_IGNORE_LOCAL_APERTURE': False
+}
+
+c_header_flag_mapping = """
+#ifndef XSTUITE_TRACK_FLAGS_H
+#define XSTUITE_TRACK_FLAGS_H
+"""
+for flag_name, bit_pos in flag_mapping.items():
+    c_header_flag_mapping += f"#define {flag_name} ({bit_pos})\n"
+c_header_flag_mapping += """
+#endif // XSTUITE_TRACK_FLAGS_H
+"""
+
+class TrackFlags:
+
+    c_header_flag_mapping = c_header_flag_mapping
+
+    def __init__(self, **kwargs):
+        object.__setattr__(self, 'flags', flag_defaults.copy())
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def get_flags_register(self):
+
+        reg = 0
+        for flag_name, bit_pos in flag_mapping.items():
+            assert bit_pos >= 0
+            assert bit_pos < 64
+            if self.flags.get(flag_name, False):
+                reg |= (1 << bit_pos)
+        return np.uint64(reg)
+
+    def __getattr__(self, name):
+        if name == 'flags':
+            return object.__getattribute__(self, 'flags')
+        if name in self.flags:
+            return self.flags[name]
+        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+    def __setattr__(self, name, value):
+        if name == 'flags': # for unpickling
+            object.__setattr__(self, 'flags', value)
+            return
+        if name not in self.flags:
+            raise KeyError(f"Unknown flag '{name}'")
+        self.flags[name] = value
+
+    def __repr__(self):
+        return f"TrackFlags({self.flags})"
+
+    def print_flag_register(self):
+        reg = self.get_flags_register()
+        print(f"Flag register (binary): {bin(reg)[2:].zfill(64)}")

{xtrack-0.89.0 → xtrack-0.89.1}/xtrack/tracker.py

@@ -23,6 +23,7 @@ from .line import freeze_longitudinal as _freeze_longitudinal
 from .pipeline import PipelineStatus
 from .progress_indicator import progress
 from .tracker_data import TrackerData
+from .track_flags import TrackFlags
 
 logger = logging.getLogger(__name__)
 
@@ -73,6 +74,7 @@ class Tracker:
         self.local_particle_src = local_particle_src
         self._enable_pipeline_hold = enable_pipeline_hold
         self.use_prebuilt_kernels = use_prebuilt_kernels
+        self.track_flags = TrackFlags()
 
         # Some data for collective mode prepared also for non-collective lines
         # to allow collective actions by the tracker (e.g. time-functions on knobs)
@@ -464,6 +466,7 @@ class Tracker:
 
         headers.extend(self.extra_headers)
         headers.append(_pkg_root.joinpath("headers/constants.h"))
+        headers.append(self.track_flags.c_header_flag_mapping)
 
         src_lines = []
         src_lines.append(
@@ -483,7 +486,9 @@ class Tracker:
                              double line_length,
                 /*gpuglmem*/ int8_t* buffer_tbt_monitor,
                              int64_t offset_tbt_monitor,
-                /*gpuglmem*/ int8_t* io_buffer){
+                /*gpuglmem*/ int8_t* io_buffer,
+                             uint64_t track_flags
+                             ){
 
             #define CONTEXT_OPENMP  //only_for_context cpu_openmp
             #ifdef CONTEXT_OPENMP
@@ -492,7 +497,7 @@ class Tracker:
 
                 #ifndef XT_OMP_SKIP_REORGANIZE
                     const int64_t num_particles_to_track = ParticlesData_get__num_active_particles(particles);
-                    
+
                     {
                         LocalParticle lpart;
                         lpart.io_buffer = io_buffer;
@@ -504,10 +509,10 @@ class Tracker:
                 #else // When we skip reorganize, we cannot just batch active particles
                     const int64_t num_particles_to_track = capacity;
                 #endif
-                
+
                 const int64_t chunk_size = (num_particles_to_track + num_threads - 1)/num_threads; // ceil division
             #endif // CONTEXT_OPENMP
-            
+
             #pragma omp parallel for                                                           //only_for_context cpu_openmp
             for (int chunk = 0; chunk < num_threads; chunk++) {                                //only_for_context cpu_openmp
             int64_t part_id = chunk * chunk_size;                                              //only_for_context cpu_openmp
@@ -521,6 +526,7 @@ class Tracker:
 
             LocalParticle lpart;
             lpart.io_buffer = io_buffer;
+            lpart.track_flags = track_flags;
 
             /*gpuglmem*/ int8_t* tbt_mon_pointer =
                             buffer_tbt_monitor + offset_tbt_monitor;
@@ -706,6 +712,7 @@ class Tracker:
                     xo.Arg(xo.Int8, pointer=True, name="buffer_tbt_monitor"),
                     xo.Arg(xo.Int64, name="offset_tbt_monitor"),
                     xo.Arg(xo.Int8, pointer=True, name="io_buffer"),
+                    xo.Arg(xo.UInt64, name="track_flags"),
                 ],
             )
         }
@@ -1148,6 +1155,8 @@ class Tracker:
         assert ele_start >= 0
         assert ele_start <= self.num_elements
 
+        track_flags = self.track_flags.get_flags_register()
+
         # Logic to split the tracking turns:
         # Case 1: 0 <= start < stop <= L
         #      Track first turn from start until stop (with num_elements_first_turn=stop-start)
@@ -1255,6 +1264,7 @@ class Tracker:
             buffer_tbt_monitor=buffer_monitor,
             offset_tbt_monitor=offset_monitor,
             io_buffer=self.io_buffer.buffer,
+            track_flags=track_flags
         )
 
         # Middle turns
@@ -1275,6 +1285,7 @@ class Tracker:
                 buffer_tbt_monitor=buffer_monitor,
                 offset_tbt_monitor=offset_monitor,
                 io_buffer=self.io_buffer.buffer,
+                track_flags=track_flags
             )
 
         # Last turn, only if incomplete
@@ -1295,6 +1306,7 @@ class Tracker:
                 buffer_tbt_monitor=buffer_monitor,
                 offset_tbt_monitor=offset_monitor,
                 io_buffer=self.io_buffer.buffer,
+                track_flags=track_flags
             )
 
         self.record_last_track = monitor