xcoll 0.5.8__py3-none-any.whl → 0.5.9__py3-none-any.whl

xcoll/beam_elements/__init__.py.orig

@@ -0,0 +1,32 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+from .base import BaseBlock, BaseCollimator, BaseCrystal
+from .absorber import BlackAbsorber, BlackCrystal
+from .everest import EverestBlock, EverestCollimator, EverestCrystal
+from .blowup import BlowUp
+from .monitor import EmittanceMonitor
+
+block_classes = tuple(v for v in globals().values()
+                      if isinstance(v, type) and issubclass(v, BaseBlock) and v != BaseBlock
+                      and v != BaseCollimator and v != BaseCrystal)
+# Includes crystals
+collimator_classes = tuple(v for v in globals().values()
+                           if isinstance(v, type) and (issubclass(v, BaseCollimator) or issubclass(v, BaseCrystal))
+                           and v != BaseCollimator and v != BaseCrystal)
+crystal_classes = tuple(v for v in globals().values()
+                        if isinstance(v, type) and issubclass(v, BaseCrystal) and v != BaseCrystal)
+
+element_classes = block_classes + (BlowUp, EmittanceMonitor)
+<<<<<<< HEAD
+
+# These should not go into any of the classes lists
+# but are added for compatibility with _K2
+from .k2 import _K2Collimator, _K2Crystal
+_all_block_classes      = block_classes + (_K2Collimator, _K2Crystal)
+_all_collimator_classes = collimator_classes + (_K2Collimator, _K2Crystal)
+_all_crystal_classes    = crystal_classes + (_K2Crystal,)
+=======
+>>>>>>> release/v0.6.0

xcoll/colldb.py.orig

