NV hitlet simulator plugin

fuse/common.py

@@ -26,44 +26,84 @@
 )
 
 
-@numba.njit()
-def dynamic_chunking(data, scale, n_min):
-    idx_sort = stable_argsort(data)
-    idx_undo_sort = stable_argsort(idx_sort)
+@numba.njit(cache=True)
+def dynamic_chunking(
+    time_gaps,
+    file_size_limit,
+    min_gap_length,
+    n_bytes_per_interaction,
+):
+
+    data_size_mb = 0
+    clusters_index = []
+
+    running_index = 0
+
+    for g in time_gaps:
+
+        data_size_mb += n_bytes_per_interaction / 1e6
+
+        if data_size_mb < file_size_limit:
+            clusters_index.append(running_index)
+            continue
+
+        if g >= min_gap_length:
+            running_index += 1
+            data_size_mb = 0
+            clusters_index.append(running_index)
+        else:
+            clusters_index.append(running_index)
+
+    return np.array(clusters_index)
+
+
+@numba.njit(cache=True)
+def dynamic_chunking_two_outputs(
+    combined_time_gaps,
+    combined_types,
+    file_size_limit,
+    min_gap_length,
+    n_bytes_per_interaction_TPC,
+    n_bytes_per_interaction_NV,
+):
+    """Function to split the TPC and NV data into chunks based on the time gaps
+    between the interactions."""
+
+    data_size_mb_tpc = 0
+    data_size_mb_nv = 0
 
-    data_sorted = data[idx_sort]
+    combined_cluster_index = []
+    running_index = 0
 
-    diff = data_sorted[1:] - data_sorted[:-1]
+    for i, (interaction_type, delta_t) in enumerate(zip(combined_types, combined_time_gaps)):
 
-    clusters = [0]
-    c = 0
-    n_cluster = 0
-    for value in diff:
-        if value <= scale:
-            clusters.append(c)
-            n_cluster += 1
-        elif n_cluster + 1 < n_min:
-            clusters.append(c)
-            n_cluster += 1
-        elif value > scale:
-            c = c + 1
-            clusters.append(c)
-            n_cluster = 0
+        if interaction_type == 0:
+            # TPC interaction
+            data_size_mb_tpc += n_bytes_per_interaction_TPC / 1e6
+        elif interaction_type == 1:
+            # NV interaction
+            data_size_mb_nv += n_bytes_per_interaction_NV / 1e6
 
-    clusters = np.array(clusters)
-    clusters_undo_sort = clusters[idx_undo_sort]
+        if (data_size_mb_tpc < file_size_limit) & (data_size_mb_nv < file_size_limit):
+            combined_cluster_index.append(running_index)
+            continue
 
-    return clusters_undo_sort
+        if delta_t >= min_gap_length:
+            running_index += 1
+            data_size_mb_tpc = data_size_mb_nv = 0
+            combined_cluster_index.append(running_index)
+        else:
+            combined_cluster_index.append(running_index)
+    return np.array(combined_cluster_index)
 
 
 def full_array_to_numpy(array, dtype):
-    len_output = len(awkward_to_flat_numpy(array["x"]))
+    len_output = len(awkward_to_flat_numpy(array[array.fields[0]]))
 
     numpy_data = np.zeros(len_output, dtype=dtype)
 
     for field in array.fields:
         numpy_data[field] = awkward_to_flat_numpy(array[field])
-
     return numpy_data
 
 

fuse/context_utils.py

@@ -3,6 +3,7 @@
 import straxen
 from straxen import URLConfig
 from copy import deepcopy
+from scipy import interpolate
 import logging
 
 logging.basicConfig(handlers=[logging.StreamHandler()])
@@ -200,3 +201,36 @@ def apply_mc_overrides(context, config_file):
             log.debug(f"[mc_overrides] Set '{key}' to '{value}'")
     except Exception as e:
         raise ValueError(f"[mc_overrides] Failed to apply overrides from {config_file}: {e}") from e
+
+
+@URLConfig.register("nveto_pmt_qe")
+def nveto_pmt_qe_dict(data):
+    """Get NV PMT quantum efficiecny values and interpolate."""
+
+    QE_array_n = []
+    # nVeto
+    for i in np.arange(2000, 2120):
+        QE_array_n.append(
+            interpolate.interp1d(
+                data["nv_pmt_qe_wavelength"],
+                data["nv_pmt_qe"][str(i)],
+                bounds_error=False,
+                fill_value=0,
+            )
+        )
+
+    pmt_id = list(np.arange(2000, 2120))
+    QE_array = QE_array_n
+
+    return QE_array
+    # pd_dict = {"pmt_id": pmt_id, "QE": QE_array}
+
+    # return pd_dict
+
+
+@URLConfig.register("nveto_spe")
+def nveto_spe_dict(data_spe):
+    """Get dictionary with NV SPE parameters."""
+
+    data_dict = {entry["pmtID"]: entry for entry in data_spe}
+    return data_dict

fuse/dtypes.py

@@ -77,3 +77,26 @@
     (("ID of the cluster creating the photon", "cluster_id"), np.int32),
     (("Type of the photon. S1 (1), S2 (2) or PMT AP (0)", "photon_type"), np.int8),
 ]
+
+
+# @ Experts: Please add a description of the fields in the dtype
+# See fields above for inspiration.
+# Is float64 necessary? Switch to float32 if possible to save space.
+neutron_veto_hitlet_dtype = [
+    ("area", np.float64),
+    ("amplitude", np.float64),
+    ("time_amplitude", np.int16),
+    ("entropy", np.float64),
+    ("range_50p_area", np.float64),
+    ("range_80p_area", np.float64),
+    ("left_area", np.float64),
+    ("low_left_area", np.float64),
+    ("range_hdr_50p_area", np.float64),
+    ("range_hdr_80p_area", np.float64),
+    ("left_hdr", np.float64),
+    ("low_left_hdr", np.float64),
+    ("fwhm", np.float64),
+    ("left", np.float64),
+    ("fwtm", np.float64),
+    ("low_left", np.float64),
+]

fuse/plugins/__init__.py

@@ -7,6 +7,9 @@
 from . import pmt_and_daq
 from .pmt_and_daq import *
 
+from . import neutron_veto
+from .neutron_veto import *
+
 from . import truth_information
 from .truth_information import *