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

xcoll/initial_distribution.py.orig

@@ -0,0 +1,256 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+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 collimator_classes, EverestCrystal, Geant4Collimator
+=======
+from .beam_elements import collimator_classes, BaseCrystal
+>>>>>>> main
+
+
+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, _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 len(line.get_elements_of_type(Geant4Collimator)[0]) > 0:
+        kwargs.setdefault('_capacity', 2*num_particles)
+
+    # 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 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:
+            twiss = tw
+
+    if twiss is None:
+        twiss = line.twiss()
+
+    # Longitudinal plane
+    if _longitudinal_coords:
+        zeta = _longitudinal_coords[0]
+        delta = _longitudinal_coords[1]
+
+    else:
+        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':
+        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,
+                                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, 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, 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, **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, 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]

xcoll/interaction_record/interaction_record.py

@@ -8,6 +8,7 @@ import xtrack as xt
 
 from .interaction_types import source, interactions, shortcuts, is_point
 from ..general import _pkg_root
+from ..headers.particle_states import particle_states_src
 
 import numpy as np
 import pandas as pd
@@ -53,7 +54,7 @@ class InteractionRecord(xt.BeamElement):
 
     _extra_c_sources = [
         source,
-        _pkg_root.joinpath('headers','particle_states.h'),
+        particle_states_src,
         _pkg_root.joinpath('interaction_record','interaction_record_src','interaction_record.h')
     ]
 

xcoll/interaction_record/interaction_types.py

