xcoll 0.4.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

xcoll/scattering_routines/geant4/collimasim/tests/test_xtrack_tracking.py

@@ -0,0 +1,97 @@
+import numpy as np
+import xobjects as xo
+import xtrack as xt
+import xpart as xp
+
+import collimasim as cs
+
+import multiprocessing
+
+def make_particles():
+    np.random.seed(seed=1994)
+    n_part = 5000
+    particles = xp.Particles(
+        _capacity=n_part,
+		p0c = 185.e9,
+        mass0 = xp.ELECTRON_MASS_EV,
+        x=np.random.uniform(-1e-3, 1e-3, n_part),
+        px=np.random.uniform(-1e-4, 1e-4, n_part),
+        y=np.random.uniform(-1e-3, 1e-3, n_part),
+        py=np.random.uniform(-1e-4, 1e-4, n_part),
+        zeta=np.random.uniform(-0.1, 0.1, n_part),
+        delta=np.random.uniform(-0.9, 0.9, n_part),
+    )
+
+    return particles
+
+
+def run_g4_logitudinal(particles):
+    np.random.seed(seed=1994)
+
+    g4man = cs.Geant4CollimationManager(collimator_file="resources/CollDB_testing.dat",
+                                        bdsim_config_file="resources/settings_black_absorber.gmad",
+                                        tfs_file="resources/twiss_file_testing.tfs",
+                                        reference_pdg_id=-11,
+                                        reference_kinetic_energy=182.5e9,
+                                        emittance_norm=(1.e-3, 1.e-6),
+                                        relative_energy_cut=0.001,
+                                        seed=1993,
+                                        # batchMode=False
+                                        batchMode=True
+                                        )
+
+    coll1 = g4man.make_xtg4_collimator("coll_open") # Use new convenience method
+
+    # Generate a simple sequence
+    line = xt.Line(
+        elements=[coll1,
+                  ])
+
+    context = xo.ContextCpu() 
+    line.build_tracker(_context=context)
+
+    part_copy = particles.copy()
+    line.track(part_copy, num_turns=1)
+    part_copy.remove_unused_space()
+
+    return part_copy
+
+def run_drift_logitudinal(particles):
+
+    line = xt.Line(
+        elements=[xt.Drift(length=10),
+                  ])
+
+    context = xo.ContextCpu() 
+    line.config.XTRACK_USE_EXACT_DRIFTS = True
+    line.build_tracker(_context=context)
+
+    part_copy = particles.copy()
+    line.track(part_copy, num_turns=1)
+    part_copy.remove_unused_space()
+
+    return part_copy
+
+
+def test_xtrack_tracking():
+    particles = make_particles()
+    part_g4 = run_g4_logitudinal(particles)
+    part_dr = run_drift_logitudinal(particles)
+
+    # BDSIM tracks trough some geometry padding now, which is not corrected for now
+    # so use a fairly lenient tolerance of 1e-8 for all coordinates over 10 m of drift
+    all_close = True
+    for var in ['x', 'px', 'y', 'py', 'zeta', 'delta']:
+        atol = 1e-8
+        rtol = 1e-5
+        var_close = np.all(np.isclose(getattr(part_g4, var), getattr(part_dr, var), atol=atol, rtol=rtol))
+
+        all_close &= var_close
+        print("Var {}: {}".format(var, var_close))
+        
+    assert all_close
+
+
+
+if __name__ == '__main__':
+    test_xtrack_tracking()

xcoll/scattering_routines/geant4/collimasim/tests/test_xtrack_tracking_ions.py

