xcoll 0.5.3__tar.gz → 0.5.4__tar.gz

{xcoll-0.5.3 → xcoll-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xcoll
-Version: 0.5.3
+Version: 0.5.4
 Summary: Xsuite collimation package
 Home-page: https://github.com/xsuite/xcoll
 License: Apache 2.0

{xcoll-0.5.3 → xcoll-0.5.4}/pyproject.toml

@@ -1,18 +1,21 @@
 [tool.poetry]
 name = "xcoll"
-version = "0.5.3"
+version = "0.5.4"
 description = "Xsuite collimation package"
 homepage = "https://github.com/xsuite/xcoll"
 repository = "https://github.com/xsuite/xcoll"
 authors = [
            "Frederik F. Van der Veken <frederik@cern.ch>",
+           "Simone O. Solstrand <simone.otelie.solstrand@cern.ch>",
+           "Bjorn Lindstrom",
+           "Andre Donadon Servelle",
            "Despina Demetriadou <despina.demetriadou@cern.ch>",
            "Andrey Abramov <andrey.abramov@cern.ch>",
            "Giovanni Iadorala <giovanni.iadorala@cern.ch>"
 ]
 readme = "README.md"
 license = "Apache 2.0"
-include = [ "LICENSE", "NOTICE" ]
+include = ["LICENSE", "NOTICE"]
 
 
 [tool.poetry.dependencies]
@@ -36,3 +39,14 @@ tests = ["pytest", "ruamel-yaml"]
 # Needed for pip install -e (BTW: need pip version 22)
 requires = ["poetry-core>=1.0.8"]
 build-backend = "poetry.core.masonry.api"
+
+# pyproject.toml
+[tool.pytest.ini_options]
+addopts = "-ra --durations=10 --durations-min=1 --junitxml=pytest_results.xml"
+junit_suite_name = "Xcoll"
+junit_duration_report = "call"
+python_functions = ["test_"]
+testpaths = [
+    "tests",
+]
+

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/__init__.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
@@ -17,7 +17,7 @@ from .initial_distribution import *
 from .lossmap import LossMap
 
 # Deprecated
-from .manager import CollimatorManager
+from ._manager import CollimatorManager
 
 # print("If you use Xcoll in your simulations, please cite us :-)")
 # print(citation)

xcoll-0.5.4/xcoll/_manager.py

@@ -0,0 +1,22 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+
+class CollimatorManager:
+
+    def __init__(self, **kwargs):
+        raise ValueError("`CollimatorManager` is deprecated. Use `CollimatorDatabase` instead.")
+
+    @classmethod
+    def from_yaml(cls, file, **kwargs):
+        raise ValueError("`CollimatorManager` is deprecated. Use `CollimatorDatabase` instead.")
+
+    @classmethod
+    def from_json(cls, file, **kwargs):
+        raise ValueError("`CollimatorManager` is deprecated. Use `CollimatorDatabase` instead.")
+
+    @classmethod
+    def from_dict(cls, file, **kwargs):
+        raise ValueError("`CollimatorManager` is deprecated. Use `CollimatorDatabase` instead.")

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/__init__.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/absorber.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/base.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/blowup.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/elements_src/black_absorber.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2024.                 #
 // ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/elements_src/black_crystal.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2024.                 #
 // ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/elements_src/blowup.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2024.                 #
 // ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/elements_src/emittance_monitor.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2024.                 #
 // ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/elements_src/everest_block.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2024.                 #
 // ######################################### #
 
@@ -59,7 +59,6 @@ EverestCollData EverestBlock_init(EverestBlockData el, LocalParticle* part0, int
             coll->record_scatterings = EverestBlockData_get_record_scatterings(el);
         }
     }
-
     return coll;
 }
 
