xcoll 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

xcoll/beam_elements/__init__.py

@@ -1,3 +1,8 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
 from .base import BaseBlock, BaseCollimator
 from .absorber import BlackAbsorber
 from .everest import EverestBlock, EverestCollimator, EverestCrystal

xcoll/beam_elements/absorber.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.  #
+# This file is part of the Xcoll Package.   #
 # Copyright (c) CERN, 2023.                 #
 # ######################################### #
 
@@ -19,6 +19,7 @@ class BlackAbsorber(BaseCollimator):
     ]
 
     isthick = True
+    allow_track = True
     behaves_like_drift = True
     skip_in_loss_location_refinement = True
 

xcoll/beam_elements/base.py

@@ -20,15 +20,14 @@ class InvalidXcoll(xt.BeamElement):
 
     isthick = True
     behaves_like_drift = True
-#     allow_track = False   # Need to wait for xtrack release to implement
+    allow_track = False
     skip_in_loss_location_refinement = True
     allow_backtrack = True
 
     # InvalidXcoll catches unallowed cases, like backtracking through a collimator
     _extra_c_sources = [
         _pkg_root.joinpath('headers','particle_states.h'),
-        _pkg_root.joinpath('headers','checks.h'),
-        _pkg_root.joinpath('beam_elements','collimators_src','invalid_collimator.h')
+        _pkg_root.joinpath('headers','checks.h')
     ]
 
     _depends_on = [xt.RandomRutherford]  # Needed for checks
@@ -47,18 +46,27 @@ class BaseBlock(xt.BeamElement):
     }
 
     isthick = True
+    allow_track = False
     behaves_like_drift = True
-#     allow_track = False   # Need to wait for xtrack release to implement
     skip_in_loss_location_refinement = True
 
-    _extra_c_sources = [
-        _pkg_root.joinpath('beam_elements','collimators_src','base_block.h')
-    ]
     _depends_on = [InvalidXcoll]
     _internal_record_class = CollimatorImpacts
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+    # This is an abstract class and cannot be instantiated
+    def __new__(cls, *args, **kwargs):
+        if cls == BaseBlock:
+            raise Exception("Abstract class `BaseBlock` cannot be instantiated!")
+        instance = super().__new__(cls)
+        return instance
+
+    def enable_scattering(self):
+        if hasattr(self, '_tracking'):
+            self._tracking = True
+
+    def disable_scattering(self):
+        if hasattr(self, '_tracking'):
+            self._tracking = False
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
         return InvalidXcoll(length=-self.length,
@@ -89,8 +97,8 @@ class BaseCollimator(xt.BeamElement):
     }
 
     isthick = True
+    allow_track = False
     behaves_like_drift = True