@@ -20,8 +20,8 @@ source = r'''
 #define  XC_PP_ELASTIC                     15        // PP   // point (no children)
 #define  XC_SINGLE_DIFFRACTIVE             16        // SD   // point (no children)
 #define  XC_COULOMB                        17        // C    // point (no children)
-#define  XC_CHANNELING                    100        // CH   // continuous
-#define  XC_DECHANNELING                  101        // DCH  // point (no children)
+#define  XC_CHANNELLING                    100       // CH   // continuous
+#define  XC_DECHANNELLING                  101       // DCH  // point (no children)
 #define  XC_VOLUME_REFLECTION_TRANS_CH    102        // VRCH // point (no children)    Transition region around +-xpcrit
 #define  XC_VOLUME_REFLECTION             103        // VR   // point (no children)
 #define  XC_VOLUME_REFLECTION_TRANS_MCS   104        // VRAM // point (no children)    Transition region around t_P

xcoll/line_tools.py

@@ -11,22 +11,108 @@ import xtrack as xt
 from .beam_elements import element_classes, collimator_classes, block_classes
 
 
-class XcollScatteringAPI:
-    def __init__(self, line):
-        self._line = line
+class XcollLineAccessor:
+    _typename = 'element'
+
+    def __init__(self, line, names=None):
+        super().__setattr__('_line', line)
+        if names:
+            super().__setattr__('names', names)
+            # self.names = names
 
     @property
     def line(self):
         return self._line
 
+    @property
+    def _coll_dict(self):
+        return {name: self.line.get(name) for name in self.names}
+
+    def keys(self):
+        return self._coll_dict.keys()
+
+    def values(self):
+        return self._coll_dict.values()
+
+    def items(self):
+        return self._coll_dict.items()
+
+    def __iter__(self):
+        super().__setattr__('_iter_names', iter(self.names))
+        return self
+
+    def __next__(self):
+        try:
+            name = next(self._iter_names)
+        except StopIteration:
+            raise StopIteration
+        else:
+            return self.line[name]
+
+    def __len__(self):
+        return len(self.names)
+
+    def __contains__(self, key):
+        return key in self.names
+
+    def __getattr__(self, attr):
+        properties = {}
+        for name, el in self.items():
+            if hasattr(el, attr):
+                properties[name] = getattr(el, attr)
+        if len(properties) == 0:
+            raise AttributeError(f"Attribute `{attr}` not found.")
+        if len({tuple(ii) if isinstance(ii, list) else ii for ii in properties.values()}) == 1:
+            # If all values are the same, return a single value
+            return next(iter(properties.values()))
+        return properties
+
+    def __setattr__(self, attr, value):
+        if isinstance(value, dict):
+            for name, el in self.items():
+                if name in value:
+                    if not hasattr(el, attr):
+                        raise AttributeError(f"Attribute `{attr}` not found in "
+                                           + f"{self._typename} `{name}`.")
+                    setattr(el, attr, value[name])
+        else:
+            # If value is not a dict, we assume it is a single value to set for all collimators
+            for name, el in self.items():
+                if hasattr(el, attr):
+                    setattr(el, attr, value)
+
+    def __getitem__(self, name):
+        if name in self.names:
+            return self.line[name]
+        else:
+            raise ValueError(f"{self._typename.capitalize()} `{name}` not found in line!")
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} at {hex(id(self))}>"
+
+    def __str__(self):
+        if len(self.names) == 0:
+            return ''
+        res = [f'{self._typename.capitalize()}s:']
+        for name in self.names:
+            res.append(f"    {name:<16} ({self.line[name].__class__.__name__})")
+        return "\n".join(res)
+
+
+class XcollScatteringAPI(XcollLineAccessor):
+
+    @property
+    def names(self):
+        # This makes sure the accessor can access the names of the collimators dynamically
+        return self.line.get_elements_of_type(element_classes)[1]
+
     def enable(self):
-        elements = self.line.get_elements_of_type(element_classes)[0]
-        if len(elements) == 0:
+        if len(self) == 0:
             print("No xcoll elements found in line.")
         else:
             nemitt_x = None
             nemitt_y = None
-            for el in elements:
+            for el in self:
                 if hasattr(el, 'optics') and el.optics is not None:
                     if nemitt_x is None:
                         nemitt_x = el.nemitt_x
@@ -40,22 +126,78 @@ class XcollScatteringAPI:
                     el.enable_scattering()
 
     def disable(self):
-        elements = self.line.get_elements_of_type(element_classes)[0]
-        if len(elements) == 0:
+        if len(self) == 0:
             print("No xcoll elements found in line.")
         else:
-            for el in elements:
+            for el in self:
                 if hasattr(el, 'disable_scattering'):
                     el.disable_scattering()
 
 
-class XcollCollimatorAPI:
-    def __init__(self, line):
-        self._line = line
+class XcollCollimatorAPI(XcollLineAccessor):
+    _typename = 'collimator'
+
+    def __str__(self):
+        res = []
+        if len(self.families) > 0:
+            res.append('Families:')
+            for family, names in self.families.items():
+                res.append(f"    {family:8}: {', '.join(names)}")
+            res.append('')
+        res.append(super().__str__())
+        return "\n".join(res)
 
     @property
-    def line(self):
-        return self._line
+    def names(self):
+        # This makes sure the accessor can access the names of the collimators dynamically
+        return self.line.get_elements_of_type(collimator_classes)[1]
+
+    @property
+    def families(self):
+        families = {}
+        try:
+            prop_families = self.family
+        except AttributeError:
+            return families
+        else:
+            for name in self.names:
+                if name in prop_families:
+                    if prop_families[name] not in families:
+                        families[prop_families[name]] = []
+                    families[prop_families[name]].append(name)
+                else:
+                    if 'no family' not in families:
+                        families['no family'] = []
+                    families['no family'].append(name)
+            return families
+
+    def __getitem__(self, name):
+        # We can getitem by name or family, so we overwrite the super method
+        if name in self.families:
+            return XcollLineAccessor(line=self.line, names=self.families[name])
+        elif name in self.names:
+            return self.line[name]
+        else:
+            raise ValueError(f"Neither family nor collimator `{name}` found in line!")
+
+
+    def open(self, names=None):
+        if names is None:
+            names = self.names
+        if len(names) == 0:
+            print("No collimators found in line.")
+        else:
+            for coll in names:
+                self.line[coll].open_jaws(keep_tilts=False)
+                self.line[coll].gap = None
+
+    def to_parking(self, names=None):
+        if names is None:
+            names = self.names
+        if len(names) == 0:
+            print("No collimators found in line.")
+        else:
+            raise NotImplementedError("Need to move this to new type manager or so.")
 
     def install(self, names, elements, *, at_s=None, apertures=None, need_apertures=False, s_tol=1.e-6):
         if self.line._has_valid_tracker():
@@ -77,9 +219,11 @@ class XcollCollimatorAPI:
         assert len(apertures) == len(names)
 
         # Verify elements
-        for el in elements:
+        for name, el in zip(names, elements):
             assert isinstance(el, block_classes)
             el._tracking = False
+            if el.name is None:
+                el.name = name
 
         # Get positions
         tab = self.line.get_table()
@@ -238,30 +382,11 @@ class XcollCollimatorAPI:
                 twiss = tw
         if not self.line._has_valid_tracker():
             raise Exception("Please build tracker before setting the openings!")
-        names = self.line.get_elements_of_type(collimator_classes)[1]
-        tw_upstream, tw_downstream = self.get_optics_at(names, twiss=twiss)
+        tw_upstream, tw_downstream = self.get_optics_at(self.names, twiss=twiss)
         beta_gamma_rel = self.line.particle_ref._xobject.gamma0[0]*self.line.particle_ref._xobject.beta0[0]
-        for coll in names:
-            self.line[coll].assign_optics(name=coll, nemitt_x=nemitt_x, nemitt_y=nemitt_y, twiss_upstream=tw_upstream,
-                                    twiss_downstream=tw_downstream, beta_gamma_rel=beta_gamma_rel)
-
-    def open(self, names=None):
-        if names is None:
-            names = self.line.get_elements_of_type(collimator_classes)[1]
-        if len(names) == 0:
-            print("No collimators found in line.")
-        else:
-            for coll in names:
-                self.line[coll].open_jaws(keep_tilts=False)
-                self.line[coll].gap = None
-
-    def to_parking(self, names=None):
-        if names is None:
-            names = self.line.get_elements_of_type(collimator_classes)[1]
-        if len(names) == 0:
-            print("No collimators found in line.")
-        else:
-            raise NotImplementedError("Need to move this to new type manager or so.")
+        for name, coll in self.items():
+            coll.assign_optics(name=name, nemitt_x=nemitt_x, nemitt_y=nemitt_y, twiss_upstream=tw_upstream,
+                               twiss_downstream=tw_downstream, beta_gamma_rel=beta_gamma_rel)
 
     def _get_s_start(self, name, *, length, table=None):
         if table is None:
@@ -273,38 +398,3 @@ class XcollCollimatorAPI:
         if name not in table.name:
             raise ValueError(f"Element {name} not found in line. Need to manually provide `at_s`.")
         return table.rows[name].s[0] + existing_length/2. - length/2
-
-
-
-# Deprecated; to be removed
-# -------------------------
-
-def assign_optics_to_collimators(line, nemitt_x=None, nemitt_y=None, twiss=None):
-    warn("The function xcoll.assign_optics_to_collimators() is deprecated and will be "
-       + "removed in the future. Please use line.collimators.assign_optics() instead.", FutureWarning)
-    line.collimators.assign_optics(nemitt_x=nemitt_x, nemitt_y=nemitt_y, twiss=twiss)
-
-def get_optics_at(names, *, twiss=None, line=None):
-    warn("The function xcoll.get_optics_at() is deprecated and will be "
-       + "removed in the future. Please use line.collimators.get_optics_at() instead.", FutureWarning)
-    return line.collimators.get_optics_at(names=names, twiss=twiss)
-
-def open_collimators(line, names=None):
-    warn("The function xcoll.open_collimators() is deprecated and will be "
-       + "removed in the future. Please use line.collimators.open_collimators() instead.", FutureWarning)
-    line.collimators.open(names=names)
-
-def send_to_parking(line, names=None):
-    warn("The function xcoll.send_to_parking() is deprecated and will be "
-       + "removed in the future. Please use line.collimators.send_to_parking() instead.", FutureWarning)
-    line.collimators.to_parking(names=names)
-
-def enable_scattering(line):
-    warn("The function xcoll.enable_scattering() is deprecated and will be "
-       + "removed in the future. Please use line.scattering.enable() instead.", FutureWarning)
-    line.scattering.enable()
-
-def disable_scattering(line):
-    warn("The function xcoll.disable_scattering() is deprecated and will be "
-       + "removed in the future. Please use line.scattering.disable() instead.", FutureWarning)
-    line.scattering.disable()