@@ -0,0 +1,96 @@
+import numpy as np
+import xobjects as xo
+import xtrack as xt
+import xpart as xp
+
+import collimasim as cs
+
+import multiprocessing
+
+def make_particles():
+    np.random.seed(seed=1994)
+    n_part = 5000
+    particles = xp.Particles(
+        _capacity=n_part,
+		p0c = 36900e9,
+        mass0 = 193687676197.13638,
+        x=np.random.uniform(-1e-3, 1e-3, n_part),
+        px=np.random.uniform(-1e-4, 1e-4, n_part),
+        y=np.random.uniform(-1e-3, 1e-3, n_part),
+        py=np.random.uniform(-1e-4, 1e-4, n_part),
+        zeta=np.random.uniform(-0.1, 0.1, n_part),
+        delta=np.random.uniform(-0.4, 0.4, n_part),
+    )
+
+    return particles
+
+
+def run_g4_logitudinal(particles):
+    np.random.seed(seed=1994)
+
+    g4man = cs.Geant4CollimationManager(collimator_file="resources/CollDB_testing.dat",
+                                        bdsim_config_file="resources/settings_black_absorber.gmad",
+                                        tfs_file="resources/twiss_file_testing.tfs",
+                                        reference_pdg_id=1000822080,
+                                        reference_kinetic_energy=36706820652500.0,
+                                        emittance_norm=(5.e-5, 5.e-5),
+                                        relative_energy_cut=0.001,
+                                        seed=1993,
+                                        # batchMode=False
+                                        batchMode=True
+                                        )
+
+    coll1 = g4man.make_xtg4_collimator("coll_open") # Use new convenience method
+
+    # Generate a simple sequence
+    line = xt.Line(
+        elements=[coll1,
+                  ])
+
+    context = xo.ContextCpu() 
+    line.build_tracker(_context=context)
+
+    part_copy = particles.copy()
+    line.track(part_copy, num_turns=1)
+    part_copy.remove_unused_space()
+
+    return part_copy
+
+def run_drift_logitudinal(particles):
+
+    line = xt.Line(
+        elements=[xt.Drift(length=10),
+                  ])
+
+    context = xo.ContextCpu() 
+    line.config.XTRACK_USE_EXACT_DRIFTS = True
+    line.build_tracker(_context=context)
+
+    part_copy = particles.copy()
+    line.track(part_copy, num_turns=1)
+    part_copy.remove_unused_space()
+
+    return part_copy
+
+
+def test_xtrack_tracking():
+    particles = make_particles()
+    part_g4 = run_g4_logitudinal(particles)
+    part_dr = run_drift_logitudinal(particles)
+
+    # BDSIM tracks trough some geometry padding now, which is not corrected for now
+    # so use a fairly lenient tolerance of 1e-8 for all coordinates over 10 m of drift
+    all_close = True
+    for var in ['x', 'px', 'y', 'py', 'zeta', 'delta']:
+        atol = 1e-6 # For ions, the zeta variable doesn't seem to agree so well, e.g it is up to 1e-7 m off
+        rtol = 1e-5
+        var_close = np.all(np.isclose(getattr(part_g4, var), getattr(part_dr, var), atol=atol, rtol=rtol))
+
+        all_close &= var_close
+        print("Var {}: {}".format(var, var_close))
+    
+    assert all_close
+
+
+if __name__ == '__main__':
+    test_xtrack_tracking()

xcoll/scattering_routines/geometry/collimator_geometry.h

@@ -33,7 +33,8 @@ typedef struct CollimatorGeometry_ {
     // Impact table
     InteractionRecordData record;
     RecordIndex record_index;
-    int8_t record_touches;
+    int8_t record_impacts;
+    int8_t record_exits;
 } CollimatorGeometry_;
 typedef CollimatorGeometry_* CollimatorGeometry;
 
