xcoll 0.3.0__tar.gz → 0.3.2__tar.gz

{xcoll-0.3.0 → xcoll-0.3.2}/PKG-INFO

@@ -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 → xcoll-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "xcoll"
-version = "0.3.0"
+version = "0.3.2"
 description = "Xsuite collimation package"
 homepage = "https://github.com/xsuite/xcoll"
 repository = "https://github.com/xsuite/xcoll"

{xcoll-0.3.0 → xcoll-0.3.2}/setup.py

@@ -68,7 +68,7 @@ extras_require = \
 
 setup_kwargs = {
     'name': 'xcoll',
-    'version': '0.3.0',
+    'version': '0.3.2',
     'description': 'Xsuite collimation package',
     'long_description': '# xcoll\n\n<!---![PyPI - Python Version](https://img.shields.io/pypi/pyversions/xcoll?logo=PyPI?style=plastic) ![PyPI - Wheel](https://img.shields.io/pypi/wheel/xcoll?logo=PyPI?style=plastic)-->\n\n![GitHub release (latest by date)](https://img.shields.io/github/v/release/xsuite/xcoll?style=plastic)\n![GitHub](https://img.shields.io/github/license/xsuite/xcoll?style=plastic)\n![PyPi](https://img.shields.io/pypi/dm/xcoll?logo=PyPI&style=plastic)\n![GitHub all releases](https://img.shields.io/github/downloads/xsuite/xcoll/total?logo=GitHub&style=plastic)\n\n![GitHub pull requests](https://img.shields.io/github/issues-pr/xsuite/xcoll?logo=GitHub&style=plastic)\n![GitHub issues](https://img.shields.io/github/issues/xsuite/xcoll?logo=GitHub&style=plastic)\n![GitHub repo size](https://img.shields.io/github/repo-size/xsuite/xcoll?logo=GitHub&style=plastic)\n\nCollimation in xtrack simulations\n\n## Description\n\n## Getting Started\n\n### Dependencies\n\n* python >= 3.8\n    * numpy\n    * pandas\n    * xsuite (in particular xobjects, xdeps, xtrack, xpart)\n\n### Installing\n`xcoll` is packaged using `poetry`, and can be easily installed with `pip`:\n```bash\npip install xcoll\n```\nFor a local installation, clone and install in editable mode (need to have `pip` >22):\n```bash\ngit clone git@github.com:xsuite/xcoll.git\npip install -e xcoll\n```\n\n### Example\n\n## Features\n\n## Authors\n\n* [Frederik Van der Veken](https://github.com/freddieknets) (frederik@cern.ch)\n* [Despina Demetriadou](https://github.com/ddemetriadou)\n* [Andrey Abramov](https://github.com/anabramo)\n* [Giovanni Iadarola](https://github.com/giadarol)\n\n\n## Version History\n\n* 0.1\n    * Initial Release\n\n## License\n\nThis project is [Apache 2.0 licensed](./LICENSE).\n',
     'author': 'Frederik F. Van der Veken',

{xcoll-0.3.0 → xcoll-0.3.2}/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-0.3.2/xcoll/beam_elements/__init__.py

@@ -0,0 +1,11 @@
+from .base import BaseBlock, BaseCollimator
+from .absorber import BlackAbsorber
+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-0.3.0 → xcoll-0.3.2}/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-0.3.0/xcoll/beam_elements/base_collimator.py → xcoll-0.3.2/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, 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-0.3.0 → xcoll-0.3.2}/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-0.3.0 → xcoll-0.3.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-0.3.0 → xcoll-0.3.2}/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-0.3.0/xcoll/beam_elements/everest_collimator.py → xcoll-0.3.2/xcoll/beam_elements/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-0.3.0 → xcoll-0.3.2}/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-0.3.0 → xcoll-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-0.3.0 → xcoll-0.3.2}/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-0.3.0 → xcoll-0.3.2}/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;