xcoll 0.6.2__tar.gz → 0.6.3__tar.gz

{xcoll-0.6.2 → xcoll-0.6.3}/PKG-INFO

@@ -1,15 +1,13 @@
 Metadata-Version: 2.3
 Name: xcoll
-Version: 0.6.2
+Version: 0.6.3
 Summary: Xsuite collimation package
-License: Apache 2.0
+License: Apache-2.0
 Author: Frederik F. Van der Veken
 Author-email: frederik@cern.ch
-Requires-Python: >=3.8
-Classifier: License :: Other/Proprietary License
+Requires-Python: >=3.10
+Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -17,13 +15,19 @@ Classifier: Programming Language :: Python :: 3.13
 Provides-Extra: tests
 Requires-Dist: numpy (>=1.0)
 Requires-Dist: pandas (>=1.4)
-Requires-Dist: ruamel-yaml (>=0.17.31,<0.18.0) ; extra == "tests"
-Requires-Dist: xdeps (>=0.10.5)
-Requires-Dist: xobjects (>=0.5.0)
-Requires-Dist: xpart (>=0.23.0)
-Requires-Dist: xtrack (>=0.84.8)
+Requires-Dist: pytest ; extra == "tests"
+Requires-Dist: pytest-html ; extra == "tests"
+Requires-Dist: pytest-xdist ; extra == "tests"
+Requires-Dist: ruamel-yaml (>=0.17.31,<0.18.15)
+Requires-Dist: xdeps (>=0.10.6)
+Requires-Dist: xobjects (>=0.5.2)
+Requires-Dist: xpart (>=0.23.1)
+Requires-Dist: xtrack (>=0.89.3)
+Project-URL: Bug Tracker, https://github.com/xsuite/xsuite/issues
+Project-URL: Documentation, https://xsuite.readthedocs.io/
 Project-URL: Homepage, https://github.com/xsuite/xcoll
 Project-URL: Repository, https://github.com/xsuite/xcoll
+Project-URL: download, https://pypi.python.org/pypi/xcoll
 Description-Content-Type: text/markdown
 
 # xcoll

xcoll-0.6.3/pyproject.toml

@@ -0,0 +1,63 @@
+[project]
+name = "xcoll"
+version = "0.6.3"
+description = "Xsuite collimation package"
+authors = [
+    {name="Frederik F. Van der Veken", email="frederik@cern.ch"},
+    {name="Simone O. Solstrand",       email="simone.otelie.solstrand@cern.ch"},
+    {name="Björn Lindstrom",           email="bjorn.lindstrom@cern.ch"},
+    {name="Giacomo Broggi",            email="giacomo.broggi@cern.ch"},
+    {name="André Donadon Servelle",    email="andre.donadon.servelle@cern.ch"},
+    {name="Dora E. Veres",             email="dora.erzsebet.veres@cern.ch"},
+    {name="Despina Demetriadou",       email="despina.demetriadou@cern.ch"},
+    {name="Andrey Abramov",            email="andrey.abramov@cern.ch"},
+    {name="Giovanni Iadarola",         email="giovanni.iadarola@cern.ch"}
+]
+readme = "README.md"
+license = {text="Apache-2.0"}
+
+requires-python = ">=3.10"
+
+dependencies = [
+    "ruamel-yaml (>=0.17.31,<0.18.15)",
+    "numpy>=1.0",
+    "pandas>=1.4",
+    "xobjects>=0.5.2",
+    "xdeps>=0.10.6",
+    "xpart>=0.23.1",
+    "xtrack>=0.89.3",
+]
+
+[project.optional-dependencies]
+tests = ["pytest", "pytest-html", "pytest-xdist"]
+
+[project.urls]
+homepage = "https://github.com/xsuite/xcoll"
+repository = "https://github.com/xsuite/xcoll"
+documentation = "https://xsuite.readthedocs.io/"
+"Bug Tracker" = "https://github.com/xsuite/xsuite/issues"
+download = "https://pypi.python.org/pypi/xcoll"
+
+[project.entry-points.xobjects]
+include = "xcoll"
+
+[tool.poetry]
+exclude = ["xcoll/lib", "xcoll/config"]
+
+[poetry.group.dev.dependencies]
+pytest = ">=7.3"
+xaux = ">=0.3.5"
+
+[build-system]
+# Needed for pip install -e (BTW: need pip version 22)
+requires = ["poetry-core>=1.0.8"]
+build-backend = "poetry.core.masonry.api"
+
+# pyproject.toml
+#[tool.pytest.ini_options]
+#addopts = "-ra --durations=10 --durations-min=1"
+#python_functions = ["test_"]
+#testpaths = [
+#    "tests",
+#]
+

