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

xcoll/interaction_record/interaction_record.py

@@ -0,0 +1,298 @@
+# 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[:],
+        '_inter':            xo.Int64[:],
+        'id_before':         xo.Int64[:],
+        's_before':          xo.Float64[:],
+        'x_before':          xo.Float64[:],
+        'px_before':         xo.Float64[:],
+        'y_before':          xo.Float64[:],
+        'py_before':         xo.Float64[:],
+        'zeta_before':       xo.Float64[:],
+        'delta_before':      xo.Float64[:],
+        'energy_before':     xo.Float64[:],
+        'mass_before':       xo.Float64[:],
+        'charge_before':     xo.Int64[:],
+        'z_before':          xo.Int64[:],
+        'a_before':          xo.Int64[:],
+        'pdgid_before':      xo.Int64[:],
+        'id_after':          xo.Int64[:],
+        's_after':           xo.Float64[:],
+        'x_after':           xo.Float64[:],
+        'px_after':          xo.Float64[:],
+        'y_after':           xo.Float64[:],
+        'py_after':          xo.Float64[:],
+        'zeta_after':        xo.Float64[:],
+        'delta_after':       xo.Float64[:],
+        'energy_after':      xo.Float64[:],
+        'mass_after':        xo.Float64[:],
+        'charge_after':      xo.Int64[:],
+        'z_after':           xo.Int64[:],
+        'a_after':           xo.Int64[:],
+        'pdgid_after':       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=None, elements=None, names=None, record_impacts=None, record_exits=None,
+              record_scatterings=None, capacity=1e6, io_buffer=None, coll_ids=None):
+        elements, names = _get_xcoll_elements(line, elements, names)
+        if len(names) == 0:
+            return
+        capacity = int(capacity)
+
+        if getattr(line, 'tracker', None) is None \
+        or getattr(line.tracker, 'io_buffer', None) is None:
+            if io_buffer is None:
+                io_buffer = xt.new_io_buffer(capacity=capacity)
+        elif io_buffer is not None:
+            raise ValueError("Cannot provide io_buffer when tracker already built!")
+        else:
+            io_buffer = line.tracker.io_buffer
+        if capacity > io_buffer.capacity:
+            io_buffer.grow(capacity - io_buffer.capacity)
+        if record_impacts is None and record_scatterings is None:
+            record_impacts = True
+            record_scatterings = True
+        elif record_impacts is None:
+            record_impacts = not record_scatterings
+        elif record_scatterings is None:
+            record_scatterings = not record_impacts
+        if record_exits is None:
+            # record_exits defaults to True only if the other two are True
+            record_exits = record_impacts and record_scatterings
+        assert record_impacts is True or record_impacts is False
+        assert record_exits is True or record_exits is False
+        assert record_scatterings is True or record_scatterings is False
+        for el in elements:
+            if not el.record_impacts and not el.record_exits and not el.record_scatterings:
+                el.record_impacts = record_impacts
+                el.record_exits = record_exits
+                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
+        recording_elements = names if len(names) > 0 else elements
+        record._recording_elements = recording_elements
+        if coll_ids is None:
+            if line is None:
+                if len(names) > 0:
+                    record._coll_ids = {name: idx for idx, name in enumerate(names)}
+            else:
+                record._coll_ids = {name: line.element_names.index(name) for name in names}
+        else:
+            assert len(coll_ids) == len(names)
+            record._coll_ids = {name: idx for name, idx in zip(names, coll_ids)}
+        if hasattr(record, '_coll_ids'):
+            record._coll_names = {vv: kk for kk, vv in record._coll_ids.items()}
+        return record
+
+    def stop(self, *, elements=False, names=False):
+        self.assert_class_init()
+        elements, names = _get_xcoll_elements(self.line, elements, names)
+        if self.line is not None and 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
+        stopping_elements = names if len(names) > 0 else elements
+        self._recording_elements = list(set(self._recording_elements) - set(stopping_elements))
+
+
+    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 = [self.line[name] for name in record_start]
+        for el in elements:  # TODO: this should be smarter
+            if not el.record_impacts and not el.record_scatterings:
+                el.record_impacts = 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_name
+        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, frame=None):
+        if frame is None:
+            frame = 'jaw'
+        frame = frame.lower()
+        if frame not in ['jaw', 'collimator', 'lab']:
+            raise ValueError(f"Invalid frame {frame}. Must be 'jaw', 'collimator', or 'lab'!")
+        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]],
+                **{
+                    f'{val}_{p}': getattr(self, f'{val}_{p}')[:n_rows]
+                    for p in ['before', 'after']
+                    for val in ['id', 's', '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.id_before[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.id_before[mask]
+                })
+            return df.groupby(['pid', 'turn'], sort=False)['int'].apply(list)
+
+    def first_touch_per_turn(self, frame=None):
+        n_rows = self._index.num_recorded
+        df = pd.DataFrame({'id_before': self.id_before[: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', 'id_before'], sort=False)]
+        df_first = self.to_pandas(frame=frame).loc[idx_first]
+        df_first.insert(2, "jaw", df_first.interaction_type.astype(str).str[-1])
+        to_drop = ['interaction_type',
+                   *[col for col in df_first.columns if col.endswith('_after')]]
+        to_rename = {col: col.replace('_before', '') for col in df_first.columns if col.endswith('before')}
+        to_rename['id_before'] = 'pid'
+        return df_first.drop(columns=to_drop).rename(columns=to_rename)
+
+
+def _get_xcoll_elements(line=None, elements=None, names=None):
+    from xcoll import element_classes
+    if names is not None and names is not False and \
+    (not hasattr(names, '__iter__') or isinstance(names, str)):
+        names = [names]
+    if elements is not None and elements is not False and \
+    (not hasattr(elements, '__iter__') or isinstance(elements, str)):
+        elements = [elements]
+    if line is None:
+        if elements is None:
+            raise ValueError("No line nor elements provided!")
+    else:
+        if elements is not None and elements is not False:
+            raise ValueError("Cannot provide both line and elements!")
+        if names is None or names is True:
+            elements, names = line.get_elements_of_type(element_classes)
+            if len(names) == 0:
+                raise ValueError("No Xcoll elements in line!")
+        elif names is False:
+            names = []
+            elements = []
+        else:
+            assert elements is not False
+            for name in names:
+                if name not in line.element_names:
+                    raise ValueError(f"Element {name} not found in line!")
+            elements = [line[name] for name in names]
+    for idx, element in enumerate(elements):
+        if not isinstance(element, element_classes):
+            name = name[idx] if names is not None else element.__class__.__name__
+            raise ValueError(f"Element {name} not an Xcoll element!")
+    return elements, names
+

xcoll/interaction_record/interaction_record_src/interaction_record.h

@@ -0,0 +1,98 @@
+// copyright ############################### #
+// This file is part of the Xcoll Package.   #
+// Copyright (c) CERN, 2024.                 #
+// ######################################### #
+
+
+#ifndef XCOLL_IMPACTS_H
+#define XCOLL_IMPACTS_H
+
+// TODO: do we need to pass RecordIndex?
+// probably can do RecordIndex record_index = InteractionRecordData_getp__index(record);  ?
+/*gpufun*/
+int64_t InteractionRecordData_log(InteractionRecordData record, RecordIndex record_index, LocalParticle* parent,
+                                  int64_t interaction){
+    // This can be used for a point-like interaction where there is no child (or because it's equal to the parent)
+    // or to log the parent first, to be followed up with InteractionRecordData_log_child on the same slot
+
+    int64_t i_slot = -1;
+    if (record){
+        // Get a slot in the record (this is thread safe)
+        i_slot = RecordIndex_get_slot(record_index);
+        // The returned slot id is negative if record is NULL or if record is full
+
+        if (i_slot>=0){
+            InteractionRecordData_set_at_element(record, i_slot, LocalParticle_get_at_element(parent));
+            InteractionRecordData_set_at_turn(record, i_slot, LocalParticle_get_at_turn(parent));
+            InteractionRecordData_set__inter(record, i_slot, interaction);
+
+            double charge_ratio = LocalParticle_get_charge_ratio(parent);
+            double mass_ratio = charge_ratio / LocalParticle_get_chi(parent);
+            double energy = ( LocalParticle_get_ptau(parent) + 1 / LocalParticle_get_beta0(parent)
+                             ) * mass_ratio * LocalParticle_get_p0c(parent);
+            // All fields have to be written, or the arrays will not have the same length
+            // TODO: maybe this is not true, as we are setting by slot index? Don't the arrays come pre-initialised?
+            InteractionRecordData_set_id_before(record, i_slot, LocalParticle_get_particle_id(parent));
+            InteractionRecordData_set_s_before(record,  i_slot, LocalParticle_get_s(parent));
+            InteractionRecordData_set_x_before(record,  i_slot, LocalParticle_get_x(parent));
+            InteractionRecordData_set_px_before(record, i_slot, LocalParticle_get_px(parent));
+            InteractionRecordData_set_y_before(record,  i_slot, LocalParticle_get_y(parent));
+            InteractionRecordData_set_py_before(record, i_slot, LocalParticle_get_py(parent));
+            InteractionRecordData_set_zeta_before(record,   i_slot, LocalParticle_get_zeta(parent));
+            InteractionRecordData_set_delta_before(record,  i_slot, LocalParticle_get_delta(parent));
+            InteractionRecordData_set_energy_before(record, i_slot, energy);
+            InteractionRecordData_set_mass_before(record,   i_slot, mass_ratio*LocalParticle_get_mass0(parent));
+            InteractionRecordData_set_charge_before(record, i_slot, charge_ratio*LocalParticle_get_q0(parent));
+            // TODO: particle info
+            InteractionRecordData_set_z_before(record, i_slot, -1);
+            InteractionRecordData_set_a_before(record, i_slot, -1);
+            InteractionRecordData_set_pdgid_before(record, i_slot, -1);
+
+            // TODO: maybe this is not needed
+            InteractionRecordData_set_id_after(record, i_slot, -1);
+            InteractionRecordData_set_s_after(record, i_slot, -1);
+            InteractionRecordData_set_x_after(record, i_slot, -1);
+            InteractionRecordData_set_px_after(record, i_slot, -1);
+            InteractionRecordData_set_y_after(record, i_slot, -1);
+            InteractionRecordData_set_py_after(record, i_slot, -1);
+            InteractionRecordData_set_zeta_after(record, i_slot, -1);
+            InteractionRecordData_set_delta_after(record, i_slot, -1);
+            InteractionRecordData_set_energy_after(record, i_slot, -1);
+            InteractionRecordData_set_mass_after(record, i_slot, -1);
+            InteractionRecordData_set_charge_after(record, i_slot, -1);
+            InteractionRecordData_set_z_after(record, i_slot, -1);
+            InteractionRecordData_set_a_after(record, i_slot, -1);
+            InteractionRecordData_set_pdgid_after(record, i_slot, -1);
+        }
+    }
+//     printf("Logging %i in slot %i\n", interaction, i_slot);
+    return i_slot;
+}
+
+/*gpufun*/
+void InteractionRecordData_log_child(InteractionRecordData record, int64_t i_slot, LocalParticle* child){
+    if (record && i_slot>=0){
+        double charge_ratio = LocalParticle_get_charge_ratio(child);
+        double mass_ratio = charge_ratio / LocalParticle_get_chi(child);
+        double energy = ( LocalParticle_get_ptau(child) + 1 / LocalParticle_get_beta0(child)
+                         ) * mass_ratio * LocalParticle_get_p0c(child);
+        InteractionRecordData_set_id_after(record, i_slot, LocalParticle_get_particle_id(child));
+        InteractionRecordData_set_s_after(record,  i_slot, LocalParticle_get_s(child));
+        InteractionRecordData_set_x_after(record,  i_slot, LocalParticle_get_x(child));
+        InteractionRecordData_set_px_after(record, i_slot, LocalParticle_get_px(child));
+        InteractionRecordData_set_y_after(record,  i_slot, LocalParticle_get_y(child));
+        InteractionRecordData_set_py_after(record, i_slot, LocalParticle_get_py(child));
+        InteractionRecordData_set_zeta_after(record,  i_slot, LocalParticle_get_zeta(child));
+        InteractionRecordData_set_delta_after(record, i_slot, LocalParticle_get_delta(child));
+        InteractionRecordData_set_energy_after(record, i_slot, energy);
+        InteractionRecordData_set_mass_after(record,   i_slot, mass_ratio*LocalParticle_get_mass0(child));
+        InteractionRecordData_set_charge_after(record, i_slot, charge_ratio*LocalParticle_get_q0(child));
+        // TODO: particle info
+        InteractionRecordData_set_z_after(record, i_slot, -1);
+        InteractionRecordData_set_a_after(record, i_slot, -1);
+        InteractionRecordData_set_pdgid_after(record, i_slot, -1);
+//     printf("Slot %i: length %f\n", i_slot, ds);
+    }
+}
+
+#endif /* XCOLL_IMPACTS_H */

xcoll/{impacts → interaction_record}/interaction_types.py

@@ -1,6 +1,6 @@
 # copyright ############################### #
 # This file is part of the Xcoll Package.   #
-# Copyright (c) CERN, 2023.                 #
+# Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
 
@@ -10,15 +10,16 @@ source = r'''
 
 #define  XC_UNITIALISED                     0        // NAN  // Do not use
 
-#define  XC_ENTER_JAW                      -1        // JI   // point (no children)    Set ds > 0 if entering later
-#define  XC_EXIT_JAW                       -2        // JO   // point (no children)
+#define  XC_ENTER_JAW_L                    -1        // JI   // point (no children)    Set ds > 0 if entering later
+#define  XC_ENTER_JAW_R                    -2        // JI   // point (no children)    Set ds > 0 if entering later
+#define  XC_EXIT_JAW                       -3        // JO   // point (no children)
+#define  XC_ENTER_JAW                      -4        // JI   // point (no children)    Set ds > 0 if entering later    still here for compatibility
 #define  XC_ABSORBED                        1        // A    // point (no children)    Don't use 0 (is default for unitialised)
 #define  XC_MULTIPLE_COULOMB_SCATTERING    13        // MCS  // continuous
 #define  XC_PN_ELASTIC                     14        // PN   // point (no children)
 #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_RUTHERFORD                     18        // RU   // point (no children)
 #define  XC_CHANNELING                    100        // CH   // continuous
 #define  XC_DECHANNELING                  101        // DCH  // point (no children)
 #define  XC_VOLUME_REFLECTION_TRANS_CH    102        // VRCH // point (no children)    Transition region around +-xpcrit
@@ -43,3 +44,9 @@ shortcuts = {
     for line in source.split('\n')
     if len(line.split()) > 1 and line.split()[1][:3] == 'XC_' # select the source lines with the definitions
 }
+
+is_point = {
+    int(line.split()[2]): line.split()[6].lower() == 'point'
+    for line in source.split('\n')
+    if len(line.split()) > 1 and line.split()[1][:3] == 'XC_' # select the source lines with the definitions
+}

xcoll/line_tools.py

@@ -0,0 +1,82 @@
+# 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, _all_collimator_classes
+
+
+def assign_optics_to_collimators(line, nemitt_x=None, nemitt_y=None, twiss=None):
+    if not line._has_valid_tracker():
+        raise Exception("Please build tracker before setting the openings!")
+    names = line.get_elements_of_type(_all_collimator_classes)[1]
+    tw_upstream, tw_downstream = get_optics_at(names, twiss=twiss, line=line)
+    beta_gamma_rel = line.particle_ref._xobject.gamma0[0]*line.particle_ref._xobject.beta0[0]
+    for coll in names:
+        line[coll].assign_optics(name=coll, nemitt_x=nemitt_x, nemitt_y=nemitt_x, twiss_upstream=tw_upstream,
+                                 twiss_downstream=tw_downstream, beta_gamma_rel=beta_gamma_rel)
+
+def get_optics_at(names, *, twiss=None, line=None):
+    if twiss is None:
+        if not line._has_valid_tracker():
+            raise Exception("Please pass a line and build tracker before computing the optics for the openings!")
+        twiss = line.twiss()
+    if not hasattr(names, '__iter__') and not isinstance(names, str):
+        names = [names]
+    coll_entry_mask = twiss.mask[names]
+    tw_entry = twiss.rows[coll_entry_mask]
+    tw_exit = twiss.rows[coll_entry_mask+1]
+    tw_exit.name = tw_entry.name
+    return tw_entry, tw_exit
+
+
+def open_collimators(line, names=None):
+    if names is None:
+        names = line.get_elements_of_type(_all_collimator_classes)[1]
+    if len(names) == 0:
+        print("No collimators found in line.")
+    else:
+        for coll in names:
+            line[coll].open_jaws(keep_tilts=False)
+            line[coll].gap = None
+
+def send_to_parking(line, names=None):
+    if names is None:
+        names = line.get_elements_of_type(_all_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.")
+
+
+def enable_scattering(line):
+    elements = line.get_elements_of_type(element_classes)[0]
+    if len(elements) == 0:
+        print("No xcoll elements found in line.")
+    else:
+        nemitt_x = None
+        nemitt_y = None
+        for el in elements:
+            if hasattr(el, 'optics') and el.optics is not None:
+                if nemitt_x is None:
+                    nemitt_x = el.nemitt_x
+                if nemitt_y is None:
+                    nemitt_y = el.nemitt_y
+                if not np.isclose(el.nemitt_x, nemitt_x) \
+                or not np.isclose(el.nemitt_x, nemitt_x):
+                    raise ValueError("Not all collimators have the same "
+                                   + "emittance. This is not supported.")
+            if hasattr(el, 'enable_scattering'):
+                el.enable_scattering()
+
+def disable_scattering(line):
+    elements = line.get_elements_of_type(element_classes)[0]
+    if len(elements) == 0:
+        print("No xcoll elements found in line.")
+    else:
+        for el in elements:
+            if hasattr(el, 'disable_scattering'):
+                el.disable_scattering()