Merge pull request #66 from xsuite/release/v0.3.5

Release 0.3.5

Author: freddieknets

examples/colldb/lhc_run3.yaml

@@ -42,7 +42,7 @@ families:
   # Dump protection
   - &TCDQ    { <<: *ALL,  gap: 7.3,   stage: tertiary,   material: C,     length: 3,     angle: 0,  side: left }
   - &TCSP    { <<: *ALL,  gap: 7.3,   stage: secondary,  material: C,     length: 1,     angle: 0 }
-  # Physics background minimalisation
+  # Triplet protection
   - &TCT15   { <<: *ALL,  gap: 8.5,   stage: tertiary,   material: Iner,  length: 1,     parking: 0.020 }
   - &TCT2    { <<: *ALL,  gap: 37,    stage: tertiary,   material: Iner,  length: 1,     }
   - &TCT8    { <<: *ALL,  gap: 11.5,  stage: tertiary,   material: Iner,  length: 1,     }

examples/colldb/lhc_run3_crystals.yaml

@@ -42,7 +42,7 @@ families:
   # Dump protection
   - &TCDQ    { <<: *ALL,  gap: 7.3,   stage: tertiary,   material: C,     length: 3,     angle: 0,  side: left }
   - &TCSP    { <<: *ALL,  gap: 7.3,   stage: secondary,  material: C,     length: 1,     angle: 0 }
-  # Physics background minimalisation
+  # Triplet protection
   - &TCT15   { <<: *ALL,  gap: 9.5,   stage: tertiary,   material: Iner,  length: 1,     parking: 0.020 }
   - &TCT2    { <<: *ALL,  gap: 37,    stage: tertiary,   material: Iner,  length: 1,     }
   - &TCT8    { <<: *ALL,  gap: 11.5,  stage: tertiary,   material: Iner,  length: 1,     }

examples/lhc_run3_lossmap.py

@@ -7,38 +7,34 @@
 from pathlib import Path
 import time
 start_time = time.time()
-import sys, os, contextlib
 
 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
-beam          = '1'
-plane         = 'H'
+# We do the majority of the script on the default context to be able to use prebuilt kernels
+context = xo.ContextCpu()
 
-num_turns     = 200
+
+# This script takes around 3 minutes on a modern CPU (90s preparation+interpolation, 90s tracking)
+beam = 1
+plane = 'H'
+
+num_turns = 200
 num_particles = 50000
-engine        = 'everest'
 
 path_in  = xc._pkg_root.parent / 'examples'
 path_out = Path.cwd()
 
 
 # Load from json
-with open(os.devnull, 'w') as fid:
-    with contextlib.redirect_stdout(fid):
-        line = xt.Line.from_json(path_in / 'machines' / f'lhc_run3_b{beam}.json')
+line = xt.Line.from_json(path_in / 'machines' / f'lhc_run3_b{beam}.json')
 
 
 # Aperture model check
 print('\nAperture model check on imported model:')
-with open(os.devnull, 'w') as fid:
-    with contextlib.redirect_stdout(fid):
-        df_imported = line.check_aperture()
+df_imported = line.check_aperture()
 assert not np.any(df_imported.has_aperture_problem)
 
 
@@ -47,17 +43,12 @@
 
 
 # Install collimators into line
-if engine == 'everest':
-    coll_manager.install_everest_collimators(verbose=True)
-else:
-    raise ValueError(f"Unknown scattering engine {engine}!")
+coll_manager.install_everest_collimators(verbose=True)
 
 
 # Aperture model check
 print('\nAperture model check after introducing collimators:')
-with open(os.devnull, 'w') as fid:
-    with contextlib.redirect_stdout(fid):
-        df_with_coll = line.check_aperture()
+df_with_coll = line.check_aperture()
 assert not np.any(df_with_coll.has_aperture_problem)
 
 
@@ -70,33 +61,41 @@
 coll_manager.set_openings()
 
 
+# Optimise the line
+line.optimize_for_tracking()
+
+
 # Generate initial pencil distribution on horizontal collimator
-tcp  = f"tcp.{'c' if plane=='H' else 'd'}6{'l' if beam=='1' else 'r'}7.b{beam}"
+tcp  = f"tcp.{'c' if plane=='H' else 'd'}6{'l' if f'{beam}'=='1' else 'r'}7.b{beam}"
 part = coll_manager.generate_pencil_on_collimator(tcp, num_particles=num_particles)
 
 
