Merge pull request #29 from xsuite/release0.2.4

Release 0.2.4

examples/lhc_run3_lossmap.py

@@ -1,9 +1,12 @@
 import json
 import numpy as np
 from pathlib import Path
+
+import xobjects as xo
 import xtrack as xt
 import xcoll as xc
 
+context = xo.ContextCpu(omp_num_threads='auto')
 
 # On a modern CPU, we get ~5000 particle*turns/s
 # So this script should take around half an hour
@@ -29,7 +32,7 @@
 
 
 # Initialise collmanager
-coll_manager = xc.CollimatorManager.from_yaml(path_in / 'colldb' / f'lhc_run3.yaml', line=line, beam=beam)
+coll_manager = xc.CollimatorManager.from_yaml(path_in / 'colldb' / f'lhc_run3.yaml', line=line, beam=beam, _context=context)
 
 
 # Install collimators into line

pyproject.toml

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

tests/test_everest.py

@@ -3,6 +3,8 @@
 import numpy as np
 import xpart as xp
 import xcoll as xc
+from xobjects.test_helpers import for_all_test_contexts
+
 
 materials_b1 = {
   'BE':   'tcl.4r1.b1',
@@ -83,56 +85,88 @@
 
 path = Path(__file__).parent / 'data_test_everest'
 
-def test_primaries():
-    _track_collimator('tcp.c6l7.b1')
-    _track_collimator('tcp.c6r7.b2')
-
-def test_materials_b1():
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_primaries(test_context):
+    _track_collimator('tcp.c6l7.b1', _context=test_context)
+    _track_collimator('tcp.c6r7.b2', _context=test_context)
+
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_materials_b1(test_context):
     for key, name in materials_b1.items():
-        _track_collimator(name)
+        _track_collimator(name, _context=test_context)
 
-def test_materials_b2():
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_materials_b2(test_context):
     for key, name in materials_b2.items():
-        _track_collimator(name)
+        _track_collimator(name, _context=test_context)
 
-def test_angles_b1():
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_angles_b1(test_context):
     for key, name in angles_b1.items():
-        _track_collimator(name)
+        _track_collimator(name, _context=test_context)
 
-def test_angles_b2():
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_angles_b2(test_context):
     for key, name in angles_b2.items():
-        _track_collimator(name)
+        _track_collimator(name, _context=test_context)
 
-def test_lengths_b1():
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_lengths_b1(test_context):
     for key, name in lengths_b1.items():
-        _track_collimator(name)
+        _track_collimator(name, _context=test_context)
 
-def test_lengths_b2():
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_lengths_b2(test_context):
     for key, name in lengths_b2.items():
-        _track_collimator(name)
+        _track_collimator(name, _context=test_context)
 
-def test_offsets_b1():
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_offsets_b1(test_context):
     for key, name in offsets_b1.items():
-        _track_collimator(name)
+        _track_collimator(name, _context=test_context)
 
-def test_offsets_b2():
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_offsets_b2(test_context):
     for key, name in offsets_b2.items():
-        _track_collimator(name)
+        _track_collimator(name, _context=test_context)
 
-def test_crystals():
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
+def test_crystals(test_context):
     for name in crystals_b1 + crystals_b2:
-        _track_collimator(name)
+        _track_collimator(name, _context=test_context)
 
 
-def _track_collimator(name, atolx=3e-9, atoly=3e-10, atolpx=5e-9, atolpy=2e-9, atolz=1e-11, atold=2e-8):
+def _track_collimator(name, atolx=3e-9, atoly=3e-10, atolpx=5e-9, atolpy=2e-9, atolz=1e-11, atold=2e-8, _context=None):
+    if _context is None:
+        _context = xo.ContextCpu()
     with open(Path(path, 'initial.json'), 'r') as fid:
-        part = xp.Particles.from_dict(json.load(fid))
+        part = xp.Particles.from_dict(json.load(fid), _context=_context)
     with open(Path(path, 'Collimators', name+'.json'), 'r') as fid:
         colldict = json.load(fid)
     if colldict['__class__'] == 'EverestCollimator':
-        coll = xc.EverestCollimator.from_dict(colldict)
+        coll = xc.EverestCollimator.from_dict(colldict, _context=_context)
     elif colldict['__class__'] == 'EverestCrystal':
-        coll = xc.EverestCrystal.from_dict(colldict)
+        coll = xc.EverestCrystal.from_dict(colldict, _context=_context)
     coll.track(part)
     part.sort(interleave_lost_particles=True)
     with open(Path(path, 'Ref',name+'.json'), 'r') as fid:

tests/test_lossmap.py

@@ -5,22 +5,27 @@
 import xcoll as xc
 import pytest
 from xpart.test_helpers import flaky_assertions, retry
+from xobjects.test_helpers import for_all_test_contexts
 
 path = Path(__file__).parent / 'data'
 
 # https://github.com/xsuite/xtrack/blob/18b1ac33d6a9d87a156e87bfb71cb2c8011085f6/tests/test_radiation.py#LL138C5-L138C29
+@for_all_test_contexts(
+    excluding=('ContextCupy', 'ContextPyopencl')  # Rutherford RNG not on GPU
+)
 @retry()
 @pytest.mark.parametrize("beam, plane, npart, interpolation, ignore_crystals", [
                             [1, 'H', 25000, 0.2, True],
                             [2, 'V', 25000, 0.3, True],
                             [1, 'V', 35000, 0.1, False],
                             [2, 'H', 30000, 0.15, False]
                         ], ids=["B1", "B2V", "B1V_crystals", "B2H_crystals"])
-def test_run_lossmap(beam, plane, npart, interpolation, ignore_crystals):
+def test_run_lossmap(beam, plane, npart, interpolation, ignore_crystals, test_context):
 
     line = xt.Line.from_json(path / f'sequence_lhc_run3_b{beam}.json')
 
-    coll_manager = xc.CollimatorManager.from_yaml(path / f'colldb_lhc_run3_ir7.yaml', line=line, beam=beam, ignore_crystals=ignore_crystals)
+    coll_manager = xc.CollimatorManager.from_yaml(path / f'colldb_lhc_run3_ir7.yaml', line=line, beam=beam,
+                                                  ignore_crystals=ignore_crystals, _context=test_context)
 
     coll_manager.install_everest_collimators()
     coll_manager.build_tracker()

tests/test_version.py

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

xcoll/__init__.py

@@ -5,4 +5,4 @@
 from .manager import CollimatorManager
 from .colldb import CollimatorDatabase, load_SixTrack_colldb
 
-__version__ = '0.2.3'
+__version__ = '0.2.4'

xcoll/manager.py

@@ -84,7 +84,7 @@ def __init__(self, **kwargs):
             raise ValueError("The variable 'line' needs to be an xtrack Line object!")
         else:
             self.line = line
-        self.line._needs_rng = True
+        self.line._needs_rng = True  # TODO not needed if only BlackAbsorbers
         if beam is not None and beam > 1:
             self._line_is_reversed = True
         else:
@@ -652,13 +652,15 @@ def generate_pencil_on_collimator(self, collimator, num_particles, *, side='+-',
             part = xp.build_particles(
                     x=pencil, px=p_pencil, y_norm=transverse_norm, py_norm=p_transverse_norm,
                     zeta=zeta, delta=delta, nemitt_x=nemitt_x, nemitt_y=nemitt_y,
-                    line=self.line, at_element=collimator, match_at_s=match_at_s
+                    line=self.line, at_element=collimator, match_at_s=match_at_s,
+                    _context=self._buffer.context
             )
         else:
             part = xp.build_particles(
                     x_norm=transverse_norm, px_norm=p_transverse_norm, y=pencil, py=p_pencil, 
                     zeta=zeta, delta=delta, nemitt_x=nemitt_x, nemitt_y=nemitt_y,
-                    line=self.line, at_element=collimator, match_at_s=match_at_s
+                    line=self.line, at_element=collimator, match_at_s=match_at_s,
+                    _context=self._buffer.context
             )
 
         part._init_random_number_generator()

xcoll/scattering_routines/everest/everest.h

@@ -6,6 +6,13 @@
 #ifndef XCOLL_EVEREST_ENGINE_H
 #define XCOLL_EVEREST_ENGINE_H
 
+/*gpufun*/
+void Drift_single_particle_4d(LocalParticle* part, double length){
+    double zeta = LocalParticle_get_zeta(part);
+    Drift_single_particle(part, length);
+    LocalParticle_set_zeta(part, zeta);
+}
+
 /*gpufun*/
 double drift_zeta_single(double rvv, double xp, double yp, double length){
     double const rv0v = 1./rvv;

xcoll/scattering_routines/everest/scatter.h

@@ -226,6 +226,8 @@ void scatter(LocalParticle* part, double length, MaterialData material, RandomRu
         y_out  =  z*cRRot -  x*sRRot;
         xp_out = xp*cRRot + zp*sRRot;
         yp_out = zp*cRRot - xp*sRRot;
+    } else {
+        Drift_single_particle_4d(part, length);
     }
 
     LocalParticle_set_x(part, x_out);