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

xcoll/interaction_record/interaction_record.py

@@ -18,34 +18,35 @@ class InteractionRecord(xt.BeamElement):
         '_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[:],
+        '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
@@ -58,45 +59,69 @@ class InteractionRecord(xt.BeamElement):
 
 
     @classmethod
-    def start(cls, line, names=None, *, record_touches=None, record_scatterings=None, capacity=1e6, io_buffer=None):
-        names = _get_xcoll_elements(line, names)
+    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
-        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
+
+        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_touches is None:
-            record_touches = not record_scatterings
+        elif record_impacts is None:
+            record_impacts = not record_scatterings
         elif record_scatterings is None:
-            record_scatterings = not record_touches
-        assert record_touches is True or record_touches is False
+            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_touches and not el.record_scatterings:
-                el.record_touches = record_touches
+            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
-        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()}
+        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, names=None):
+    def stop(self, *, elements=False, names=False):
         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:
+        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
-        self._recording_elements = list(set(self._recording_elements) - set(names))
+        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):
@@ -145,10 +170,10 @@ class InteractionRecord(xt.BeamElement):
             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
+        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)
@@ -172,25 +197,29 @@ class InteractionRecord(xt.BeamElement):
 
     def _collimator_id(self, element_name):
         if not hasattr(self, '_coll_ids'):
-            return element_id
+            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):
+    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]],
-                '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']
+                    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
@@ -207,46 +236,63 @@ class InteractionRecord(xt.BeamElement):
             mask = mask & (self.at_turn == turn)
             df = pd.DataFrame({
                     'int':  [shortcuts[inter] for inter in self._inter[mask]],
-                    'pid':  self.parent_id[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.parent_id[mask]
+                    'pid':   self.id_before[mask]
                 })
             return df.groupby(['pid', 'turn'], sort=False)['int'].apply(list)
 
-    def first_touch_per_turn(self):
+    def first_touch_per_turn(self, frame=None):
         n_rows = self._index.num_recorded
-        df = pd.DataFrame({'parent_id': self.parent_id[:n_rows],
+        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', 'parent_id'], sort=False)]
-        df_first = self.to_pandas().loc[idx_first]
+        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 = ['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_')}
+        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, 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):
+def _get_xcoll_elements(line=None, elements=None, names=None):
+    from xcoll.beam_elements import block_classes
+    if names is not None and names is not False and \
+    (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):
+    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(block_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, block_classes):
+            name = name[idx] if names is not None else element.__class__.__name__
             raise ValueError(f"Element {name} not an Xcoll element!")
-    return names
+    return elements, names
 

xcoll/interaction_record/interaction_record_src/interaction_record.h

@@ -24,7 +24,6 @@ int64_t InteractionRecordData_log(InteractionRecordData record, RecordIndex reco
         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);
@@ -33,35 +32,37 @@ int64_t InteractionRecordData_log(InteractionRecordData record, RecordIndex reco
                              ) * 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));
+            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_parent_z(record, i_slot, -1);
-            InteractionRecordData_set_parent_a(record, i_slot, -1);
-            InteractionRecordData_set_parent_pdgid(record, i_slot, -1);
+            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_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);
+            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);
@@ -69,28 +70,27 @@ int64_t InteractionRecordData_log(InteractionRecordData record, RecordIndex reco
 }
 
 /*gpufun*/
-void InteractionRecordData_log_child(InteractionRecordData record, int64_t i_slot, LocalParticle* child, double ds){
+void InteractionRecordData_log_child(InteractionRecordData record, int64_t i_slot, LocalParticle* child){
     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));
+        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_child_z(record, i_slot, -1);
-        InteractionRecordData_set_child_a(record, i_slot, -1);
-        InteractionRecordData_set_child_pdgid(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("Slot %i: length %f\n", i_slot, ds);
     }
 }

xcoll/line_tools.py

@@ -6,17 +6,16 @@
 import numpy as np
 import xtrack as xt
 
-from .beam_elements import element_classes, collimator_classes
+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(collimator_classes)[1]
+    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:
-        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)
 
@@ -36,7 +35,7 @@ def get_optics_at(names, *, twiss=None, line=None):
 
 def open_collimators(line, names=None):
     if names is None:
-        names = line.get_elements_of_type(collimator_classes)[1]
+        names = line.get_elements_of_type(_all_collimator_classes)[1]
     if len(names) == 0:
         print("No collimators found in line.")
     else:
@@ -46,7 +45,7 @@ def open_collimators(line, names=None):
 
 def send_to_parking(line, names=None):
     if names is None:
-        names = line.get_elements_of_type(collimator_classes)[1]
+        names = line.get_elements_of_type(_all_collimator_classes)[1]
     if len(names) == 0:
         print("No collimators found in line.")
     else:
@@ -70,9 +69,8 @@ def enable_scattering(line):
                 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()
+            if hasattr(el, 'enable_scattering'):
+                el.enable_scattering()
 
 def disable_scattering(line):
     elements = line.get_elements_of_type(element_classes)[0]
@@ -80,4 +78,5 @@ def disable_scattering(line):
         print("No xcoll elements found in line.")
     else:
         for el in elements:
