xcoll 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

xcoll/colldb.py

@@ -4,13 +4,16 @@
 # ######################################### #
 
 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, EverestCollimator, EverestCrystal, collimator_classes, element_classes
+from .beam_elements import BlackAbsorber, BlackCrystal, EverestCollimator, EverestCrystal, \
+                           BaseCollimator, BaseCrystal, _all_collimator_classes
 from .install import install_elements
 from .scattering_routines.everest.materials import SixTrack_to_xcoll
 
@@ -19,7 +22,7 @@ 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': 0, 'bending_angle': 0, 'width': 0, 'height': 0, 'miscut': 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
@@ -76,7 +79,7 @@ class CollimatorDatabase:
                          kwargs.get('ignore_crystals', True))
         self.nemitt_x = kwargs['nemitt_x']
         self.nemitt_y = kwargs['nemitt_y']
-        self.line = None
+        self._elements = {}
 
 
     def _parse_dict(self, coll, fam, beam=None, _yaml_merged=False, ignore_crystals=True):
@@ -147,7 +150,6 @@ class CollimatorDatabase:
 
     @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
@@ -155,7 +157,8 @@ class CollimatorDatabase:
         if isinstance(file, io.IOBase):
             dct = yaml.load(file)
         else:
-            with open(file, 'r') as fid:
+            file = Path(file).resolve()
+            with file.open('r') as fid:
                 dct = yaml.load(fid)
         dct = _dict_keys_to_lower(dct)
 
@@ -214,7 +217,8 @@ class CollimatorDatabase:
         if isinstance(file, io.IOBase):
             dct = json.load(file)
         else:
-            with open(file, 'r') as fid:
+            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)
@@ -259,17 +263,16 @@ class CollimatorDatabase:
 
     @classmethod
     def from_SixTrack(cls, file, ignore_crystals=True, **kwargs):
-        # only import regex here
-        import re
-        with open(file, 'r') as infile:
+        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', 'miscut', 'crystal']
+            cry_fields = ['bending_radius', 'width', 'height', 'thick', 'tilt', 'miscut', 'crystal']
             cry = {}
 
-            for line in infile:
+            for line in fp:
                 if line.startswith('#'):
                     continue # Comment
                 sline = line.split()
@@ -282,8 +285,8 @@ class CollimatorDatabase:
                     elif sline[0].lower() == "crystal":
                         cry[sline[1]] = {}
                         for i, key in enumerate(cry_fields):
-                            idx = i+2 if i < 4 else i+3  # we skip "tilt"
-                            if i < 5:
+                            idx = i+2
+                            if i < 6:
                                 cry[sline[1]][key] = float(sline[idx])
                             else:
                                 cry[sline[1]][key] = int(sline[idx])
@@ -507,77 +510,85 @@ class CollimatorDatabase:
                 _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 install_black_absorbers(self, line, names=None, *, verbose=False, need_apertures=True):
-        self.line = line
-        elements = []
-        if names is None:
+    def _get_names_from_line(self, line, names, families):
+        if names is None and families is None:
             names = self.collimator_names
