Merge pull request #26 from xsuite/release0.2.3

Release 0.2.3

Author: freddieknets

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "xcoll"
-version = "0.2.2"
+version = "0.2.3"
 description = "Xsuite collimation package"
 homepage = "https://github.com/xsuite/xcoll"
 repository = "https://github.com/xsuite/xcoll"
@@ -25,7 +25,6 @@ xdeps = ">=0.1.1"
 xpart = ">=0.15.0"
 xtrack = ">=0.36.5"
 xfields = ">=0.12.1"
-xcoll = ">=0.2.1"
 
 [tool.poetry.dev-dependencies]
 pytest = ">=7.3"

tests/test_version.py

@@ -1,5 +1,5 @@
 from xcoll import __version__
 
 def test_version():
-    assert __version__ == '0.2.2'
+    assert __version__ == '0.2.3'
 

xcoll/__init__.py

@@ -1,9 +1,8 @@
 from .general import _pkg_root
 
 from .beam_elements import BaseCollimator, BlackAbsorber, EverestCollimator, EverestCrystal
-from .scattering_routines.everest import materials
-from .scattering_routines.everest.materials import GeneralMaterial, Material, CrystalMaterial
+from .scattering_routines.everest import materials, Material, CrystalMaterial
 from .manager import CollimatorManager
 from .colldb import CollimatorDatabase, load_SixTrack_colldb
 
-__version__ = '0.2.2'
+__version__ = '0.2.3'

xcoll/beam_elements/collimators_src/everest_collimator.h

@@ -8,6 +8,12 @@
 #include <math.h>
 #include <stdio.h>
 