-#     allow_track = False   # Need to wait for xtrack release to implement
     skip_in_loss_location_refinement = True
 
     _skip_in_to_dict  = ['jaw_L', 'jaw_R', 'ref_x', 'ref_y',
@@ -99,18 +107,17 @@ class BaseCollimator(xt.BeamElement):
     _store_in_to_dict = ['angle', 'tilt', 'jaw', 'reference_center', 'side']
     # Extra fields (only in Python): angle_L, angle_R, tilt_L, tilt_R, jaw_LU, jaw_LD, jaw_RU, jaw_RD
 
-    _extra_c_sources = [
-        _pkg_root.joinpath('beam_elements','collimators_src','base_collimator.h')
-    ]
     _depends_on = [InvalidXcoll, xt.Drift, xt.XYShift, xt.SRotation, xt.YRotation]
     _internal_record_class = CollimatorImpacts
 
+    # This is an abstract class and cannot be instantiated
+    def __new__(cls, *args, **kwargs):
+        if cls == BaseCollimator:
+            raise Exception("Abstract class `BaseCollimator` cannot be instantiated!")
+        instance = super().__new__(cls)
+        return instance
 
     def __init__(self, **kwargs):
-        # TODO: quick hack to avoid instantiation; did not manage to get it to work correclty with ABC
-        if self.__class__.__name__ == 'BaseCollimator':
-            raise Exception("Abstract class 'BaseCollimator' cannot be instantiated!")
-
         if '_xobject' not in kwargs:
             # Set jaw
             if 'jaw' in kwargs:
@@ -197,6 +204,14 @@ class BaseCollimator(xt.BeamElement):
         super().__init__(**kwargs)
 
 
+    def enable_scattering(self):
+        if hasattr(self, '_tracking'):
+            self._tracking = True
+
+    def disable_scattering(self):
+        if hasattr(self, '_tracking'):
+            self._tracking = False
+
     @property
     def jaw(self):
         return _get_LR(self, 'jaw', neg=True)

xcoll/beam_elements/collimators_src/absorber.h

@@ -1,3 +1,8 @@
+// copyright ############################### #
+// This file is part of the Xcoll Package.   #
+// Copyright (c) CERN, 2024.                 #
+// ######################################### #
+
 #ifndef XCOLL_ABSORBER_H
 #define XCOLL_ABSORBER_H
 

xcoll/beam_elements/everest.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.  #
+# This file is part of the Xcoll Package.   #
 # Copyright (c) CERN, 2023.                 #
 # ######################################### #
 
@@ -31,6 +31,8 @@ class EverestBlock(BaseBlock):
     }
 
     isthick = True
+    needs_rng = True
+    allow_track = True
     behaves_like_drift = True
     skip_in_loss_location_refinement = True
 
@@ -67,12 +69,7 @@ class EverestBlock(BaseBlock):
             use_prebuilt_kernels = kwargs.pop('use_prebuilt_kernels', True)
         super().__init__(**kwargs)
         if '_xobject' not in kwargs:
-            try:   # TODO: small workaround until PR
-                self.compile_kernels(use_prebuilt_kernels=use_prebuilt_kernels,
-                                     only_if_needed=True)
-            except TypeError:
-                self.compile_kernels(only_if_needed=True)
-            self._context.kernels.EverestBlock_set_material(el=self)
+            self.EverestBlock_set_material(el=self)
 
 
     @property
@@ -86,8 +83,7 @@ class EverestBlock(BaseBlock):
                 raise ValueError("Invalid material!")
         if not xt.line._dicts_equal(self.material.to_dict(), material.to_dict()):
             self._material = material
-            self.compile_kernels(only_if_needed=True)
-            self._context.kernels.EverestBlock_set_material(el=self)
+            self.EverestBlock_set_material(el=self)
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
         return InvalidXcoll(length=-self.length, _context=_context,
@@ -102,6 +98,8 @@ class EverestCollimator(BaseCollimator):
     }
 
     isthick = True
+    needs_rng = True
+    allow_track = True
     behaves_like_drift = True
     skip_in_loss_location_refinement = True
 
@@ -137,12 +135,7 @@ class EverestCollimator(BaseCollimator):
             use_prebuilt_kernels = kwargs.pop('use_prebuilt_kernels', True)
         super().__init__(**kwargs)
         if '_xobject' not in kwargs:
-            try:   # TODO: small workaround until PR
-                self.compile_kernels(use_prebuilt_kernels=use_prebuilt_kernels,
-                                     only_if_needed=True)
-            except TypeError:
-                self.compile_kernels(only_if_needed=True)
-            self._context.kernels.EverestCollimator_set_material(el=self)
+            self.EverestCollimator_set_material(el=self)
 
     @property
     def material(self):
@@ -155,8 +148,7 @@ class EverestCollimator(BaseCollimator):
                 raise ValueError("Invalid material!")
         if not xt.line._dicts_equal(self.material.to_dict(), material.to_dict()):
             self._material = material