@@ -0,0 +1,654 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import io
+import re
+import json
+import numpy as np
+import pandas as pd
+from pathlib import Path
+
+import xtrack as xt
+
+from .beam_elements import BlackAbsorber, BlackCrystal, EverestCollimator, EverestCrystal, \
+<<<<<<< HEAD
+                           FlukaCollimator, BaseCollimator, BaseCrystal, _all_collimator_classes
+=======
+                           BaseCollimator, BaseCrystal, collimator_classes
+>>>>>>> release/v0.6.0
+from .scattering_routines.everest.materials import SixTrack_to_xcoll
+from .scattering_routines.fluka import FlukaEngine
+
+
+def _initialise_None(dct):
+    fields = {'gap': None, 'angle': 0, 'offset': 0, 'parking': 1, 'jaw': None, 'family': None}
+    fields.update({'overwritten_keys': [], 'side': 'both', 'material': None, 'stage': None})
+    fields.update({'length': 0, 'collimator_type': None, 'active': True, 'crystal': None, 'tilt': 0})
+    fields.update({'bending_radius': None, 'bending_angle': None, 'width': 0, 'height': 0, 'miscut': 0})
+    for f, val in fields.items():
+        if f not in dct.keys():
+            dct[f] = val
+    for key in dct.keys():
+        if key not in fields.keys():
+            raise ValueError(f"Illegal setting {key} in collimator!")
+
+
+def _dict_keys_to_lower(dct):
+    if isinstance(dct, dict):
+        return {k.lower(): _dict_keys_to_lower(v) for k,v in dct.items()}
+    else:
+        return dct
+
+
+def _get_coll_dct_by_beam(coll, beam):
+    # The dictionary can be a CollimatorDatabase for a single beam (beam=None)
+    # or for both beams (beam='b1' or beam='b2)
+    if beam is not None:
+        if isinstance(beam, int) or len(beam) == 1:
+            beam = f'b{beam}'
+        beam = beam.lower()
+    beam_in_db = list(coll.keys())
+
+    if beam_in_db == ['b1','b2']:
+        if beam is None:
+            raise ValueError("Need to specify a beam, because the given dict is for both beams!")
+        return coll[beam]
+
+    elif len(beam_in_db) == 1:
+        if beam is None:
+            beam = beam_in_db[0].lower()
+        elif beam != beam_in_db[0].lower():
+            raise ValueError("Variable 'beam' does not match beam specified in CollimatorDatabase!")
+        return coll[beam]
+
+    elif beam is not None:
+        print("Warning: Specified a beam, but the CollimatorDatabase is for a single beam only!")
+    return coll
+
+
+class CollimatorDatabase:
+
+    def __init__(self, **kwargs):
+        # Get all arguments
+        for key in ['collimator_dict', 'nemitt_x', 'nemitt_y']:
+            if key not in kwargs.keys():
+                raise ValueError(f"CollimatorDatabase is missing required argument '{key}'!")
+
+        self._parse_dict(kwargs['collimator_dict'],
+                         kwargs.get('family_dict', {}),
+                         kwargs.get('beam', None),
+                         kwargs.get('_yaml_merged', False),
+                         kwargs.get('ignore_crystals', True))
+        self.nemitt_x = kwargs['nemitt_x']
+        self.nemitt_y = kwargs['nemitt_y']
+        self._elements = {}
+
+
+    def _parse_dict(self, coll, fam, beam=None, _yaml_merged=False, ignore_crystals=True):
+        # We make all keys case-insensitive to avoid confusion between different conventions
+        coll = _dict_keys_to_lower(coll)
+        fam  = _dict_keys_to_lower(fam)
+
+        # The dictionary can be a CollimatorDatabase for a single beam (beam=None)
+        # or for both beams (beam='b1' or beam='b2)
+        coll = _get_coll_dct_by_beam(coll, beam)
+
+        # Apply family settings
+        crystals = []
+        for thiscoll, settings in coll.items():
+            settings = {k.lower(): v for k,v in settings.items()}
+            if 'family' in settings.keys() and settings['family'] is not None:
+                settings['family'] = settings['family'].lower()
+                thisfam = settings['family']
+                if thisfam not in fam.keys():
+                    raise ValueError(f"Collimator {thiscoll} depends on family {thisfam}, "
+                                   + f"but the latter is not defined!")
+
+                # Check if some family settings are overwritten for this collimator
+                # 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.lower() for key in settings.keys() if key in fam[thisfam]]
+                    if len(overwritten_keys) > 0:
+                        settings['overwritten_keys'] = overwritten_keys
+
+                # Load family settings, potentially overwriting settings for this collimator
+                settings = {**fam[thisfam], **settings}
+
+            else:
+                settings['family'] = None
+            coll[thiscoll] = settings
+
+            # Save list of crystals
+            if 'crystal' in settings:
+                if settings['crystal'] != 0.0:
+                    crystals += [thiscoll]
+                else:
+                    settings['crystal'] = None
+
+        # Remove crystals from colldb
+        if ignore_crystals:
+            for thiscoll in crystals:
+                del coll[thiscoll]
+
+        # Check that all collimators have gap settings
+        if not np.all(['gap' in val.keys() or 'jaw' in val.keys() for val in coll.values()]):
+            raise ValueError("Ill-defined CollimatorDatabase: Not all collimators have a gap or "
+                           + "jaw setting, (or the keys / structure of the dictionary is wrong)!")
+
+        # Update collimators with default values for missing keys
+        for name, collimator in coll.items():
+            # Change all values to lower case
+            for key, val in collimator.items():
+                collimator[key] = val.lower() if isinstance(val, str) else val
+            _initialise_None(collimator)
+
+        self._collimator_dict = coll
+        self._family_dict = fam
+
+    # =======================================
+    # ====== Loading/dumping functions ======
+    # =======================================
+
+    @classmethod
+    def from_yaml(cls, file, **kwargs):
+        # Only do the import here, as to not force people to install
+        # ruamel if they don't load CollimatorDatabase yaml's
+        from ruamel.yaml import YAML
+        yaml = YAML(typ='safe')
+        if isinstance(file, io.IOBase):
+            dct = yaml.load(file)
+        else:
+            file = Path(file).resolve()
+            with file.open('r') as fid:
+                dct = yaml.load(fid)
+        dct = _dict_keys_to_lower(dct)
+
+        # If the CollimatorDatabase uses YAML merging, we need a bit of hackery to get the
+        # family names from the tags (anchors/aliases)
+        _yaml_merged = False
+        if 'families' in dct.keys() and not isinstance(dct['families'], dict):
+            _yaml_merged = True
+
+            # First we load a round-trip yaml
+            yaml = YAML()
+            if isinstance(file, io.IOBase):
+                full_dct = yaml.load(file)
+            else:
+                with open(file, 'r') as fid:
+                    full_dct = yaml.load(fid)
+            famkey = [key for key in full_dct.keys() if key.lower() == 'families'][0]
+            collkey = [key for key in full_dct.keys() if key.lower() == 'collimators'][0]
+            families = {}
+
+            # We loop over family names to get the name tag ('anchor') of each family
+            for fam, full_fam in zip(dct['families'], full_dct[famkey]):
+                if full_fam.anchor.value is None:
+                    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()] = {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
+            beam = kwargs.get('beam', None)
+            coll_dct = _get_coll_dct_by_beam(dct['collimators'], beam)
+            full_coll_dct = _get_coll_dct_by_beam(full_dct[collkey], beam)
+            for coll, full_coll in zip(coll_dct.values(), full_coll_dct.values()):
+                if not isinstance(coll['gap'], (int,float)):
+                    coll['gap'] = None
+                if 'family' in coll.keys():
+                    raise ValueError(f"Error in {coll}: Cannot use merging for families "
+                                    + "and manually specify family as well!")
+                elif len(full_coll.merge) > 0:
+                    coll['family'] = full_coll.merge[0][1].anchor.value.lower()
+                    # Check if some family settings are overwritten for this collimator
+                    overwritten_keys = [key.lower() for key in full_coll.keys()
+                                        if full_coll._unmerged_contains(key)
+                                        and key.lower() in families[coll['family']].keys()]
+                    if len(overwritten_keys) > 0:
+                        coll['overwritten_keys'] = overwritten_keys
+                else:
+                    coll['family'] = None
+
+        return cls.from_dict(dct, _yaml_merged=_yaml_merged, **kwargs)
+
+
+    @classmethod
+    def from_json(cls, file, **kwargs):
+        if isinstance(file, io.IOBase):
+            dct = json.load(file)
+        else:
+            file = Path(file).resolve()
+            with file.open('r') as fid:
+                dct = json.load(fid)
+        dct = _dict_keys_to_lower(dct)
+        return cls.from_dict(dct, **kwargs)
+
+
+    @classmethod
+    def from_dict(cls, dct, beam=None, _yaml_merged=False, nemitt_x=None, nemitt_y=None, ignore_crystals=True):
+        # We make all keys case-insensitive to avoid confusion between different conventions
+        # The families are optional
+        fam = {}
+        dct = _dict_keys_to_lower(dct)
+
+        # Get the emittance
+        if nemitt_x is None and nemitt_y is None:
+            if 'emittance' not in dct.keys():
+                raise ValueError("Missing emittance info! Add 'emittance' as a key to "
+                               + "the CollimatorDatabase file, or specify it as 'nemitt_x' "
+                               + "and 'nemitt_y' to the loader!")
+            nemitt_x = dct['emittance']['x']
+            nemitt_y = dct['emittance']['y']
+        elif nemitt_x is None or nemitt_y is None:
+            raise ValueError("Need to provide both 'nemitt_x' and 'nemitt_y'!")
+        elif 'emittance' in dct.keys():
+            if dct['emittance']['x'] != nemitt_x or dct['emittance']['y'] != nemitt_y:
+                raise ValueError("Emittance in CollimatorDatabase file different from "
+                               + "'nemitt_x' and 'nemitt_y'!")
+
+        # Get family and collimator dicts
+        if 'families' in dct.keys():
+            if not 'collimators' in dct.keys():
+                raise ValueError("Could not find 'collimators' dict in CollimatorDatabase!")
+            fam  = dct['families']
+            coll = dct['collimators']
+        elif 'collimators' in dct.keys():
+            coll = dct['collimators']
+        else:
+            coll = dct
+
+        return cls(collimator_dict=coll, family_dict=fam, nemitt_x=nemitt_x, nemitt_y=nemitt_y,
+                   beam=beam, _yaml_merged=_yaml_merged, ignore_crystals=ignore_crystals)
+
+
+    @classmethod
+    def from_SixTrack(cls, file, ignore_crystals=True, **kwargs):
+        file = Path(file).resolve()
+        with file.open('r') as fp:
+            coll_data_string = ''
+            family_settings = {}
+            family_types = {}
+            side = {}
+            cry_fields = ['bending_radius', 'width', 'height', 'thick', 'tilt', 'miscut', 'crystal']
+            cry = {}
+
+            for line in fp:
+                if line.startswith('#'):
+                    continue # Comment
+                sline = line.split()
+                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 < 6:
+                                cry[sline[1]][key] = float(sline[idx])
+                            else:
+                                cry[sline[1]][key] = int(sline[idx])
+                    elif sline[0].lower() == 'settings':
+                        pass # Acknowledge and ignore this line
+                    elif len(sline) == 6:
+                        # Standard collimator definition
+                        coll_data_string += line
+                    else:
+                        print(f"Unknown setting {line}")
+
+        defaults = {}
+        _initialise_None(defaults)
+        defaults['thick'] = 0
+
+        famdct = {key: {'gap': None if family_settings[key] > 900 else family_settings[key],
+                        'stage': family_types[key]} for key in family_settings}
+        names = ['name', 'gap', 'material', 'length', 'angle', 'offset']
+
+        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: None if s in family_types else 'UNKNOWN'))
+
+        df['gap'] = df['gap'].apply(lambda s: None if not isinstance(s, str) and s > 900 else s)
+        df['gap'] = df['gap'].apply(lambda s: None if isinstance(s, str) else s)
+
+        # 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']]
+            if not np.allclose(np.unique(df.thick.values), 0):
+                print("Warning: Keyword 'thick' is currently not supported in xcoll! Ignoring.")
+            df = df.drop('thick', axis=1)
+            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']]
+        if not np.allclose(np.unique(df.offset.values), 0):
+            print("Warning: Keyword 'offset' is currently not supported in xcoll! Ignoring.")
+        df = df.drop('offset', axis=1)
+        df = df.set_index('name')
+        df = df.replace(np.nan, None)
+
+        colldict = df.transpose().to_dict()
+        # Remove Nonetype families
+        colldict = {coll: {kk: vv for kk, vv in coll_settings.items()
+                           if kk != 'family' or vv is not None}
+                    for coll, coll_settings in colldict.items()}
+        # Remove None gaps and stages if the collimator is assigned a family (they will be set in _parse_dict)
+        colldict = {coll: {kk: vv for kk, vv in coll_settings.items()
+                           if kk != 'gap' or vv is not None or 'family' not in coll_settings}
+                    for coll, coll_settings in colldict.items()}
+        colldict = {coll: {kk: vv for kk, vv in coll_settings.items()
+                           if kk != 'stage' or vv is not None or 'family' not in coll_settings}
+                    for coll, coll_settings in colldict.items()}
+
+        return cls.from_dict({'collimators': colldict, 'families': famdct}, \
+                             ignore_crystals=ignore_crystals, **kwargs)
+
+    def to_pandas(self):
+        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.')
+
+
+    # ====================================
+    # ====== Installing collimators ======
+    # ====================================
+
+    def _get_names_from_line(self, line, names, families):
+        if names is None and families is None:
+            names = self.collimator_names
+        elif names is None:
+            names = self.get_collimators_from_family(families)
+        elif families is not None:
+            names.append(self.get_collimators_from_family(families))
+        return list(set(names)) # Remove duplicates
+
+    def _check_installed(self, line, name, collimator_class):
+            # Check that collimator is not installed as different type
+            # TODO: automatically replace collimator type and print warning
+            if isinstance(line[name], collimator_classes):
+                raise ValueError(f"Trying to install {name} as {collimator_class.__name__}, "
+                               + f"but it is already installed as {line[name].__class__.__name__}!\n"
+                               + f"Please reconstruct the line.")
+            # TODO: only allow Marker elements, no Drifts!!
+            #       How to do this with importing a line for MAD-X or SixTrack...?
+            #       Maybe we want a DriftCollimator type in Xtrack as a general placeholder
+            elif not isinstance(line[name], (xt.Marker, xt.Drift)):
+                raise ValueError(f"Trying to install {name} as {collimator_class.__name__}, "
+                               + f"but the line element to replace is not an xtrack.Marker "
+                               + 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)
+        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)
+        el.emittance = [self.nemitt_x, self.nemitt_y]
+        self._elements[name] = el
+
+    def install_black_absorbers(self, line, *, names=None, families=None, verbose=False, need_apertures=True):
+        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)
+            else:
+                self._create_crystal(line, BlackCrystal, name, verbose=verbose)
+        elements = [self._elements[name] for name in names]
+        line.collimators.install(names, elements, need_apertures=need_apertures)
+
+    def install_everest_collimators(self, line, *, names=None, families=None, verbose=False, need_apertures=True):
+        names = self._get_names_from_line(line, names, families)
+        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],
+                                        verbose=verbose)
+            else:
+                self._create_crystal(line, EverestCrystal, name, material=mat[1],
+                                     lattice=self[name]['crystal'], verbose=verbose,
+                                     miscut=self[name]['miscut'])
+        elements = [self._elements[name] for name in names]
+        line.collimators.install(names, elements, need_apertures=need_apertures)
+
+    def install_fluka_collimators(self, line, *, names=None, families=None, verbose=False, need_apertures=True,
+                                  fluka_input_file=None, remove_missing=True):
+        # Check server
+        if FlukaEngine.is_running():
+            print("Warning: FLUKA server is already running. Stopping server to install collimators.")
+            FlukaEngine.stop_server(fluka_input_file)
+        names = self._get_names_from_line(line, names, families)
+        for name in names:
+            self._create_collimator(line, FlukaCollimator, name, verbose=verbose)
+        elements = [self._elements[name] for name in names]
+        line.collimators.install(names, elements, need_apertures=need_apertures)
+
+    # ==================================
+    # ====== Accessing attributes ======
+    # ==================================
+
+    @property
+    def collimator_names(self):
+        return list(self._collimator_dict.keys())
+
+    @property
+    def collimator_families(self):
+        families = {fam: [] for fam in self._family_dict.keys()}
+        families["no family"] = []
+        for name in self.collimator_names:
+            if 'family' not in self[name] or self[name]['family'].lower() == 'unknown':
+                families["no family"].append(name)
+            else:
+                families[self[name]['family']].append(name)
+        return families
+
+    def get_collimators_from_family(self, family):
+        if not hasattr(family, '__iter__') and not isinstance(family, str):
+            family = [family]
+        result = []
+        for fam in family:
+            if fam not in self.collimator_families:
+                raise ValueError(f"Family '{fam}' not found in CollimatorDatabase!")
+            result += self.collimator_families[fam]
+        return result
+
+    @property
+    def properties(self):
+        return {attr for d in self._collimator_dict.values() for attr in d.keys()}
+
+    def __getattr__(self, attr):
+        if attr in self.properties:
+            # TODO: include families
+            return {kk: vv.get(attr, None) for kk, vv in self._collimator_dict.items()}
+        else:
+            raise ValueError(f"Property `{attr}` not present in CollimatorDatabase!")
+
+    def __getitem__(self, name):
+        if name in self._family_dict:
+            return self._family_dict[name]
+        elif name in self._collimator_dict:
+            return self._collimator_dict[name]
+        else:
+            raise ValueError(f"Family nor collimator `{name}` found in CollimatorDatabase!")
+

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.8'
+__version__ = '0.5.9'
 # ======================

