Temp branch ads

examples/colldb/lhc_run3_crystals.yaml

@@ -92,11 +92,11 @@ collimators:
     tcp.c6l7.b1:        { <<: *TCP7,   angle:   0,    material: MoGR,  assembly: 'hilumi_tcppm' }
     tcp.b6l7.b1:        { <<: *TCP7,   angle: 127.5   }
     tcsg.a6l7.b1:       { <<: *TCSG7,  angle: 141.1   }
-    tcpcv.a6l7.b1:      { <<: *CRY7,   angle:  90,    bending_radius: 85.10,  width: 5.0e-3,  height: 30.0e-3,  assembly: 'hilumi_TCPCVB1' }
+    tcpcv.a6l7.b1:      { <<: *CRY7,   angle:  90,    bending_radius: 85.10,  width: 5.0e-3,  height: 30.0e-3 } #,  assembly: 'hilumi_TCPCVB1' }
     tcsg.b5l7.b1:       { <<: *TCSG7,  angle: 143.5   }
     tcsg.a5l7.b1:       { <<: *TCSG7,  angle:  40.7   }
     tcsg.d4l7.b1:       { <<: *TCSG7,  angle:  90,    material: MoGR,  assembly: 'hilumi_tcspm' }
-    tcpch.a4l7.b1:      { <<: *CRY7,   angle:   0,    bending_radius: 61.54,  width: 2.0e-3,  height: 50.0e-3,  assembly: 'hilumi_TCPCHB1' }
+    tcpch.a4l7.b1:      { <<: *CRY7,   angle:   0,    bending_radius: 61.54,  width: 2.0e-3,  height: 50.0e-3 } #,  assembly: 'hilumi_TCPCHB1' }
     tcsg.b4l7.b1:       { <<: *TCSG7,  angle:   0,    active: false }
     tcspm.b4l7.b1:      { <<: *TCSG7,  angle:   0,    material: MoGR,  assembly: 'hilumi_tcspm' }
     tcsg.a4l7.b1:       { <<: *TCSG7,  angle: 134.6   }

examples/fluka_crystal.py

@@ -18,10 +18,10 @@
 # context = xo.ContextCupy()      # For CUDA GPUs
 # context = xo.ContextPyopencl()  # For OpenCL GPUs
 
-num_part = 100
+num_part = int(10_000)
 
 coll = xc.FlukaCrystal(length=0.002, material='si', bending_angle=149e-6,
-                       width=0.002, height=0.05, side='+', jaw=1.e-12,  # Hack to make it work for now; jaw should be almost zero
+                       width=0.002, height=0.05, side='+', jaw=1e-3,  # Hack to make it work for now; jaw should be almost zero
                          _context=context)
 
 # Connect to FLUKA
@@ -34,10 +34,9 @@
 px_init  = np.random.uniform(low=-50.e-6, high=250.e-6, size=num_part)
 y_init   = np.random.normal(loc=0., scale=1e-3, size=num_part)
 py_init  = np.random.normal(loc=0., scale=5.e-6, size=num_part)
-part = xp.build_particles(x=x_init, px=px_init, y=y_init, py=py_init, particle_ref=xc.fluka.engine.particle_ref, _context=context)
+part = xp.build_particles(x=x_init, px=px_init, y=y_init, py=py_init, particle_ref=xc.fluka.engine.particle_ref, _context=context, _capacity=xc.fluka.engine.capacity)
 part_init = part.copy()
 
-
 coll.track(part)
 
 xc.fluka.engine.stop()

examples/lhc_run3_lossmap_with_crystals_fluka.py

