xcoll 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

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/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/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.6'
 # ===================

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-0.3.4.dist-info → xcoll-0.3.6.dist-info}/METADATA

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

{xcoll-0.3.4.dist-info → xcoll-0.3.6.dist-info}/RECORD

@@ -3,23 +3,23 @@ NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
 xcoll/__init__.py,sha256=XO53i_KKg6kaJfTCpmNJLy14tou4RMF9MDuXlH70iTQ,642
 xcoll/beam_elements/__init__.py,sha256=4ekQv_1Kxwqd0VSaPzMpVQ1V9RMX5tr-ycBBiTIFyRw,776
 xcoll/beam_elements/absorber.py,sha256=yxF-NhnQs_l7s-M_0AhdKzx0YVgvhujOAwQnAbVTSOs,1260
-xcoll/beam_elements/base.py,sha256=93pTQ1vwYaZ1uSfRiwWx-yr6pp7EGUfRPloUZYHrXDk,15353
+xcoll/beam_elements/base.py,sha256=Nwgg2v1RdEEpByFG8AZgL7E6ecESbYm3mqnIfoh5SFQ,15785
 xcoll/beam_elements/collimators_src/absorber.h,sha256=ZsAzIZ3uWP15unMJyuxbuD5lVgFeWAwQsuyOj_ZSkjI,5502
 xcoll/beam_elements/collimators_src/everest_block.h,sha256=Gpr5wyFLtcaNgVj3auKgFWB-2KPP5cjbEWTfC3f9q1U,5596
 xcoll/beam_elements/collimators_src/everest_collimator.h,sha256=wsGIFBb-VMCRhxwK2G5k9SEX-6SzIgRX4XE5g_I2okg,6296
 xcoll/beam_elements/collimators_src/everest_crystal.h,sha256=pX9cmW5lfsUoszR14tXWU4Emu-0VwX4bbfg6Q6pxXR4,5164
 xcoll/beam_elements/everest.py,sha256=JxmREXln3p3c9XEcHp-DgDbXFzc9txiGrxvjOrDUacs,11369
-xcoll/colldb.py,sha256=vXF-9MgpF9iKfm5oOQKm8Q0mfzdplPb1rdd8j_HUxfQ,49527
+xcoll/colldb.py,sha256=_MywNPz9cDocGPsfXpBVkjADyAGGVFQQRUI4kxOQaOg,49476
 xcoll/collimator_settings.py,sha256=uj25cpSaM-EY7VM65OUq3Cxb8gedSmi-r9mpkP-ikrg,17015
-xcoll/general.py,sha256=knZwyYB7rrpcteXV15Hh8ItXCNWxnBf8jpH6YvRQ_ac,525
+xcoll/general.py,sha256=Pcqa-FWVvvs3X5gBEhnNw2qmw-Rl9IdCMX2Mww33wt4,525
 xcoll/headers/checks.h,sha256=3CLIa2etrJVKBfxdsHcS8NirqzVDHG9aCbrlMwkWuHw,1586
 xcoll/headers/particle_states.h,sha256=YxpS56_ZD8mhrqKTohcyVe3hbQZ6H9vspYHSVM6a5-U,844
 xcoll/impacts/__init__.py,sha256=yKQdIeLG2CT3qo-iZB6MODm-WHzkTmulgUvOAS41w7E,39
 xcoll/impacts/impacts.py,sha256=J7FfcxIprUMHyn5oZhmjhvBewXOp7kwJtErStOsAwJQ,3953
 xcoll/impacts/impacts_src/impacts.h,sha256=0eq4Z8xDPMJzN8yPF3lTkDsld6W54tbXmynZEbXCh4c,5950
 xcoll/impacts/interaction_types.py,sha256=iHbhC4beN1vZRWeHbALqSFo8lG10R7asc08BRpIqXx0,2493
-xcoll/manager.py,sha256=AF9XziOQ2JT62VSKdzOZzlNFeuaqowzWfqv6msONR9k,45001
-xcoll/rf_sweep.py,sha256=sUstmeeXmozmM4f8xUGCKLyJYOV3xD4aVZG_FCRspVA,5402
+xcoll/manager.py,sha256=A7hCuSgS7sSFKoV26Nt8xc_GXUNOm-WRqsGftRoBMzA,45052
+xcoll/rf_sweep.py,sha256=_cnDb2d0_tfVrrwuubLOSyUHpMQHMvmSzYgd_-WfJg0,8936
 xcoll/scattering_routines/everest/__init__.py,sha256=yuCzolZkiL7mHaBXCKHRl66VxjKFB0t9Vyw9TtFv1Bo,286
 xcoll/scattering_routines/everest/constants.h,sha256=9Kus-RKyxsGdlxx9tfMWUz2H5Q-qy6y_F4kddIZzO30,1390
 xcoll/scattering_routines/everest/crystal.h,sha256=WlA_e9ePuVhc9fmMVdSXFBKvHrFvCQ5aG4d67iMrtuQ,54119
@@ -104,8 +104,8 @@ xcoll/scattering_routines/fluka/flukaio/tests/fakes/FakeFlukaIOServer.h,sha256=0
 xcoll/scattering_routines/fluka/flukaio/tests/fakes/FakeFortranFlukaIO.h,sha256=WxAjG2SPlWqMG38iHBM8GQwh5yp031jAGeyPi4y9G4c,330
 xcoll/scattering_routines/fluka/flukaio/tests/fakes/FakeNetIO.c,sha256=JcgCWxbT-ZOeG0c9mz1kk3PfXjaazaUk2P_ROtoBKaY,1452
 xcoll/scattering_routines/fluka/flukaio/tests/fakes/FakeNetIO.h,sha256=JTMU9FtCBJkitsnHp7kWKUbPSN4rviRfJgVHK26AuPY,906
-xcoll-0.3.4.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-xcoll-0.3.4.dist-info/METADATA,sha256=788lvabRGh9xA1nMET6gwOk2-oWD7OKEJG5TGp764Ms,2709
-xcoll-0.3.4.dist-info/NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
-xcoll-0.3.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-xcoll-0.3.4.dist-info/RECORD,,
+xcoll-0.3.6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+xcoll-0.3.6.dist-info/METADATA,sha256=-HXnJmvM23LxNmSOOxUuLhk49SuI8DnG7uIWl7DtHPg,2709
+xcoll-0.3.6.dist-info/NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
+xcoll-0.3.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+xcoll-0.3.6.dist-info/RECORD,,