{xcoll-0.6.2 → xcoll-0.6.3}/xcoll/__init__.py

@@ -11,7 +11,7 @@ from .beam_elements import BlackAbsorber, BlackCrystal, TransparentCollimator, T
 from .scattering_routines.everest import materials, Material, CrystalMaterial
 from .colldb import CollimatorDatabase
 from .interaction_record import InteractionRecord
-from .rf_sweep import RFSweep
+from .rf_sweep import RFSweep, prepare_rf_sweep
 from .lossmap import LossMap, MultiLossMap
 
 

{xcoll-0.6.2 → xcoll-0.6.3}/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.6.2'
+__version__ = '0.6.3'
 # ======================

xcoll-0.6.3/xcoll/prebuild_kernels/__init__.py

@@ -0,0 +1,7 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+from .elements import DEFAULT_XCOLL_ELEMENTS
+from .element_inits import XCOLL_ELEMENTS_INIT_DEFAULTS

xcoll-0.6.3/xcoll/prebuild_kernels/element_inits.py

@@ -0,0 +1,6 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+XCOLL_ELEMENTS_INIT_DEFAULTS = {}

xcoll-0.6.3/xcoll/prebuild_kernels/elements.py

@@ -0,0 +1,20 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+from ..beam_elements import *
+
+
+DEFAULT_XCOLL_ELEMENTS = [
+    BlackAbsorber,
+    BlackCrystal,
+    TransparentCollimator,
+    TransparentCrystal,
+    EverestBlock,
+    EverestCollimator,
+    EverestCrystal,
+    BlowUp,
+    # ParticleStatsMonitor,
+    EmittanceMonitor
+]

xcoll-0.6.3/xcoll/rf_sweep.py

