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

xcoll/general.py

@@ -1,6 +1,6 @@
 # copyright ############################### #
 # This file is part of the Xcoll Package.   #
-# Copyright (c) CERN, 2023.                 #
+# Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
 from pathlib import Path
@@ -12,5 +12,5 @@ citation = "F.F. Van der Veken, et al.: Recent Developments with the New Tools f
 # ===================
 # Do not change
 # ===================
-__version__ = '0.3.5'
+__version__ = '0.4.0'
 # ===================

xcoll/headers/checks.h

@@ -31,5 +31,5 @@ int8_t xcoll_check_particle_init(RandomRutherfordData rng, LocalParticle* part)
     }
     return is_tracking*rng_is_set*ruth_is_set;
 }
-    
+
 #endif /* XCOLL_CHECKS_H */

xcoll/headers/particle_states.h

@@ -1,6 +1,6 @@
 // copyright ############################### #
 // This file is part of the Xcoll Package.   #
-// Copyright (c) CERN, 2023.                 #
+// Copyright (c) CERN, 2024.                 #
 // ######################################### #
 
 #ifndef XCOLL_STATES_H
@@ -21,5 +21,5 @@
 #define  XC_ERR_NOT_IMPLEMENTED     -391
 #define  XC_ERR_INVALID_XOFIELD     -392
 #define  XC_ERR                     -399
-    
+
 #endif /* XCOLL_STATES_H */

xcoll/initial_distribution.py

@@ -0,0 +1,195 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import numpy as np
+
+import xtrack as xt
+import xobjects as xo
+import xpart as xp
+
+from .beam_elements import collimator_classes
+
+
+def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', pencil_spread=1e-6,
+                                  impact_parameter=0, sigma_z=7.61e-2, tw=None, longitudinal=None,
+                                  longitudinal_betatron_cut=None):
+    """
+    Generate a pencil beam on a collimator.
+    """
+
+    if not line._has_valid_tracker():
+        raise Exception("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 Exception("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:
+        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 None:
+        tw = line.twiss()    # TODO: can we do this smarter by caching?
+
+    # Is it converging or diverging?    # TODO: This might change with a tilt!!!!!!
+    s_front = line.get_s_position(name)
+    s_back  = s_front + coll.length
+    is_converging  = tw[f'alf{plane}', name] > 0
+    print(f"Collimator {name} is {'con' if is_converging else 'di'}verging.")
+
+    beam_sizes = tw.get_beam_covariance(nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y)
+    if is_converging:
+        # pencil at front of jaw
+        match_at_s = s_front
+        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]
+    else:
+        # pencil at back of jaw
+        match_at_s = s_back
+        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]
+    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, match_at_s, is_converging)
+        coords_min  = _generate_4D_pencil_one_jaw(line, name, num_min,  plane, '-', impact_parameter, dr_sigmas, match_at_s, is_converging)
+        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, match_at_s, is_converging)
+    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=name, match_at_s=match_at_s,
+                _context=coll._buffer.context
+        )
+    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=name, match_at_s=match_at_s,
+                _context=coll._buffer.context
+        )
+
+    part._init_random_number_generator()
+
+    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, match_at_s, is_converging):
+    coll = line[name]
+
+    if side == '+':
+        if is_converging:
+            pencil_pos = coll.jaw_LU + impact_parameter
+        else:
+            pencil_pos = coll.jaw_LD + impact_parameter
+    elif side == '-':
+        if is_converging:
+            pencil_pos = coll.jaw_RU - 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=name, side=side,match_at_s=match_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/install.py