xcoll/initial_distribution.py.orig

@@ -0,0 +1,224 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import numpy as np
+from warnings import warn
+
+import xtrack as xt
+import xobjects as xo
+import xpart as xp
+
+<<<<<<< HEAD
+from .beam_elements import _all_collimator_classes, EverestCrystal, FlukaCollimator
+=======
+from .beam_elements import collimator_classes, EverestCrystal
+>>>>>>> release/v0.6.0
+
+
+def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', pencil_spread=1e-6,
+                                  impact_parameter=0, sigma_z=7.61e-2, twiss=None, longitudinal=None,
+                                  longitudinal_betatron_cut=None, tw=None, **kwargs):
+    """
+    Generate a pencil beam on a collimator.
+    """
+
+    if not line._has_valid_tracker():
+        raise ValueError("Please build tracker before generating pencil distribution!")
+
+    coll = line[name]
+
+    if not isinstance(coll, tuple(collimator_classes)):
+        raise ValueError("Need to provide a valid collimator!")
+
+    if coll.optics is None:
+        raise ValueError("Need to assign optics to collimators before generating pencil distribution!")
+
+    num_particles = int(num_particles)
+    if len(line.get_elements_of_type(FlukaCollimator)[0]) > 0:
+        kwargs.setdefault('_capacity', 2*num_particles)
+
+    if coll.side == 'left':
+        side = '+'
+    if coll.side == 'right':
+        side = '-'
+
+    # Define the plane
+    angle = coll.angle
+    if abs(np.mod(angle-90,180)-90) < 1e-6:
+        plane = 'x'
+        transv_plane = 'y'
+    elif abs(np.mod(angle,180)-90) < 1e-6:
+        plane = 'y'
+        transv_plane = 'x'
+    else:
+        raise NotImplementedError("Pencil beam on a skew collimator not yet supported!")
+
+    if tw is not None:
+        warn("The argument tw is deprecated. Please use twiss instead.", FutureWarning)
+        if twiss is None:
+            twiss = tw
+
+    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
+    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)
+    if is_converging:
+        # pencil at front of jaw
+        sigma = beam_sizes.rows[name:f'{name}>>1'][f'sigma_{plane}'][0]
+        sigma_transv = beam_sizes.rows[name:f'{name}>>1'][f'sigma_{transv_plane}'][0]
+        tw_at_s = twiss.rows[name]
+        at_element = name
+    else:
+        # pencil at back of jaw
+        sigma = beam_sizes.rows[name:f'{name}>>1'][f'sigma_{plane}'][1]
+        sigma_transv = beam_sizes.rows[name:f'{name}>>1'][f'sigma_{transv_plane}'][1]
+        tw_at_s = twiss.rows[f'{name}>>1']
+        at_element = line.element_names[line.element_names.index(name)+1]
+
+    dr_sigmas = pencil_spread/sigma
+
+    # Generate 4D coordinates
+    # TODO: there is some looping in the calculation here and in xpart. Can it be improved?
+    if side == '+-':
+        num_plus = int(num_particles/2)
+        num_min  = int(num_particles - num_plus)
+        coords_plus = _generate_4D_pencil_one_jaw(line, name, num_plus, plane, '+', impact_parameter, dr_sigmas, at_element, is_converging, tw_at_s)
+        coords_min  = _generate_4D_pencil_one_jaw(line, name, num_min,  plane, '-', impact_parameter, dr_sigmas, at_element, is_converging, tw_at_s)
+        coords      = [ [*c_plus, *c_min] for c_plus, c_min in zip(coords_plus, coords_min)]
+    else:
+        coords      = _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter, dr_sigmas, at_element, is_converging, tw_at_s)
+    pencil            = coords[0]
+    p_pencil          = coords[1]
+    transverse_norm   = coords[2]
+    p_transverse_norm = coords[3]
+
+    # Longitudinal plane
+    # TODO: make this more general, make this better
+    if longitudinal is None:
+        delta = 0
+        zeta  = 0
+    elif longitudinal == 'matched_dispersion':
+        raise NotImplementedError
+        # if longitudinal_betatron_cut is None:
+        #     cut = 0
+        # else:
+        #     cut = np.random.uniform(-longitudinal_betatron_cut, longitudinal_betatron_cut,
+        #                             num_particles)
+        # delta = generate_delta_from_dispersion(line, name, plane=plane, position_mm=pencil,
+        #                                        nemitt_x=nemitt_x, nemitt_y=nemitt_y, twiss=tw,
+        #                                        betatron_cut=cut, match_at_front=is_converging)
+        # zeta  = 0
+    elif longitudinal == 'bucket':
+        zeta, delta = xp.generate_longitudinal_coordinates(
+                num_particles=num_particles, distribution='gaussian', sigma_z=sigma_z, line=line
+        )
+    elif not hasattr(longitudinal, '__iter__'):
+        raise ValueError
+    elif len(longitudinal) != 2:
+        raise ValueError
+    elif isinstance(longitudinal, str):
+        raise ValueError
+    elif isinstance(longitudinal, dict):
+        zeta = longitudinal['zeta']
+        delta = longitudinal['delta']
+    else:
+        zeta = longitudinal[0]
+        delta = longitudinal[1]
+
+    # Build the particles
+    if plane == 'x':
+        part = xp.build_particles(
+                x=pencil, px=p_pencil, y_norm=transverse_norm, py_norm=p_transverse_norm,
+                zeta=zeta, delta=delta, nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y,
+                line=line, at_element=at_element, _context=coll._buffer.context, **kwargs
+        )
+    else:
+        part = xp.build_particles(
+                x_norm=transverse_norm, px_norm=p_transverse_norm, y=pencil, py=p_pencil, 
+                zeta=zeta, delta=delta, nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y,
+                line=line, at_element=at_element, _context=coll._buffer.context, **kwargs
+        )
+
+    part._init_random_number_generator()
+
+    if not is_converging:
+        dri = xt.Drift(length=-coll.length)
+        dri.track(part)
+        part.start_tracking_at_element -= 1
+        part.at_element -= 1
+
+    return part
+
+
+def generate_delta_from_dispersion(line, at_element, *, plane, position_mm, nemitt_x, nemitt_y,
+                                   twiss=None, betatron_cut=0, match_at_front=True):
+    if line.tracker is None:
+        raise ValueError("Need to build tracker first!")
+    if not hasattr(betatron_cut, '__iter__'):
+        if hasattr(position_mm, '__iter__'):
+            betatron_cut = np.full_like(position_mm, betatron_cut)
+    elif not hasattr(position_mm, '__iter__'):
+        position_mm = np.full_like(betatron_cut, position_mm)
+    elif len(position_mm) != len(betatron_cut):
+        raise ValueError
+    if plane not in ['x', 'y']:
+        raise ValueError("The variable 'plane' needs to be either 'x' or 'y'!")
+
+    if twiss is None:
+        twiss = line.twiss()
+
+    beam_sizes = twiss.get_beam_covariance(nemitt_x=nemitt_x, nemitt_y=nemitt_y)
+    beam_sizes = beam_sizes.rows[at_element:f'{at_element}>>1'][f'sigma_{plane}']
+    sigma = beam_sizes[0] if match_at_front else beam_sizes[1]
+    delta  = (position_mm - betatron_cut*sigma - twiss.rows[at_element][plane])
+    delta /= twiss.rows[at_element][f'd{plane}']
+
+    return delta
+
+
+def _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter,
+                                dr_sigmas, at_element, is_converging, tw_at_s=None):
+    coll = line[name]
+
+    if side == '+':
+        if is_converging:
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_U + impact_parameter
+            else:
+                pencil_pos = coll.jaw_LU + impact_parameter
+        else:
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_D - impact_parameter
+            else:
+                pencil_pos = coll.jaw_LD + impact_parameter
+    elif side == '-':
+        if is_converging:
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_U - impact_parameter
+            else:
+                pencil_pos = coll.jaw_RU - impact_parameter
+        else:
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_D + impact_parameter
+            else:
+                pencil_pos = coll.jaw_RD - impact_parameter
+
+    # Collimator plane: generate pencil distribution
+    pencil, p_pencil = xp.generate_2D_pencil_with_absolute_cut(
+                        num_particles, plane=plane, absolute_cut=pencil_pos, line=line,
+                        dr_sigmas=dr_sigmas, nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y,
+                        at_element=at_element, side=side, twiss=tw_at_s
+                       )
+
+    # Other plane: generate gaussian distribution in normalized coordinates
+    transverse_norm   = np.random.normal(size=num_particles)
+    p_transverse_norm = np.random.normal(size=num_particles)
+
+    return [pencil, p_pencil, transverse_norm, p_transverse_norm]
+