-            self.compile_kernels(only_if_needed=True)
-            self._context.kernels.EverestCollimator_set_material(el=self)
+            self.EverestCollimator_set_material(el=self)
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
         return InvalidXcoll(length=-self.length, _context=_context,
@@ -181,6 +173,8 @@ class EverestCrystal(BaseCollimator):
     }
 
     isthick = True
+    needs_rng = True
+    allow_track = True
     behaves_like_drift = True
     skip_in_loss_location_refinement = True
 
@@ -236,12 +230,7 @@ class EverestCrystal(BaseCollimator):
                 self._bending_angle = np.arcsin(self.active_length/bending_radius)
             if bending_angle:
                 self._bending_radius = self.active_length / np.sin(bending_angle)
-            try:   # TODO: small workaround until PR
-                self.compile_kernels(use_prebuilt_kernels=use_prebuilt_kernels,
-                                     only_if_needed=True)
-            except TypeError:
-                self.compile_kernels(only_if_needed=True)
-            self._context.kernels.EverestCrystal_set_material(el=self)
+            self.EverestCrystal_set_material(el=self)
 
 
     @property
@@ -290,8 +279,7 @@ class EverestCrystal(BaseCollimator):
                 raise ValueError("Invalid material!")
         if not xt.line._dicts_equal(self.material.to_dict(), material.to_dict()):
             self._material = material
-            self.compile_kernels(only_if_needed=True)
-            self._context.kernels.EverestCrystal_set_material(el=self)
+            self.EverestCrystal_set_material(el=self)
 
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
@@ -307,3 +295,4 @@ def _lattice_setter(lattice):
     else:
         raise ValueError(f"Illegal value {lattice} for 'lattice'! "
                         + "Only use 'strip' (110) or 'quasi-mosaic' (111).")
+

xcoll/colldb.py

@@ -111,7 +111,7 @@ class CollimatorDatabase:
                     raise ValueError("Missing name tag / anchor in "
                                    + "CollimatorDatabase['families']!")
                 # We get the anchor from the rt yaml, and use it as key in the families dict
-                families[full_fam.anchor.value.lower()] = fam
+                families[full_fam.anchor.value.lower()] = {f.lower(): v for f, v in fam.items()}
             dct['families'] = families
 
             # Now we need to loop over each collimator, and verify which family was used
@@ -184,10 +184,8 @@ class CollimatorDatabase:
                    beam=beam, _yaml_merged=_yaml_merged, ignore_crystals=ignore_crystals)
 
 
-    # TODO: load crystals with SixTrack loader
-    # TODO: load families with SixTrack loader
     @classmethod
-    def from_SixTrack(cls, file, **kwargs):
+    def from_SixTrack(cls, file, ignore_crystals=True, **kwargs):
         # only import regex here
         import re
         with open(file, 'r') as infile:
@@ -195,49 +193,69 @@ class CollimatorDatabase:
             family_settings = {}
             family_types = {}
             side = {}
+            cry_fields = ['bending_radius', 'xdim', 'ydim', 'thick', 'miscut', 'crystal']
+            cry = {}
 
             for l_no, line in enumerate(infile):
                 if line.startswith('#'):
                     continue # Comment
-
                 sline = line.split()
-                if len(sline) > 0 and len(sline) < 6:
+                if len(sline) > 0:
                     if sline[0].lower() == 'nsig_fam':
                         family_settings[sline[1]] = float(sline[2])
                         family_types[sline[1]] = sline[3]
                     elif sline[0].lower() == 'onesided':
                         side[sline[1]] = int(sline[2])
+                    elif sline[0].lower() == "crystal":
+                        cry[sline[1]] = {}
+                        for i, key in enumerate(cry_fields):
+                            idx = i+2 if i < 4 else i+3  # we skip "tilt"
+                            if i < 5:
+                                cry[sline[1]][key] = float(sline[idx])
+                            else:
+                                cry[sline[1]][key] = int(sline[idx])
                     elif sline[0].lower() == 'settings':
-                        # TODO CRYSTAL
                         pass # Acknowledge and ignore this line
