xcoll 0.5.12__py3-none-any.whl → 0.6.1__py3-none-any.whl

xcoll/beam_elements/monitor.py

@@ -57,6 +57,7 @@ class EmittanceMonitor(xt.BeamElement):
     behaves_like_drift = True
     allow_loss_refinement = True
 
+    _noexpr_fields   = {'name', 'line'}
     _extra_c_sources = [
         xt._pkg_root.joinpath('headers/atomicadd.h'),
         _pkg_root.joinpath('beam_elements/elements_src/emittance_monitor.h')

xcoll/beam_elements/transparent.py

@@ -0,0 +1,83 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import xtrack as xt
+import xobjects as xo
+from .base import BaseCollimator, BaseCrystal, InvalidXcoll
+from ..scattering_routines.geometry import XcollGeometry
+from ..general import _pkg_root
+
+
+class TransparentCollimator(BaseCollimator):
+    _xofields = { **BaseCollimator._xofields,
+        '_tracking':        xo.Int8
+    }
+
+    isthick = True
+    needs_rng = False
+    allow_track = True
+    allow_double_sided = True
+    behaves_like_drift = True
+    allow_rot_and_shift = False
+    allow_loss_refinement = True
+    skip_in_loss_location_refinement = True
+
+    _noexpr_fields         = BaseCollimator._noexpr_fields
+    _skip_in_to_dict       = BaseCollimator._skip_in_to_dict
+    _store_in_to_dict      = BaseCollimator._store_in_to_dict
+    _internal_record_class = BaseCollimator._internal_record_class
+
+    _depends_on = [BaseCollimator, XcollGeometry]
+
+    _extra_c_sources = [
+        _pkg_root.joinpath('beam_elements','elements_src','transparent_collimator.h')
+    ]
+
+    def __init__(self, **kwargs):
+        if '_xobject' not in kwargs:
+            kwargs.setdefault('_tracking', True)
+        super().__init__(**kwargs)
+        if not isinstance(self._context, xo.ContextCpu):
+            raise ValueError('TransparentCollimator is currently not supported on GPU.')
+
+    def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
+        return InvalidXcoll(length=-self.length, _context=_context, _buffer=_buffer, _offset=_offset)
+
+
+class TransparentCrystal(BaseCrystal):
+    _xofields = { **BaseCrystal._xofields,
+        '_tracking':        xo.Int8
+    }
+
+    isthick = True
+    needs_rng = False
+    allow_track = True
+    allow_double_sided = False
+    behaves_like_drift = True
+    allow_rot_and_shift = False
+    allow_loss_refinement = True
+    skip_in_loss_location_refinement = True
+
+    _noexpr_fields         = BaseCrystal._noexpr_fields
+    _skip_in_to_dict       = BaseCrystal._skip_in_to_dict
+    _store_in_to_dict      = BaseCrystal._store_in_to_dict
+    _internal_record_class = BaseCrystal._internal_record_class
+
+    _depends_on = [BaseCrystal, XcollGeometry]
+
+    _extra_c_sources = [
+        _pkg_root.joinpath('beam_elements','elements_src','transparent_crystal.h')
+    ]
+
+    def __init__(self, **kwargs):
+        if '_xobject' not in kwargs:
+            kwargs.setdefault('_tracking', True)
+        super().__init__(**kwargs)
+        if not isinstance(self._context, xo.ContextCpu):
+            raise ValueError('TransparentCrystal is currently not supported on GPU.')
+
+    def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
+        return InvalidXcoll(length=-self.length, _context=_context, _buffer=_buffer, _offset=_offset)
+

xcoll/colldb.py

@@ -547,18 +547,26 @@ class CollimatorDatabase:
             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['name'] = name
         prop_dict.update(kwargs)
         el = cls(**prop_dict)
         el.emittance = [self.nemitt_x, self.nemitt_y]
+        if 'family' in self[name] and self[name]['family'].lower() != 'unknown':
+            if self[name]['family'] == name:
+                raise ValueError(f"Collimator {name} has the same name as its family!")
+            el.family = self[name]['family']
         self._elements[name] = el
+        el.name = name
+        return 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(BlackAbsorber, line, name, verbose=verbose)
-            else:
+            if ('bending_radius' in self[name] and self[name]['bending_radius']) \
+            or ('bending_angle' in self[name] and self[name]['bending_angle']):
                 self._create_collimator(BlackCrystal, line, name, verbose=verbose)
+            else:
+                self._create_collimator(BlackAbsorber, line, name, verbose=verbose)
         elements = [self._elements[name] for name in names]
         line.collimators.install(names, elements, need_apertures=need_apertures)
 
@@ -566,11 +574,12 @@ class CollimatorDatabase:
         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(EverestCollimator, line, name, material=mat[0],
+            if ('bending_radius' in self[name] and self[name]['bending_radius']) \
+            or ('bending_angle' in self[name] and self[name]['bending_angle']):
+                self._create_collimator(EverestCrystal, line, name, material=mat[1],
                                         verbose=verbose)
             else:
-                self._create_collimator(EverestCrystal, line, name, material=mat[1],
+                self._create_collimator(EverestCollimator, line, name, material=mat[0],
                                         verbose=verbose)
         elements = [self._elements[name] for name in names]
         line.collimators.install(names, elements, need_apertures=need_apertures)

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.12'
+__version__ = '0.6.1'
 # ======================

xcoll/general.py.orig

@@ -0,0 +1,20 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+from pathlib import Path
+
+_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
+# ======================
+<<<<<<< HEAD
+__version__ = '0.5.11.dev0+geant4'
+=======
+__version__ = '0.6.0'
+>>>>>>> main
+# ======================

xcoll/headers/particle_states.py

@@ -0,0 +1,51 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+
+LOST_ON_EVEREST_BLOCK   = -330
+LOST_ON_EVEREST_COLL    = -331
+LOST_ON_EVEREST_CRYSTAL = -332
+
+LOST_ON_FLUKA_BLOCK     = -333
+LOST_ON_FLUKA_COLL      = -334
+LOST_ON_FLUKA_CRYSTAL   = -335
+
+LOST_ON_GEANT4_BLOCK    = -336
+LOST_ON_GEANT4_COLL     = -337
+LOST_ON_GEANT4_CRYSTAL  = -338
+
+LOST_ON_ABSORBER        = -340
+
+MASSLESS_OR_NEUTRAL     = -341
+ACC_IONISATION_LOSS     = -342  # Not a real particle, but accumulation of some ionisation losses that are not accounted for (per collimator)
+VIRTUAL_ENERGY          = -343  # Not a real particle, but energy that is deposited
+
+ERR_INVALID_TRACK       = -390
+ERR_NOT_IMPLEMENTED     = -391
+ERR_INVALID_XOFIELD     = -392
+ERR                     = -399
+
+particle_states_src = f"""
+#ifndef XCOLL_STATES_H
+#define XCOLL_STATES_H
+#define  XC_LOST_ON_EVEREST_BLOCK   {LOST_ON_EVEREST_BLOCK}
+#define  XC_LOST_ON_EVEREST_COLL    {LOST_ON_EVEREST_COLL}
+#define  XC_LOST_ON_EVEREST_CRYSTAL {LOST_ON_EVEREST_CRYSTAL}
+#define  XC_LOST_ON_FLUKA_BLOCK     {LOST_ON_FLUKA_BLOCK}
+#define  XC_LOST_ON_FLUKA_COLL      {LOST_ON_FLUKA_COLL}
+#define  XC_LOST_ON_FLUKA_CRYSTAL   {LOST_ON_FLUKA_CRYSTAL}
+#define  XC_LOST_ON_GEANT4_BLOCK    {LOST_ON_GEANT4_BLOCK}
+#define  XC_LOST_ON_GEANT4_COLL     {LOST_ON_GEANT4_COLL}
+#define  XC_LOST_ON_GEANT4_CRYSTAL  {LOST_ON_GEANT4_CRYSTAL}
+#define  XC_LOST_ON_ABSORBER        {LOST_ON_ABSORBER}
+#define  XC_MASSLESS_OR_NEUTRAL     {MASSLESS_OR_NEUTRAL}
+#define  XC_ACC_IONISATION_LOSS     {ACC_IONISATION_LOSS}
+#define  XC_VIRTUAL_ENERGY          {VIRTUAL_ENERGY}
+#define  XC_ERR_INVALID_TRACK       {ERR_INVALID_TRACK}
+#define  XC_ERR_NOT_IMPLEMENTED     {ERR_NOT_IMPLEMENTED}
+#define  XC_ERR_INVALID_XOFIELD     {ERR_INVALID_XOFIELD}
+#define  XC_ERR                     {ERR}
+#endif /* XCOLL_STATES_H */
+"""

xcoll/initial_distribution.py

@@ -1,6 +1,6 @@
 # copyright ############################### #
 # This file is part of the Xcoll package.   #
-# Copyright (c) CERN, 2024.                 #
+# Copyright (c) CERN, 2025.                 #
 # ######################################### #
 
 import numpy as np
@@ -10,34 +10,25 @@ import xtrack as xt
 import xobjects as xo
 import xpart as xp
 
-from .beam_elements import collimator_classes, EverestCrystal
+from .beam_elements import collimator_classes, BaseCrystal
 
 
 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.
-    """
+                                  longitudinal_betatron_cut=None, _capacity=None, tw=None,
+                                  _longitudinal_coords=None, **kwargs):
+    """Generate a pencil beam on a collimator."""
 
+    # Do some general checks
     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 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:
@@ -49,6 +40,17 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
     else:
         raise NotImplementedError("Pencil beam on a skew collimator not yet supported!")
 
+    if coll.side == 'left':
+        if side == '-':
+            raise ValueError("Cannot generate a pencil on the right jaw of a left-"
+                             "sided collimator!")
+        side = '+'
+    if coll.side == 'right':
+        if side == '+':
+            raise ValueError("Cannot generate a pencil on the left jaw of a right-"
+                             "sided collimator!")
+        side = '-'
+
     if tw is not None:
         warn("The argument tw is deprecated. Please use twiss instead.", FutureWarning)
         if twiss is None:
@@ -57,79 +59,48 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
     if twiss is None:
         twiss = line.twiss()
 
-    # 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)
-    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']
+    if _longitudinal_coords:
+        zeta = _longitudinal_coords[0]
+        delta = _longitudinal_coords[1]
+
     else:
-        zeta = longitudinal[0]
-        delta = longitudinal[1]
+        zeta, delta = _generate_longitudinal_dist(line, num_particles, sigma_z, longitudinal)
+
+        # 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)
+            zeta_plus = zeta[:num_plus] if hasattr(zeta, '__iter__') else zeta
+            zeta_min  = zeta[num_plus:] if hasattr(zeta, '__iter__') else zeta
+            delta_plus = delta[:num_plus] if hasattr(delta, '__iter__') else delta
+            delta_min  = delta[num_plus:] if hasattr(delta, '__iter__') else delta
+            if _capacity:
+                _capacity_plus = int(_capacity/2)
+                _capacity_min  = int(_capacity - _capacity_plus)
+            else:
+                _capacity_plus = None
+                _capacity_min  = None
+            part_plus = generate_pencil_on_collimator(line=line, name=name, num_particles=num_plus,
+                                impact_parameter=impact_parameter, _capacity=_capacity_plus,
+                                side='+', pencil_spread=pencil_spread, twiss=twiss, 
+                                _longitudinal_coords=[zeta_plus, delta_plus],
+                                **kwargs)
+            part_min = generate_pencil_on_collimator(line=line, name=name, num_particles=num_min,
+                                impact_parameter=impact_parameter, _capacity=_capacity_min,
+                                side='-', pencil_spread=pencil_spread, twiss=twiss, 
+                                _longitudinal_coords=[zeta_min, delta_min],
+                                **kwargs)
+
+            part = xt.Particles.merge([part_plus, part_min])
+            part.start_tracking_at_element = part_plus.start_tracking_at_element
+            assert part.start_tracking_at_element == part_min.start_tracking_at_element
+            return part
+
+    pencil, p_pencil, transverse_norm, p_transverse_norm, is_converging, at_element = \
+                    _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side,
+                                                impact_parameter, pencil_spread, twiss, **kwargs)
 
     # Build the particles
     if plane == 'x':
@@ -183,38 +154,97 @@ def generate_delta_from_dispersion(line, at_element, *, plane, position_mm, nemi
 
 
 def _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter,
-                                dr_sigmas, at_element, is_converging, tw_at_s=None):
+                                pencil_spread, twiss=None, _capacity=None, **kwargs):
     coll = line[name]
+    beam_sizes = twiss.get_beam_covariance(nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y)
 
+    # Is it converging or diverging?
+    # TODO: dispersion might change this
+    # TODO: skew collimators
+    tolerance_tilt = 1e-12 # 0.1 urad tolerance on jaw tilt  =>  we prioritise converging
+    divergence = coll.divergence
     if side == '+':
-        if isinstance(coll, EverestCrystal):
+        if isinstance(coll, BaseCrystal):
             # A pencil on the crystal should always be upstream
+            is_converging = True
             pencil_pos = coll.jaw_U + impact_parameter
         else:
+            betatron_angle = coll.gap_L * divergence
+            is_converging = coll.tilt_L + tolerance_tilt >= betatron_angle
+            print(f"Left jaw of collimator {name} is {'con' if is_converging else 'di'}verging.")
             if is_converging:
                 pencil_pos = coll.jaw_LU + impact_parameter
             else:
                 pencil_pos = coll.jaw_LD + impact_parameter
     elif side == '-':
-        if isinstance(coll, EverestCrystal):
+        if isinstance(coll, BaseCrystal):
             # A pencil on the crystal should always be upstream
+            is_converging = True
             pencil_pos = coll.jaw_U - impact_parameter
         else:
+            betatron_angle = coll.gap_R * divergence
+            is_converging = coll.tilt_R - tolerance_tilt <= betatron_angle
+            print(f"Right jaw of collimator {name} is {'con' if is_converging else 'di'}verging.")
             if is_converging:
                 pencil_pos = coll.jaw_RU - impact_parameter
             else:
                 pencil_pos = coll.jaw_RD - impact_parameter
+    else:
+        raise ValueError(f"Sinde {side} not supported in  _generate_4D_pencil_one_jaw!")
+
+    if is_converging:
+        # pencil at front of jaw
+        sigma = beam_sizes.rows[name:f'{name}>>1'][f'sigma_{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]
+        tw_at_s = twiss.rows[f'{name}>>1']
+        at_element = line.element_names[line.element_names.index(name)+1]
+
+    dr_sigmas = pencil_spread/sigma
 
     # 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
+                        at_element=at_element, side=side, twiss=tw_at_s, **kwargs
                        )
 
     # 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]
+    return pencil, p_pencil, transverse_norm, p_transverse_norm, is_converging, at_element
+
 
+def _generate_longitudinal_dist(line, num_particles, sigma_z, longitudinal):
+    # TODO: make this more general, make this better
+    if longitudinal is None:
+        return 0, 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':
+        return 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):
+        return longitudinal['zeta'], longitudinal['delta']
+    else:
+        return longitudinal[0], longitudinal[1]