xcoll 0.5.9__tar.gz → 0.5.11__tar.gz

{xcoll-0.5.9 → xcoll-0.5.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xcoll
-Version: 0.5.9
+Version: 0.5.11
 Summary: Xsuite collimation package
 Home-page: https://github.com/xsuite/xcoll
 License: Apache 2.0
@@ -18,10 +18,10 @@ 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.7.3)
+Requires-Dist: xdeps (>=0.8.1)
 Requires-Dist: xobjects (>=0.4.5)
-Requires-Dist: xpart (>=0.19.1)
-Requires-Dist: xtrack (>=0.69.7)
+Requires-Dist: xpart (>=0.19.3)
+Requires-Dist: xtrack (>=0.70.3)
 Project-URL: Repository, https://github.com/xsuite/xcoll
 Description-Content-Type: text/markdown
 

{xcoll-0.5.9 → xcoll-0.5.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "xcoll"
-version = "0.5.9"
+version = "0.5.11"
 description = "Xsuite collimation package"
 homepage = "https://github.com/xsuite/xcoll"
 repository = "https://github.com/xsuite/xcoll"
@@ -26,16 +26,16 @@ ruamel-yaml = { version = "^0.17.31", optional = true }
 numpy = ">=1.0"
 pandas = ">=1.4"
 xobjects = ">=0.4.5"
-xdeps = ">=0.7.3"
-xpart = ">=0.19.1"
-xtrack = ">=0.69.7"
+xdeps = ">=0.8.1"
+xpart = ">=0.19.3"
+xtrack = ">=0.70.3"
 
 [tool.poetry.dev-dependencies]
 pytest = ">=7.3"
 xaux = ">=0.2.1"
 
 [tool.poetry.extras]
-tests = ["pytest", "ruamel-yaml"]
+tests = ["pytest", "ruamel-yaml", "pytest-html", "pytest-xdist"]
 
 [build-system]
 # Needed for pip install -e (BTW: need pip version 22)

{xcoll-0.5.9 → xcoll-0.5.11}/xcoll/beam_elements/base.py

@@ -12,8 +12,8 @@ from ..interaction_record import InteractionRecord
 from ..general import _pkg_root
 
 
-OPEN_JAW = 3.
-OPEN_GAP = 999.
+OPEN_JAW = 3
+OPEN_GAP = 999
 
 
 class InvalidXcoll(xt.BeamElement):
@@ -198,6 +198,10 @@ class BaseCollimator(BaseBlock):
                 to_assign['angle_L'] = kwargs.pop('angle_L', 0)
                 to_assign['angle_R'] = kwargs.pop('angle_R', 0)
 
+            # We do not allow any combination of jaw_ and gap_ attributes
+            # (except when jaw=..., gap=None or jaw=None, gap=... is used, as this is how the colldb installs it)
+            kwargs = {kk: vv for kk, vv in kwargs.items() if not vv is None}
+
             # Set jaw
             if 'jaw' in kwargs:
                 for key in ['jaw_L', 'jaw_R', 'jaw_LU', 'jaw_LD', 'jaw_RU', 'jaw_RD', 'gap', 'gap_L', 'gap_R']:
@@ -333,7 +337,7 @@ class BaseCollimator(BaseBlock):
         or   (self.tilt_L == 0 and self.tilt_R == 0):
             return [self.jaw_L, self.jaw_R]
         else:
-            return [[self.jaw_LU, self.jaw_RU], [self.jaw_LD, self.jaw_RD]]
+            return [[self.jaw_LU, self.jaw_LD], [self.jaw_RU, self.jaw_RD]]
 
     @jaw.setter   # Keeps the tilts unless all 4 corners are specified
     def jaw(self, val):
@@ -352,8 +356,8 @@ class BaseCollimator(BaseBlock):
             if hasattr(val[0], '__iter__'):
                 if hasattr(val[1], '__iter__') and len(val[0]) == 2 and len(val[1]) == 2:
                     self.jaw_LU = val[0][0]
-                    self.jaw_RU = val[0][1]
-                    self.jaw_LD = val[1][0]
+                    self.jaw_LD = val[0][1]
+                    self.jaw_RU = val[1][0]
                     self.jaw_RD = val[1][1]
                     return
             else:
@@ -556,6 +560,8 @@ class BaseCollimator(BaseBlock):
             print("Warning: Setting a tilt does not preserve the hierarchy, as there "
                 + "will always be one corner that tightens (the tilt is applied at "
                 + "the centre of the jaw).")
+            if val > np.pi/2 or val < -np.pi/2:
+                raise ValueError("Tilts larger than 90 degrees are not supported.")
         self._sin_yL = np.sin(val)
         self._cos_yL = np.cos(val)
         self._tan_yL = np.tan(val)
@@ -577,6 +583,8 @@ class BaseCollimator(BaseBlock):
             print("Warning: Setting a tilt does not preserve the hierarchy, as there "
                 + "will always be one corner that tightens (the tilt is applied at "
                 + "the centre of the jaw).")
+            if val > np.pi/2 or val < -np.pi/2:
+                raise ValueError("Tilts larger than 90 degrees are not supported.")
         self._sin_yR = np.sin(val)
         self._cos_yR = np.cos(val)
         self._tan_yR = np.tan(val)
@@ -651,6 +659,11 @@ class BaseCollimator(BaseBlock):
         self._nemitt_x = val
         self._apply_optics()
 
+    @property
+    def gemitt_x(self):
+        if self.nemitt_x is not None and self.optics_ready():
+            return self.nemitt_x / self.optics['beta_gamma_rel']
+
     @property
     def nemitt_y(self):
         if self._nemitt_y == 0:
@@ -666,6 +679,11 @@ class BaseCollimator(BaseBlock):
         self._nemitt_y = val
         self._apply_optics()
 
+    @property
+    def gemitt_y(self):
+        if self.nemitt_y is not None and self.optics_ready():
+            return self.nemitt_y / self.optics['beta_gamma_rel']
+
     @property
     def emittance(self):
         if self.nemitt_x is not None and self.nemitt_y is not None:
@@ -726,8 +744,8 @@ class BaseCollimator(BaseBlock):
             alfy = self.optics[self.align]['alfy'][0]
             betx = self.optics[self.align]['betx'][0]
             bety = self.optics[self.align]['bety'][0]
-            divx = -np.sqrt(self.nemitt_x/self.optics['beta_gamma_rel']/betx)*alfx
-            divy = -np.sqrt(self.nemitt_y/self.optics['beta_gamma_rel']/bety)*alfy
+            divx = -np.sqrt(self.gemitt_x/betx)*alfx
+            divy = -np.sqrt(self.gemitt_y/bety)*alfy
             if hasattr(self, '_cos_zL'):
                 if self.side != 'right':
                     return divx if abs(self.angle_L) < 1e-6 else divy
@@ -965,8 +983,6 @@ class BaseCollimator(BaseBlock):
             assert self._jaws_parallel == False
             assert np.isclose(self._sin_zDiff, self._cos_zL*self._sin_zR - self._sin_zL*self._cos_zR)
             assert np.isclose(self._cos_zDiff, self._cos_zL*self._cos_zR + self._sin_zL*self._sin_zR)
-        if self.side == 'both' and abs(self.tilt_L - self.tilt_R) >= 90.:
-            raise ValueError("Tilts of both jaws differ more than 90 degrees!")
         if self.side != 'right':
             ang = abs(np.arccos(self._cos_yL))
             ang = np.pi - ang if ang > np.pi/2 else ang
@@ -1058,6 +1074,11 @@ class BaseCrystal(BaseBlock):
             # Set angle
             to_assign['angle'] = kwargs.pop('angle', 0)
 
+            # We do not allow any combination of jaw_ and gap_ attributes
+            # (except when jaw=..., gap=None or jaw=None, gap=... is used, as this is how the colldb installs it)
+            kwargs = {kk: vv for kk, vv in kwargs.items() if not vv is None}
+
+
             # Set jaw
             if 'jaw' in kwargs:
                 for key in ['jaw_U', 'jaw_D', 'gap']:
@@ -1214,6 +1235,8 @@ class BaseCrystal(BaseBlock):
             if val > min(0, -self.bending_angle/2):
                 print("Warning: Setting a positive tilt does not preserve the hierarchy, as the "
                     + "crystal tightens towards the beam.")
+        if val > np.pi/2 or val < -np.pi/2:
+            raise ValueError("Tilts larger than 90 degrees are not supported.")
         self._sin_y = np.sin(val)
         self._cos_y = np.cos(val)
         self._tan_y = np.tan(val)
@@ -1243,6 +1266,10 @@ class BaseCrystal(BaseBlock):
     def nemitt_x(self, val):
         BaseCollimator.nemitt_x.fset(self, val)
 
+    @property
+    def gemitt_x(self):
+        return BaseCollimator.gemitt_x.fget(self)
+
     @property
     def nemitt_y(self):
         return BaseCollimator.nemitt_y.fget(self)
@@ -1251,6 +1278,10 @@ class BaseCrystal(BaseBlock):
     def nemitt_y(self, val):
         BaseCollimator.nemitt_y.fset(self, val)
 
+    @property
+    def gemitt_y(self):
+        return BaseCollimator.gemitt_y.fget(self)
+
     @property
     def emittance(self):
         return BaseCollimator.emittance.fget(self)

{xcoll-0.5.9 → xcoll-0.5.11}/xcoll/colldb.py

@@ -358,156 +358,158 @@ class CollimatorDatabase:
         return pd.DataFrame(self._collimator_dict).transpose()
 
     def to_yaml(self, out, lhc_style=True):
-        """
-        Writes a colldb in memory to disk in the yaml format.
-
-        > colldb_object.write_to_yaml(<path+name>, lhc_style=Bool)
-
-        if lhc_style == True, it will add comments assuming that the collimators are named
-        as in the lhc.
-
-        The function can dump b1, b2 and a general bx, however multi-beam functionality is not yet
-        added to the collmanager. TODO
-
-        If any of the dumped keys contains capital letters (e.g. gap_L), it will not be possible
-        to load it back into xcoll, since all keys are set to lowercase when importing TODO
-        """
-        # Dumps collimator database to a YAML file with optional LHC style formatting
-        import re
-
-        # Local helper functions
-        def _format_dict_entry(key, value, spacing='', mapping=False, key_width=15):
-            # Formats a dictionary entry into a string for YAML output
-            formatted_values = ',    '.join(f"{k}: {v}" for k, v in value.items())
-            formatted_values = re.sub(r'none', 'null', formatted_values, flags=re.IGNORECASE)
-            # Ensure key has a fixed width for alignment
-            if mapping:
-                formatted_key = f'{key}'.ljust(key_width)
-            else:
-                formatted_key = f'{key}:'.ljust(key_width)
-            #formatted_values = formatted_values.ljust(key_width)
-            return f"{spacing}{formatted_key} {{ {formatted_values} }}\n"
-
-        def _print_values(keys, dct, file, spacing='', mapping=False):
-            # Writes formatted dictionary entries to a file
-            for key in keys:
-                file.write(_format_dict_entry(key, dct[key], spacing=spacing, mapping=mapping))
-
-        def _print_colls(colls, dcts, beam, file):
-            # Filters and formats collimator data, then writes to a file
-            coll_items_to_print = ['<<','gap','angle','material','active','length','side']
-            file.write(f'  {beam}:\n')
-            for coll in colls:
-                coll_dict = dcts.to_pandas().transpose().to_dict()[coll]
-                fam = coll_dict['family']
-                fam_keys = []
-                if fam is not None:
-                    fam_keys = dcts._family_dict[fam].keys()
-                    coll_dict = {**{'<<': '*'+fam}, **coll_dict}
-                temp_items_to_print = []
-                if coll_dict['crystal'] and str(coll_dict['crystal'])!='nan':
-                    temp_items_to_print = ['bending_radius','width','height','miscut','crystal']
-                # if 'angle_L' in coll_dict and coll_dict['angle_L'] == coll_dict['angle_R']:
-                #     coll_dict.update({'angle': coll_dict['angle_L']})
-                # else:
-                #     temp_items_to_print = temp_items_to_print + ['angle_L','angle_R']
-                # if coll_dict['gap_L'] == coll_dict['gap_R']:
-                #     coll_dict.update({'gap': coll_dict['gap_L']})
-                # elif coll_dict['gap_L'] is None and coll_dict['gap_R'] is not None:
-                #     coll_dict.update({'gap': coll_dict['gap_R']})
-                # elif coll_dict['gap_L'] is not None and coll_dict['gap_R'] is None:
-                #     coll_dict.update({'gap': coll_dict['gap_L']})
-                # else:
-                #     temp_items_to_print = temp_items_to_print + ['gap_L','gap_R']
-                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='    '))
-            file.write('\n')
-
-        LHC_families = ['tcp3', 'tcsg3', 'tcsm3', 'tcla3', 'tcp7', 'tcsg7', 'tcsm7', 'tcla7', 'tcli', 'tdi', 'tcdq', 'tcstcdq', 'tcth1', 'tcth2', 'tcth5', 'tcth8', 'tctv1', 'tctv2', 'tctv5', 'tctv8', 'tclp', 'tcxrp', 'tcryo', 'tcl4', 'tcl5', 'tcl6', 'tct15', 'tct2', 'tct8', 'tcsp', 'tcld']
-        with open(f'{out}.yaml', 'w') as file:
-            if '_family_dict' in self.__dict__.keys():
-                file.write('families:\n')
-                if lhc_style:
-                    printed_families = []
-                    fams_in_dict = self._family_dict.keys()
-
-                    # Momentum cleaning
-                    file.write('  # Momentum cleaning\n')
-                    sel_fam = [fam for fam in LHC_families if re.match('.*3', fam) and (fam in fams_in_dict)]
-                    printed_families += sel_fam
-                    _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
-
-                    # Betatron cleaning
-                    file.write('  # Betatron cleaning\n')
-                    sel_fam = [fam for fam in LHC_families if re.match('.*7', fam) and (fam in fams_in_dict)]
-                    printed_families += sel_fam
-                    _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
-
-                    # Injection protection
-                    file.write('  # Injection protection\n')
-                    sel_fam = [fam for fam in LHC_families if (fam in ['tcli', 'tdi']) and (fam in fams_in_dict)]
-                    printed_families += sel_fam
-                    _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
-
-                    # Dump protection
-                    file.write('  # Dump protection\n')
-                    sel_fam = [fam for fam in LHC_families if (fam in ['tcdq', 'tcsp', 'tcstcdq']) and (fam in fams_in_dict)]
-                    printed_families += sel_fam
-                    _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
-
-                    # Physics background / debris
-                    file.write('  # Physics background / debris\n')
-                    sel_fam = [fam for fam in LHC_families if ((re.match('tc[lt][0-9dp].*', fam)) or (fam in ['tcryo', 'tcxrp'])) and (fam in fams_in_dict)]
-                    printed_families += sel_fam
-                    _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
-
-                    # Other families
-                    if set(printed_families) != set(fams_in_dict):
-                        file.write('  # Other families\n')
-                        _print_values(set(fams_in_dict) - set(printed_families), self._family_dict, file, spacing='  - &', mapping=True)
-                else:
-                    file.write('  # Families\n')
-                    _print_values(self._family_dict.keys(), self._family_dict, file, spacing='  - &', mapping=True)
-
-            # Emittance section
-            ex = self.nemitt_x
-            ey = self.nemitt_y
-            file.write(f'\nemittance:\n  x: {ex}\n  y: {ey}\n')
-
-            # Collimators section
-            file.write('\ncollimators:\n')
-            b1_colls, b2_colls, bx_colls = [], [], []
-            for coll in self.to_pandas().index:
-                if coll == 'tclia.4r2' or coll == 'tclia.4l8':    # TODO: hardcoded!!!
-                    b1_colls.append(coll)
-                    b2_colls.append(coll)
-                elif coll[-2:] == 'b1':
-                    b1_colls.append(coll)
-                elif coll[-2:] == 'b2':
-                    b2_colls.append(coll)
-                else:
-                    bx_colls.append(coll)
-
-            # Handle special cases for collimators
-            if (('tclia.4r2' in b1_colls) or ('tclia.4l8' in b1_colls)) and (len(b1_colls) <= 2):
-                b1_colls = []
-            if (('tclia.4r2' in b2_colls) or ('tclia.4l8' in b2_colls)) and (len(b2_colls) <= 2):
-                b2_colls = []
-
-            # Print collimators for each beam
-            if len(b1_colls) > 0:
-                _print_colls(b1_colls, self, 'b1', file)
-            if len(b2_colls) > 0:
-                _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.')
+        raise NotImplementedError("This functionality needs to be updated to be compatible with tilts and read "
+                                + "the families from the colldb instead of hard-coding them!")
+        # """
+        # Writes a colldb in memory to disk in the yaml format.
+
+        # > colldb_object.write_to_yaml(<path+name>, lhc_style=Bool)
+
+        # if lhc_style == True, it will add comments assuming that the collimators are named
+        # as in the lhc.
+
+        # The function can dump b1, b2 and a general bx, however multi-beam functionality is not yet
+        # added to the collmanager. TODO
+
+        # If any of the dumped keys contains capital letters (e.g. gap_L), it will not be possible
+        # to load it back into xcoll, since all keys are set to lowercase when importing TODO
+        # """
+        # # Dumps collimator database to a YAML file with optional LHC style formatting
+        # import re
+
+        # # Local helper functions
+        # def _format_dict_entry(key, value, spacing='', mapping=False, key_width=15):
+        #     # Formats a dictionary entry into a string for YAML output
+        #     formatted_values = ',    '.join(f"{k}: {v}" for k, v in value.items())
+        #     formatted_values = re.sub(r'none', 'null', formatted_values, flags=re.IGNORECASE)
+        #     # Ensure key has a fixed width for alignment
+        #     if mapping:
+        #         formatted_key = f'{key}'.ljust(key_width)
+        #     else:
+        #         formatted_key = f'{key}:'.ljust(key_width)
+        #     #formatted_values = formatted_values.ljust(key_width)
+        #     return f"{spacing}{formatted_key} {{ {formatted_values} }}\n"
+
+        # def _print_values(keys, dct, file, spacing='', mapping=False):
+        #     # Writes formatted dictionary entries to a file
+        #     for key in keys:
+        #         file.write(_format_dict_entry(key, dct[key], spacing=spacing, mapping=mapping))
+
+        # def _print_colls(colls, dcts, beam, file):
+        #     # Filters and formats collimator data, then writes to a file
+        #     coll_items_to_print = ['<<','gap','angle','material','active','length','side']
+        #     file.write(f'  {beam}:\n')
+        #     for coll in colls:
+        #         coll_dict = dcts.to_pandas().transpose().to_dict()[coll]
+        #         fam = coll_dict['family']
+        #         fam_keys = []
+        #         if fam is not None:
+        #             fam_keys = dcts._family_dict[fam].keys()
+        #             coll_dict = {**{'<<': '*'+fam}, **coll_dict}
+        #         temp_items_to_print = []
+        #         if coll_dict['crystal'] and str(coll_dict['crystal'])!='nan':
+        #             temp_items_to_print = ['bending_radius','width','height','miscut','crystal']
+        #         # if 'angle_L' in coll_dict and coll_dict['angle_L'] == coll_dict['angle_R']:
+        #         #     coll_dict.update({'angle': coll_dict['angle_L']})
+        #         # else:
+        #         #     temp_items_to_print = temp_items_to_print + ['angle_L','angle_R']
+        #         # if coll_dict['gap_L'] == coll_dict['gap_R']:
+        #         #     coll_dict.update({'gap': coll_dict['gap_L']})
+        #         # elif coll_dict['gap_L'] is None and coll_dict['gap_R'] is not None:
+        #         #     coll_dict.update({'gap': coll_dict['gap_R']})
+        #         # elif coll_dict['gap_L'] is not None and coll_dict['gap_R'] is None:
+        #         #     coll_dict.update({'gap': coll_dict['gap_L']})
+        #         # else:
+        #         #     temp_items_to_print = temp_items_to_print + ['gap_L','gap_R']
+        #         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='    '))
+        #     file.write('\n')
+
+        # LHC_families = ['tcp3', 'tcsg3', 'tcsm3', 'tcla3', 'tcp7', 'tcsg7', 'tcsm7', 'tcla7', 'tcli', 'tdi', 'tcdq', 'tcstcdq', 'tcth1', 'tcth2', 'tcth5', 'tcth8', 'tctv1', 'tctv2', 'tctv5', 'tctv8', 'tclp', 'tcxrp', 'tcryo', 'tcl4', 'tcl5', 'tcl6', 'tct15', 'tct2', 'tct8', 'tcsp', 'tcld']
+        # with open(f'{out}.yaml', 'w') as file:
+        #     if '_family_dict' in self.__dict__.keys():
+        #         file.write('families:\n')
+        #         if lhc_style:
+        #             printed_families = []
+        #             fams_in_dict = self._family_dict.keys()
+
+        #             # Momentum cleaning
+        #             file.write('  # Momentum cleaning\n')
+        #             sel_fam = [fam for fam in LHC_families if re.match('.*3', fam) and (fam in fams_in_dict)]
+        #             printed_families += sel_fam
+        #             _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
+
+        #             # Betatron cleaning
+        #             file.write('  # Betatron cleaning\n')
+        #             sel_fam = [fam for fam in LHC_families if re.match('.*7', fam) and (fam in fams_in_dict)]
+        #             printed_families += sel_fam
+        #             _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
+
+        #             # Injection protection
+        #             file.write('  # Injection protection\n')
+        #             sel_fam = [fam for fam in LHC_families if (fam in ['tcli', 'tdi']) and (fam in fams_in_dict)]
+        #             printed_families += sel_fam
+        #             _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
+
+        #             # Dump protection
+        #             file.write('  # Dump protection\n')
+        #             sel_fam = [fam for fam in LHC_families if (fam in ['tcdq', 'tcsp', 'tcstcdq']) and (fam in fams_in_dict)]
+        #             printed_families += sel_fam
+        #             _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
+
+        #             # Physics background / debris
+        #             file.write('  # Physics background / debris\n')
+        #             sel_fam = [fam for fam in LHC_families if ((re.match('tc[lt][0-9dp].*', fam)) or (fam in ['tcryo', 'tcxrp'])) and (fam in fams_in_dict)]
+        #             printed_families += sel_fam
+        #             _print_values(sel_fam, self._family_dict, file, spacing='  - &', mapping=True)
+
+        #             # Other families
+        #             if set(printed_families) != set(fams_in_dict):
+        #                 file.write('  # Other families\n')
+        #                 _print_values(set(fams_in_dict) - set(printed_families), self._family_dict, file, spacing='  - &', mapping=True)
+        #         else:
+        #             file.write('  # Families\n')
+        #             _print_values(self._family_dict.keys(), self._family_dict, file, spacing='  - &', mapping=True)
+
+        #     # Emittance section
+        #     ex = self.nemitt_x
+        #     ey = self.nemitt_y
+        #     file.write(f'\nemittance:\n  x: {ex}\n  y: {ey}\n')
+
+        #     # Collimators section
+        #     file.write('\ncollimators:\n')
+        #     b1_colls, b2_colls, bx_colls = [], [], []
+        #     for coll in self.to_pandas().index:
+        #         if coll == 'tclia.4r2' or coll == 'tclia.4l8':    # TODO: hardcoded!!!
+        #             b1_colls.append(coll)
+        #             b2_colls.append(coll)
+        #         elif coll[-2:] == 'b1':
+        #             b1_colls.append(coll)
+        #         elif coll[-2:] == 'b2':
+        #             b2_colls.append(coll)
+        #         else:
+        #             bx_colls.append(coll)
+
+        #     # Handle special cases for collimators
+        #     if (('tclia.4r2' in b1_colls) or ('tclia.4l8' in b1_colls)) and (len(b1_colls) <= 2):
+        #         b1_colls = []
+        #     if (('tclia.4r2' in b2_colls) or ('tclia.4l8' in b2_colls)) and (len(b2_colls) <= 2):
+        #         b2_colls = []
+
+        #     # Print collimators for each beam
+        #     if len(b1_colls) > 0:
+        #         _print_colls(b1_colls, self, 'b1', file)
+        #     if len(b2_colls) > 0:
+        #         _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.')
 
 
     # ====================================
