cleanup a bit and changed everest to blackabsorber

Author: andredonadon

xcoll/scattering_routines/fluka/track.py

@@ -59,82 +59,51 @@ def track(coll, particles):
 
     start = time.time()
 
-    coll_everest = xc.EverestCollimator(length=coll.length, material=xc.materials.MolybdenumGraphite)
-    coll_everest.jaw = coll.jaw
-
-    part_everest = particles.copy()
-    coll_everest.track(part_everest)
-
-    mask_hit_base = (
-        (abs(particles.px - part_everest.px) > 1e-12) |
-        (abs(particles.py - part_everest.py) > 1e-12) |
-        (particles.pdg_id == -999999999)
-    )
+    # Checking particles that hit the collimator
+    coll_blackabsorber = xc.BlackAbsorber(length=coll.length+coll.length_front+coll.length_back)
+    part_blackabsorber = particles.copy()
+    coll_blackabsorber.track(part_blackabsorber)
 
     end = time.time()
-    print(f"Tracking Everest collimator took {end - start:.4f} seconds")
+    print(f"Tracking black absorber collimator took {end - start:.4f} seconds")
 
+    mask_capacity = particles.pdg_id != -999999999
 
-    start = time.time()
+    # Mask for particles that hit the collimator
+    mask_hit = ((part_blackabsorber.state != 1) & (particles.state == 1)) | (particles.pdg_id == -999999999)
+
+    # Mask for particles already lost
     mask_lost = (particles.state != 1) & (particles.pdg_id != -999999999)
 
     # Final hit mask includes both
-    mask_hit = mask_hit_base # | mask_lost
-
-
-#     mask_hit = ((abs(particles.px - part_everest.px) > 1e-12) | (abs(particles.py - part_everest.py) >1e-12) | (particles.pdg_id == -999999999))
-# 
-    # mask_lost = ((particles.state != 1) & (particles.pdg_id != -999999999))
-    mask_miss_no_lost = ~((abs(particles.px - part_everest.px) > 1e-12) |  (abs(particles.py - part_everest.py) > 1e-12)) & (particles.pdg_id != -999999999)
+    mask_miss_no_lost = ~mask_hit
     mask_miss = mask_miss_no_lost | mask_lost
 
 
     hit_id = particles.particle_id[mask_hit]
-    miss_id = particles.particle_id[mask_miss]
     is_hit = np.isin(particles.particle_id, hit_id)
+    miss_id = particles.particle_id[mask_miss]
     is_miss = np.isin(particles.particle_id, miss_id)
 
-    if mask_lost.any():
-        particles_lost = particles.filter(mask_lost)
-    else:
-        particles_lost = None 
-
     particles_hit = 0
     particles_miss = 0
     particles_hit = particles.filter(is_hit)
-    end = time.time()
-    print(f"Tracking Everest collimator took {end - start:.4f} seconds")
 
     print(f"Tracking {particles._num_active_particles} particles (FLUKA)...     ", end='')
     print(f"Hit: {particles_hit._num_active_particles}")
     print(f"Collimator: {coll.name}")
-
-
-    # if ((particles_hit._num_active_particles == 0) & (particles_miss == 0)):
-    #     # track_core(coll, particles)
-    #     _drift(coll, particles, coll.length)
-    #     return
-
 
     # check if ~is_hit has any true value
     if np.any(~is_hit):
-        particles_miss = particles.filter(~is_hit)
+        miss_id = particles.particle_id[mask_miss]
+        is_miss = np.isin(particles.particle_id, miss_id)
+        particles_miss = particles.filter(is_miss)
     elif (particles_hit._num_active_particles == 0) and (particles._num_active_particles > 0):
-        # import pdb; pdb.set_trace()
         print("Could not find any particles that hit the collimator, but some particles are still active. ")
         print("This is likely due to a bug in the collimator tracking code. Please report this issue.")
 
-    # particles_miss = particles.filter(~is_hit) if np.any(is_miss) else 0
-#     particles_miss = particles.filter(is_miss) if np.any(is_miss) else 0
-#     particles_miss = particles_miss.filter(particles_miss.state == 1) if np.any(is_miss) else 0
-
-    # missing_capacity = 0
-    # if particles_miss:
-    #     miss_part = particles_miss._num_active_particles
-    #     more_capacity = particles.filter([False] * (particles._capacity - miss_part) + [True] * miss_part)
     start = time.time()
 
-
     if particles_miss:
         _drift(coll, particles_miss, coll.length)
     if particles_hit._num_active_particles > 0: