xcoll 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

xcoll/__init__.py

@@ -3,11 +3,13 @@
 # Copyright (c) CERN, 2023.                 #
 # ######################################### #
 
-from .general import _pkg_root, __version__
+from .general import _pkg_root, __version__, citation
 
-from .beam_elements import BlackAbsorber, EverestBlock, EverestCollimator, EverestCrystal
+from .beam_elements import BlackAbsorber, EverestBlock, EverestCollimator, EverestCrystal, element_classes
 from .scattering_routines.everest import materials, Material, CrystalMaterial
 from .manager import CollimatorManager
 from .colldb import CollimatorDatabase, load_SixTrack_colldb
 from .rf_sweep import RFSweep
 
+# print("If you use Xcoll in your simulations, please cite us :-)")
+# print(citation)

xcoll/beam_elements/__init__.py

@@ -1,5 +1,11 @@
-from .base_collimator import BaseCollimator
+from .base import BaseBlock, BaseCollimator
 from .absorber import BlackAbsorber
-from .everest_collimator import EverestBlock, EverestCollimator, EverestCrystal
+from .everest import EverestBlock, EverestCollimator, EverestCrystal
 
 _all_collimator_types = {BlackAbsorber, EverestCollimator, EverestCrystal}
+
+block_classes = tuple(v for v in globals().values()
+                      if isinstance(v, type) and issubclass(v, BaseBlock) and v != BaseBlock)
+collimator_classes = tuple(v for v in globals().values()
+                           if isinstance(v, type) and issubclass(v, BaseCollimator) and v != BaseCollimator)
+element_classes = block_classes + collimator_classes

xcoll/beam_elements/absorber.py

@@ -5,7 +5,7 @@
 
 import xtrack as xt
 import xobjects as xo
-from .base_collimator import BaseCollimator, InvalidXcoll
+from .base import BaseCollimator, InvalidXcoll
 from ..general import _pkg_root
 
 

