Merge pull request #58 from xsuite/release/v0.3.2

Release 0.3.2

Author: freddieknets

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"

tests/test__regenerate_kernels.py

@@ -0,0 +1,7 @@
+import xtrack as xt
+
+def test_init():
+    from xtrack.prebuilt_kernels.kernel_definitions import kernel_definitions
+    xcoll_kernels = [ker for ker in kernel_definitions if 'xcoll' in ker]
+    if len(xcoll_kernels) > 0:
+        xt.prebuild_kernels.regenerate_kernels(kernels=xcoll_kernels)

tests/test_rf_sweep.py

@@ -1,17 +1,20 @@
 from pathlib import Path
 import numpy as np
+import pytest
+
 import xtrack as xt
 import xcoll as xc
 from xobjects.test_helpers import for_all_test_contexts
 
 path = Path(__file__).parent / 'data'
 
 @for_all_test_contexts
-def test_dp_pos(test_context):
-    sweep = -300   # negative sweep => positive off-momentum
+@pytest.mark.parametrize("sweep, beam", [[-300, 1], [300, 2]],
+                         ids=["DP pos", "DP neg"])
+def test_rf_sweep(sweep, beam, test_context):
     num_turns = 6000
     num_particles = 5
-    line = xt.Line.from_json(path / f'sequence_lhc_run3_b1.json')
+    line = xt.Line.from_json(path / f'sequence_lhc_run3_b{beam}.json')
 
     line.build_tracker()
 
@@ -21,22 +24,9 @@ def test_dp_pos(test_context):
     rf_sweep = xc.RFSweep(line)
     # This sweep is 3.5 buckets, so check that all particles are at least 3 buckets away
     rf_sweep.track(sweep=sweep, num_turns=num_turns, particles=part)
-    assert np.all(part.delta > 1.5e-3)
-
-
-@for_all_test_contexts
-def test_dp_neg(test_context):
-    sweep = 300   # positive sweep => negative off-momentum
-    num_turns = 6000
-    num_particles = 5
-    line = xt.Line.from_json(path / f'sequence_lhc_run3_b2.json')
-
-    line.build_tracker()
-
-    part = line.build_particles(delta=np.linspace(-2e-4, 2e-4, num_particles),
-                                x_norm=0, px_norm=0, y_norm=0, py_norm=0)
 
-    rf_sweep = xc.RFSweep(line)
-    # This sweep is 3.5 buckets, so check that all particles are at least 3 buckets away
-    rf_sweep.track(sweep=sweep, num_turns=num_turns, particles=part)
-    assert np.all(part.delta < -1.5e-3)
+    # negative sweep => positive off-momentum etc
+    if sweep < 0:
+        assert np.all(part.delta > 1.5e-3)
+    else:
+        assert np.all(part.delta < -1.5e-3)

tests/test_version.py

@@ -1,5 +1,5 @@
 from xcoll import __version__
 
 def test_version():
-    assert __version__ == '0.3.1'
+    assert __version__ == '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/beam_elements/__init__.py

@@ -3,3 +3,9 @@
 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/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);