@@ -0,0 +1,182 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+import numpy as np
+import scipy.constants as sc
+
+import xtrack as xt
+
+
+def prepare_rf_sweep(line, cavities=None, sweep=None, sweep_per_turn=None, num_turns=None):
+    rf_sweep = RFSweep(line=line, cavities=cavities)
+    rf_sweep.prepare(sweep=sweep, sweep_per_turn=sweep_per_turn, num_turns=num_turns)
+
+
+class RFSweep:
+    def __init__(self, line, cavities=None):
+        self.line = line
+        if line._var_management is None:
+            raise ValueError("Line must have a `var_management` to use RFSweep!" \
+                           + "Do not use `optimize_for_tracking` as it will "
+                           + "disable expressions, which are needed for RFSweep.")
+        self.env = line.env
+        self.cavities = cavities
+        self.env['rf_sweep_df'] = 0
+
+    @property
+    def current_sweep_value(self):
+        return self.env['rf_sweep_df']
+
+    def prepare(self, sweep=None, sweep_per_turn=None, num_turns=None):
+        if sweep_per_turn is not None:
+            if np.isclose(sweep_per_turn, 0):
+                raise ValueError("Variable `sweep_per_turn` must be non-zero!")
+            if sweep is not None:
+                raise ValueError("Provide either `sweep` or `sweep_per_turn`, not both.")
+            if num_turns is not None:
+                raise ValueError("Variable `num_turns` cannot be set when using `sweep_per_turn`.")
+            self.sweep_per_turn = sweep_per_turn
+        elif sweep is not None:
+            if np.isclose(sweep, 0):
+                raise ValueError("Variable `sweep` must be non-zero!")
+            if num_turns is None:
+                num_turns = int(abs(sweep))
+            if num_turns <= 0:
+                raise ValueError("When using `sweep`, `num_turns` must be a positive integer.")
+            self.sweep_per_turn = sweep / num_turns
+        else:
+            raise ValueError("Either `sweep` or `sweep_per_turn` must be provided.")
+        self._get_cavity_data(self.sweep_per_turn)
+        self._install_zeta_shift()
+        self.tw = self.line.twiss()
+        self.reset() # Initialize rf_sweep_df
+        self.env['rf_sweep'].dzeta = f"{self.L} * rf_sweep_df / ({self.f_RF} + rf_sweep_df)"
+        for cavs in self.cavities:
+            for cav in cavs['names']:
+                scale_factor = int(np.round(cavs['freq'] / self.f_RF))
+                if scale_factor == 1:
+                    self.env.ref[cav].frequency += self.env.ref["rf_sweep_df"]
+                else:
+                    self.env.ref[cav].frequency += scale_factor * self.env.ref["rf_sweep_df"]
+        self.line.enable_time_dependent_vars = True
+        print("Enabled time-dependent variables in the line.")
+
+
+    def reset(self):
+        if self.sweep_per_turn is None:
+            raise ValueError("RFSweep not prepared. Call `prepare` first.")
+        t_turn = self.tw.T_rev0
+        current_time = self.env['t_turn_s']
+        if current_time == 0:
+            self.env['rf_sweep_df'] = f"(t_turn_s + {t_turn}) / {t_turn} * {self.sweep_per_turn}"
+        else:
+            self.env['rf_sweep_df'] = f"(t_turn_s - {current_time} + {t_turn}) / {t_turn} * {self.sweep_per_turn}"
+
+
+    def info(self, sweep=None, num_turns=None):
+        if sweep is not None:
+            raise DeprecationWarning("The `sweep` argument is deprecated.")
+        if num_turns is not None:
+            raise DeprecationWarning("The `num_turns` argument is deprecated.")
+        if self.sweep_per_turn is None:
+            raise ValueError("RFSweep not prepared. Call `prepare` first.")
+        existing_sweep = self.current_sweep_value
+
+        beta0 = self.line.particle_ref.beta0[0]
+        E = self.line.particle_ref.energy0[0]
+        q = self.line.particle_ref.q0
+        h = self.f_RF * self.L / beta0 / sc.c
+        eta = self.tw.slip_factor
+
+        bucket_height = np.sqrt(abs(q*self.V*beta0**2 / (np.pi*h*eta*E) * (
+                            2*np.cos(self.phi) + (2*self.phi-np.pi)*np.sin(self.phi)
+                        )))
+        delta_shift = -self.sweep_per_turn / self.f_RF / eta
+        bucket_turns = bucket_height * 2 / abs(delta_shift)
+        print(f"The current frequency is {self.f_RF + existing_sweep}Hz, adding "
+            + f"{self.sweep_per_turn}Hz per turn.")
+        print(f"This sweep will move the center of the bucket with \u0394\u03B4 = "
+            + f"{delta_shift} per turn.")
+        print(f"This implies one bucket shift every {bucket_turns} turns.")
+        if self.tw.qs * bucket_turns < 3:
+            print(f"Warning: This is a very fast sweep, moving ~1 bucket in "
+                + f"~{round(self.tw.qs * bucket_turns, 2)} synchrotron oscillations "
+                + f"(on average). If the bucket moves faster than a particle "
+                + f"can follow, that particle will move out of the bucket and "
+                + f"remain uncaptured.")
+
+
+    # For backward compatibility
+    def track(self, sweep=0, particles=None, num_turns=0, verbose=True, *args, **kwargs):
+        self.prepare(sweep, num_turns)
+        self.line.track(particles=particles, *args, **kwargs)
+
+
+    def _get_cavity_data(self, sweep):
+        tt = self.line.get_table()
+        tt_c = tt.rows[[nn for nn, ttt in zip(tt.name, tt.element_type)
+                        if 'Cavity' in ttt and 'CrabCavity' not in ttt]]
+        mask = tt_c.parent_name == None  # Unsliced cavities
+        if self.cavities is None:
+            self.cavities = np.unique(list(tt_c.name[mask]) + list(tt_c.parent_name[~mask]))
+        else:
+            if isinstance(self.cavities, str):
+                self.cavities = [self.cavities]
+            for cav in self.cavities:
+                if cav not in self.env.elements:
+                    raise ValueError(f"Cavity `{cav}` not found in environment!")
+        if len(self.cavities) == 0:
+            raise ValueError("No cavities found in the line!")
+        freq = np.array([self.env[cav].frequency for cav in self.cavities])
+        volt = np.array([self.env[cav].voltage for cav in self.cavities])
+        lag  = np.array([self.env[cav].lag for cav in self.cavities])
+        s_pos = []
+        name_first_in_line = []
+        for cav in self.cavities:
+            if cav in tt_c.name:
+                s_pos.append(tt_c.rows[cav].s_start[0])
+                name_first_in_line.append(cav)
+            elif cav in tt_c.parent_name:
+                s_pos.append(tt_c.rows[tt_c.parent_name == cav].s_start.min())
+                name_first_in_line.append(tt_c.rows[tt_c.parent_name == cav].name[0])
+            else:
+                raise ValueError(f"Cavity `{cav}` not found in the line!")
+        idx  = np.argsort(freq)
+        boundaries = np.where(~np.isclose(freq[idx][1:], freq[idx][:-1],
+                                          rtol=1e-8, atol=1e-12)
+                              )[0] + 1
+        groups = []
+        for gidx in np.split(idx, boundaries):
+            groups.append({
+                'names': self.cavities[gidx],
+                'freq': np.mean(freq[gidx]),
+                'voltage': volt[gidx].sum(),
+                's_pos': np.array(s_pos)[gidx],
+                'lag': lag[gidx],
+                'name_first_in_line': np.array(name_first_in_line)[gidx],
+            })
+        self.cavities = sorted(groups, key=lambda g: (g['voltage'], g['freq']), reverse=True)
+        if np.isclose(self.cavities[0]['voltage'], 0):
+            raise ValueError("No active cavity found!")
+        self.f_RF = self.cavities[0]['freq']
+        self.phi = np.deg2rad(self.cavities[0]['lag'].mean())
+        self.V = self.cavities[0]['voltage']
+        self.L = self.line.get_length()
+        if len(self.cavities) > 1:
+            print("Found multiple cavities with different frequencies:")
+            for g in self.cavities:
+                print(f"{g['freq']}Hz  at {g['voltage']}V: {g['names']}")
+            print(r"The sweep ($\Delta f$ = " + f"{sweep}Hz per turn) will be "
+                + f"performed with respect to the highest voltage cavity at "
+                + f"{self.f_RF}Hz. The other cavities will be shifted "
+                + f"accordingly.")
+
+
+    def _install_zeta_shift(self):
+        if 'rf_sweep' not in self.env.elements:
+            idx = np.argsort(self.cavities[0]['s_pos'])
+            first_cavity = self.cavities[0]['name_first_in_line'][idx[0]]
+            self.env.elements['rf_sweep'] = xt.ZetaShift(dzeta=0)
+            self.line.insert('rf_sweep', at=f'{first_cavity}@start')