+/*gpufun*/
+void EverestCollimator_set_material(EverestCollimatorData el, LocalParticle* part0){
+    MaterialData material = EverestCollimatorData_getp__material(el);
+    RandomRutherfordData rng = EverestCollimatorData_getp_rutherford_rng(el);
+    RandomRutherford_set_by_xcoll_material(rng, (GeneralMaterialData) material);
+}
 
 /*gpufun*/
 void EverestCollimator_track_local_particle(EverestCollimatorData el, LocalParticle* part0) {
@@ -19,7 +25,6 @@ void EverestCollimator_track_local_particle(EverestCollimatorData el, LocalParti
 
     MaterialData material = EverestCollimatorData_getp__material(el);
     RandomRutherfordData rng = EverestCollimatorData_getp_rutherford_rng(el);
-    RandomRutherford_set_by_xcoll_material(rng, (GeneralMaterialData) material);
 
     // Collimator properties
     double const length     = EverestCollimatorData_get_active_length(el);
@@ -48,10 +53,10 @@ void EverestCollimator_track_local_particle(EverestCollimatorData el, LocalParti
         } else {
             // Check collimator initialisation
             int8_t is_tracking = assert_tracking(part, XC_ERR_INVALID_TRACK);
-            int8_t rng_set     = assert_rng_set(part, RNG_ERR_SEEDS_NOT_SET);
-            int8_t ruth_set    = assert_rutherford_set(rng, part, RNG_ERR_RUTH_NOT_SET);
+            int8_t rng_is_set  = assert_rng_set(part, RNG_ERR_SEEDS_NOT_SET);
+            int8_t ruth_is_set = assert_rutherford_set(rng, part, RNG_ERR_RUTH_NOT_SET);
 
-            if (is_tracking && rng_set && ruth_set) {
+            if (is_tracking && rng_is_set && ruth_is_set) {
                 // Drift inactive front
                 Drift_single_particle(part, inactive_front);
 

xcoll/beam_elements/collimators_src/everest_crystal.h

@@ -9,6 +9,14 @@
 #include <stdio.h>
 
 
+
+/*gpufun*/
+void EverestCrystal_set_material(EverestCrystalData el, LocalParticle* part0){
+    CrystalMaterialData material = EverestCrystalData_getp__material(el);
+    RandomRutherfordData rng = EverestCrystalData_getp_rutherford_rng(el);
+    RandomRutherford_set_by_xcoll_material(rng, (GeneralMaterialData) material);
+}
+
 /*gpufun*/
 void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* part0) {
     int8_t active      = EverestCrystalData_get_active(el);
@@ -19,7 +27,6 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
 
     CrystalMaterialData material = EverestCrystalData_getp__material(el);
     RandomRutherfordData rng = EverestCrystalData_getp_rutherford_rng(el);
-    RandomRutherford_set_by_xcoll_material(rng, (GeneralMaterialData) material);
 
     // Crystal properties
     double length  = EverestCrystalData_get_active_length(el);
@@ -60,6 +67,13 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
     double const cry_orient = EverestCrystalData_get__orient(el);
     double const cry_miscut = EverestCrystalData_get_miscut(el);
 
+    // Impact table
+    CollimatorImpactsData record = EverestCrystalData_getp_internal_record(el, part0);
+    RecordIndex record_index = NULL;
+    if (record){
+        record_index = CollimatorImpactsData_getp__index(record);
+    }
+
     //start_per_particle_block (part0->part)
         if (!active){
             // Drift full length
@@ -68,10 +82,10 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
         } else {
             // Check collimator initialisation
             int8_t is_tracking = assert_tracking(part, XC_ERR_INVALID_TRACK);
-            int8_t rng_set     = assert_rng_set(part, RNG_ERR_SEEDS_NOT_SET);
-            int8_t ruth_set    = assert_rutherford_set(rng, part, RNG_ERR_RUTH_NOT_SET);
+            int8_t rng_is_set  = assert_rng_set(part, RNG_ERR_SEEDS_NOT_SET);
+            int8_t ruth_is_set = assert_rutherford_set(rng, part, RNG_ERR_RUTH_NOT_SET);
 
-            if (is_tracking && rng_set && ruth_set) {
+            if (is_tracking && rng_is_set && ruth_is_set) {
                 // Drift inactive front
                 Drift_single_particle(part, inactive_front);
 
@@ -83,7 +97,7 @@ void EverestCrystal_track_local_particle(EverestCrystalData el, LocalParticle* p
 
                 scatter_cry(part, length, material, rng, cos_zL, sin_zL, c_aperture, c_offset,
                             side, cry_tilt, cry_rcurv, cry_bend, cry_alayer, cry_xmax, cry_ymax, cry_orient, 
-                            cry_miscut);
+                            cry_miscut, record, record_index);
 
                 // Return from closed orbit
                 LocalParticle_add_to_x(part, co_x);

xcoll/beam_elements/everest_collimator.py

@@ -8,7 +8,7 @@
 import xtrack as xt
 
 from .base_collimator import BaseCollimator, InvalidCollimator
-from ..scattering_routines.everest import Material, CrystalMaterial, EverestEngine
+from ..scattering_routines.everest import GeneralMaterial, Material, CrystalMaterial, EverestEngine
 from ..general import _pkg_root
 
 
@@ -20,6 +20,8 @@
 #      only activated around the track command. Furthermore, because of 'iscollective = False' we need to specify 
 #      get_backtrack_element. We want it nicer..
 
+# TODO: _per_particle_kernels should be a normal kernel (such that we don't need to pass a dummy Particles() )
+
 class EverestCollimator(BaseCollimator):
     _xofields = { **BaseCollimator._xofields,
         '_material':        Material,
@@ -41,6 +43,13 @@ class EverestCollimator(BaseCollimator):
         _pkg_root.joinpath('beam_elements','collimators_src','everest_collimator.h')
     ]
 
+    _per_particle_kernels = {
+        '_EverestCollimator_set_material': xo.Kernel(
+                c_name='EverestCollimator_set_material',
+                args=[]
+            )
+        }
+
 
     def __init__(self, **kwargs):
         if '_xobject' not in kwargs:
@@ -54,8 +63,8 @@ def __init__(self, **kwargs):
             kwargs.setdefault('rutherford_rng', xt.RandomRutherford())
             kwargs.setdefault('_tracking', True)
         super().__init__(**kwargs)
-#         if '_xobject' not in kwargs:
-#             self.random_generator.set_rutherford_by_xcoll_material(self.material)
+        if '_xobject' not in kwargs:
+            self._EverestCollimator_set_material(xp.Particles())
 
     @property
     def material(self):
@@ -67,7 +76,7 @@ def material(self, material):
             if not isinstance('material', dict) or material['__class__'] != "Material":
                 raise ValueError("Invalid material!")
         self._material = material
-#         self.random_generator.set_rutherford_by_xcoll_material(material)
+        self._EverestCollimator_set_material(xp.Particles())
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
         # TODO: this should be an InvalidCollimator
@@ -103,6 +112,13 @@ class EverestCrystal(BaseCollimator):
         _pkg_root.joinpath('beam_elements','collimators_src','everest_crystal.h')
     ]
 
+    _per_particle_kernels = {
+        '_EverestCrystal_set_material': xo.Kernel(
+                c_name='EverestCrystal_set_material',
+                args=[]
+            )
+        }
+
 
     def __init__(self, **kwargs):
         if '_xobject' not in kwargs:
@@ -122,8 +138,8 @@ def __init__(self, **kwargs):
             kwargs.setdefault('rutherford_rng', xt.RandomRutherford())
             kwargs.setdefault('_tracking', True)
         super().__init__(**kwargs)
-#         if '_xobject' not in kwargs:
-#             self.random_generator.set_rutherford_by_xcoll_material(self.material)
+        if '_xobject' not in kwargs:
+            self._EverestCrystal_set_material(xp.Particles())
 
     @property
     def lattice(self):
@@ -148,7 +164,7 @@ def material(self, material):
             if not isinstance(material, dict) or material['__class__'] != "CrystalMaterial":
                 raise ValueError("Invalid material!")
         self._material = material
-#         self.random_generator.set_rutherford_by_xcoll_material(material)
+        self._EverestCrystal_set_material(xp.Particles())
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
         # TODO: this should be an InvalidCollimator

xcoll/scattering_routines/everest/__init__.py

@@ -1,2 +1,2 @@
-from .materials import Material, CrystalMaterial
+from .materials import GeneralMaterial, Material, CrystalMaterial
 from .everest import EverestEngine