@@ -539,27 +541,14 @@ class CollimatorDatabase:
                                + f"(or xtrack.Drift)!\nPlease check the name, or correct the "
                                + f"element.")
 
-    def _create_collimator(self, line, collimator_class, name, **kwargs):
-        assert issubclass(collimator_class, BaseCollimator)
-        self._check_installed(line, name, collimator_class)
+    def _create_collimator(self, cls, line, name, **kwargs):
+        self._check_installed(line, name, cls)
         if kwargs.pop('verbose', False):
-            print(f"Installing {name:20} as {collimator_class.__name__}")
-        el = collimator_class(gap=self[name]['gap'], angle=self[name]['angle'],
-                              length=self[name]['length'], side=self[name]['side'],
-                              _tracking=False, **kwargs)
-        el.emittance = [self.nemitt_x, self.nemitt_y]
-        self._elements[name] = el
-
-    def _create_crystal(self, line, crystal_class, name, **kwargs):
-        assert issubclass(crystal_class, BaseCrystal)
-        self._check_installed(line, name, crystal_class)
-        if kwargs.pop('verbose', False):
-            print(f"Installing {name:20} as {crystal_class.__name__}")
-        el = crystal_class(gap=self[name]['gap'], angle=self[name]['angle'],
-                           length=self[name]['length'], side=self[name]['side'],
-                           bending_radius=self[name]['bending_radius'],
-                           width=self[name]['width'], height=self[name]['height'],
-                           _tracking=False, **kwargs)
+            print(f"Installing {name:20} as {cls.__name__}")
+        prop_dict = {kk: vv for kk, vv in self[name].items() \
+                     if kk in cls._xofields or kk in cls._store_in_to_dict}
+        prop_dict.update(kwargs)
+        el = cls(**prop_dict)
         el.emittance = [self.nemitt_x, self.nemitt_y]
         self._elements[name] = el
 