xcoll-0.6.2/pyproject.toml

@@ -1,52 +0,0 @@
-[tool.poetry]
-name = "xcoll"
-version = "0.6.2"
-description = "Xsuite collimation package"
-homepage = "https://github.com/xsuite/xcoll"
-repository = "https://github.com/xsuite/xcoll"
-authors = [
-           "Frederik F. Van der Veken <frederik@cern.ch>",
-           "Simone O. Solstrand <simone.otelie.solstrand@cern.ch>",
-           "Björn Lindstrom <bjorn.lindstrom@cern.ch>",
-           "Giacomo Broggi <giacomo.broggi@cern.ch>",
-           "André Donadon Servelle <andre.donadon.servelle@cern.ch>",
-           "Dora E. Veres <dora.erzsebet.veres@cern.ch>",
-           "Despina Demetriadou <despina.demetriadou@cern.ch>",
-           "Andrey Abramov <andrey.abramov@cern.ch>",
-           "Giovanni Iadarola <giovanni.iadarola@cern.ch>"
-]
-readme = "README.md"
-license = "Apache 2.0"
-exclude = ["xcoll/lib", "xcoll/config"]
-
-
-[tool.poetry.dependencies]
-python = ">=3.8"
-ruamel-yaml = { version = "^0.17.31", optional = true }
-numpy = ">=1.0"
-pandas = ">=1.4"
-xobjects = ">=0.5.0"
-xdeps = ">=0.10.5"
-xpart = ">=0.23.0"
-xtrack = ">=0.84.8"
-
-[poetry.group.dev.dependencies]
-pytest = ">=7.3"
-xaux = ">=0.3.5"
-
-[tool.poetry.extras]
-tests = ["pytest", "ruamel-yaml", "pytest-html", "pytest-xdist"]
-
-[build-system]
-# Needed for pip install -e (BTW: need pip version 22)
-requires = ["poetry-core>=1.0.8"]
-build-backend = "poetry.core.masonry.api"
-
-# pyproject.toml
-[tool.pytest.ini_options]
-addopts = "-ra --durations=10 --durations-min=1"
-python_functions = ["test_"]
-testpaths = [
-    "tests",
-]
-

xcoll-0.6.2/xcoll/rf_sweep.py