{xcoll-0.5.8.dist-info → xcoll-0.5.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xcoll
-Version: 0.5.8
+Version: 0.5.9
 Summary: Xsuite collimation package
 Home-page: https://github.com/xsuite/xcoll
 License: Apache 2.0
@@ -18,11 +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.1.1)
-Requires-Dist: xfields (>=0.12.1)
-Requires-Dist: xobjects (>=0.2.6)
-Requires-Dist: xpart (>=0.15.0)
-Requires-Dist: xtrack (>=0.36.5)
+Requires-Dist: xdeps (>=0.7.3)
+Requires-Dist: xobjects (>=0.4.5)
+Requires-Dist: xpart (>=0.19.1)
+Requires-Dist: xtrack (>=0.69.7)
 Project-URL: Repository, https://github.com/xsuite/xcoll
 Description-Content-Type: text/markdown
 

{xcoll-0.5.8.dist-info → xcoll-0.5.9.dist-info}/RECORD

@@ -3,6 +3,7 @@ NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
 xcoll/__init__.py,sha256=b_61vh5irhf5fPmqTFJlyhNSt4rmftXg9uXPIEpgVB4,1612
 xcoll/_manager.py,sha256=9NQKaNxZR2I1ChMVBeKQc0A8h6W8gVgRRg72a5NgbXU,808
 xcoll/beam_elements/__init__.py,sha256=06bU8rzvlUPhcvwpaUippddm5IChpcCHBvpmvXJQU74,1122