@@ -567,9 +556,9 @@ class CollimatorDatabase:
         names = self._get_names_from_line(line, names, families)
         for name in names:
             if self[name]['bending_radius'] is None:
-                self._create_collimator(line, BlackAbsorber, name, verbose=verbose)
+                self._create_collimator(BlackAbsorber, line, name, verbose=verbose)
             else:
-                self._create_crystal(line, BlackCrystal, name, verbose=verbose)
+                self._create_collimator(BlackCrystal, line, name, verbose=verbose)
         elements = [self._elements[name] for name in names]
         line.collimators.install(names, elements, need_apertures=need_apertures)
 
@@ -578,12 +567,11 @@ class CollimatorDatabase:
         for name in names:
             mat = SixTrack_to_xcoll(self[name]['material'])
             if self[name]['bending_radius'] is None:
-                self._create_collimator(line, EverestCollimator, name, material=mat[0],
+                self._create_collimator(EverestCollimator, line, name, material=mat[0],
                                         verbose=verbose)
             else:
-                self._create_crystal(line, EverestCrystal, name, material=mat[1],
-                                     lattice=self[name]['crystal'], verbose=verbose,
-                                     miscut=self[name]['miscut'])
+                self._create_collimator(EverestCrystal, line, name, material=mat[1],
+                                        verbose=verbose)
         elements = [self._elements[name] for name in names]
         line.collimators.install(names, elements, need_apertures=need_apertures)
 

{xcoll-0.5.9 → xcoll-0.5.11}/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.5.9'
+__version__ = '0.5.11'
 # ======================

{xcoll-0.5.9 → xcoll-0.5.11}/xcoll/initial_distribution.py

@@ -57,8 +57,14 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
     if twiss is None:
         twiss = line.twiss()
 
-    # Is it converging or diverging?    # TODO: This might change with a tilt!!!!!!
-    is_converging  = twiss[f'alf{plane}', name] > 0
+    # Is it converging or diverging?
+    # TODO: dispersion might change this
+    # TODO: this should be checked jaw by jaw (we are currently checking the left jaw - watch out for sign of tilt of right jaw)
+    # TODO: skew collimators
+    tilt = coll.tilt[0] if isinstance(coll.tilt, list) else coll.tilt
+    betatron_angle = coll.gap * coll.divergence
+    tolerance_tilt = 1e-12 # 0.1 urad tolerance on jaw tilt  =>  we prioritise converging
+    is_converging = tilt + tolerance_tilt >= betatron_angle
     print(f"Collimator {name} is {'con' if is_converging else 'di'}verging.")
 
     beam_sizes = twiss.get_beam_covariance(nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y)