@@ -1,197 +0,0 @@
-# copyright ############################### #
-# This file is part of the Xcoll package.   #
-# Copyright (c) CERN, 2024.                 #
-# ######################################### #
-
-import numpy as np
-import scipy.constants as sc
-from time import perf_counter
-
-import xtrack as xt
-from xtrack.progress_indicator import progress
-
-class RFSweep:
-
-    def __init__(self, line):
-        self.line = line
-        self._get_base_frequency()
-        self._install_zeta_shift()
-        self.reset()
-
-    def _get_base_frequency(self):
-        self.cavities = self.line.get_elements_of_type(xt.Cavity)[1]
-        freq = np.unique([round(self.line[cav].frequency, 9) for cav in self.cavities])
-        if len(freq) > 1:
-            raise NotImplementedError(f"Cannot sweep multiple cavities with different frequencies!")
-        elif abs(freq[0]) < 1e-9 :
-            raise ValueError(f"Cavity frequency not set!")
-        self.f_RF = freq[0]
-        phi = np.unique(np.array([self.line[cav].lag for cav in self.cavities])*np.pi/180.)
-        if len(phi) > 1:
-            raise NotImplementedError(f"Cannot sweep multiple cavities with different phases!")
-        self.phi = phi[0]
-        self.V = np.array([self.line[cav].voltage for cav in self.cavities]).sum()
-        self.L = self.line.get_length()
-
-    def _install_zeta_shift(self):
-        if 'rf_sweep' in self.line.element_names:
-            print(f"Found existing RF sweep in line.")
-        else:
-            s_cav = min([self.line.get_s_position(cav) for cav in self.cavities])
-            if self.line.tracker is not None:
-                self.line.unfreeze()
-                line_was_built = True
-            else:
-                line_was_built = False
-            self.line.insert_element(element=xt.ZetaShift(dzeta=0), name='rf_sweep', at_s=s_cav)
-            if line_was_built:
-                self.line.build_tracker()
-
-    @property
-    def current_sweep_value(self):
-        dzeta = self.line['rf_sweep'].dzeta
-        return round(self.f_RF * dzeta/(self.L-dzeta), 6)
-
-
-    def info(self, sweep=0, num_turns=0):
-        if abs(sweep) > 1e-9:
-            existing_sweep = self.current_sweep_value
-
-            beta0 = self.line.particle_ref.beta0[0]
-            E = self.line.particle_ref.energy0[0]
-            q = self.line.particle_ref.q0
-            h = self.f_RF * self.L / beta0 / sc.c
-            tw = self.line.twiss()
-            eta = tw.slip_factor
-
-            bucket_height = np.sqrt(abs(q*self.V*beta0**2 / (np.pi*h*eta*E) * (
-                                2*np.cos(self.phi) + (2*self.phi-np.pi)*np.sin(self.phi)
-                            )))
-            delta_shift = -sweep / self.f_RF / eta
-            bucket_shift = delta_shift / bucket_height / 2
-            if num_turns > 0:
-                rf_shift_per_turn = sweep / num_turns
-                print(f"The current frequency is {self.f_RF + existing_sweep}Hz, adding {rf_shift_per_turn}Hz "
-                    + f"per turn until {self.f_RF + existing_sweep + sweep} (for {num_turns} turns).")
-            print(f"This sweep will move the center of the bucket with \u0394\u03B4 = "
-                + f"{delta_shift} ({bucket_shift} buckets).")
-            if num_turns > 0 and num_turns < 3*bucket_shift/tw.qs:
-                print(f"Warning: This is a very fast sweep, moving ~{round(bucket_shift, 2)} buckets in "
-                    + f"~{round(num_turns*tw.qs, 2)} synchrotron oscillations (on average). If the "
-                    + f"bucket moves faster than a particle can follow, that particle will move out of "
-                    + f"the bucket and remain uncaptured.")
-
-
-    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
-        existing_sweep = self.current_sweep_value
-
-        # Just set the new RF frequency, no tracking
-        if num_turns == 0:
-            sweep += existing_sweep
-            if verbose:
-                print(f"The current frequency is {self.f_RF + existing_sweep}Hz, moving to {self.f_RF + sweep}Hz."
-                     + "No tracking performed.")
-            self.line['rf_sweep'].dzeta = self.L * sweep / (self.f_RF + sweep)
-
-        # Sweep and track
-        else:
-            if self.line.tracker is None:
-                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
-
-            # 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 reset(self):
-        self.line['rf_sweep'].dzeta = 0
-
-
-    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