@@ -0,0 +1,177 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import numpy as np
+import xtrack as xt
+
+from .beam_elements import element_classes
+
+def install_elements(line, names, elements, *, at_s=None, apertures=None, need_apertures=False, s_tol=1.e-6):
+    if line._has_valid_tracker():
+        raise Exception("Tracker already built!\nPlease install collimators before building "
+                      + "tracker!")
+
+    if not hasattr(names, '__iter__') or isinstance(names, str):
+        names = [names]
+    if not hasattr(elements, '__iter__') or isinstance(elements, str):
+        elements = [elements]
+    names = np.array(names)
+    length = np.array([coll.length for coll in elements])
+    assert len(length) == len(names)
+    if not hasattr(at_s, '__iter__'):
+        at_s = [at_s for _ in range(len(names))]
+    assert len(at_s) == len(names)
+    if isinstance(apertures, str) or not hasattr(apertures, '__iter__'):
+        apertures = [apertures for _ in range(len(names))]
+    assert len(apertures) == len(names)
+
+    # Verify elements
+    for el in elements:
+        assert isinstance(el, element_classes)
+        el._tracking = False
+
+    # Get positions
+    tab = line.get_table()
+    tt = tab.rows[[name for name in names if name in line.element_names]]
+    s_start = []
+    for name, s, l in zip(names, at_s, length):
+        if s is None:
+            s_start.append(_get_s_start(line, name, l, tt))
+        else:
+            s_start.append(s)
+    s_start = np.array(s_start)
+    s_end = s_start + length
+
+    # Check positions
+    l_line = line.get_length() 
+    for s1, s2, name, s3 in zip(s_start, s_end, names, at_s):
+        check_element_position(line, name, s1, s2, s3, l_line, s_tol=s_tol)
+
+    # Look for apertures
+    aper_upstream   = []
+    aper_downstream = []
+    for s1, s2, name, aper in zip(s_start, s_end, names, apertures):
+        if not need_apertures:
+            aper_upstream.append(None)
+            aper_downstream.append(None)
+        else:
+            aper1, aper2 = get_aperture_for_element(line, name, s1, s2, aper, tab, s_tol=s_tol)
+            aper_upstream.append(aper1)
+            aper_downstream.append(aper2)
+
+    # Remove elements at location of collimator (by changing them into markers)
+    for s1, s2, name in zip(s_start, s_end, names):
+        prepare_space_for_element(line, name, s1, s2, tab=tab, s_tol=s_tol)
+
+    # Install
+    line._insert_thick_elements_at_s(element_names=list(names), elements=elements, at_s=s_start, s_tol=s_tol)
+
+    # Install apertures
+    if need_apertures:
+        for s1, name, aper1, aper2 in zip(s_start, names, aper_upstream, aper_downstream):
+            line.insert_element(element=aper1, name=f'{name}_aper_upstream', at=name, s_tol=s_tol)
+            idx = line.element_names.index(name) + 1
+            line.insert_element(element=aper2, name=f'{name}_aper_downstream', at=idx, s_tol=s_tol)
+
+
+def _get_s_start(line, name, length, tab=None):
+    if tab is None:
+        tab = line.get_table()
+    if name in line.element_names and hasattr(line[name], 'length'):
+        existing_length = line[name].length
+    else:
+        existing_length = 0
+    return tab.rows[name].s[0] + existing_length/2. - length/2
+
+
+def check_element_position(line, name, s_start, s_end, at_s, length=None, s_tol=1.e-6):
+    if at_s is None:
+        if name not in line.element_names:
+            raise ValueError(f"Element {name} not found in line. Provide `at_s`.")
+    elif name in line.element_names:
+        if at_s < s_start or at_s > s_end:
+            raise ValueError(f"Element {name} already exists in line at different "
+                           + f"location: at_s = {at_s}, exists at [{s_start}, {s_end}].")
+    if length is None:
+        length = line.get_length() 
+    if s_start <= s_tol:
+        raise ValueError(f"Position of {name} too close to start of line. Please cycle.")
+    if s_end >= length - s_tol:
+        raise ValueError(f"Position of {name} too close to end of line. Please cycle.")
+
+
+def get_apertures_at_s(tab, s, s_tol=1.e-6):
+    tab_s = tab.rows[s-s_tol:s+s_tol:'s']
+    aper = tab_s.rows[[cls.startswith('Limit') for cls in tab_s.element_type]]
+    if len(aper) == 0:
+        return None
+    elif len(aper) == 1:
+        return aper.name[0]
+    else:
+        raise ValueError(f"Multiple apertures found at location {s} with "
+                       + f"tolerance {s_tol}: {aper.name}. Not supported.")
+
+
+def get_aperture_for_element(line, name, s_start, s_end, aperture=None, tab=None, s_tol=1.e-6):
+    if aperture is not None:
+        if isinstance(aperture, str):
+            aper1 = line[aperture]
+            aper2 = line[aperture]
+        elif hasattr(aperture, '__iter__'):
+            if len(aperture) != 2:
+                raise ValueError(f"The value `aperture` should be None or a list "
+                               + f"[upstream, downstream].")
+            assert aperture[0] is not None and aperture[1] is not None
+            if isinstance(aperture[0], str):
+                aper1 = line[aperture[0]]
+            if isinstance(aperture[1], str):
+                aper2 = line[aperture[1]]
+        else:
+            aper1 = aperture
+            aper2 = aperture
+        if not xt.line._is_aperture(aper1, line):
+            raise ValueError(f"Not a valid aperture: {aper1}")
+        if not xt.line._is_aperture(aper2, line):
+            raise ValueError(f"Not a valid aperture: {aper2}")
+        return aper1.copy(), aper2.copy()
+    else:
+        if tab is None:
+            tab = line.get_table()
+        aper1 = get_apertures_at_s(tab, s_start, s_tol=s_tol)
+        aper2 = get_apertures_at_s(tab, s_end, s_tol=s_tol)
+        if aper1 is None and aper2 is not None:
+            aper1 = aper2
+            print(f"Warning: Could not find upstream aperture for {name}! "
+                + f"Used copy of downstream aperture. Proceed with caution.")
+        elif aper2 is None and aper1 is not None:
+            aper2 = aper1
+            print(f"Warning: Could not find downstream aperture for {name}! "
+                + f"Used copy of upstream aperture. Proceed with caution.")
+        elif aper1 is None and aper2 is None:
+            aper_mid = get_apertures_at_s(tab, (s_start+s_end)/2, s_tol=s_tol)
+            if aper_mid is None:
+                raise ValueError(f"No aperture found for {name}! Please provide one.")
+            if line[aper_mid].allow_rot_and_shift \
+            and xt.base_element._tranformations_active(line[aper_mid]):
+                print(f"Warning: Using the centre aperture for {name}, but "
+                    + f"transformations are present. Proceed with caution.")
+            aper1 = aper_mid
+            aper2 = aper_mid
+        return line[aper1].copy(), line[aper2].copy()
+
+
+def prepare_space_for_element(line, name, s_start, s_end, tab=None, s_tol=1.e-6):
+    if tab is None:
+        tab = line.get_table()
+    tt = tab.rows[s_start-s_tol:s_end+s_tol:'s']
+    for element_name, element_type in zip(tt.name[:-1], tt.element_type[:-1]):
+        if element_type == 'Marker' or element_type.startswith('Drift'):
+            continue
+        if not element_type.startswith('Limit'):
+            print(f"Warning: Removed active element {element_name} "
+                + f"at location inside collimator!")
+        length = line[element_name].length if hasattr(line[element_name], 'length') else 0
+        line.element_dict[element_name] = xt.Drift(length=length)
+