@@ -109,7 +110,6 @@ int8_t hit_jaws_check_and_transform(LocalParticle* part, CollimatorGeometry rest
     // if bothsided  and no hit: lab  frame
     // if bothsided  and hit:    hit   frame
 
-
     // Drift to the impact position or end, and move to jaw frame if relevant
     if (is_hit == 1){
         // Move to the impact position
@@ -124,7 +124,7 @@ int8_t hit_jaws_check_and_transform(LocalParticle* part, CollimatorGeometry rest
         // Rotate the reference frame to tilt
         double new_s = YRotation_single_particle_rotate_only(part, LocalParticle_get_s(part), asin(cg->sin_yL));
         LocalParticle_set_s(part, new_s);
-        if (cg->record_touches){
+        if (cg->record_impacts){
             InteractionRecordData_log(cg->record, cg->record_index, part, XC_ENTER_JAW_L);
         }
 
@@ -148,7 +148,7 @@ int8_t hit_jaws_check_and_transform(LocalParticle* part, CollimatorGeometry rest
 #else
         LocalParticle_scale_xp(part, -1);
 #endif
-        if (cg->record_touches){
+        if (cg->record_impacts){
             InteractionRecordData_log(cg->record, cg->record_index, part, XC_ENTER_JAW_R);
         }
 
@@ -166,6 +166,11 @@ int8_t hit_jaws_check_and_transform(LocalParticle* part, CollimatorGeometry rest
 
 /*gpufun*/
 void hit_jaws_transform_back(int8_t is_hit, LocalParticle* part, CollimatorGeometry restrict cg){
+    if (is_hit != 0 && LocalParticle_get_state(part) > 0){
+        if (cg->record_exits){
+            InteractionRecordData_log(cg->record, cg->record_index, part, XC_EXIT_JAW);
+        }
+    }
     if (is_hit == 1){
         // Rotate back from tilt
         double new_s = YRotation_single_particle_rotate_only(part, LocalParticle_get_s(part), -asin(cg->sin_yL));
@@ -175,9 +180,6 @@ void hit_jaws_transform_back(int8_t is_hit, LocalParticle* part, CollimatorGeome
         LocalParticle_add_to_s(part, cg->length/2*(1 - cg->cos_yL));
         // If particle survived, drift to end of element
         if (LocalParticle_get_state(part) > 0){
-            if (cg->record_touches){
-                InteractionRecordData_log(cg->record, cg->record_index, part, XC_EXIT_JAW);
-            }
 #ifdef XCOLL_USE_EXACT
             Drift_single_particle_exact(part, cg->length - LocalParticle_get_s(part));
 #else
@@ -202,9 +204,6 @@ void hit_jaws_transform_back(int8_t is_hit, LocalParticle* part, CollimatorGeome
         LocalParticle_add_to_s(part, cg->length/2*(1 - cg->cos_yR));
         // If particle survived, drift to end of element
         if (LocalParticle_get_state(part) > 0){
-            if (cg->record_touches){
-                InteractionRecordData_log(cg->record, cg->record_index, part, XC_EXIT_JAW);
-            }
 #ifdef XCOLL_USE_EXACT
             Drift_single_particle_exact(part, cg->length - LocalParticle_get_s(part));
 #else

xcoll/scattering_routines/geometry/crystal_geometry.h

@@ -35,7 +35,8 @@ typedef struct CrystalGeometry_ {
     // Impact table
     InteractionRecordData record;
     RecordIndex record_index;
-    int8_t record_touches;
+    int8_t record_impacts;
+    int8_t record_exits;
 } CrystalGeometry_;
 typedef CrystalGeometry_* CrystalGeometry;
 
@@ -82,7 +83,7 @@ int8_t hit_crystal_check_and_transform(LocalParticle* part, CrystalGeometry rest
         double new_s = YRotation_single_particle_rotate_only(part, LocalParticle_get_s(part), asin(cg->sin_y));
         LocalParticle_set_s(part, new_s);
         if (cg->side == 1){
-            if (cg->record_touches){
+            if (cg->record_impacts){
                 InteractionRecordData_log(cg->record, cg->record_index, part, XC_ENTER_JAW_L);
             }
 
@@ -94,7 +95,7 @@ int8_t hit_crystal_check_and_transform(LocalParticle* part, CrystalGeometry rest
 #else
             LocalParticle_scale_xp(part, -1);
 #endif
-            if (cg->record_touches){
+            if (cg->record_impacts){
                 InteractionRecordData_log(cg->record, cg->record_index, part, XC_ENTER_JAW_R);
             }
         }
@@ -117,6 +118,11 @@ int8_t hit_crystal_check_and_transform(LocalParticle* part, CrystalGeometry rest
 /*gpufun*/
 void hit_crystal_transform_back(int8_t is_hit, LocalParticle* part, CrystalGeometry restrict cg){
     if (is_hit != 0){
+        if (LocalParticle_get_state(part) > 0){
+            if (cg->record_exits){
+                InteractionRecordData_log(cg->record, cg->record_index, part, XC_EXIT_JAW);
+            }
+        }
         if (cg->side == -1){
             // Mirror back
             LocalParticle_scale_x(part, -1);
@@ -133,9 +139,6 @@ void hit_crystal_transform_back(int8_t is_hit, LocalParticle* part, CrystalGeome
         XYShift_single_particle(part, -cg->jaw_U, 0);
         // If particle survived, drift to end of element
         if (LocalParticle_get_state(part) > 0){
-            if (cg->record_touches){
-                InteractionRecordData_log(cg->record, cg->record_index, part, XC_EXIT_JAW);
-            }
 #ifdef XCOLL_USE_EXACT
             Drift_single_particle_exact(part, cg->length - LocalParticle_get_s(part));
 #else

{xcoll-0.4.0.dist-info → xcoll-0.5.1.dist-info}/METADATA

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