@@ -69,15 +68,10 @@ EverestData EverestBlock_init_data(LocalParticle* part, EverestCollData coll){
     EverestData everest = (EverestData) malloc(sizeof(EverestData_));
     everest->coll = coll;
     everest->rescale_scattering = 1;
-#ifndef XCOLL_REFINE_ENERGY
     // Preinitialise scattering parameters
-    double charge_ratio = LocalParticle_get_charge_ratio(part);
-    double mass_ratio = charge_ratio / LocalParticle_get_chi(part);
-    double energy = ( LocalParticle_get_ptau(part) + 1 / LocalParticle_get_beta0(part)
-                     ) * mass_ratio * LocalParticle_get_p0c(part) / 1e9; // energy in GeV
+    double energy = LocalParticle_get_energy(part) / 1e9; // energy in GeV
     calculate_scattering(everest, energy);
     calculate_ionisation_properties(everest, energy);
-#endif
     return everest;
 }
 
@@ -107,9 +101,6 @@ void EverestBlock_track_local_particle(EverestBlockData el, LocalParticle* part0
                 LocalParticle_set_s(part, 0);
 
                 // Store initial coordinates for updating later
-                double const e0         = LocalParticle_get_energy0(part);
-                double const p0         = LocalParticle_get_p0c(part);
-                double const ptau_in    = LocalParticle_get_ptau(part);
                 double const rvv_in     = LocalParticle_get_rvv(part);
 #ifdef XCOLL_USE_EXACT
                 double const xp_in      = LocalParticle_get_exact_xp(part);
@@ -119,20 +110,19 @@ void EverestBlock_track_local_particle(EverestBlockData el, LocalParticle* part0
                 double const yp_in      = LocalParticle_get_yp(part);
 #endif
                 double const zeta_in    = LocalParticle_get_zeta(part);
-                double const mass_ratio = LocalParticle_get_charge_ratio(part) / LocalParticle_get_chi(part);   // m/m0
-                double energy           = (p0*ptau_in + e0) * mass_ratio;
+                double const energy_in  = LocalParticle_get_energy(part);
+                double energy_out;
 
                 EverestData everest = EverestBlock_init_data(part, coll);
-                energy = jaw(everest, part, energy, length, 0);
+                energy_out = jaw(everest, part, energy_in, length, 0);
                 free(everest);
                 LocalParticle_add_to_s(part, s_block);
 
                 LocalParticle_set_zeta(part, zeta_in);
                 // Survived particles need correcting:
                 if (LocalParticle_get_state(part)>0){
-                    // Update energy
-                    double ptau_out = (energy/mass_ratio - e0) / p0;
-                    LocalParticle_update_ptau(part, ptau_out);
+                    // Update energy; the last flag keeps angles constant (even valid for exact angles!)
+                    LocalParticle_add_to_energy(part, energy_out - energy_in, 0);
                     // Update zeta
 #ifdef XCOLL_USE_EXACT
                     double xp  = LocalParticle_get_exact_xp(part);
@@ -155,4 +145,4 @@ void EverestBlock_track_local_particle(EverestBlockData el, LocalParticle* part0
 
 
 
-#endif /* XCOLL_EVEREST_BLOCK_H */
+#endif /* XCOLL_EVEREST_BLOCK_H */

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/elements_src/everest_collimator.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2024.                 #
 // ######################################### #
 
@@ -125,7 +125,6 @@ EverestCollData EverestCollimator_init(EverestCollimatorData el, LocalParticle*
             coll->record_scatterings = EverestCollimatorData_get_record_scatterings(el);
         }
     }
-
     return coll;
 }
 
@@ -135,16 +134,10 @@ EverestData EverestCollimator_init_data(LocalParticle* part, EverestCollData col
     EverestData everest = (EverestData) malloc(sizeof(EverestData_));
     everest->coll = coll;
     everest->rescale_scattering = 1;
-#ifndef XCOLL_REFINE_ENERGY
     // Preinitialise scattering parameters
-    double charge_ratio = LocalParticle_get_charge_ratio(part);
-    double mass_ratio = charge_ratio / LocalParticle_get_chi(part);
-    double energy = ( LocalParticle_get_ptau(part) + 1 / LocalParticle_get_beta0(part)
-                     ) * mass_ratio * LocalParticle_get_p0c(part) / 1e9; // energy in GeV
-    energy = LocalParticle_get_energy0(part) / 1e9;
+    double energy = LocalParticle_get_energy(part) / 1e9; // energy in GeV
     calculate_scattering(everest, energy);
     calculate_ionisation_properties(everest, energy);
-#endif
     return everest;
 }
 
@@ -175,9 +168,6 @@ void EverestCollimator_track_local_particle(EverestCollimatorData el, LocalParti
                 LocalParticle_set_s(part, 0);
 
                 // Store initial coordinates for updating later
-                double const e0         = LocalParticle_get_energy0(part);
-                double const p0         = LocalParticle_get_p0c(part);
-                double const ptau_in    = LocalParticle_get_ptau(part);
                 double const rvv_in     = LocalParticle_get_rvv(part);
 #ifdef XCOLL_USE_EXACT
                 double const xp_in      = LocalParticle_get_exact_xp(part);
@@ -187,8 +177,8 @@ void EverestCollimator_track_local_particle(EverestCollimatorData el, LocalParti
                 double const yp_in      = LocalParticle_get_yp(part);
 #endif
                 double const zeta_in    = LocalParticle_get_zeta(part);
-                double const mass_ratio = LocalParticle_get_charge_ratio(part) / LocalParticle_get_chi(part);   // m/m0
-                double energy           = (p0*ptau_in + e0) * mass_ratio;
+                double const energy_in  = LocalParticle_get_energy(part);
+                double energy_out;
 
                 // Check if hit on jaws
                 int8_t is_hit = hit_jaws_check_and_transform(part, cg);
@@ -198,7 +188,7 @@ void EverestCollimator_track_local_particle(EverestCollimatorData el, LocalParti
                     double remaining_length = length - LocalParticle_get_s(part);
                     // Scatter
                     EverestData everest = EverestCollimator_init_data(part, coll);
-                    energy = jaw(everest, part, energy, remaining_length, 1);
+                    energy_out = jaw(everest, part, energy_in, remaining_length, 1);
                     free(everest);
                 }
 
@@ -209,9 +199,8 @@ void EverestCollimator_track_local_particle(EverestCollimatorData el, LocalParti
                 LocalParticle_set_zeta(part, zeta_in);
                 // Hit and survived particles need correcting:
                 if (is_hit!=0 && LocalParticle_get_state(part)>0){
-                    // Update energy
-                    double ptau_out = (energy/mass_ratio - e0) / p0;
-                    LocalParticle_update_ptau(part, ptau_out);
+                    // Update energy; the last flag keeps angles constant (even valid for exact angles!)
+                    LocalParticle_add_to_energy(part, energy_out - energy_in, 0);
                     // Update zeta
 #ifdef XCOLL_USE_EXACT
                     double xp  = LocalParticle_get_exact_xp(part);

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/elements_src/everest_crystal.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2024.                 #
 // ######################################### #
 
@@ -162,17 +162,12 @@ EverestData EverestCrystal_init_data(LocalParticle* part, EverestCollData restri
     EverestData everest = (EverestData) malloc(sizeof(EverestData_));
     everest->coll = coll;
     everest->rescale_scattering = 1;
-#ifndef XCOLL_REFINE_ENERGY
     // Preinitialise scattering parameters
-    double charge_ratio = LocalParticle_get_charge_ratio(part);
-    double mass_ratio = charge_ratio / LocalParticle_get_chi(part);
-    double energy = ( LocalParticle_get_ptau(part) + 1 / LocalParticle_get_beta0(part)
-                     ) * mass_ratio * LocalParticle_get_p0c(part) / 1e9; // energy in GeV
+    double energy = LocalParticle_get_energy(part) / 1e9; // energy in GeV
     calculate_scattering(everest, energy);
     calculate_ionisation_properties(everest, energy);
     calculate_critical_angle(everest, part, cg, energy);
     calculate_VI_parameters(everest, part, energy);
-#endif
     return everest;
 }
 
@@ -188,7 +183,13 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
     EverestCollData coll = EverestCrystal_init(el, part0, active);
     CrystalGeometry cg   = EverestCrystal_init_geometry(el, part0, active);
 
-    double t_c = 0;
+    // For info
+    double const e0 = LocalParticle_get_energy0(part0);
+    double t_c0  = _critical_angle0(coll, e0);
+    double Rcrit = _critical_radius(coll, e0);
+    double t_c = _critical_angle(coll, t_c0, Rcrit / fabs(cg->bending_radius));
+    EverestCrystalData_set__critical_radius(el, Rcrit);
+    EverestCrystalData_set__critical_angle(el, t_c);
 
     //start_per_particle_block (part0->part)
         if (!active){
@@ -200,13 +201,11 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
             int8_t is_valid = xcoll_check_particle_init(coll->rng, part);
 
             if (is_valid) {
+                // Store s-location of start of crystal
                 double const s_coll = LocalParticle_get_s(part);
                 LocalParticle_set_s(part, 0);
 
                 // Store initial coordinates for updating later
-                double const e0         = LocalParticle_get_energy0(part);
-                double const p0         = LocalParticle_get_p0c(part);
-                double const ptau_in    = LocalParticle_get_ptau(part);
                 double const rvv_in     = LocalParticle_get_rvv(part);
 #ifdef XCOLL_USE_EXACT
                 double const xp_in      = LocalParticle_get_exact_xp(part);
@@ -216,8 +215,8 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
                 double const yp_in      = LocalParticle_get_yp(part);
 #endif
                 double const zeta_in    = LocalParticle_get_zeta(part);
-                double const mass_ratio = LocalParticle_get_charge_ratio(part) / LocalParticle_get_chi(part);   // m/m0
-                double energy           = (p0*ptau_in + e0) * mass_ratio;
+                double const energy_in  = LocalParticle_get_energy(part);
+                double energy_out;
 
                 // Check if hit on jaws
                 int8_t is_hit = hit_crystal_check_and_transform(part, cg);
@@ -226,7 +225,7 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
                     // Hit one of the jaws, so scatter
                     double remaining_length = length - LocalParticle_get_s(part);
                     // Scatter
-                    EverestData everest = EverestCrystal_init_data(part0, coll, cg);
+                    EverestData everest = EverestCrystal_init_data(part, coll, cg);
                     calculate_initial_angle(everest, part, cg);
 #ifdef XCOLL_USE_EXACT
                     double const xp = LocalParticle_get_exact_xp(part);
@@ -234,13 +233,10 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
                     double const xp = LocalParticle_get_xp(part);
 #endif
                     if (fabs(xp - everest->t_I) < everest->t_c) {
-                        energy = Channel(everest, part, cg, energy/1.e9, remaining_length)*1.e9;
+                        energy_out = Channel(everest, part, cg, energy_in/1.e9, remaining_length)*1.e9;
                     } else {
-                        energy = Amorphous(everest, part, cg, energy/1.e9, remaining_length)*1.e9;
+                        energy_out = Amorphous(everest, part, cg, energy_in/1.e9, remaining_length, 1)*1.e9;
                     }
-                    // Temporary workaround to store the critical angle for use later
-                    calculate_critical_angle(everest, part, cg, e0/1.e9);
-                    t_c = everest->t_c;
                     free(everest);
                 }
 
@@ -251,9 +247,8 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
                 LocalParticle_set_zeta(part, zeta_in);
                 // Hit and survived particles need correcting:
                 if (is_hit!=0 && LocalParticle_get_state(part)>0){
-                    // Update energy
-                    double ptau_out = (energy/mass_ratio - e0) / p0;
-                    LocalParticle_update_ptau(part, ptau_out);
+                    // Update energy; the last flag keeps angles constant (even valid for exact angles!)
+                    LocalParticle_add_to_energy(part, energy_out - energy_in, 0);
                     // Update zeta
 #ifdef XCOLL_USE_EXACT
                     double xp  = LocalParticle_get_exact_xp(part);
@@ -271,10 +266,9 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
             }
         }
     //end_per_particle_block
-    EverestCrystalData_set__critical_angle(el, t_c);
     EverestCrystal_free(cg, active);
     free(coll);
 }
 
 
-#endif /* XCOLL_EVEREST_CRYSTAL_H */
+#endif /* XCOLL_EVEREST_CRYSTAL_H */

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/everest.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
@@ -154,6 +154,7 @@ class EverestCrystal(BaseCrystal):
         'miscut':             xo.Float64,
         '_orient':            xo.Int8,
         '_critical_angle':    xo.Float64,
+        '_critical_radius':   xo.Float64,
         '_material':          CrystalMaterial,
         'rutherford_rng':     xt.RandomRutherford,
         '_tracking':          xo.Int8
@@ -215,7 +216,11 @@ class EverestCrystal(BaseCrystal):
 
     @property
     def critical_angle(self):
-        return self._critical_angle if abs(self._critical_angle) > 1.e-10 else None
+        return self._critical_angle if abs(self._critical_angle) > 1.e-12 else None
+
+    @property
+    def critical_radius(self):
+        return self._critical_radius if abs(self._critical_radius) > 1.e-10 else None
 
     @property
     def lattice(self):

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/beam_elements/monitor.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/colldb.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/general.py

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

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/headers/checks.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2023.                 #
 // ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/headers/particle_states.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2024.                 #
 // ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/initial_distribution.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
@@ -13,14 +13,14 @@ from .beam_elements import _all_collimator_classes, EverestCrystal
 
 
 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):
+                                  impact_parameter=0, sigma_z=7.61e-2, twiss=None, longitudinal=None,
+                                  longitudinal_betatron_cut=None, tw=None):
     """
     Generate a pencil beam on a collimator.
     """
 
     if not line._has_valid_tracker():
-        raise Exception("Please build tracker before generating pencil distribution!")
+        raise ValueError("Please build tracker before generating pencil distribution!")
 
     coll = line[name]
 
@@ -28,7 +28,7 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
         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!")
+        raise ValueError("Need to assign optics to collimators before generating pencil distribution!")
 
     num_particles = int(num_particles)
 
@@ -48,26 +48,32 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
     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?
+    if tw is not None:
+        print("The argument tw is deprecated. Please use twiss instead.")
+        if twiss is None:
+            twiss = tw
+
+    if twiss is None:
+        twiss = line.twiss()
 
     # 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
+    is_converging  = twiss[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)
+    beam_sizes = twiss.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]
+        tw_at_s = twiss.rows[name]
+        at_element = name
     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]
+        tw_at_s = twiss.rows[f'{name}>>1']
+        at_element = line.element_names[line.element_names.index(name)+1]
+
     dr_sigmas = pencil_spread/sigma
 
     # Generate 4D coordinates
@@ -75,11 +81,11 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
     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_plus = _generate_4D_pencil_one_jaw(line, name, num_plus, plane, '+', impact_parameter, dr_sigmas, at_element, is_converging, tw_at_s)
+        coords_min  = _generate_4D_pencil_one_jaw(line, name, num_min,  plane, '-', impact_parameter, dr_sigmas, at_element, is_converging, tw_at_s)
         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)