xcoll/interaction_record/__init__.py

@@ -0,0 +1 @@
+from .interaction_record import InteractionRecord

xcoll/interaction_record/interaction_record.py

@@ -0,0 +1,252 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import xobjects as xo
+import xtrack as xt
+
+from .interaction_types import source, interactions, shortcuts, is_point
+from ..general import _pkg_root
+
+import numpy as np
+import pandas as pd
+
+
+class InteractionRecord(xt.BeamElement):
+    _xofields = {
+        '_index':            xt.RecordIndex,
+        'at_element':        xo.Int64[:],
+        'at_turn':           xo.Int64[:],
+        'ds':                xo.Float64[:],
+        '_inter':            xo.Int64[:],
+        'parent_id':         xo.Int64[:],
+        'parent_x':          xo.Float64[:],
+        'parent_px':         xo.Float64[:],
+        'parent_y':          xo.Float64[:],
+        'parent_py':         xo.Float64[:],
+        'parent_zeta':       xo.Float64[:],
+        'parent_delta':      xo.Float64[:],
+        'parent_energy':     xo.Float64[:],
+        'parent_mass':       xo.Float64[:],
+        'parent_charge':     xo.Int64[:],
+        'parent_z':          xo.Int64[:],
+        'parent_a':          xo.Int64[:],
+        'parent_pdgid':      xo.Int64[:],
+        'child_id':          xo.Int64[:],
+        'child_x':           xo.Float64[:],
+        'child_px':          xo.Float64[:],
+        'child_y':           xo.Float64[:],
+        'child_py':          xo.Float64[:],
+        'child_zeta':        xo.Float64[:],
+        'child_delta':       xo.Float64[:],
+        'child_energy':      xo.Float64[:],
+        'child_mass':        xo.Float64[:],
+        'child_charge':      xo.Int64[:],
+        'child_z':           xo.Int64[:],
+        'child_a':           xo.Int64[:],
+        'child_pdgid':       xo.Int64[:],
+    }
+
+    allow_track = False
+
+    _extra_c_sources = [
+        source,
+        _pkg_root.joinpath('headers','particle_states.h'),
+        _pkg_root.joinpath('interaction_record','interaction_record_src','interaction_record.h')
+    ]
+
+
+    @classmethod
+    def start(cls, line, names=None, *, record_touches=None, record_scatterings=None, capacity=1e6, io_buffer=None):
+        names = _get_xcoll_elements(line, names)
+        if len(names) == 0:
+            return
+        elements = [line[name] for name in names]
+        capacity = int(capacity)
+        if io_buffer is None:
+            io_buffer = xt.new_io_buffer(capacity=capacity)
+        if record_touches is None and record_scatterings is None:
+            record_touches = True
+            record_scatterings = True
+        elif record_touches is None:
+            record_touches = not record_scatterings
+        elif record_scatterings is None:
+            record_scatterings = not record_touches
+        assert record_touches is True or record_touches is False
+        assert record_scatterings is True or record_scatterings is False
+        for el in elements:
+            if not el.record_touches and not el.record_scatterings:
+                el.record_touches = record_touches
+                el.record_scatterings = record_scatterings
+        record = xt.start_internal_logging(io_buffer=io_buffer, capacity=capacity, \
+                                           elements=elements)
+        record._line = line
+        record._io_buffer = io_buffer
+        record._recording_elements = names
+        record._coll_ids = {name: line.element_names.index(name) for name in names}
+        record._coll_names = {vv: kk for kk, vv in record._coll_ids.items()}
+        return record
+
+    def stop(self, names=None):
+        self.assert_class_init()
+        names = _get_xcoll_elements(self.line, names)
+        elements = [self.line[name] for name in names]
+        if self.line.tracker is not None:
+            self.line.tracker._check_invalidated()
+        xt.stop_internal_logging(elements=elements)
+        # Removed the stopped collimators from list of logged elements
+        self._recording_elements = list(set(self._recording_elements) - set(names))
+
+
+    def assert_class_init(self):
+        if not hasattr(self, '_io_buffer') or not hasattr(self, '_line') \
+        or not hasattr(self, '_recording_elements'):
+            raise ValueError("This InteractionRecord has been manually instantiated, "
+                           + "hence the expanded API is not available. Use "
+                           + "InteractionRecord.start() to initialise with extended API.")
+
+    @property
+    def line(self):
+        if hasattr(self, '_line'):
+            return self._line
+
+    @property
+    def io_buffer(self):
+        if hasattr(self, '_io_buffer'):
+            return self._io_buffer
+
+    @property
+    def capacity(self):
+        if hasattr(self, '_io_buffer'):
+            return self.io_buffer.capacity
+
+    # @capacity.setter
+    # def capacity(self, val):
+    #     if hasattr(self, '_io_buffer'):
+    #         capacity = int(capacity)
+    #         if capacity < self.capacity:
+    #             raise NotImplementedError("Shrinking of capacity not yet implemented!")
+    #         elif capacity == self.capacity:
+    #             return
+    #         else:
+    #             self.io_buffer.grow(capacity - self.capacity)
+    #             # TODO: increase capacity of iobuffer AND of fields in record table
+
+    @property
+    def recording_elements(self):
+        if hasattr(self, '_recording_elements'):
+            return self._recording_elements
+
+    @recording_elements.setter
+    def recording_elements(self, val):
+        self.assert_class_init()
+        if val is None:
+            val = []
+        record_start = _get_xcoll_elements(self.line, val)
+        self.stop(set(self.recording_elements) - set(record_start))
+        elements = [line[name] for name in record_start]
+        for el in elements:
+            if not el.record_touches and not el.record_scatterings:
+                el.record_touches = True
+                el.record_scatterings = True
+        xt.start_internal_logging(io_buffer=self.io_buffer, capacity=self.capacity, \
+                                  record=self, elements=elements)
+        self._recording_elements = record_start
+        # Updating coll IDs: careful to correctly overwrite existing values
+        self._coll_ids.update({name: self.line.element_names.index(name) for name in record_start})
+        self._coll_names = {vv: kk for kk, vv in self._coll_ids.items()}
+
+    @property
+    def interaction_type(self):
+        return np.array([interactions[inter] for inter in self._inter])
+
+    def _collimator_name(self, element_id):
+        if not hasattr(self, '_coll_names'):
+            return element_id
+        elif element_id not in self._coll_names:
+            raise ValueError(f"Element {element_id} not found in list of collimators of this record table! "
+                           + f"Did the line change without updating the list in the table?")
+        else:
+            return self._coll_names[element_id]
+
+    def _collimator_id(self, element_name):
+        if not hasattr(self, '_coll_ids'):
+            return element_id
+        elif element_name not in self._coll_ids:
+            raise ValueError(f"Element {element_name} not found in list of collimators of this record table! "
+                           + f"Did the line change without updating the list in the table?")
+        else:
+            return self._coll_ids[element_name]
+
+    def to_pandas(self):
+        n_rows = self._index.num_recorded
+        coll_header = 'collimator' if hasattr(self, '_coll_names') else 'collimator_id'
+        df = pd.DataFrame({
+                'turn':              self.at_turn[:n_rows],
+                coll_header:         [self._collimator_name(element_id) for element_id in self.at_element[:n_rows]],
+                'interaction_type':  [interactions[inter] for inter in self._inter[:n_rows]],
+                'ds':                self.ds[:n_rows],
+                **{
+                    f'{p}_{val}': getattr(self, f'{p}_{val}')[:n_rows]
+                    for p in ['parent', 'child']
+                    for val in ['id', 'x', 'px', 'y', 'py', 'zeta', 'delta', 'energy', 'mass', 'charge', 'z', 'a', 'pdgid']
+                }
+            })
+        return df
+
+    # TODO: list of impacted collimators
+
+
+    # TODO: does not work when multiple children
+    def interactions_per_collimator(self, collimator=0, *, turn=None):
+        if isinstance(collimator, str):
+            collimator = self._collimator_id(collimator)
+        mask = (self._inter > 0) & (self.at_element == collimator)
+        if turn is not None:
+            mask = mask & (self.at_turn == turn)
+            df = pd.DataFrame({
+                    'int':  [shortcuts[inter] for inter in self._inter[mask]],
+                    'pid':  self.parent_id[mask]
+                })
+            return df.groupby('pid', sort=False)['int'].agg(list)
+        else:
+            df = pd.DataFrame({
+                    'int':   [shortcuts[inter] for inter in self._inter[mask]],
+                    'turn':  self.at_turn[mask],
+                    'pid':   self.parent_id[mask]
+                })
+            return df.groupby(['pid', 'turn'], sort=False)['int'].apply(list)
+
+    def first_touch_per_turn(self):
+        n_rows = self._index.num_recorded
+        df = pd.DataFrame({'parent_id': self.parent_id[:n_rows],
+                           'at_turn': self.at_turn[:n_rows],
+                           'at_element': self.at_element[:n_rows]})
+        mask = np.char.startswith(self.interaction_type[:n_rows], 'Enter Jaw')
+        idx_first = [group.at_element.idxmin() for _, group in df[mask].groupby(['at_turn', 'parent_id'], sort=False)]
+        df_first = self.to_pandas().loc[idx_first]
+        df_first.insert(2, "jaw", df_first.interaction_type.astype(str).str[-1])
+        to_drop = ['ds', 'interaction_type',
+                   *[col for col in df_first.columns if col.startswith('child_')]]
+        to_rename = {col: col.replace('parent_', '') for col in df_first.columns if col.startswith('parent_')}
+        return df_first.drop(columns=to_drop).rename(columns=to_rename)
+
+
+def _get_xcoll_elements(line, names):
+    from xcoll import element_classes
+    if names is None or names is True:
+        names = line.get_elements_of_type(element_classes)[1]
+        if len(names) == 0:
+            raise ValueError("No Xcoll elements in line!")
+    if names is False:
+        names = []
+    if not hasattr(names, '__iter__') or isinstance(names, str):
+        names = [names]
+    for name in names:
+        if name not in line.element_names:
+            raise ValueError(f"Element {name} not found in line!")
+        if not isinstance(line[name], element_classes):
+            raise ValueError(f"Element {name} not an Xcoll element!")
+    return names
+