@@ -0,0 +1,168 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import numpy as np
+from pathlib import Path
+import time
+start_time = time.time()
+import matplotlib.pyplot as plt
+
+import xobjects as xo
+import xpart as xp
+import xtrack as xt
+import xcoll as xc
+
+
+# We do the majority of the script on the default context to be able to use prebuilt kernels
+context = xo.ContextCpu()
+
+
+# This script takes around 3 minutes on a modern CPU (90s preparation+interpolation, 90s tracking)
+beam = 1
+plane = 'H'
+tilt = 0  # rad !!
+
+num_turns = 2 # 00
+num_particles = 50 #000
+
+path_in  = xc._pkg_root.parent / 'examples'
+path_out = Path.cwd()
+
+
+# Load from json
+line = xt.Line.from_json(path_in / 'machines' / f'lhc_run3_b{beam}.json')
+
+
+# Initialise colldb
+colldb = xc.CollimatorDatabase.from_yaml(path_in / 'colldb' / f'lhc_run3_crystals.yaml',
+                                               beam=beam, ignore_crystals=False)
+
+
+# Install collimators into line
+colldb.install_fluka_collimators(line=line, verbose=True)
+
+# Aperture model check
+print('\nAperture model check after introducing collimators:')
+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 f'{beam}'=='1' else 5}{'l' if f'{beam}'=='1' else 'r'}7.b{beam}"
+
+tw = line.twiss4d()
+
+alfa_x = tw.rows[tcpc].alfx[0]
+alfa_y = tw.rows[tcpc].alfy[0]
+beta_x  = tw.rows[tcpc].betx[0]
+beta_y  = tw.rows[tcpc].bety[0]
+gap_tcpc = colldb[tcpc]['gap']
+gamma0 = line.particle_ref.gamma0
+beta0 = line.particle_ref.beta0
+
+# import pdb; pdb.set_trace()
+
+geo_emitt_x = colldb.nemitt_x/(gamma0*beta0)
+geo_emitt_y = colldb.nemitt_y/(gamma0*beta0)
+
+ch_x = -gap_tcpc * alfa_x * np.sqrt(geo_emitt_x/beta_x)
+ch_y = -gap_tcpc * alfa_y * np.sqrt(geo_emitt_y/beta_y)
+
+# Apply settings
+line[tcpc].bending_angle = ch_x # 40.e-6
+line[tcpc].width         = 0.002
+line[tcpc].height        = 0.05
+line[tcpc].length         = 0.004
+# line[tcpc].align_to_beam_divergence()
+
+# Build the tracker
+line.build_tracker()
+
+
+# Assign the optics to deduce the gap settings
+line.collimators.assign_optics()
+
+# Connect to FLUKA
+xc.fluka.engine.particle_ref = xt.Particles.reference_from_pdg_id(pdg_id='Pb208', p0c=6.8e12*82)
+xc.fluka.engine.capacity = 20*num_particles
+xc.fluka.engine.seed = 5656565
+xc.fluka.engine.start(line=line, capacity=xc.fluka.engine.capacity, cwd='run_fluka_temp', clean=False, verbose=True, return_ions=True)
+
+
+# Optimise the line
+# line.optimize_for_tracking()
+
+
+# # Generate initial pencil distribution on crystal
+# part = line[tcp].generate_pencil(num_particles)
+# Generate initial halo
+x_norm, px_norm, _, _ = xp.generate_2D_uniform_circular_sector(r_range=(5, 5.04), num_particles=num_particles)
+y_norm  = np.random.normal(scale=0.01, size=num_particles)
+py_norm = np.random.normal(scale=0.01, size=num_particles)
+part = line.build_particles(
+            x_norm=x_norm, px_norm=px_norm, y_norm=y_norm, py_norm=py_norm,
+            nemitt_x=line[tcpc].nemitt_x, nemitt_y=line[tcpc].nemitt_y,
+            at_element=tcpc, particle_ref=xc.fluka.engine.particle_ref, _context=context, _capacity=xc.fluka.engine.capacity)
+
+
+# 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!
+line.scattering.enable()
+line.track(part, num_turns=num_turns, time=True, with_progress=1)
+line.scattering.disable()
+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())
+
+# Stop the FLUKA connection (and return to the previous directory)
+xc.fluka.engine.stop(clean=False)
+
+# Save lossmap to json, which can be loaded, combined (for more statistics),
+# and plotted with the 'lossmaps' package
+line_is_reversed = True if f'{beam}' == '2' else False
+ThisLM = xc.LossMap(line, line_is_reversed=line_is_reversed, part=part)
+ThisLM.to_json(file=Path(path_out, f'lossmap_B{beam}{plane}.json'))
+
+# Save a summary of the collimator losses to a text file
+ThisLM.save_summary(file=Path(path_out, f'coll_summary_B{beam}{plane}.out'))
+print(ThisLM.summary)
+
+
+# Impacts
+# =======
+#nabs = summary.loc[summary.collname==tcpc, 'nabs'].values[0]
+#imp = colldb.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 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)
+
+#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")
+
+ThisLM.plot()
+plt.show()