-# Optimise the line
-line.optimize_for_tracking()
-idx = line.element_names.index(tcp)
-part.at_element = idx
-part.start_tracking_at_element = idx
+
+# Move the line to an OpenMP context to be able to use all cores
+line.discard_tracker()
+line.build_tracker(_context=xo.ContextCpu(omp_num_threads='auto'))
+# Should move iobuffer as well in case of impacts
 
 
-# Track
+# Track!
 coll_manager.enable_scattering()
-line.track(part, num_turns=num_turns, time=True)
+line.track(part, num_turns=num_turns, time=True, with_progress=1)
 coll_manager.disable_scattering()
 print(f"Done tracking in {line.time_last_track:.1f}s.")
 
 
+# Move the line back to the default context to be able to use all prebuilt kernels for the aperture interpolation
+line.discard_tracker()
+line.build_tracker(_context=xo.ContextCpu())
+
+
 # Save lossmap to json, which can be loaded, combined (for more statistics),
 # and plotted with the 'lossmaps' package
-_ = coll_manager.lossmap(part, file=Path(path_out,f'lossmap_B{beam}{plane}.json'))
+coll_manager.lossmap(part, file=Path(path_out,f'lossmap_B{beam}{plane}.json'))
 
 
 # Save a summary of the collimator losses to a text file
 summary = coll_manager.summary(part, file=Path(path_out,f'coll_summary_B{beam}{plane}.out'))
 print(summary)
 print(f"Total calculation time {time.time()-start_time}s")
 
-exit()

examples/lhc_run3_lossmap_with_crystals.py

@@ -3,74 +3,60 @@
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
-import json
 import numpy as np
 from pathlib import Path
 import time
 start_time = time.time()
-import sys, os, contextlib
 
+import xobjects as xo
 import xpart as xp
 import xtrack as xt
 import xcoll as xc
 
-if len(sys.argv) < 4:
-    raise ValueError("Need three arguments: beam, plane, and tilt.")
-beam  = int(sys.argv[1])
-plane = str(sys.argv[2])
-tilt  = int(sys.argv[3])*1.e-6
 
+# We do the majority of the script on the default context to be able to use prebuilt kernels
+context = xo.ContextCpu()
 
-# On a modern CPU, we get ~5000 particle*turns/s
-# So this script should take a bit less than 40 minutes
-# beam          =  1
-# plane         = 'H'
-# tilt          = 0         # rad !!
 
-num_turns     = 200
-num_particles = 5000
-engine        = 'everest'
+# This script takes around 3 minutes on a modern CPU (90s preparation+interpolation, 90s tracking)
+beam = 1
+plane = 'H'
+tilt = 0  # rad !!
+
+num_turns = 200
+num_particles = 50000
 
 path_in  = xc._pkg_root.parent / 'examples'
 path_out = Path.cwd()
 
 
 # Load from json
-with open(os.devnull, 'w') as fid:
-    with contextlib.redirect_stdout(fid):
-        line = xt.Line.from_json(path_in / 'machines' / f'lhc_run3_b{beam}.json')
+line = xt.Line.from_json(path_in / 'machines' / f'lhc_run3_b{beam}.json')
 
 
 # Aperture model check
 print('\nAperture model check on imported model:')
-with open(os.devnull, 'w') as fid:
-    with contextlib.redirect_stdout(fid):
-        df_imported = line.check_aperture()
+df_imported = line.check_aperture()
 assert not np.any(df_imported.has_aperture_problem)
 
 
 # Initialise collmanager
 coll_manager = xc.CollimatorManager.from_yaml(path_in / 'colldb' / f'lhc_run3_crystals.yaml', line=line,
-                                              beam=beam, ignore_crystals=False, record_impacts=True)
+                                              beam=beam, ignore_crystals=False, _context=context)
 
 
 # Install collimators into line
-if engine == 'everest':
-    coll_manager.install_everest_collimators(verbose=True)
-else:
-    raise ValueError(f"Unknown scattering engine {engine}!")
+coll_manager.install_everest_collimators(verbose=True)
 
 
 # Aperture model check
 print('\nAperture model check after introducing collimators:')
-with open(os.devnull, 'w') as fid:
-    with contextlib.redirect_stdout(fid):
-        df_with_coll = line.check_aperture()
+df_with_coll = line.check_aperture()
 assert not np.any(df_with_coll.has_aperture_problem)
 
 
 # Primary crystal
-tcpc = f"tcpc{plane.lower()}.a{6 if plane=='V' else 4 if beam==1 else 5}{'l' if beam==1 else 'r'}7.b{beam}"
+tcpc = f"tcpc{plane.lower()}.a{6 if plane=='V' else 4 if f'{beam}'=='1' else 5}{'l' if f'{beam}'=='1' else 'r'}7.b{beam}"
 
 
 # Build the tracker
