xcoll 0.3.1__tar.gz → 0.3.2__tar.gz

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

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

{xcoll-0.3.1 → xcoll-0.3.2}/pyproject.toml

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

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

@@ -68,7 +68,7 @@ extras_require = \
 
 setup_kwargs = {
     'name': 'xcoll',
-    'version': '0.3.1',
+    '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.1 → 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.1 → 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.1 → 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.1 → 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.1 → 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.1 → xcoll-0.3.2}/xcoll/beam_elements/everest.py

@@ -14,7 +14,6 @@ from ..scattering_routines.everest import GeneralMaterial, Material, CrystalMate
 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.1 → xcoll-0.3.2}/xcoll/general.py

@@ -7,10 +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."
+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.1'
+__version__ = '0.3.2'
 # ===================

{xcoll-0.3.1 → 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.1 → xcoll-0.3.2}/xcoll/manager.py

@@ -913,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.1 → 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;
 

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

@@ -1,5 +0,0 @@
-from .base import BaseBlock, BaseCollimator
-from .absorber import BlackAbsorber
-from .everest import EverestBlock, EverestCollimator, EverestCrystal
-
-_all_collimator_types = {BlackAbsorber, EverestCollimator, EverestCrystal}