xcoll 0.3.4__tar.gz → 0.3.5__tar.gz

{xcoll-0.3.4 → xcoll-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xcoll
-Version: 0.3.4
+Version: 0.3.5
 Summary: Xsuite collimation package
 Home-page: https://github.com/xsuite/xcoll
 License: Apache 2.0

{xcoll-0.3.4 → xcoll-0.3.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "xcoll"
-version = "0.3.4"
+version = "0.3.5"
 description = "Xsuite collimation package"
 homepage = "https://github.com/xsuite/xcoll"
 repository = "https://github.com/xsuite/xcoll"

{xcoll-0.3.4 → xcoll-0.3.5}/xcoll/beam_elements/base.py

@@ -60,6 +60,14 @@ class BaseBlock(xt.BeamElement):
         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,
                               _context=_context, _buffer=_buffer, _offset=_offset)
@@ -196,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-0.3.4 → xcoll-0.3.5}/xcoll/colldb.py

@@ -214,9 +214,6 @@ class CollimatorDatabase:
                                 cry[sline[1]][key] = float(sline[idx])
                             else:
                                 cry[sline[1]][key] = int(sline[idx])
-                    elif sline[0].lower() == "target":
-                        variables['xdim'][sline[1]] = float(sline[2])
-                        variables['ydim'][sline[1]] = float(sline[3])
                     elif sline[0].lower() == 'settings':
                         pass # Acknowledge and ignore this line
                     elif len(sline) == 6:
@@ -231,7 +228,7 @@ class CollimatorDatabase:
         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), sep='\s+', 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')))
@@ -307,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:
@@ -324,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='    '))

{xcoll-0.3.4 → xcoll-0.3.5}/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.4'
+__version__ = '0.3.5'
 # ===================

{xcoll-0.3.4 → xcoll-0.3.5}/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-0.3.5/xcoll/rf_sweep.py

@@ -0,0 +1,191 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2023.                 #
+# ######################################### #
+
+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()
+
+    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 _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-0.3.4/xcoll/rf_sweep.py

@@ -1,120 +0,0 @@
-# copyright ############################### #
-# This file is part of the Xcoll Package.   #
-# Copyright (c) CERN, 2023.                 #
-# ######################################### #
-
-import numpy as np
-import pandas as pd
-import scipy.constants as sc
-
-import xtrack as xt
-
-
-class RFSweep:
-
-    def __init__(self, line):
-        self.line = line
-        self._get_base_frequency()
-        self._install_zeta_shift()
-
-    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, num_turns=0, particles=None, 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!")
-            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