xcoll/beam_elements/{base_collimator.py → 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, 2023.                 #
 # ######################################### #
 
@@ -20,6 +20,7 @@ class InvalidXcoll(xt.BeamElement):
 
     isthick = True
     behaves_like_drift = True
+#     allow_track = False   # Need to wait for xtrack release to implement
     skip_in_loss_location_refinement = True
     allow_backtrack = True
 
@@ -47,6 +48,7 @@ class BaseBlock(xt.BeamElement):
 
     isthick = True
     behaves_like_drift = True
+#     allow_track = False   # Need to wait for xtrack release to implement
     skip_in_loss_location_refinement = True
 
     _extra_c_sources = [
@@ -88,6 +90,7 @@ class BaseCollimator(xt.BeamElement):
 
     isthick = True
     behaves_like_drift = True
+#     allow_track = False   # Need to wait for xtrack release to implement
     skip_in_loss_location_refinement = True
 
     _skip_in_to_dict  = ['jaw_L', 'jaw_R', 'ref_x', 'ref_y',

xcoll/beam_elements/collimators_src/everest_block.h

@@ -10,12 +10,13 @@
 
 
 /*gpufun*/
-void EverestBlock_set_material(EverestBlockData el, LocalParticle* part0){
+void EverestBlock_set_material(EverestBlockData el){
     MaterialData material = EverestBlockData_getp__material(el);
     RandomRutherfordData rng = EverestBlockData_getp_rutherford_rng(el);
     RandomRutherford_set_by_xcoll_material(rng, (GeneralMaterialData) material);
 }
 
+
 /*gpufun*/
 EverestCollData EverestBlock_init(EverestBlockData el, LocalParticle* part0, int8_t active){
     EverestCollData coll = (EverestCollData) malloc(sizeof(EverestCollData_));
@@ -32,6 +33,7 @@ EverestCollData EverestBlock_init(EverestBlockData el, LocalParticle* part0, int
         coll->csref[0] = MaterialData_get_cross_section(material, 0);
         coll->csref[1] = MaterialData_get_cross_section(material, 1);
         coll->csref[5] = MaterialData_get_cross_section(material, 5);
+        coll->only_mcs = MaterialData_get__only_mcs(material);
 
         // Impact table
         coll->record = EverestBlockData_getp_internal_record(el, part0);
@@ -46,29 +48,27 @@ EverestCollData EverestBlock_init(EverestBlockData el, LocalParticle* part0, int
 
 
 /*gpufun*/
-EverestData EverestBlock_init_data(LocalParticle* part, EverestCollData coll, int8_t only_mcs){
+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
-    if (!only_mcs){
-        // 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
-        calculate_scattering(everest, energy);
-        calculate_ionisation_properties(everest, 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
+    calculate_scattering(everest, energy);
+    calculate_ionisation_properties(everest, energy);
 #endif
     return everest;
 }
 
+
 /*gpufun*/
 void EverestBlock_track_local_particle(EverestBlockData el, LocalParticle* part0) {
     int8_t active = EverestBlockData_get__tracking(el);
     double const length   = EverestBlockData_get_length(el);
-    int8_t const only_mcs = EverestBlockData_get__only_mcs(el);
 
     // Initialise collimator data
     // TODO: we want this to happen before tracking (instead of every turn), as a separate kernel
@@ -84,7 +84,7 @@ void EverestBlock_track_local_particle(EverestBlockData el, LocalParticle* part0
             int8_t is_valid = xcoll_check_particle_init(coll->rng, part);
 
             if (is_valid) {
-                EverestData everest     = EverestBlock_init_data(part, coll, only_mcs);
+                EverestData everest     = EverestBlock_init_data(part, coll);
                 double const e0         = LocalParticle_get_energy0(part) / 1.e9;        // Reference energy in GeV
                 double const p0         = LocalParticle_get_p0c(part) / 1e9;            // Reference momentum in GeV
                 double const mass_ratio = LocalParticle_get_charge_ratio(part) / LocalParticle_get_chi(part);   // m/m0
@@ -95,7 +95,7 @@ void EverestBlock_track_local_particle(EverestBlockData el, LocalParticle* part0
                 double const px_in   = LocalParticle_get_px(part);
                 double const py_in   = LocalParticle_get_py(part);
 
-                double* result = jaw(everest, part, energy, length, only_mcs, 0);
+                double* result = jaw(everest, part, energy, length, 0);
                 energy = result[0];
                 if (result[1] == 1){ is_abs = 1; }
                 double s_out = result[2];

xcoll/beam_elements/collimators_src/everest_collimator.h

@@ -9,8 +9,7 @@
 #include <stdio.h>
 
 
-/*gpufun*/
-void EverestCollimator_set_material(EverestCollimatorData el, LocalParticle* part0){
+void EverestCollimator_set_material(EverestCollimatorData el){
     MaterialData material = EverestCollimatorData_getp__material(el);
     RandomRutherfordData rng = EverestCollimatorData_getp_rutherford_rng(el);
     RandomRutherford_set_by_xcoll_material(rng, (GeneralMaterialData) material);
@@ -36,6 +35,7 @@ EverestCollData EverestCollimator_init(EverestCollimatorData el, LocalParticle*
         coll->csref[0] = MaterialData_get_cross_section(material, 0);
         coll->csref[1] = MaterialData_get_cross_section(material, 1);
         coll->csref[5] = MaterialData_get_cross_section(material, 5);
+        coll->only_mcs = MaterialData_get__only_mcs(material);
 
         // Impact table
         coll->record = EverestCollimatorData_getp_internal_record(el, part0);
@@ -93,8 +93,8 @@ void EverestCollimator_track_local_particle(EverestCollimatorData el, LocalParti
     double const sin_zR     = EverestCollimatorData_get_sin_zR(el);
     double const cos_zR     = EverestCollimatorData_get_cos_zR(el);
     if (fabs(sin_zL-sin_zR) > 1.e-10 || fabs(cos_zL-cos_zR) > 1.e-10 ){
-        printf("Jaws with different angles not yet implemented!");
-        fflush(stdout);
+        printf("Jaws with different angles not yet implemented!");  //only_for_context cpu_serial
+        fflush(stdout);                                             //only_for_context cpu_serial
         kill_all_particles(part0, XC_ERR_NOT_IMPLEMENTED);
     };
 

xcoll/beam_elements/collimators_src/everest_crystal.h

@@ -9,9 +9,8 @@
 #include <stdio.h>
 
 
-
 /*gpufun*/
-void EverestCrystal_set_material(EverestCrystalData el, LocalParticle* part0){
+void EverestCrystal_set_material(EverestCrystalData el){
     CrystalMaterialData material = EverestCrystalData_getp__material(el);
     RandomRutherfordData rng = EverestCrystalData_getp_rutherford_rng(el);
     RandomRutherford_set_by_xcoll_material(rng, (GeneralMaterialData) material);

xcoll/beam_elements/{everest_collimator.py → everest.py}

@@ -9,12 +9,11 @@ import xobjects as xo
 import xpart as xp
 import xtrack as xt
 
-from .base_collimator import BaseBlock, BaseCollimator, InvalidXcoll
+from .base import BaseBlock, BaseCollimator, InvalidXcoll
 from ..scattering_routines.everest import GeneralMaterial, Material, CrystalMaterial, EverestEngine
 from ..general import _pkg_root
 
 
-
 # TODO:
 #      We want these elements to behave as if 'iscollective = True' when doing twiss etc (because they would ruin the CO),
 #      but as if 'iscollective = False' for normal tracking as it is natively in C...
@@ -28,8 +27,7 @@ class EverestBlock(BaseBlock):
     _xofields = { **BaseBlock._xofields,
         '_material':        Material,
         'rutherford_rng':   xt.RandomRutherford,
-        '_tracking':        xo.Int8,
-        '_only_mcs':        xo.Int8
+        '_tracking':        xo.Int8
     }
 
     isthick = True
@@ -46,10 +44,10 @@ class EverestBlock(BaseBlock):
         _pkg_root.joinpath('beam_elements','collimators_src','everest_block.h')
     ]
 
-    _per_particle_kernels = {
-        '_EverestBlock_set_material': xo.Kernel(
+    _kernels = {
+        'EverestBlock_set_material': xo.Kernel(
                 c_name='EverestBlock_set_material',
-                args=[]
+                args=[xo.Arg(xo.ThisClass, name='el')]
             )
         }
 
@@ -64,20 +62,20 @@ class EverestBlock(BaseBlock):
                 or mat['__class__'] != "Material":
                     raise ValueError("Invalid material!")
             kwargs['_material'] = mat
-            # TODO: this should be better
-            if np.allclose(mat.Z, 0.) or np.allclose(mat.A, 0.) \
-            or np.allclose(mat.density, 0.) \
-            or np.allclose(mat.excitation_energy, 0.) \
-            or np.allclose(mat.nuclear_radius, 0.) \
-            or np.allclose(mat.nuclear_elastic_slope, 0.):
-                kwargs['_only_mcs'] = True
-            else:
-                kwargs['_only_mcs'] = False
             kwargs.setdefault('rutherford_rng', xt.RandomRutherford())
             kwargs.setdefault('_tracking', True)
+            use_prebuilt_kernels = kwargs.pop('use_prebuilt_kernels', True)
         super().__init__(**kwargs)
         if '_xobject' not in kwargs:
-            self._EverestBlock_set_material(xp.Particles())
+            try:   # TODO: small workaround until PR
+                self.compile_kernels(use_prebuilt_kernels=use_prebuilt_kernels,
+                                 particles_class=xp.Particles,
+                                 only_if_needed=True)
+            except TypeError:
+                self.compile_kernels(particles_class=xp.Particles,
+                                 only_if_needed=True)
+            self._context.kernels.EverestBlock_set_material(el=self)
+
 
     @property
     def material(self):
@@ -88,8 +86,10 @@ class EverestBlock(BaseBlock):
         if not isinstance(material, Material):
             if not isinstance('material', dict) or material['__class__'] != "Material":
                 raise ValueError("Invalid material!")
-        self._material = material
-        self._EverestBlock_set_material(xp.Particles())
+        if not xt.line._dicts_equal(self.material.to_dict(), material.to_dict()):
+            self._material = material
+            self.compile_kernels(particles_class=xp.Particles, only_if_needed=True)
+            self._context.kernels.EverestBlock_set_material(el=self)
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
         return InvalidXcoll(length=-self.length, _context=_context,
@@ -117,10 +117,10 @@ class EverestCollimator(BaseCollimator):
         _pkg_root.joinpath('beam_elements','collimators_src','everest_collimator.h')
     ]
 
-    _per_particle_kernels = {
-        '_EverestCollimator_set_material': xo.Kernel(
+    _kernels = {
+        'EverestCollimator_set_material': xo.Kernel(
                 c_name='EverestCollimator_set_material',
-                args=[]
+                args=[xo.Arg(xo.ThisClass, name='el')]
             )
         }
 
@@ -136,9 +136,17 @@ class EverestCollimator(BaseCollimator):
             kwargs['_material'] = kwargs.pop('material')
             kwargs.setdefault('rutherford_rng', xt.RandomRutherford())
             kwargs.setdefault('_tracking', True)
+            use_prebuilt_kernels = kwargs.pop('use_prebuilt_kernels', True)
         super().__init__(**kwargs)
         if '_xobject' not in kwargs:
-            self._EverestCollimator_set_material(xp.Particles())
+            try:   # TODO: small workaround until PR
+                self.compile_kernels(use_prebuilt_kernels=use_prebuilt_kernels,
+                                 particles_class=xp.Particles,
+                                 only_if_needed=True)
+            except TypeError:
+                self.compile_kernels(particles_class=xp.Particles,
+                                 only_if_needed=True)
+            self._context.kernels.EverestCollimator_set_material(el=self)
 
     @property
     def material(self):
@@ -149,8 +157,10 @@ class EverestCollimator(BaseCollimator):
         if not isinstance(material, Material):
             if not isinstance('material', dict) or material['__class__'] != "Material":
                 raise ValueError("Invalid material!")
-        self._material = material
-        self._EverestCollimator_set_material(xp.Particles())
+        if not xt.line._dicts_equal(self.material.to_dict(), material.to_dict()):
+            self._material = material
+            self.compile_kernels(particles_class=xp.Particles, only_if_needed=True)
+            self._context.kernels.EverestCollimator_set_material(el=self)
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
         return InvalidXcoll(length=-self.length, _context=_context,
@@ -189,10 +199,10 @@ class EverestCrystal(BaseCollimator):
         _pkg_root.joinpath('beam_elements','collimators_src','everest_crystal.h')
     ]
 
-    _per_particle_kernels = {
-        '_EverestCrystal_set_material': xo.Kernel(
+    _kernels = {
+        'EverestCrystal_set_material': xo.Kernel(
                 c_name='EverestCrystal_set_material',
-                args=[]
+                args=[xo.Arg(xo.ThisClass, name='el')]
             )
         }
 
@@ -223,13 +233,22 @@ class EverestCrystal(BaseCollimator):
             kwargs['_orient'] = _lattice_setter(kwargs.pop('lattice', 'strip'))
             kwargs.setdefault('rutherford_rng', xt.RandomRutherford())
             kwargs.setdefault('_tracking', True)
+            use_prebuilt_kernels = kwargs.pop('use_prebuilt_kernels', True)
         super().__init__(**kwargs)
         if '_xobject' not in kwargs:
             if bending_radius:
                 self._bending_angle = np.arcsin(self.active_length/bending_radius)
             if bending_angle:
                 self._bending_radius = self.active_length / np.sin(bending_angle)
-            self._EverestCrystal_set_material(xp.Particles())
+            try:   # TODO: small workaround until PR
+                self.compile_kernels(use_prebuilt_kernels=use_prebuilt_kernels,
+                                 particles_class=xp.Particles,
+                                 only_if_needed=True)
+            except TypeError:
+                self.compile_kernels(particles_class=xp.Particles,
+                                 only_if_needed=True)
+            self._context.kernels.EverestCrystal_set_material(el=self)
+
 
     @property
     def critical_angle(self):
@@ -275,8 +294,10 @@ class EverestCrystal(BaseCollimator):
         if not isinstance(material, CrystalMaterial):
             if not isinstance(material, dict) or material['__class__'] != "CrystalMaterial":
                 raise ValueError("Invalid material!")
-        self._material = material
-        self._EverestCrystal_set_material(xp.Particles())
+        if not xt.line._dicts_equal(self.material.to_dict(), material.to_dict()):
+            self._material = material
+            self.compile_kernels(particles_class=xp.Particles, only_if_needed=True)
+            self._context.kernels.EverestCrystal_set_material(el=self)
 
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
@@ -293,3 +314,17 @@ def _lattice_setter(lattice):
         raise ValueError(f"Illegal value {lattice} for 'lattice'! "
                         + "Only use 'strip' (110) or 'quasi-mosaic' (111).")
 
+
+# TODO: We want this in the HybridClass to get Kernels attached automatically,
+#       like the PerParticlePyMethod in BeamElement
+# def _exec_kernel(el, kernel_name, **kwargs):
+# #     context = el._context
+# #     desired_classes  = tuple(a.atype for a in el._kernels[kernel_name].args)
+# #     if (kernel_name, desired_classes) not in context.kernels:
+# #         el.compile_kernels(particles_class=xp.Particles)
+# #     kern = context.kernels[(kernel_name, desired_classes)]
+# #     return kern(el=el._xobject, **kwargs)
+#     el.compile_kernels(particles_class=xp.Particles, only_if_needed=True)
+#     return getattr(el._context.kernels, kernel_name)(el=el, **kwargs)
+
+

xcoll/general.py

@@ -7,8 +7,10 @@ from pathlib import Path
 
 _pkg_root = Path(__file__).parent.absolute()
 
+citation = "F.F. Van der Veken, et al.: Recent Developments with the New Tools for Collimation Simulations in Xsuite, Proceedings of HB2023, Geneva, Switzerland."
+
 # ===================
 # Do not change
 # ===================
-__version__ = '0.3.0'
+__version__ = '0.3.2'
 # ===================

xcoll/headers/checks.h

@@ -10,18 +10,18 @@
 int8_t xcoll_check_particle_init(RandomRutherfordData rng, LocalParticle* part) {
     int8_t is_tracking = assert_tracking(part, XC_ERR_INVALID_TRACK);
     if (!is_tracking){
-        printf("Collimator tracking code is called, but we are not supposed to be tracking!");
-        fflush(stdout);
+        printf("Collimator tracking code is called, but we are not supposed to be tracking!"); //only_for_context cpu_serial
+        fflush(stdout);                                                                        //only_for_context cpu_serial
     }
     int8_t rng_is_set  = assert_rng_set(part, RNG_ERR_SEEDS_NOT_SET);
     if (!rng_is_set){
-        printf("Random generator seeds in particles object are not set!");
-        fflush(stdout);
+        printf("Random generator seeds in particles object are not set!"); //only_for_context cpu_serial
+        fflush(stdout);                                                    //only_for_context cpu_serial
     }
     int8_t ruth_is_set = assert_rutherford_set(rng, part, RNG_ERR_RUTH_NOT_SET);
     if (!ruth_is_set){
-        printf("Rutherford random generator not initialised!");
-        fflush(stdout);
+        printf("Rutherford random generator not initialised!"); //only_for_context cpu_serial
+        fflush(stdout);                                         //only_for_context cpu_serial
     }
     return is_tracking*rng_is_set*ruth_is_set;
 }

xcoll/manager.py

@@ -332,22 +332,19 @@ class CollimatorManager:
             elif name not in self.collimator_names:
                 raise Exception(f"Warning: Collimator {name} not found in CollimatorDatabase!...")
         if self.tracker_ready:
-            raise Exception("Tracker already built!\nPlease install collimators before building tracker!")
-
-        # Get collimator centers
-        positions = dict(zip(names,line.get_s_position(names)))
+            raise Exception("Tracker already built!\nPlease install collimators before building "
+                          + "tracker!")
 
         # Loop over collimators to install
         for name in names:
- 
-            # Get the settings from the CollimatorDatabase
-            thiscoll = df.loc[name]
-            # Create the collimator element
-            newcoll = install_func(thiscoll, name)
-            collimator_class = newcoll.__class__
+
+            # Get s positions
+            # This cannot go outside the loop as the indices will change!
+            ss = line.get_s_position()
 
             # Get the settings from the CollimatorDatabase
             thiscoll = df.loc[name]
+            idx = line.element_names.index(name)
             # Create the collimator element
             newcoll = install_func(thiscoll, name)
             collimator_class = newcoll.__class__
@@ -356,15 +353,16 @@ class CollimatorManager:
             # TODO: automatically replace collimator type and print warning
             if isinstance(line[name], tuple(_all_collimator_types - {collimator_class})):
                 raise ValueError(f"Trying to install {name} as {collimator_class.__name__},"
-                                 + f" but it is already installed as {type(line[name]).__name__}!\n"
-                                 + "Please reconstruct the line.")
+                               + f" but it is already installed as {type(line[name]).__name__}!\n"
+                               + f"Please reconstruct the line.")
 
             # Check that collimator is not installed previously
             elif isinstance(line[name], collimator_class):
                 if df.loc[name,'collimator_type'] != collimator_class.__name__:
-                    raise Exception(f"Something is wrong: Collimator {name} already installed in line "
-                                    + f"as {collimator_class.__name__} element, but registered in CollimatorDatabase "
-                                    + f"as {df.loc[name, 'collimator_type']}. Please reconstruct the line.")
+                    raise Exception(f"Something is wrong: Collimator {name} already installed in "
+                                  + f"line as {collimator_class.__name__} element, but registered "
+                                  + f"in CollimatorDatabase as {df.loc[name, 'collimator_type']}. "
+                                  + f"Please reconstruct the line.")
                 if verbose: print(f"Collimator {name} already installed. Skipping...")
                 continue
 
@@ -372,37 +370,98 @@ class CollimatorManager:
             #       How to do this with importing a line for MAD-X or SixTrack...?
             elif not isinstance(line[name], (xt.Marker, xt.Drift)) and not support_legacy_elements:
                 raise ValueError(f"Trying to install {name} as {collimator_class.__name__},"
-                                 + f" but the line element to replace is not an xtrack.Marker (or xtrack.Drift)!\n"
-                                 + "Please check the name, or correct the element.")
+                               + f" but the line element to replace is not an xtrack.Marker "
+                               + f"(or xtrack.Drift)!\nPlease check the name, or correct the "
+                               + f"element.")
 
-            if verbose: print(f"Installing {name:16} as {collimator_class.__name__}")
+            if verbose: print(f"Installing {name:20} as {collimator_class.__name__}")
             # Update the position and type in the CollimatorDatabase
-            df.loc[name,'s_center'] = positions[name]
+            df.loc[name,'s_center'] = ss[idx]
             df.loc[name,'collimator_type'] = collimator_class.__name__
-            # Do the installation
-            s_install = df.loc[name,'s_center'] - thiscoll['active_length']/2 - thiscoll['inactive_front']
-            has_apertures = np.unique([nn for nn in line.element_names if name + '_aper' in nn])
-            if len(has_apertures) > 0:
-                if len(has_apertures) > 1:
-                    # Choose the aperture closest to the element
-                    has_apertures = [aa for aa in has_apertures
-                                     if line.element_names.index(aa) < line.element_names.index(name)]
-                    has_apertures.sort(key=lambda nn: line.element_names.index(nn))
-                coll_aper = line[has_apertures[-1]]
-                assert coll_aper.__class__.__name__.startswith('Limit')
-                if np.any([name + '_aper_tilt_' in nn for nn in line.element_names]):
-                    raise NotImplementedError("Collimator apertures with tilt not implemented!")
-                if np.any([name + '_aper_offset_' in nn for nn in line.element_names]):
-                    raise NotImplementedError("Collimator apertures with offset not implemented!")
-            else:
-                coll_aper = None
 
+            # Find apertures and store them
+            # TODO: same with cryotanks for FLUKA
+            # TODO: use compound info  ->  need full collimator info from MADX
+            # TODO: this is all very hacky....
+            aper_before = {}
+            aper_after = {}
+            if f'{name}_mken' in line.element_names\
+            and f'{name}_mkex'in line.element_names:
+                # TODO what with transformations? How to shift them in s if different?
+                aper_before = {nn.replace('mken', 'upstream'): line[nn].copy()
+                               for nn in line.element_names if nn.startswith(f'{name}_mken_aper')}
+                aper_after  = {nn.replace('mkex', 'downstream'): line[nn].copy()
+                               for nn in line.element_names if nn.startswith(f'{name}_mkex_aper')}
+            if len(aper_before) == 0:
+                # TODO what with transformations? How to shift them in s from centre to start/end?
+                aper_before = {nn.replace('_aper', '_upstream_aper'): line[nn].copy()
+                               for nn in line.element_names if nn.startswith(f'{name}_aper')}
+            if len(aper_after) == 0:
+                aper_after  = {nn.replace('_aper', '_downstream_aper'): line[nn].copy()
+                               for nn in line.element_names if nn.startswith(f'{name}_aper')}
+            if len(aper_before) == 0 or len(aper_after) == 0:
+                print(f"Warning: No aperture found for collimator {name}!")
+
+            # Remove stuff at location of collimator
+            l = thiscoll['active_length']
+            to_remove = []
+            i = idx - 1
+            # We remove everything between the beginning and end of the collimator except drifts
+            while ss[i] >= ss[idx] - l/2:
+                el = line[i]
+                nn = line.element_names[i]
+                if el.__class__.__name__ == 'Drift':
+                    i -= 1
+                    continue
+                if hasattr(el, 'length') and el.length > 0:
+                    raise ValueError(f"Found active element {nn} with length "
+                                   + f"{el.length} at location inside collimator!")
+                # I don't like this class selection...
+                if not el.__class__.__name__ in ['Marker', 'SRotation', \
+                                       'YRotation', 'XRotation', 'XYShift'] \
+                and not el.__class__.__name__.startswith('Limit'):
+                    print(f"Warning: Removed active element {nn} "
+                        + f"at location inside collimator!")
+                to_remove.append(nn)
+                i -= 1
+            i = idx + 1
+            while ss[i] <= ss[idx] + l/2:
+                el = line[i]
+                nn = line.element_names[i]
+                if el.__class__.__name__ == 'Drift':
+                    i += 1
+                    continue
+                if hasattr(el, 'length') and el.length > 0:
+                    raise ValueError(f"Found active element {nn} with length "
+                                   + f"{el.length} at location inside collimator!")
+                # I don't like this class selection...
+                if not el.__class__.__name__ in ['Marker', 'SRotation', \
+                                       'YRotation', 'XRotation', 'XYShift'] \
+                and not el.__class__.__name__.startswith('Limit'):
+                    print(f"Warning: Removed active element {line.element_names[i]} "
+                        + f"at location inside collimator!")
+                to_remove.append(nn)
+                i += 1
+            for nn in to_remove:
+                # TODO: need to update Compounds
+                line.element_names.remove(nn)
+                line.element_dict.pop(nn)
+
+            # Do the installation
+            s_install = df.loc[name,'s_center'] - thiscoll['active_length']/2 \
+                            - thiscoll['inactive_front']
             line.insert_element(element=newcoll, name=name, at_s=s_install)
 
-            if coll_aper is not None:
-                line.insert_element(element=coll_aper, name=name+'_aper_front', index=name)
-                line.insert_element(element=coll_aper, name=name+'_aper_back',
+            # Reinstall apertures
+            for aper, el in aper_before.items():
+                # TODO: need to update Compounds
+                line.insert_element(element=el, name=aper, index=name)
+            for aper, el in reversed(aper_after.items()):
+                # Reversed because of index+1
+                # TODO: need to update Compounds
+                line.insert_element(element=el, name=aper,
                                     index=line.element_names.index(name)+1)
+
         self._set_record_impacts()
 
 
@@ -854,7 +913,7 @@ class CollimatorManager:
 
         if file is not None:
             with open(Path(file), 'w') as fid:
-                json.dump(self._lossmap, fid, indent=True)
+                json.dump(self._lossmap, fid, cls=xo.JEncoder, indent=True)
     
         return self._lossmap
 

xcoll/scattering_routines/everest/everest.h

@@ -43,6 +43,7 @@ typedef struct EverestCollData_ {
     double ai;
     double eum;
     double collnt;
+    int8_t only_mcs;
 } EverestCollData_;
 typedef EverestCollData_ *EverestCollData;
 

xcoll/scattering_routines/everest/jaw.h

@@ -90,7 +90,7 @@ int ichoix(EverestData restrict everest, LocalParticle* part) {
 
 
 /*gpufun*/
-double* jaw(EverestData restrict everest, LocalParticle* part, double p, double zlm, int only_mcs, int edge_check) {
+double* jaw(EverestData restrict everest, LocalParticle* part, double p, double zlm, int edge_check) {
 
     double* result = (double*)malloc(3 * sizeof(double));
 
@@ -108,7 +108,7 @@ double* jaw(EverestData restrict everest, LocalParticle* part, double p, double
     double z  = LocalParticle_get_y(part);
     double zp = LocalParticle_get_py(part)*rpp_in;
 
-    if (only_mcs) {
+    if (everest->coll->only_mcs) {
         double* res = mcs(everest, part, zlm, p, x, xp, z, zp, edge_check);
         s = res[0];
         x = res[1];
@@ -174,7 +174,7 @@ double* jaw(EverestData restrict everest, LocalParticle* part, double p, double
             int inter = ichoix(everest, part);
             nabs = inter;
             if (inter == 1) {
-                s = (zlm-rlen)+zlm1;
+                s = zlm - rlen + zlm1;
                 m_dpodx = calcionloss(everest, part, rlen);
                 p = p-m_dpodx*s;
                 break;
@@ -208,14 +208,14 @@ double* jaw(EverestData restrict everest, LocalParticle* part, double p, double
             xp = xp + tx;
             zp = zp + tz;
 
-            // Treat single-diffractive scattering.
+            // Treat single-diffractive scattering.    TODO: this does nothing??
             if(inter == 4) {
                 // added update for s
-                s    = (zlm-rlen)+zlm1;
+                s = zlm - rlen + zlm1;
             }
 
             // Calculate the remaining interaction length and close the iteration loop.
-            rlen = rlen-zlm1;
+            rlen = rlen - zlm1;
         }
     }
 
@@ -223,6 +223,7 @@ double* jaw(EverestData restrict everest, LocalParticle* part, double p, double
     LocalParticle_set_px(part, xp/rpp_in);
     LocalParticle_set_y(part, z);
     LocalParticle_set_py(part, zp/rpp_in);
+    LocalParticle_add_to_s(part, s);  // TODO: is this correct with tilt etc?
 
     result[0] = p;
     result[1] = nabs;

xcoll/scattering_routines/everest/materials.py

@@ -1,6 +1,6 @@
 # copyright ############################### #
-# This file is part of the Xcoll Package.  #
-# Copyright (c) CERN, 2023.                 #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
 import xobjects as xo
@@ -31,7 +31,8 @@ class GeneralMaterial(xo.HybridClass):
                 # Index 0:Total, 1:absorption, 2:nuclear elastic, 3:pp or pn elastic
                 #       4:Single Diffractive pp or pn, 5:Coulomb for t above mcs
         'hcut':                     xo.Float64,
-        'name':                     xo.String
+        'name':                     xo.String,
+        '_only_mcs':                xo.Int8
     }
 
     def __init__(self, **kwargs):
@@ -40,6 +41,15 @@ class GeneralMaterial(xo.HybridClass):
             kwargs.setdefault('cross_section', [0., 0., 0., 0., 0., 0.])
             kwargs.setdefault('name', "NO NAME")
             kwargs['name'] = kwargs['name'].ljust(55)  # Pre-allocate 64 byte using whitespace
+            # TODO: this should be better
+            if kwargs.get('Z', None) is None or kwargs.get('A', None) is None \
+            or kwargs.get('density', None) is None \
+            or kwargs.get('excitation_energy', None) is None \
+            or kwargs.get('nuclear_radius', None) is None \
+            or kwargs.get('nuclear_elastic_slope', None) is None:
+                kwargs['_only_mcs'] = True
+            else:
+                kwargs['_only_mcs'] = False
         super().__init__(**kwargs)
         self.name = self.name.strip()
 

xcoll/scattering_routines/everest/multiple_coulomb_scattering.h

@@ -145,10 +145,10 @@ double* mcs(EverestData restrict everest, LocalParticle* part, double zlm1, doub
             xp = res[1];
             free(res);
             if (x <= 0) {
-                s = (rlen0-rlen)+ s;
+                s = rlen0 - rlen + s;
                 break; // go to 20
             }
-            if ((s + dh) >= rlen) {
+            if (s + dh >= rlen) {
                 s = rlen0;
                 break; // go to 20
             }
@@ -157,11 +157,11 @@ double* mcs(EverestData restrict everest, LocalParticle* part, double zlm1, doub
         }
 
     } else {
-        s = rlen0;
-        double* res = scamcs(part, x, xp, s);
+        double* res = scamcs(part, x, xp, rlen0);
         x  = res[0];
         xp = res[1];
         free(res);
+        s = rlen0;
     }
 
     double* res = scamcs(part, z, zp, s);

xcoll/scattering_routines/everest/scatter.h

@@ -89,7 +89,7 @@ void scatter(EverestData restrict everest, LocalParticle* part, double length){
         if (zlm > 0.) {
             is_hit = 1;
 
-            double* jaw_result = jaw(everest, part, energy, zlm, 0, 1);
+            double* jaw_result = jaw(everest, part, energy, zlm, 1);
 
             energy = jaw_result[0];
             if (jaw_result[1] == 1){
@@ -99,7 +99,7 @@ void scatter(EverestData restrict everest, LocalParticle* part, double length){
             free(jaw_result);
 
             if (is_abs != 1) {
-            // Do the rest drift, if particle left collimator early
+               // Do the rest drift, if particle left collimator early
                 Drift_single_particle_4d(part, zlm-s_out);
             }
         }

{xcoll-0.3.0.dist-info → xcoll-0.3.2.dist-info}/METADATA

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

{xcoll-0.3.0.dist-info → xcoll-0.3.2.dist-info}/RECORD

@@ -1,39 +1,39 @@
 LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
-xcoll/__init__.py,sha256=aR1cGO47NuYwAZp1w4xBlnk2lTCxff6aG3IeXeJCO9M,529
-xcoll/beam_elements/__init__.py,sha256=C7MOF6VkGc04b_cadW6pz-nBqTaW9Jmi4RHNA7f8dC8,236
-xcoll/beam_elements/absorber.py,sha256=pTs2PL8L8ewg0AJ-HLDQtDuaIsTX3pJ5Xbq4g4A7KFs,1247
-xcoll/beam_elements/base_collimator.py,sha256=Rfh6ocJKxHJHhxElDDvvTFV5Pz4XgTld796seYMjOl8,15347
+xcoll/__init__.py,sha256=XO53i_KKg6kaJfTCpmNJLy14tou4RMF9MDuXlH70iTQ,642
+xcoll/beam_elements/__init__.py,sha256=9rKZ5O2jRpjkBlU457ii27eiVGAauCM1M-ALDgJ4f0g,591
+xcoll/beam_elements/absorber.py,sha256=qWBEQA7fEJkLffbr2p0Kx5WYJOcoF7Ikz_A_hL1hdAo,1236
+xcoll/beam_elements/base.py,sha256=jeGSn1SXAtzhSpLDXt9sQO6AkeFrmEE80Q5NENtzCDM,15573
 xcoll/beam_elements/collimators_src/absorber.h,sha256=rW8ODVDsKb1VR0Xkrhk2i5AQ1F616ugOfrcI9JPK9Lw,5313
 xcoll/beam_elements/collimators_src/base_block.h,sha256=KlxxqKVNbkoqVvgHFqoPGtFaD82rL6lFzkzkD7G0jsY,421
 xcoll/beam_elements/collimators_src/base_collimator.h,sha256=zUkv-wczcN8xGJqC_c35QhHTTMSTTEy0eG_URUclUrk,428
-xcoll/beam_elements/collimators_src/everest_block.h,sha256=qKPdU24nLMhBbnJRGCrLNcDezE1v-_I_ceeMGMS21uc,5708
-xcoll/beam_elements/collimators_src/everest_collimator.h,sha256=981zTeOmDKdQxmTVyuqwdTcu4_OYkaMlEfEemYrfXIo,6161
-xcoll/beam_elements/collimators_src/everest_crystal.h,sha256=1L9njfnC-2wT6_o3sJujXWtF9bUXcx1OXismkQZWBH0,5187
+xcoll/beam_elements/collimators_src/everest_block.h,sha256=Gpr5wyFLtcaNgVj3auKgFWB-2KPP5cjbEWTfC3f9q1U,5596
+xcoll/beam_elements/collimators_src/everest_collimator.h,sha256=wsGIFBb-VMCRhxwK2G5k9SEX-6SzIgRX4XE5g_I2okg,6296
+xcoll/beam_elements/collimators_src/everest_crystal.h,sha256=pX9cmW5lfsUoszR14tXWU4Emu-0VwX4bbfg6Q6pxXR4,5164
 xcoll/beam_elements/collimators_src/invalid_collimator.h,sha256=xwZEQ1STOS0_a3Rso-03Wp0L7wmC5YVy88xtOt1IyEk,423
-xcoll/beam_elements/everest_collimator.py,sha256=5ipSK_2zBX7BZ4nT7c1HUFuPVtV6PfhAS1ndm0DBk2g,11209
+xcoll/beam_elements/everest.py,sha256=Thh5rk9rulieW6JiiEnq4C9Q3LbOyUg7NIu_dKi7Rko,13469
 xcoll/colldb.py,sha256=e2VK1TE2zOjVLwWD-kSs8g0OEGXDVDQvJKrku7OxCwI,48253
 xcoll/collimator_settings.py,sha256=uj25cpSaM-EY7VM65OUq3Cxb8gedSmi-r9mpkP-ikrg,17015
-xcoll/general.py,sha256=ZMCCv5sIZBnQo76NZ3L4IlkeDlDC2NLHwN9HPp8i4-E,361
-xcoll/headers/checks.h,sha256=xEhZQHIvgl3_sh9nkD-yISbX1tKMXTSAyi81pLnvfJk,1023
+xcoll/general.py,sha256=6Yxa4-rchtnWttv0NdskFDH7tRMFNeAdQa3DG85Vqm4,525
+xcoll/headers/checks.h,sha256=nTo2bMdtfEXdODcImLwkBPfIcdyf4lBn-CrYXlDiT48,1365
 xcoll/headers/particle_states.h,sha256=YxpS56_ZD8mhrqKTohcyVe3hbQZ6H9vspYHSVM6a5-U,844
 xcoll/impacts/__init__.py,sha256=yKQdIeLG2CT3qo-iZB6MODm-WHzkTmulgUvOAS41w7E,39
 xcoll/impacts/impacts.py,sha256=ShyKrN8nmCG9UWS40V5zDeW2W0VTUHZaTKR3FIycOkc,3929
 xcoll/impacts/impacts_src/impacts.h,sha256=w_hDpuffO2zrc_iJQkdoVZwozGDN8QqMLQ1IqgxUtRA,6110
 xcoll/impacts/interaction_types.py,sha256=iHbhC4beN1vZRWeHbALqSFo8lG10R7asc08BRpIqXx0,2493
 xcoll/install_collimators.py,sha256=wJZ7QelXokWuM7kOySg8LMELMCcTmvWWetTOjUjTU04,9003
-xcoll/manager.py,sha256=pI60QiMmGKYcqRdLL7HM3giUJOR5TCYic0JYJ-xVaEk,42067
+xcoll/manager.py,sha256=AF9XziOQ2JT62VSKdzOZzlNFeuaqowzWfqv6msONR9k,45001
 xcoll/rf_sweep.py,sha256=sUstmeeXmozmM4f8xUGCKLyJYOV3xD4aVZG_FCRspVA,5402
 xcoll/scattering_routines/everest/__init__.py,sha256=JVJQnyNeUiA3bTNz-rFp-RYxIMBiKEdBhcokk8aj-dg,101
 xcoll/scattering_routines/everest/constants.h,sha256=9Kus-RKyxsGdlxx9tfMWUz2H5Q-qy6y_F4kddIZzO30,1390
 xcoll/scattering_routines/everest/crystal.h,sha256=WlA_e9ePuVhc9fmMVdSXFBKvHrFvCQ5aG4d67iMrtuQ,54119
-xcoll/scattering_routines/everest/everest.h,sha256=rOlEGmxemOOl7GVv4rn8Ago5G60zqp0ytLPY7BRrbJ8,3722
+xcoll/scattering_routines/everest/everest.h,sha256=p4yfzxRvh19N-V1vorf_CCIBotodV40Lcrzzy7339WM,3743
 xcoll/scattering_routines/everest/everest.py,sha256=dP4br8KN1vJ7IkvQ92KedI4zXyxOvNvOWnS6vAAYUy8,1305
-xcoll/scattering_routines/everest/jaw.h,sha256=iQ1cDyErjLllrkc5uSWtZtEssaXI2q00DRqNkQFV9Ms,7197
-xcoll/scattering_routines/everest/materials.py,sha256=Ck4EC9mDnqtDRnVJIAvNnxwtq0FZNZps8Cm7QBxcnBQ,9150
-xcoll/scattering_routines/everest/multiple_coulomb_scattering.h,sha256=hW1CnnvQYAYUGk5Z9jGqQXKsbBnYUVirVVx6H5RJGlg,4515
+xcoll/scattering_routines/everest/jaw.h,sha256=lmijUU7ih4veifVho3ehGe0DVyIYkDHMVIIclDNbZds,7308
+xcoll/scattering_routines/everest/materials.py,sha256=zzQcJyTCWkzdvRffuBI8oHYeEE6T_wHt9-c9-YICZJ4,9665
+xcoll/scattering_routines/everest/multiple_coulomb_scattering.h,sha256=0Euxsnc2BU8YqxYwfXZIiplNcnDzFq9HgWAoqGwuSAg,4518
 xcoll/scattering_routines/everest/properties.h,sha256=aRqkX5ZpUeTnbygqFVsoIKwWyoB5VXzD5vq5d-xiI74,6585
-xcoll/scattering_routines/everest/scatter.h,sha256=3hQph45_sToEYLCdVxJMtcUGUXznDmxbZtLZOS6CeaA,6333
+xcoll/scattering_routines/everest/scatter.h,sha256=jl-SAIpizYCbUfy1KAhuO00sA5BBhNBCIjueQ8bCxoQ,6333
 xcoll/scattering_routines/everest/scatter_crystal.h,sha256=BHBPNiMKynq1XNxIiWJXYxMwW2bhmmbmC8j7-Dy-SNs,9455
 xcoll/scattering_routines/fluka/flukaio/.git,sha256=5ZxurpwP1waJk2Zu2AstENGOp4rXU8-C-oU4CypmIVI,73
 xcoll/scattering_routines/fluka/flukaio/.gitignore,sha256=bhrpWiAnKChkcwIaIPX_zKxyG2nCtNzNgQAHnIJa8Fw,12
@@ -385,7 +385,7 @@ xcoll/scattering_routines/geant4/collimasim/tests/test_xtrack_colldb_load.py,sha
 xcoll/scattering_routines/geant4/collimasim/tests/test_xtrack_interaction.py,sha256=Z8-ntyXTqK-sr_1PcGtkdQpaOPNVl3xAsu6zSsH7qnQ,6015
 xcoll/scattering_routines/geant4/collimasim/tests/test_xtrack_tilt.py,sha256=q-8CpGFH_5u3r3Xa25EpMbxY1YlLrzJ7Q_i9hSwgTGg,2396
 xcoll/scattering_routines/geant4/collimasim/tests/test_xtrack_tracking.py,sha256=wgf0ncOe3miF3ZOrte_zebY0yulnFz2GFAWksF5-0z0,2940
-xcoll-0.3.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-xcoll-0.3.0.dist-info/WHEEL,sha256=vxFmldFsRN_Hx10GDvsdv1wroKq8r5Lzvjp6GZ4OO8c,88
-xcoll-0.3.0.dist-info/METADATA,sha256=rFEsb8ylvSBoh4itokxFA74Oerl92-Mq6mSIixob3PA,2657
-xcoll-0.3.0.dist-info/RECORD,,
+xcoll-0.3.2.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+xcoll-0.3.2.dist-info/WHEEL,sha256=vxFmldFsRN_Hx10GDvsdv1wroKq8r5Lzvjp6GZ4OO8c,88
+xcoll-0.3.2.dist-info/METADATA,sha256=I8dB3azZohNOKlsDDMJ1zLXjMvcGJfc4lL0_ECKndvM,2657
+xcoll-0.3.2.dist-info/RECORD,,