-            el.disable_scattering()
+            if hasattr(el, 'disable_scattering'):
+                el.disable_scattering()

xcoll/lossmap.py

@@ -12,13 +12,13 @@ import xtrack as xt
 import xpart as xp
 import xobjects as xo
 
-from .beam_elements import collimator_classes
+from .beam_elements import _all_collimator_classes, _all_crystal_classes
 
 
 class LossMap:
 
     def __init__(self, line, part, *, line_is_reversed, interpolation=0.1,
-                 weights=None, weight_function=None):
+                 weights=None, weight_function=None, verbose=True):
         self._line = line
         self._line_is_reversed = line_is_reversed
         self._machine_length = line.get_length()
@@ -34,26 +34,31 @@ class LossMap:
                 raise ValueError("Use either 'weights' or 'weight_function', not both!")
             self._weights = part.sort(interleave_lost_particles=True)
 
-        # loss location refinement
+        # Correct particles that are lost in aperture directly after collimator -> should be absorbed
+        self._correct_absorbed(verbose=verbose)
+
+        # 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()
+        aper_s, aper_names, aper_nabs, aper_energy = self._get_aperture_losses()
 
         self._lossmap = {
                 'collimator': {
                     's':      coll_summary['s'],
                     'name':   coll_summary['collname'],
                     'length': coll_summary['length'],
-                    'n':      coll_summary['nabs']
+                    'n':      coll_summary['nabs'],
+                    # 'e':      coll_summary['energy']
                 }
                 ,
                 'aperture': {
                     's':    aper_s,
                     'name': aper_names,
-                    'n':    aper_nabs
+                    'n':    aper_nabs,
+                    # 'e':    aper_energy
                 }
                 ,
                 'machine_length': self._machine_length
@@ -106,35 +111,37 @@ class LossMap:
         return self._weights
 
 
-    def _interpolate(self):
-        new_state = self._part.state.copy()
-        new_elem  = self._part.at_element.copy()
-
+    def _correct_absorbed(self, verbose=True):
         # 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
+        part = self._part
+        coll_classes = list(set(_all_collimator_classes) - set(_all_crystal_classes))
+        coll_elements = self._line.get_elements_of_type(coll_classes)[1]
+        for idx, elem in enumerate(part.at_element):
+            if part.state[idx] == 0:
+                if elem == 0:
+                    prev_elem = len(self._line.element_names) - 1
                 else:
-                    raise ValueError(f"Unknown collimator type {what_type}")
-        self._part.state      = new_state
-        self._part.at_element = new_elem
+                    prev_elem = elem - 1
+                if self._line.element_names[prev_elem] in coll_elements:
+                    if verbose:
+                        print(f"Found at {self._line.element_names[elem]}, "
+                            + f"moved to {self._line.element_names[elem-1]}")
+                    part.at_element[idx] = elem - 1
+                    what_type = self._line[elem-1].__class__.__name__
+                    if what_type == 'EverestCollimator':
+                        part.state[idx] = -331
+                    elif what_type == 'EverestCrystal':
+                        part.state[idx] = -332
+                    elif what_type == 'FlukaCollimator':
+                        part.state[idx] = -334   # TODO: what if crystal?
+                    elif what_type == 'Geant4Collimator':
+                        part.state[idx] = -337   # TODO: what if crystal?
+                    elif what_type == 'BlackAbsorber':
+                        part.state[idx] = -340
+                    else:
+                        raise ValueError(f"Unknown collimator type {what_type}")
 
-        # do the interpolation
+    def _interpolate(self):
         aper_s = list(self._part.s[self._part.state==0])
         if len(aper_s) > 0:
             print("Performing the aperture losses refinement.")
@@ -149,7 +156,7 @@ class LossMap:
 
 
     def _make_coll_summary(self):
-        collimator_names = self._line.get_elements_of_type(collimator_classes)[1]
+        collimator_names = self._line.get_elements_of_type(_all_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]])
@@ -157,7 +164,7 @@ class LossMap:
         coll_pos      = [(self._line.get_s_position(i) + self._line[i].length/2)
                          for i in collimator_names]
 
-        if self._line is reversed:
+        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]  
@@ -165,10 +172,11 @@ class LossMap:
         nabs          = [coll_weights[coll_losses == j].sum() for j in collimator_names]
 
         self._summary = pd.DataFrame({
-            'collname': collimator_names, 
+            'collname': collimator_names,
             'nabs':     nabs,    # of particles lost on collimators
+            # 'energy':   energy,
             'length':   coll_lengths,
-            's':        coll_pos,    
+            's':        coll_pos,
             'type':     coll_types
         })
 
@@ -179,7 +187,7 @@ class LossMap:
         aper_s       = list(self._part.s[aper_mask])
 
         if len(aper_s) == 0:
-            return [], [], []
+            return [], [], [], []
         if self._line_is_reversed:
             aper_s   = [ self._machine_length - s for s in aper_s ]
 
@@ -193,7 +201,9 @@ class LossMap:
         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
+        aper_energy = 0
+
+        return aper_pos, aper_names, aper_nabs, aper_energy
 
 
 def _create_weights_from_initial_state(part, function):