-        for name in names:
-            if verbose: print(f"Installing {name:20} as BlackAbsorber")
-            el = BlackAbsorber(gap=self[name]['gap'], angle=self[name]['angle'],
-                               length=self[name]['length'], side=self[name]['side'],
-                               _tracking=False)
+        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], tuple(collimator_classes)):
-                raise ValueError(f"Trying to install {name} as {el.__class__.__name__},"
-                               + f" but it is already installed as {line[name].__class__.__name__}!\n"
+            if isinstance(line[name], _all_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 {el.__class__.__name__},"
-                               + f" but the line element to replace is not an xtrack.Marker "
+                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.")
-            el.emittance = [self.nemitt_x, self.nemitt_y]
-            elements.append(el)
-        install_elements(line, names, elements, need_apertures=need_apertures)
 
-    def install_everest_collimators(self, line, names=None, *, verbose=False, need_apertures=True):
-        self.line = line
-        elements = []
-        if names is None:
-            names = self.collimator_names
+    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:
-            mat = SixTrack_to_xcoll[self[name]['material']]
-            if self[name]['crystal'] is None:
-                if verbose: print(f"Installing {name:20} as EverestCollimator")
-                el = EverestCollimator(gap=self[name]['gap'], angle=self[name]['angle'],
-                                       length=self[name]['length'], side=self[name]['side'],
-                                       material=mat[0], _tracking=False)
+            if self[name]['bending_radius'] is None:
+                self._create_collimator(line, BlackAbsorber, name, verbose=verbose)
             else:
-                if verbose: print(f"Installing {name:20} as EverestCrystal")
-                el = EverestCrystal(gap=self[name]['gap'], angle=self[name]['angle'],
-                                    length=self[name]['length'], side=self[name]['side'], material=mat[1],
-                                    lattice=self[name]['crystal'], bending_radius=self[name]['bending_radius'],
-                                    width=self[name]['width'], height=self[name]['height'],
-                                    miscut=self[name]['miscut'], _tracking=False)
-
-            # Check that collimator is not installed as different type
-            # TODO: automatically replace collimator type and print warning
-            if isinstance(line[name], tuple(collimator_classes)):
-                raise ValueError(f"Trying to install {name} as {el.__class__.__name__},"
-                               + f" but it is already installed as {line[name].__class__.__name__}!\n"
-                               + f"Please reconstruct the line.")
+                self._create_crystal(line, BlackCrystal, name, verbose=verbose)
+        elements = [self._elements[name] for name in names]
+        install_elements(line, names, elements, need_apertures=need_apertures)
 
-            # TODO: only allow Marker elements, no Drifts!!
-            #       How to do this with importing a line for MAD-X or SixTrack...?
-            elif not isinstance(line[name], (xt.Marker, xt.Drift)):
-                raise ValueError(f"Trying to install {name} as {el.__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.")
-            el.emittance = [self.nemitt_x, self.nemitt_y]
-            elements.append(el)
+    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]
         install_elements(line, names, elements, need_apertures=need_apertures)
 
+
     # ==================================
     # ====== Accessing attributes ======
     # ==================================
@@ -588,15 +599,35 @@ class CollimatorDatabase:
 
     @property
     def collimator_families(self):
-        return self._family_dict.keys()
+        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
 
-    def __getattr__(self, name):
-        if name in self._family_dict:
-            return self._family_dict[name]
-        elif name in self._collimator_dict:
-            return self._collimator_dict[name]
+    @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"Family nor collimator '{name}' found in CollimatorDatabase!")
+            raise ValueError(f"Property `{attr}` not present in CollimatorDatabase!")
 
     def __getitem__(self, name):
         if name in self._family_dict:
@@ -604,7 +635,5 @@ class CollimatorDatabase:
         elif name in self._collimator_dict:
             return self._collimator_dict[name]
         else:
-            raise ValueError(f"Family nor collimator '{name}' found in CollimatorDatabase!")
-        
-        
+            raise ValueError(f"Family nor collimator `{name}` found in CollimatorDatabase!")
 

xcoll/general.py

@@ -9,8 +9,8 @@ _pkg_root = Path(__file__).parent.absolute()
 
 citation = "F.F. Van der Veken, et al.: Recent Developments with the New Tools for Collimation Simulations in Xsuite, Proceedings of HB2023, Geneva, Switzerland."
 
-# ===================
+# ======================
 # Do not change
-# ===================
-__version__ = '0.4.0'
-# ===================
+# ======================
+__version__ = '0.5.0'
+# ======================

xcoll/initial_distribution.py

@@ -9,7 +9,7 @@ import xtrack as xt
 import xobjects as xo
 import xpart as xp
 
-from .beam_elements import collimator_classes
+from .beam_elements import _all_collimator_classes, EverestCrystal
 
 
 def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', pencil_spread=1e-6,
@@ -24,7 +24,7 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
 
     coll = line[name]
 
-    if not isinstance(coll, tuple(collimator_classes)):
+    if not isinstance(coll, tuple(_all_collimator_classes)):
         raise ValueError("Need to provide a valid collimator!")
 
     if coll.optics is None:
@@ -171,14 +171,26 @@ def _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_p
 
     if side == '+':
         if is_converging:
-            pencil_pos = coll.jaw_LU + impact_parameter
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_U + impact_parameter
+            else:
+                pencil_pos = coll.jaw_LU + impact_parameter
         else:
-            pencil_pos = coll.jaw_LD + impact_parameter
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_D - impact_parameter
+            else:
+                pencil_pos = coll.jaw_LD + impact_parameter
     elif side == '-':
         if is_converging:
-            pencil_pos = coll.jaw_RU - impact_parameter
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_U - impact_parameter
+            else:
+                pencil_pos = coll.jaw_RU - impact_parameter
         else:
-            pencil_pos = coll.jaw_RD - impact_parameter
+            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(

xcoll/install.py

@@ -83,6 +83,8 @@ def _get_s_start(line, name, length, tab=None):
         existing_length = line[name].length
     else:
         existing_length = 0
+    if name not in tab.name:
+        raise ValueError(f"Element {name} not found in line. Need to mmanually provide `at_s`.")
     return tab.rows[name].s[0] + existing_length/2. - length/2
 
 
@@ -171,7 +173,7 @@ def prepare_space_for_element(line, name, s_start, s_end, tab=None, s_tol=1.e-6)
             continue
         if not element_type.startswith('Limit'):
             print(f"Warning: Removed active element {element_name} "
-                + f"at location inside collimator!")
+                + f"at location inside collimator {name}!")
         length = line[element_name].length if hasattr(line[element_name], 'length') else 0
         line.element_dict[element_name] = xt.Drift(length=length)