1+ # copyright ############################### #
2+ # This file is part of the Xcoll package. #
3+ # Copyright (c) CERN, 2024. #
4+ # ######################################### #
5+
6+ import numpy as np
7+ import matplotlib as mpl
8+ import matplotlib .pyplot as plt
9+
10+ import xobjects as xo
11+ import xtrack as xt
12+ import xpart as xp
13+ import xcoll as xc
14+
15+ # ============================================
16+ # With collimators
17+ # ============================================
18+
19+ # Get line and collimators
20+ line = xt .Line .from_json (xc ._pkg_root / '..' / 'examples' / 'machines' / 'lhc_run3_b1.json' )
21+
22+ colldb = xc .CollimatorDatabase .from_yaml (xc ._pkg_root / '..' / 'examples' / 'colldb' / 'lhc_run3.yaml' , beam = 1 )
23+ colldb .install_everest_collimators (verbose = True , line = line )
24+ df_with_coll = line .check_aperture ()
25+ assert not np .any (df_with_coll .has_aperture_problem )
26+
27+ # Start interaction record
28+ impacts = xc .InteractionRecord .start (line = line )
29+
30+ # Build tracker, assign optics and generate particles
31+ line .build_tracker ()
32+ xc .assign_optics_to_collimators (line = line )
33+ part = xc .generate_pencil_on_collimator (line , 'tcp.d6l7.b1' , 50000 )
34+
35+ # This is not needed, but is done here so that we can track with 12 treads.
36+ line .discard_tracker ()
37+ line .build_tracker (_context = xo .ContextCpu (omp_num_threads = 12 ))
38+
39+ # Track
40+ xc .enable_scattering (line )
41+ line .track (part , num_turns = 20 , time = True , with_progress = 1 )
42+ xc .disable_scattering (line )
43+ line .discard_tracker ()
44+
45+ df = impacts .to_pandas ()
46+ df [df .interaction_type == 'Enter Jaw L' ]
47+ df .to_csv ('impacts.csv' , index = False )
48+
49+ # ============================================
50+ # With crystal
51+ # ============================================
52+ coll = xc .EverestCrystal (length = 0.002 , material = xc .materials .SiliconCrystal , bending_angle = 149e-6 ,
53+ width = 0.002 , height = 0.05 , side = '+' , lattice = 'strip' , jaw = 0.001 )
54+
55+ num_part = int (50000 )
56+ x_init = np .random .normal (loc = 1.5e-3 , scale = 75.e-6 , size = num_part )
57+ px_init = np .random .uniform (low = - 50.e-6 , high = 250.e-6 , size = num_part )
58+ y_init = np .random .normal (loc = 0. , scale = 1e-3 , size = num_part )
59+ py_init = np .random .normal (loc = 0. , scale = 5.e-6 , size = num_part )
60+ part = xp .Particles (x = x_init , px = px_init , y = y_init , py = py_init , delta = 0 , p0c = 4e11 )
61+
62+ io_buffer = xt .new_io_buffer (capacity = int (2e7 )) # 4-5 GB of memory
63+ coll .record_scatterings = True
64+ impacts_crystal = xt .start_internal_logging (elements = [coll ], io_buffer = io_buffer , capacity = io_buffer .capacity )
65+
66+ coll .track (part )
67+ part .sort (interleave_lost_particles = True )
68+
69+ impacts_crystal .to_pandas ()
70+ df_crystal = impacts_crystal .interactions_per_collimator ()
71+ df_crystal .to_csv ('interactions_per_crystal.csv' , index = False )
0 commit comments