+xcoll/beam_elements/__init__.py.orig,sha256=XZMbC-0IzLmKN4LXXT5zm2zOXzRWvkpuElCBGE3VWtw,1507
 xcoll/beam_elements/absorber.py,sha256=efK6gyUgD4x_FnBLpMR7-5_HCdp_753nkYikcdn6ulw,2502
 xcoll/beam_elements/base.py,sha256=6scwhofls5AHPJcUyEbTJOJ8U86EQU4S7BB7NGoE_j0,52728
 xcoll/beam_elements/blowup.py,sha256=gBXdlISvoDiMjXVpA77ls5QdAU3H9krwvFt2bSW_NII,8029
@@ -16,10 +17,12 @@ xcoll/beam_elements/elements_src/everest_crystal.h,sha256=WwwNF6it7TuOimhpSXJa7U
 xcoll/beam_elements/everest.py,sha256=PA_VWpnPrIuO1xN__eKyT_ejbGZK7p93QHDVi3re7cM,8541
 xcoll/beam_elements/monitor.py,sha256=zzMdN3JMFSAs-30_ntRvd5qZGdsXfGtColhiFDuMcIk,16928
 xcoll/colldb.py,sha256=lEpkDRAT54szT9i7-jbTMRvkuW_0W2piZCdDaungcIs,30983