+        coords      = _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter, dr_sigmas, at_element, is_converging, tw_at_s)
     pencil            = coords[0]
     p_pencil          = coords[1]
     transverse_norm   = coords[2]
@@ -123,19 +129,25 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
         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,
+                line=line, at_element=at_element, #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,
+                line=line, at_element=at_element, #match_at_s=match_at_s,
                 _context=coll._buffer.context
         )
 
     part._init_random_number_generator()
 
+    if not is_converging:
+        dri = xt.Drift(length=-coll.length)
+        dri.track(part)
+        part.start_tracking_at_element -= 1
+        part.at_element -= 1
+
     return part
 
 
@@ -166,7 +178,7 @@ def generate_delta_from_dispersion(line, at_element, *, plane, position_mm, nemi
 
 
 def _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter,
-                                dr_sigmas, match_at_s, is_converging):
+                                dr_sigmas, at_element, is_converging, tw_at_s=None):
     coll = line[name]
 
     if side == '+':
@@ -196,7 +208,7 @@ def _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_p
     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
+                        at_element=at_element, side=side, twiss=tw_at_s
                        )
 
     # Other plane: generate gaussian distribution in normalized coordinates

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/install.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/interaction_record/interaction_record.py

@@ -1,5 +1,5 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.   #
+# This file is part of the Xcoll package.   #
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 

{xcoll-0.5.3 → xcoll-0.5.4}/xcoll/interaction_record/interaction_record_src/interaction_record.h

@@ -1,5 +1,5 @@
 // copyright ############################### #
-// This file is part of the Xcoll Package.   #
+// This file is part of the Xcoll package.   #
 // Copyright (c) CERN, 2024.                 #
 // ######################################### #