+                    elif len(sline) == 6:
+                        # Standard collimator definition
+                        coll_data_string += line
                     else:
                         print(f"Unknown setting {line}")
-                else:
-                    coll_data_string += line
+
+        defaults = {}
+        _initialise_None(defaults)
 
         famdct = {key: {'gap': family_settings[key], 'stage': family_types[key]} for key in family_settings}
         names = ['name', 'gap', 'material', 'active_length', 'angle', 'offset']
 
-        df = pd.read_csv(io.StringIO(coll_data_string), delim_whitespace=True,
-                        index_col=False, names=names)
+        df = pd.read_csv(io.StringIO(coll_data_string), sep=r'\s+', index_col=False, names=names)
         df['family'] = df['gap'].copy()
         df['family'] = df['family'].apply(lambda s: None if re.match(r'^-?\d+(\.\d+)?$', str(s)) else s)
         df.insert(5,'stage', df['gap'].apply(lambda s: family_types.get(s, 'UNKNOWN')))
-        sides = df['name'].apply(lambda s: side.get(s, 0))
+
         df['gap'] = df['gap'].apply(lambda s: float(family_settings.get(s, s)))
-        df['name'] = df['name'].str.lower() # Make the names lowercase for easy processing # TODO this breaks code if a key has upper case, e.g. gap_L
+        # TODO this breaks code if a key has upper case, e.g. gap_L
+        df['name'] = df['name'].str.lower() # Make the names lowercase for easy processing
         df['parking'] = 0.025
+        if ignore_crystals:
+            df = df[~df.name.isin(list(cry.keys()))]
+        else:
+            for key in cry_fields:
+                df[key] = [cry[name][key] if name in cry else defaults[key]
+                           for name in df['name']]
+            df['crystal'] = ['strip' if s==1 else s for s in df['crystal']]
+            df['crystal'] = ['quasi-mosaic' if s==2 else s for s in df['crystal']]
+        df['side'] = [side[name] if name in side else defaults['side']
+                      for name in df['name']]
+        df['side'] = ['both'  if s==0 else s for s in df['side']]
+        df['side'] = ['left'  if s==1 else s for s in df['side']]
+        df['side'] = ['right' if s==2 else s for s in df['side']]
         df = df.set_index('name')
-        df['side'] = sides.values
-        df['side'] = [ 'both'  if s==0 else s for s in df['side'] ]
-        df['side'] = [ 'left'  if s==1 else s for s in df['side'] ]
-        df['side'] = [ 'right' if s==2 else s for s in df['side'] ]
-        return cls.from_dict({'collimators': df.transpose().to_dict(), 'families': famdct}, **kwargs)
 