-xcoll/general.py,sha256=tYbVvHS7A8CxAZYyCBXWq0RaD9mPy8wBP8aLcV17TKU,534
+xcoll/colldb.py.orig,sha256=v4AXbY44o_mlealtFXzMzDNLVt-Hkt_ZkjKs5QUJlec,31898
+xcoll/general.py,sha256=-xLa-F2jd_1I3hIhlntMYsz_B-DZw_bj0qqh1SR1Cno,534
 xcoll/headers/checks.h,sha256=qdXsOTBOK1MwW6bdFF93j4yE648mcDtEv5rGN1w9sfk,1582
 xcoll/headers/particle_states.h,sha256=DZa_ZSaJrjnA8aHFriZKfRCkArQ8nK1t445MRwevDtA,840
 xcoll/initial_distribution.py,sha256=x5G4LTXn4boEg5jBFrQCk_l759h91XiAUhDTdcUvLkc,8779
+xcoll/initial_distribution.py.orig,sha256=NOtDtAoNYx8p5l-hUO9ueafl3jtbIzGU0z8-xHtTX_E,9026
 xcoll/install.py,sha256=SxEFQnhWXlsXyPBIo847q6wPgId_f5ZtFRR1awGbkjc,2108
 xcoll/interaction_record/__init__.py,sha256=UFoLiKa-z2oX7YoszP-7Vgdt1nM6kT382v1CaIu8_u0,50
 xcoll/interaction_record/interaction_record.py,sha256=ixsQtVZn71vVEuTAA27a2NWcZZZ8iAcWFOa58bcWEQU,13271
