Update yields and clustering

fuse/plugins/micro_physics/lineage_cluster.py

@@ -1,7 +1,6 @@
 import numpy as np
 import strax
 import straxen
-from numba import njit
 
 import re
 import periodictable as pt
@@ -48,9 +47,18 @@ class LineageClustering(FuseBasePlugin):
 
     # Config options
     gamma_distance_threshold = straxen.URLConfig(
-        default=0.9,
+        default=0.0,
         type=(int, float),
-        help="Distance threshold to break lineage for gamma rays [cm]. Default from NEST code",
+        help="Distance threshold to break lineage for gamma rays [cm]. \
+        Do not break if distance is smaller than threshold. \
+        Default at 0 means we always break the lineage.",
+    )
+
+    brem_distance_threshold = straxen.URLConfig(
+        default=0.1,
+        type=(int, float),
+        help="Distance threshold to break lineage for bremsstrahlung [cm]. \
+        Do not break if distance is smaller than threshold.",
     )
 
     time_threshold = straxen.URLConfig(
@@ -59,6 +67,19 @@ class LineageClustering(FuseBasePlugin):
         help="Time threshold to break the lineage [ns]",
     )
 
+    classify_ic_as_gamma = straxen.URLConfig(
+        default=True,
+        type=bool,
+        help="Classify internal conversion electrons as gamma particles",
+    )
+
+    classify_phot_as_beta = straxen.URLConfig(
+        default=True,
+        type=bool,
+        help="Classify photoabsorption electrons as beta particles \
+        (if False, classify as gamma particles)",
+    )
+
     def compute(self, geant4_interactions):
         """
         Args:
@@ -170,14 +191,21 @@ def build_lineage_for_event(self, event):
                         parent,
                         parent_lineage,
                         self.gamma_distance_threshold,
+                        self.brem_distance_threshold,
                         self.time_threshold,
                     )
 
                     if broken_lineage:
                         # The lineage is broken. We can start a new one!
                         running_lineage_index += 1
+
                         tmp_result = start_new_lineage(
-                            particle, tmp_result, i, running_lineage_index
+                            particle,
+                            tmp_result,
+                            i,
+                            running_lineage_index,
+                            self.classify_ic_as_gamma,
+                            self.classify_phot_as_beta,
                         )
 
                     else:
@@ -187,7 +215,15 @@ def build_lineage_for_event(self, event):
                 else:
                     # Particle without parent. Start a new lineage
                     running_lineage_index += 1
-                    tmp_result = start_new_lineage(particle, tmp_result, i, running_lineage_index)
+
+                    tmp_result = start_new_lineage(
+                        particle,
+                        tmp_result,
+                        i,
+                        running_lineage_index,
+                        self.classify_ic_as_gamma,
+                        self.classify_phot_as_beta,
+                    )
 
             else:
                 # We have seen this particle before. Now evaluate if we have to break the lineage
@@ -202,13 +238,20 @@ def build_lineage_for_event(self, event):
                         last_particle_interaction,
                         last_particle_lineage,
                         self.gamma_distance_threshold,
+                        self.brem_distance_threshold,
                         self.time_threshold,
                     )
                     if broken_lineage:
                         # New lineage!
                         running_lineage_index += 1
+
                         tmp_result = start_new_lineage(
-                            particle, tmp_result, i, running_lineage_index
+                            particle,
+                            tmp_result,
+                            i,
+                            running_lineage_index,
+                            self.classify_ic_as_gamma,
+                            self.classify_phot_as_beta,
                         )
 
                     else:
@@ -295,10 +338,15 @@ def num_there(s):
     return any(i.isdigit() for i in s)
 
 
-def classify_lineage(particle_interaction):
+def classify_lineage(particle_interaction, classify_ic_as_gamma, classify_phot_as_beta):
     """Function to classify a new lineage based on the particle and its parent
     information."""
 
+    # Excited states of nuclei, decaying electromagnetically
+    # this will become the lineage of internal conversion electrons
+    if "[" in particle_interaction["type"]:
+        return NEST_GAMMA if classify_ic_as_gamma else NEST_BETA
+
     # NR interactions
     if (particle_interaction["parenttype"] == "neutron") & (
         num_there(particle_interaction["type"])
@@ -317,7 +365,7 @@ def classify_lineage(particle_interaction):
         elif particle_interaction["creaproc"] == "conv":
             return NEST_BETA
         elif particle_interaction["creaproc"] == "phot":
-            return NEST_GAMMA
+            return NEST_BETA if classify_phot_as_beta else NEST_GAMMA
         else:
             # This case should not happen or? Classify it as nontype
             return NEST_NONE
@@ -333,15 +381,8 @@ def classify_lineage(particle_interaction):
         elif particle_interaction["edproc"] == "conv":
             return NEST_BETA
         elif particle_interaction["edproc"] == "phot":
-            # Not sure about this, but was looking for a way to give beta yields
-            # to gammas that are coming directly from a radioactive decay
-            if particle_interaction["creaproc"] == "RadioactiveDecayBase":
-                return NEST_BETA
-            # Need this case for custom geant4 inputs...
-            elif particle_interaction["creaproc"] == "Null":
-                return NEST_BETA
-            else:
-                return NEST_GAMMA
+            # This is gamma photoabsorption. Return gamma
+            return NEST_BETA if classify_phot_as_beta else NEST_GAMMA
         else:
             # could be rayleigh scattering or something else. Classify it as gamma...
             return NEST_BETA
@@ -351,6 +392,11 @@ def classify_lineage(particle_interaction):
         particle_interaction["parenttype"] == "none"
     ):
 
+        # If [ in type, it is a nucleus excitation
+        # we give it a beta for the possible conversion electrons
+        if "[" in particle_interaction["type"]:
+            return NEST_BETA
+
         # Alpha particles
         if particle_interaction["type"] == "alpha":
             return NEST_ALPHA
@@ -369,28 +415,37 @@ def classify_lineage(particle_interaction):
         return NEST_NONE
 
 
-@njit()
 def is_lineage_broken(
     particle,
     parent,
     parent_lineage,
     gamma_distance_threshold,
+    brem_distance_threshold,
     time_threshold,
 ):
     """Function to check if the lineage is broken."""
 
-    # In the nest code: Lineage is always broken if the parent is a ion
-    # But if it's an alpha particle, we want to keep the lineage
-    break_for_ion = parent_lineage["lineage_type"] == 6
-    break_for_ion &= parent["type"] != "alpha"
-    break_for_ion &= particle["creaproc"] != "eIoni"
+    if (
+        particle["creaproc"] == "RadioactiveDecayBase"
+        and particle["edproc"] == "RadioactiveDecayBase"
+    ):
+        # second step of a decay. We want to split the lineage
+        return True
 
-    if break_for_ion:
+    # In the nest code: Lineage is always broken if the parent is a ion
+    # this breaks the lineage for all ions, also for alpha decays (we need it)
+    # but if it's via an excited nuclear state, we want to keep the lineage
+    if (num_there(parent["type"])) and ("[" not in parent["type"]):
         return True
 
     # For gamma rays, check the distance between the parent and the particle
     if particle["type"] == "gamma":
 
+        if particle["creaproc"] == "phot" and particle["edproc"] == "phot":
+            # We do not want to split a photo absorption into two clusters
+            # The second photo absorption (that we see) could be x rays
+            return False
+
         # Break the lineage for these transportation gammas
         # Transportations is a special case. They are not real gammas.
         # They are just used to transport the energy
@@ -400,9 +455,14 @@ def is_lineage_broken(
 
         particle_position = np.array([particle["x"], particle["y"], particle["z"]])
         parent_position = np.array([parent["x"], parent["y"], parent["z"]])
-
         distance = np.sqrt(np.sum((parent_position - particle_position) ** 2, axis=0))
 
+        if particle["creaproc"] == "eBrem":
+            # we do not want to split a bremsstrahlung into two clusters
+            # if the distance is really small, it is most likely the same interaction
+            if distance < brem_distance_threshold:
+                return False
+
         if distance > gamma_distance_threshold:
             return True
 
@@ -497,9 +557,13 @@ def propagate_mega_type(mega_type):
     return main_cluster_types
 
 
-def start_new_lineage(particle, tmp_result, i, running_lineage_index):
+def start_new_lineage(
+    particle, tmp_result, i, running_lineage_index, classify_ic_as_gamma, classify_phot_as_beta
+):
 
-    lineage_class, lineage_A, lineage_Z = classify_lineage(particle)
+    lineage_class, lineage_A, lineage_Z = classify_lineage(
+        particle, classify_ic_as_gamma, classify_phot_as_beta
+    )
     tmp_result[i]["lineage_index"] = running_lineage_index
     tmp_result[i]["lineage_type"] = lineage_class
     tmp_result[i]["lineage_A"] = lineage_A