-    # -------------------------------
-    # ------ Dump to disk -----------
-    # -------------------------------
-    
+        return cls.from_dict({'collimators': df.transpose().to_dict(), 'families': famdct}, \
+                             ignore_crystals=ignore_crystals, **kwargs)
+
+
     def write_to_yaml(self, out, lhc_style=True):
         """
         Writes a colldb in memory to disk in the yaml format.
@@ -286,8 +304,8 @@ class CollimatorDatabase:
                     fam_keys = dcts._family_dict[fam].keys()
                     coll_dict = {**{'<<': '*'+fam}, **coll_dict}
                 temp_items_to_print = []
-                if coll_dict['crystal']:
-                    temp_items_to_print = ['bend','xdim','ydim','miscut','crystal']
+                if coll_dict['crystal'] and str(coll_dict['crystal'])!='nan':
+                    temp_items_to_print = ['bending_radius','xdim','ydim','miscut','crystal', 'thick']
                 if coll_dict['angle_L'] == coll_dict['angle_R']:
                     coll_dict.update({'angle': coll_dict['angle_L']})
                 else:
@@ -303,6 +321,8 @@ class CollimatorDatabase:
                 value = {}
                 overwritten_keys = coll_dict['overwritten_keys']
                 for key, val in coll_dict.items():
+                    if key == 'active_length':
+                        key = 'length' 
                     if (key in coll_items_to_print+temp_items_to_print) and (key not in (set(fam_keys)-set(overwritten_keys))) and (val != 'both'):
                         value.update({key: val})
                 file.write(_format_dict_entry(coll, value, spacing='    '))
@@ -387,7 +407,7 @@ class CollimatorDatabase:
                 _print_colls(b2_colls, self, 'b2', file)
             if len(bx_colls) > 0:
                 _print_colls(bx_colls, self, 'bx', file)
-                print('WARNING -- some collimators could not be assigned to b1 or b2. Tracking might not work with those collimators. Please manually change the output file if necessary.')    
+                print('WARNING -- some collimators could not be assigned to b1 or b2. Tracking might not work with those collimators. Please manually change the output file if necessary.')
 
 
     def __init__(self, **kwargs):
@@ -450,7 +470,7 @@ class CollimatorDatabase:
                 # Only do this check if we didn't do a YAML merge earlier (because then it
                 # is already taken care of)
                 if not _yaml_merged:
-                    overwritten_keys = [key for key in settings.keys() if key in fam[thisfam]]
+                    overwritten_keys = [key.lower() for key in settings.keys() if key in fam[thisfam]]
                     if len(overwritten_keys) > 0:
                         settings['overwritten_keys'] = overwritten_keys
 
@@ -463,7 +483,10 @@ class CollimatorDatabase:
 
             # Save list of crystals
             if 'crystal' in settings:
-                crystals += [thiscoll]
+                if settings['crystal'] != 0.0:
+                    crystals += [thiscoll]
+                else: 
+                    settings['crystal'] = None
 
         # Remove crystals from colldb
         if ignore_crystals:

xcoll/general.py

@@ -12,5 +12,5 @@ citation = "F.F. Van der Veken, et al.: Recent Developments with the New Tools f
 # ===================
 # Do not change
 # ===================
-__version__ = '0.3.2'
+__version__ = '0.3.5'
 # ===================

xcoll/headers/checks.h

@@ -6,6 +6,12 @@
 #ifndef XCOLL_CHECKS_H
 #define XCOLL_CHECKS_H
 
+// This is a quick macro to use inside a function body on a parameter that is not
+// used inside the function (this avoids throwing warnings at compilation time).
+#ifndef UNUSED
+#define UNUSED(expr) (void)(expr)
+#endif
+
 /*gpufun*/
 int8_t xcoll_check_particle_init(RandomRutherfordData rng, LocalParticle* part) {
     int8_t is_tracking = assert_tracking(part, XC_ERR_INVALID_TRACK);
@@ -26,4 +32,4 @@ int8_t xcoll_check_particle_init(RandomRutherfordData rng, LocalParticle* part)
     return is_tracking*rng_is_set*ruth_is_set;
 }
     
-#endif /* XCOLL_CHECKS_H */
+#endif /* XCOLL_CHECKS_H */

xcoll/impacts/impacts.py

@@ -14,7 +14,6 @@ import pandas as pd
 
 
 class CollimatorImpacts(xt.BeamElement):
-
     _xofields = {
         '_index':            xt.RecordIndex,
         'at_element':        xo.Int64[:],
@@ -49,6 +48,8 @@ class CollimatorImpacts(xt.BeamElement):
         'child_pdgid':       xo.Int64[:],
     }
 
+    allow_track = False
+
     _extra_c_sources = [
         source,
         _pkg_root.joinpath('headers','particle_states.h'),

xcoll/impacts/impacts_src/impacts.h

@@ -7,11 +7,6 @@
 #ifndef XCOLL_IMPACTS_H
 #define XCOLL_IMPACTS_H
 
-/*gpufun*/
-void CollimatorImpacts_track_local_particle(CollimatorImpactsData el, LocalParticle* part0) {
-    kill_all_particles(part0, XC_ERR_INVALID_TRACK);
-}
-
 // TODO: do we need to pass RecordIndex?
 // probably can do RecordIndex record_index = CollimatorImpactsData_getp__index(record);  ?
 /*gpufun*/
@@ -101,4 +96,4 @@ void CollimatorImpactsData_log_child(CollimatorImpactsData record, int64_t i_slo
     }
 }
 
-#endif /* XCOLL_IMPACTS_H */
+#endif /* XCOLL_IMPACTS_H */

xcoll/manager.py

@@ -8,7 +8,7 @@ from pathlib import Path
 import numpy as np
 import pandas as pd
 
-from .beam_elements import BaseCollimator, BlackAbsorber, EverestCollimator, EverestCrystal, _all_collimator_types
+from .beam_elements import BaseCollimator, BlackAbsorber, EverestCollimator, EverestCrystal, _all_collimator_types, element_classes
 from .colldb import CollimatorDatabase
 from .impacts import CollimatorImpacts
 from .scattering_routines.everest.materials import SixTrack_to_xcoll, CrystalMaterial
@@ -775,14 +775,14 @@ class CollimatorManager:
 
     def enable_scattering(self):
         # Prepare collimators for tracking
-        for coll in self.collimator_names:
-            self.line[coll]._tracking = True
+        for el in self.line.get_elements_of_type(element_classes)[0]:
+            el.enable_scattering()
         self._set_record_impacts()
 
     def disable_scattering(self):
         # Prepare collimators for tracking
-        for coll in self.collimator_names:
-            self.line[coll]._tracking = False
+        for el in self.line.get_elements_of_type(element_classes)[0]:
+            el.disable_scattering()
 
     @property
     def scattering_enabled(self):

xcoll/rf_sweep.py

@@ -4,11 +4,11 @@
 # ######################################### #
 
 import numpy as np
-import pandas as pd
 import scipy.constants as sc
+from time import perf_counter
 
 import xtrack as xt
-
+from xtrack.progress_indicator import progress
 
 class RFSweep:
 
@@ -81,7 +81,7 @@ class RFSweep:
                     + f"the bucket and remain uncaptured.")
 
 
-    def track(self, sweep=0, num_turns=0, particles=None, verbose=True, *args, **kwargs):
+    def track(self, sweep=0, particles=None, num_turns=0, verbose=True, *args, **kwargs):
 
         # Was there a previous sweep?
         # If yes, we do not overwrite it but continue from there
@@ -101,20 +101,91 @@ class RFSweep:
                 raise ValueError("Need to build tracker first!")
             if particles is None:
                 raise ValueError("Need particles to track!")
+            time = kwargs.pop('time', False)
+            with_progress = kwargs.pop('with_progress', False)
             rf_shift_per_turn = sweep / num_turns
-            if 'time' in kwargs and kwargs['time']:
-                self.line.tracker.time_last_track = 0
-            for i in range(num_turns):
-                sweep = existing_sweep + i*rf_shift_per_turn
-                self.line['rf_sweep'].dzeta = self.L * sweep / (self.f_RF + sweep)
-#                 for cav in cavities:
-#                     self.line[cav].frequency = freq + sweep
-                if 'time' in kwargs and kwargs['time']:
-                    prev_time = self.line.time_last_track
-                self.line.track(particles, num_turns=1, *args, **kwargs)
-                if 'time' in kwargs and kwargs['time']:
-                    self.line.tracker.time_last_track += prev_time
-                if not np.any(particles.state == 1):
-                    print(f"All particles lost at turn {i}, stopped sweep at "
-                        + f"{existing_sweep + i*rf_shift_per_turn}Hz.")
-                    break
+
+            # This is taken from xtrack.tracker.Tracker._track
+            if time:
+                t0 = perf_counter()
+            if with_progress:
+                if self.line.tracker.enable_pipeline_hold:
+                    raise ValueError("Progress indicator is not supported with pipeline hold")
+                if num_turns < 2:
+                    raise ValueError('Tracking with progress indicator is only '
+                                    'possible over more than one turn.')
+                if with_progress is True:
+                    batch_size = scaling = 100
+                else:
+                    batch_size = int(with_progress)
+                    scaling = with_progress if batch_size > 1 else None
+                if kwargs.get('turn_by_turn_monitor') is True:
+                    ele_start = kwargs.get('ele_start') or 0
+                    ele_stop = kwargs.get('ele_stop')
+                    if ele_stop is None:
+                        ele_stop = len(self.line)
+                    if ele_start >= ele_stop:
+                        # we need an additional turn and space in the monitor for
+                        # the incomplete turn
+                        num_turns += 1
+                    _, monitor, _, _ = self.line.tracker._get_monitor(particles, True, num_turns)
+                    kwargs['turn_by_turn_monitor'] = monitor
+
+                for ii in progress(
+                        range(0, num_turns, batch_size),
+                        desc='Tracking',
+                        unit_scale=scaling,
+                ):
+                    one_turn_kwargs = kwargs.copy()
+                    is_first_batch = ii == 0
+                    is_last_batch = ii + batch_size >= num_turns
+
+                    if is_first_batch and is_last_batch:
+                        # This is the only batch, we track as normal
+                        pass
+                    elif is_first_batch:
+                        # Not the last batch, so track until the last element
+                        one_turn_kwargs['ele_stop'] = None
+                        one_turn_kwargs['num_turns'] = batch_size
+                    elif is_last_batch:
+                        # Not the first batch, so track from the first element
+                        one_turn_kwargs['ele_start'] = None
+                        remaining_turns = num_turns % batch_size
+                        if remaining_turns == 0:
+                            remaining_turns = batch_size
+                        one_turn_kwargs['num_turns'] = remaining_turns
+                        one_turn_kwargs['_reset_log'] = False
+                    elif not is_first_batch and not is_last_batch:
+                        # A 'middle batch', track from first to last element
+                        one_turn_kwargs['num_turns'] = batch_size
+                        one_turn_kwargs['ele_start'] = None
+                        one_turn_kwargs['ele_stop'] = None
+                        one_turn_kwargs['_reset_log'] = False
+                    self._tracking_func(particles, rf_shift_per_turn, **one_turn_kwargs)
+                    if not np.any(particles.state == 1):
+                        break
+
+            else:
+                self._tracking_func(particles, rf_shift_per_turn, num_turns=num_turns, *args, **kwargs)
+
+            if not np.any(particles.state == 1):
+                print(f"All particles lost at turn {particles.at_turn.max()}, stopped sweep at "
+                    + f"{self.current_sweep_value}Hz.")
+
+            if time:
+                t1 = perf_counter()
+                self.line.tracker._context.synchronize()
+                self.line.tracker.time_last_track = t1 - t0
+            else:
+                self.line.tracker.time_last_track = None
+
+    def _tracking_func(self, particles, rf_shift_per_turn, num_turns=1, *args, **kwargs):
+        existing_sweep = self.current_sweep_value
+        for i in range(num_turns):
+            sweep = existing_sweep + i*rf_shift_per_turn
+            self.line['rf_sweep'].dzeta = self.L * sweep / (self.f_RF + sweep)
+#             for cav in cavities:
+#                 self.line[cav].frequency = freq + sweep
+            self.line.track(particles, num_turns=1, *args, **kwargs)
+            if not np.any(particles.state == 1):
+                break

xcoll/scattering_routines/everest/__init__.py

@@ -1,2 +1,7 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
 from .materials import GeneralMaterial, Material, CrystalMaterial
 from .everest import EverestEngine

xcoll/scattering_routines/everest/everest.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.  #
+# This file is part of the Xcoll Package.   #
 # Copyright (c) CERN, 2023.                 #
 # ######################################### #