xcoll/beam_elements/base.py

@@ -1502,8 +1502,8 @@ def _verify_consistency(self):
         if '_side' in self._xofields:
             assert self._side in [-1, 0, 1]
         # Crystal specific
-        if '_bending_radius' in self._xofields and '_bending_angle' in self._xofields:
-            assert isinstance(self._bending_radius, float) and not np.isclose(self._bending_radius, 0)
-            assert isinstance(self._bending_angle, float) and abs(self._bending_angle) <= np.pi/2
-            assert np.isclose(self._bending_angle, np.arcsin(self.length/self._bending_radius))
+        # if '_bending_radius' in self._xofields and '_bending_angle' in self._xofields:
+        #     assert isinstance(self._bending_radius, float) and not np.isclose(self._bending_radius, 0)
+        #     assert isinstance(self._bending_angle, float) and abs(self._bending_angle) <= np.pi/2
+        #     assert np.isclose(self._bending_angle, np.arcsin(self.length/self._bending_radius))
 

xcoll/beam_elements/elements_src/fluka_crystal.h

@@ -8,35 +8,35 @@
 
 
 /*gpufun*/
-int8_t FLukaCrystalData_get_record_impacts(FLukaCrystalData el){
-    return FLukaCrystalData_get__record_interactions(el) % 2;
+int8_t FlukaCrystalData_get_record_impacts(FlukaCrystalData el){
+    return FlukaCrystalData_get__record_interactions(el) % 2;
 }
 
 /*gpufun*/
-int8_t FLukaCrystalData_get_record_exits(FLukaCrystalData el){
-    return (FLukaCrystalData_get__record_interactions(el) >> 1) % 2;
+int8_t FlukaCrystalData_get_record_exits(FlukaCrystalData el){
+    return (FlukaCrystalData_get__record_interactions(el) >> 1) % 2;
 }
 
 /*gpufun*/
-int8_t FLukaCrystalData_get_record_scatterings(FLukaCrystalData el){
-    return (FLukaCrystalData_get__record_interactions(el) >> 2) % 2;
+int8_t FlukaCrystalData_get_record_scatterings(FlukaCrystalData el){
+    return (FlukaCrystalData_get__record_interactions(el) >> 2) % 2;
 }
 
 
 /*gpufun*/
-CrystalGeometry FLukaCrystal_init_geometry(FLukaCrystalData el, LocalParticle* part0){
+CrystalGeometry FlukaCrystal_init_geometry(FlukaCrystalData el, LocalParticle* part0){
     CrystalGeometry cg = (CrystalGeometry) malloc(sizeof(CrystalGeometry_));
-    cg->length = FLukaCrystalData_get_length(el);
-    cg->side   = FLukaCrystalData_get__side(el);
-    cg->bending_radius = FLukaCrystalData_get__bending_radius(el);
-    cg->bending_angle  = FLukaCrystalData_get__bending_angle(el);
-    cg->width  = FLukaCrystalData_get__width(el);
-    cg->height = FLukaCrystalData_get__height(el);
-    cg->jaw_U  = FLukaCrystalData_get__jaw_U(el);
-    cg->sin_z  = FLukaCrystalData_get__sin_z(el);
-    cg->cos_z  = FLukaCrystalData_get__cos_z(el);
-    cg->sin_y  = FLukaCrystalData_get__sin_y(el);
-    cg->cos_y  = FLukaCrystalData_get__cos_y(el);
+    cg->length = FlukaCrystalData_get_length(el);
+    cg->side   = FlukaCrystalData_get__side(el);
+    cg->bending_radius = FlukaCrystalData_get__bending_radius(el);
+    cg->bending_angle  = FlukaCrystalData_get__bending_angle(el);
+    cg->width  = FlukaCrystalData_get__width(el);
+    cg->height = FlukaCrystalData_get__height(el);
+    cg->jaw_U  = FlukaCrystalData_get__jaw_U(el);
+    cg->sin_z  = FlukaCrystalData_get__sin_z(el);
+    cg->cos_z  = FlukaCrystalData_get__cos_z(el);
+    cg->sin_y  = FlukaCrystalData_get__sin_y(el);
+    cg->cos_y  = FlukaCrystalData_get__cos_y(el);
     double jaw;
     if (cg->side == 1){
         jaw = cg->jaw_U;
@@ -48,14 +48,14 @@ CrystalGeometry FLukaCrystal_init_geometry(FLukaCrystalData el, LocalParticle* p
     }
     cg->segments = create_crystal(cg->bending_radius, cg->width, cg->length, jaw, cg->sin_y, cg->cos_y);
     // Impact table
-    cg->record = FLukaCrystalData_getp_internal_record(el, part0);
+    cg->record = FlukaCrystalData_getp_internal_record(el, part0);
     cg->record_index = NULL;
     cg->record_impacts = 0;
     cg->record_exits = 0;
     if (cg->record){
         cg->record_index = InteractionRecordData_getp__index(cg->record);
-        cg->record_impacts = FLukaCrystalData_get_record_impacts(el);
-        cg->record_exits = FLukaCrystalData_get_record_exits(el);
+        cg->record_impacts = FlukaCrystalData_get_record_impacts(el);
+        cg->record_exits = FlukaCrystalData_get_record_exits(el);
     }
     // Not needed, set to zero
     cg->miscut_angle = 0;
@@ -66,21 +66,21 @@ CrystalGeometry FLukaCrystal_init_geometry(FLukaCrystalData el, LocalParticle* p
 }
 
 /*gpufun*/
-void FLukaCrystal_free(CrystalGeometry restrict cg){
+void FlukaCrystal_free(CrystalGeometry restrict cg){
     destroy_crystal(cg->segments);
     free(cg);
 }
 
 
 /*gpufun*/
-void FLukaCrystal_track_local_particle(FLukaCrystalData el, LocalParticle* part0){
-    int8_t active = FLukaCrystalData_get_active(el);
-    active       *= FLukaCrystalData_get__tracking(el);
+void FlukaCrystal_track_local_particle(FlukaCrystalData el, LocalParticle* part0){
+    int8_t active = FlukaCrystalData_get_active(el);
+    active       *= FlukaCrystalData_get__tracking(el);
 
     // Get geometry
     CrystalGeometry cg;
     if (active){
-        cg = FLukaCrystal_init_geometry(el, part0);
+        cg = FlukaCrystal_init_geometry(el, part0);
         if (cg->width==0 || cg->height==0 || cg->bending_radius==0){
             kill_all_particles(part0, XC_ERR_INVALID_XOFIELD);
         }
@@ -159,7 +159,7 @@ void FLukaCrystal_track_local_particle(FLukaCrystalData el, LocalParticle* part0
         }
     //end_per_particle_block
     if (active){
-        FLukaCrystal_free(cg);
+        FlukaCrystal_free(cg);
     }
 }
 

xcoll/beam_elements/fluka.py

@@ -152,9 +152,13 @@ def assembly(self, val):
 
     def track(self, part):
         if track_pre(self, part):
-            super().track(part)
+            if self.material != "vacuum":
+                super().track(part)
+            else:
+                part.state[part.state == 1] = 334
             track_core(self, part)
-            track_post(self, part)
+            if self.material != "vacuum":
+                track_post(self, part)
 
     def __setattr__(self, name, value):
         import xcoll as xc
@@ -336,7 +340,10 @@ def assembly(self, assembly):
         FlukaCollimator.assembly.fset(self, assembly)
 
     def track(self, part):
-        FlukaCollimator.track(self, part)
+        if track_pre(self, part):
+            part.state[part.state == 1] = 334
+            track_core(self, part)
+            track_post(self, part)
 
     def __setattr__(self, name, value):
         import xcoll as xc

xcoll/colldb.py

@@ -14,7 +14,7 @@
 import xtrack as xt
 
 from .beam_elements import (BlackAbsorber, BlackCrystal, EverestCollimator, EverestCrystal,
-                            FlukaCollimator, Geant4Collimator, collimator_classes)
+                            FlukaCollimator, FlukaCrystal, Geant4Collimator, collimator_classes)
 
 
 def _initialise_None(dct):