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

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_ds(record, i_slot, 0);
+            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_parent_id(record, i_slot, LocalParticle_get_particle_id(parent));
+            InteractionRecordData_set_parent_x(record,  i_slot, LocalParticle_get_x(parent));
+            InteractionRecordData_set_parent_px(record, i_slot, LocalParticle_get_px(parent));
+            InteractionRecordData_set_parent_y(record,  i_slot, LocalParticle_get_y(parent));
+            InteractionRecordData_set_parent_py(record, i_slot, LocalParticle_get_py(parent));
+            InteractionRecordData_set_parent_zeta(record,   i_slot, LocalParticle_get_zeta(parent));
+            InteractionRecordData_set_parent_delta(record,  i_slot, LocalParticle_get_delta(parent));
+            InteractionRecordData_set_parent_energy(record, i_slot, energy);
+            InteractionRecordData_set_parent_mass(record,   i_slot, mass_ratio*LocalParticle_get_mass0(parent));
+            InteractionRecordData_set_parent_charge(record, i_slot, charge_ratio*LocalParticle_get_q0(parent));
+            // TODO: particle info
+            InteractionRecordData_set_parent_z(record, i_slot, -1);
+            InteractionRecordData_set_parent_a(record, i_slot, -1);
+            InteractionRecordData_set_parent_pdgid(record, i_slot, -1);
+
+            // TODO: maybe this is not needed
+            InteractionRecordData_set_child_id(record, i_slot, -1);
+            InteractionRecordData_set_child_x(record, i_slot, -1);
+            InteractionRecordData_set_child_px(record, i_slot, -1);
+            InteractionRecordData_set_child_y(record, i_slot, -1);
+            InteractionRecordData_set_child_py(record, i_slot, -1);
+            InteractionRecordData_set_child_zeta(record, i_slot, -1);
+            InteractionRecordData_set_child_delta(record, i_slot, -1);
+            InteractionRecordData_set_child_energy(record, i_slot, -1);
+            InteractionRecordData_set_child_mass(record, i_slot, -1);
+            InteractionRecordData_set_child_charge(record, i_slot, -1);
+            InteractionRecordData_set_child_z(record, i_slot, -1);
+            InteractionRecordData_set_child_a(record, i_slot, -1);
+            InteractionRecordData_set_child_pdgid(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, double ds){
+    if (record && i_slot>=0){
+        InteractionRecordData_set_ds(record, i_slot, ds);
+
+        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_child_id(record, i_slot, LocalParticle_get_particle_id(child));
+        InteractionRecordData_set_child_x(record,  i_slot, LocalParticle_get_x(child));
+        InteractionRecordData_set_child_px(record, i_slot, LocalParticle_get_px(child));
+        InteractionRecordData_set_child_y(record,  i_slot, LocalParticle_get_y(child));
+        InteractionRecordData_set_child_py(record, i_slot, LocalParticle_get_py(child));
+        InteractionRecordData_set_child_zeta(record,  i_slot, LocalParticle_get_zeta(child));
+        InteractionRecordData_set_child_delta(record, i_slot, LocalParticle_get_delta(child));
+        InteractionRecordData_set_child_energy(record, i_slot, energy);
+        InteractionRecordData_set_child_mass(record,   i_slot, mass_ratio*LocalParticle_get_mass0(child));
+        InteractionRecordData_set_child_charge(record, i_slot, charge_ratio*LocalParticle_get_q0(child));
+        // TODO: particle info
+        InteractionRecordData_set_child_z(record, i_slot, -1);
+        InteractionRecordData_set_child_a(record, i_slot, -1);
+        InteractionRecordData_set_child_pdgid(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,83 @@
+# 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, 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(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:
+        print(coll)
+        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(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(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.")
+            el.enable_scattering()
+        # self.line.tracker.io_buffer = self._io_buffer
+        # self._set_record_interaction_record()
+
+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:
+            el.disable_scattering()

xcoll/lossmap.py

@@ -0,0 +1,209 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import numpy as np
+import pandas as pd
+from pathlib import Path
+import json
+
+import xtrack as xt
+import xpart as xp
+import xobjects as xo
+
+from .beam_elements import collimator_classes
+
+
+class LossMap:
+
+    def __init__(self, line, part, *, line_is_reversed, interpolation=0.1,
+                 weights=None, weight_function=None):
+        self._line = line
+        self._line_is_reversed = line_is_reversed
+        self._machine_length = line.get_length()
+        self._part = part
+        self._interpolation = interpolation
+        if weights is None:
+            if weight_function is None:
+                self._weights = np.ones(len(part.x))
+            else:
+                self._weights = _create_weights_from_initial_state(part, weight_function)
+        else:
+            if weight_function is not None:
+                raise ValueError("Use either 'weights' or 'weight_function', not both!")
+            self._weights = part.sort(interleave_lost_particles=True)
+
+        # loss location refinement
+        if interpolation is not None:
+            self._interpolate()
+
+        self._make_coll_summary()
+        coll_summary = self._summary[self._summary.nabs > 0].to_dict('list')
+        aper_s, aper_names, aper_nabs = self._get_aperture_losses()
+
+        self._lossmap = {
+                'collimator': {
+                    's':      coll_summary['s'],
+                    'name':   coll_summary['collname'],
+                    'length': coll_summary['length'],
+                    'n':      coll_summary['nabs']
+                }
+                ,
+                'aperture': {
+                    's':    aper_s,
+                    'name': aper_names,
+                    'n':    aper_nabs
+                }
+                ,
+                'machine_length': self._machine_length
+                ,
+                'interpolation': interpolation
+                ,
+                'reversed': self._line_is_reversed
+            }
+
+
+    def to_json(self, file):
+        with open(Path(file), 'w') as fid:
+            json.dump(self._lossmap, fid, indent=True, cls=xo.JEncoder)
+
+    def save_summary(self, file):
+        with open(Path(file), 'w') as fid:
+            fid.write(self._summary.__repr__())
+
+
+    @property
+    def lossmap(self):
+        return self._lossmap
+
+    @property
+    def summary(self):
+        return self._summary
+
+    @property
+    def line(self):
+        return self._line
+
+    @property
+    def line_is_reversed(self):
+        return self._line_is_reversed
+
+    @property
+    def machine_length(self):
+        return self._machine_length
+
+    @property
+    def part(self):
+        return self._part
+
+    @property
+    def interpolation(self):
+        return self._interpolation
+
+    @property
+    def weights(self):
+        return self._weights
+
+
+    def _interpolate(self):
+        new_state = self._part.state.copy()
+        new_elem  = self._part.at_element.copy()
+
+        # Correct particles that are at an aperture directly after a collimator
+        for idx, elem in enumerate(self._part.at_element):
+            if (self._part.state[idx] == 0 and elem > 0
+            and self._line.element_names[elem-1] in
+            self._line.get_elements_of_type(collimator_classes)[1]):
+                print(f"Found at {self._line.element_names[elem]}, "
+                    + f"should be {self._line.element_names[elem-1]}")
+                new_elem[idx] = elem - 1
+                what_type = self._line[elem-1].__class__.__name__
+                if what_type == 'EverestCollimator':
+                    new_state[idx] = -331
+                elif what_type == 'EverestCrystal':
+                    new_state[idx] = -332
+                elif what_type == 'FlukaCollimator':
+                    new_state[idx] = -334   # TODO: what if crystal?
+                elif what_type == 'Geant4Collimator':
+                    new_state[idx] = -337   # TODO: what if crystal?
+                elif what_type == 'BlackAbsorber':
+                    new_state[idx] = -340
+                else:
+                    raise ValueError(f"Unknown collimator type {what_type}")
+        self._part.state      = new_state
+        self._part.at_element = new_elem
+
+        # do the interpolation
+        aper_s = list(self._part.s[self._part.state==0])
+        if len(aper_s) > 0:
+            print("Performing the aperture losses refinement.")
+            loss_loc_refinement = xt.LossLocationRefinement(
+                self._line,
+                n_theta = 360,            # Angular resolution
+                r_max = 0.5,              # Maximum transverse aperture [m]
+                dr = 50e-6,               # Transverse accuracy [m]
+                ds = self.interpolation   # Longitudinal accuracy [m]
+            )
+            loss_loc_refinement.refine_loss_location(self._part)
+
+
+    def _make_coll_summary(self):
+        collimator_names = self._line.get_elements_of_type(collimator_classes)[1]
+        coll_mask     = (self._part.state <= -330) & (self._part.state >= -340)
+        coll_losses   = np.array([self._line.element_names[i]
+                                  for i in self._part.at_element[coll_mask]])
+        coll_lengths  = [self._line[j].length for j in collimator_names] 
+        coll_pos      = [(self._line.get_s_position(i) + self._line[i].length/2)
+                         for i in collimator_names]
+
+        if self._line is reversed:
+            coll_pos  = [self._machine_length - s for s in coll_pos]
+
+        coll_types    = [self._line[i].__class__.__name__ for i in collimator_names]  
+        coll_weights  = self._weights[coll_mask]
+        nabs          = [coll_weights[coll_losses == j].sum() for j in collimator_names]
+
+        self._summary = pd.DataFrame({
+            'collname': collimator_names, 
+            'nabs':     nabs,    # of particles lost on collimators
+            'length':   coll_lengths,
+            's':        coll_pos,    
+            'type':     coll_types
+        })
+
+
+    def _get_aperture_losses(self):
+        # Get s position per particle (lost on aperture)
+        aper_mask    = self._part.state == 0
+        aper_s       = list(self._part.s[aper_mask])
+
+        if len(aper_s) == 0:
+            return [], [], []
+        if self._line_is_reversed:
+            aper_s   = [ self._machine_length - s for s in aper_s ]
+
+        # Store names of aperture markers
+        aper_names   = [self._line.element_names[i] for i in self._part.at_element[aper_mask]]
+        name_dict    = dict(zip(aper_s, aper_names)) # TODO: not floating-point-safe and slow
+
+        # Create output arrays
+        aper_pos     = np.unique(aper_s)
+        aper_weights = self._weights[aper_mask]
+        aper_nabs    = [aper_weights[aper_s == j].sum() for j in aper_pos] 
+        aper_names   = [name_dict[ss] for ss in aper_pos]
+
+        return aper_pos, aper_names, aper_nabs
+
+
+def _create_weights_from_initial_state(part, function):
+    if len(function) == 4:
+        return function[0](part.x)*function[1](part.px)*\
+               function[2](part.y)*function[3](part.py)
+    elif len(function) == 6:
+        return function[0](part.x)*function[1](part.px)*\
+               function[2](part.y)*function[3](part.py)*\
+               function[4](part.zeta)*function[5](part.delta)
+    else:
+        raise NotImplementedError
+