@@ -406,8 +409,8 @@ xcoll/scattering_routines/geometry/objects.h,sha256=A5ktGvVlSkC4hUsI_PQFsE80CuDw
 xcoll/scattering_routines/geometry/rotation.h,sha256=lO3RaQBys9r0ROMjR8T8Rr7UsIEm-H9_C_70Nwz4MXY,701
 xcoll/scattering_routines/geometry/segments.h,sha256=7nKnnin2ByxkKyaYwGvFaqgLQg5uBba4CdLHL7L3iQs,7667
 xcoll/scattering_routines/geometry/sort.h,sha256=b1MkFO2ddzv1fWGeQzsLuz46qo2pKyRSXHjoAEVU7Ts,5763
-xcoll-0.5.8.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-xcoll-0.5.8.dist-info/METADATA,sha256=ikWtL-ecHne0cta39XCx61xI5-P_DwuJY2jmHTPLmSo,2709
-xcoll-0.5.8.dist-info/NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
-xcoll-0.5.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-xcoll-0.5.8.dist-info/RECORD,,
+xcoll-0.5.9.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+xcoll-0.5.9.dist-info/METADATA,sha256=ZBJuoHNhd-lFobMiSbj1SmSZ8yyie9OC_QTUo2Af4PI,2675
+xcoll-0.5.9.dist-info/NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
+xcoll-0.5.9.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+xcoll-0.5.9.dist-info/RECORD,,