@@ -84,6 +70,10 @@
 coll_manager.set_openings()
 
 
+# Optimise the line
+line.optimize_for_tracking()
+
+
 # # Generate initial pencil distribution on crystal
 # part = coll_manager.generate_pencil_on_collimator(tcpc, num_particles=num_particles)
 # Generate initial halo
@@ -97,12 +87,6 @@
             _buffer=coll_manager._part_buffer
 )
 
-# Optimise the line
-line.optimize_for_tracking()
-idx = line.element_names.index(tcpc)
-part.at_element = idx
-part.start_tracking_at_element = idx
-
 
 # Apply settings
 line[tcpc].tilt_L        = tilt
@@ -111,16 +95,27 @@
 line[tcpc].ydim          = 0.05
 
 
-# Track
+# Move the line to an OpenMP context to be able to use all cores
+line.discard_tracker()
+line.build_tracker(_context=xo.ContextCpu(omp_num_threads='auto'))
+# Should move iobuffer as well in case of impacts
+
+
+# Track!
 coll_manager.enable_scattering()
-line.track(part, num_turns=num_turns, time=True)
+line.track(part, num_turns=num_turns, time=True, with_progress=1)
 coll_manager.disable_scattering()
 print(f"Done tracking in {line.time_last_track:.1f}s.")
 
 
+# Move the line back to the default context to be able to use all prebuilt kernels for the aperture interpolation
+line.discard_tracker()
+line.build_tracker(_context=xo.ContextCpu())
+
+
 # Save lossmap to json, which can be loaded, combined (for more statistics),
 # and plotted with the 'lossmaps' package
-# _ = coll_manager.lossmap(part, file=Path(path_out,f'lossmap_B{beam}{plane}.json'))
+coll_manager.lossmap(part, file=Path(path_out,f'lossmap_B{beam}{plane}.json'))
 
 
 # Save a summary of the collimator losses to a text file
@@ -129,28 +124,28 @@
 print(summary)
 
 
-# Impacts
-# =======
-imp = coll_manager.impacts.to_pandas()
-outfile = Path(path_out, f'cry_{round(tilt*1e6)}urad_impacts_B{beam}{plane}.json')
-imp.to_json(outfile)
+# # Impacts
+# # =======
+# imp = coll_manager.impacts.to_pandas()
+# outfile = Path(path_out, f'cry_{round(tilt*1e6)}urad_impacts_B{beam}{plane}.json')
+# imp.to_json(outfile)
+
+# # Only keep the first impact
+# imp.drop_duplicates(subset=['parent_id'], inplace=True)
+# assert np.unique(imp.interaction_type) == ['Enter Jaw']
 
-# Only keep the first impact
-imp.drop_duplicates(subset=['parent_id'], inplace=True)
-assert np.unique(imp.interaction_type) == ['Enter Jaw']
+# # Only keep those first impacts that are on the crystal
+# first_impacts = imp[imp.collimator==tcpc].parent_id
+# num_first_impacts = len(first_impacts)
+# assert num_first_impacts==len(np.unique(first_impacts))
 
-# Only keep those first impacts that are on the crystal
-first_impacts = imp[imp.collimator==tcpc].parent_id
-num_first_impacts = len(first_impacts)
-assert num_first_impacts==len(np.unique(first_impacts))
+# ineff = nabs/num_first_impacts
+# outfile = Path(path_out, f'cry_{round(tilt*1e6)}urad_ineff_B{beam}{plane}.json')
+# with outfile.open('w') as fid:
+#     json.dump({'first_impacts': num_first_impacts, 'nabs': nabs, 'ineff': ineff}, fid)
 
-ineff = nabs/num_first_impacts
-outfile = Path(path_out, f'cry_{round(tilt*1e6)}urad_ineff_B{beam}{plane}.json')
-with outfile.open('w') as fid:
-    json.dump({'first_impacts': num_first_impacts, 'nabs': nabs, 'ineff': ineff}, fid)
+# print(f"Out of {num_first_impacts} particles hitting the crystal {tcpc} (angle {round(tilt*1e6)}urad) first, {round(nabs)} are absorbed in the crystal (for an inefficiency of {ineff:5}.")
+# print(f"Critical angle is {round(line[tcpc].critical_angle*1.e6, 1)}urad.")
 
-print(f"Out of {num_first_impacts} particles hitting the crystal {tcpc} (angle {round(tilt*1e6)}urad) first, {round(nabs)} are absorbed in the crystal (for an inefficiency of {ineff:5}.")
-print(f"Critical angle is {round(line[tcpc].critical_angle*1.e6, 1)}urad.")
 print(f"Total calculation time {time.time()-start_time}s")
 
-exit()