diff --git a/fuse/common.py b/fuse/common.py index 0b79bcc2..bfe9b861 100644 --- a/fuse/common.py +++ b/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 diff --git a/fuse/context_utils.py b/fuse/context_utils.py index 64c3a6df..fd779094 100644 --- a/fuse/context_utils.py +++ b/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 diff --git a/fuse/dtypes.py b/fuse/dtypes.py index 498bf3c6..3d8e97c8 100644 --- a/fuse/dtypes.py +++ b/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), +] diff --git a/fuse/plugins/__init__.py b/fuse/plugins/__init__.py index bebf6792..e2640eac 100644 --- a/fuse/plugins/__init__.py +++ b/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 * diff --git a/fuse/plugins/micro_physics/input.py b/fuse/plugins/micro_physics/input.py index 3a33690d..ccbebfb7 100644 --- a/fuse/plugins/micro_physics/input.py +++ b/fuse/plugins/micro_physics/input.py @@ -8,6 +8,8 @@ import strax import straxen +from numba import njit + from ...dtypes import g4_fields, primary_positions_fields, deposit_positions_fields from ...common import ( stable_sort, @@ -16,6 +18,7 @@ reshape_awkward, dynamic_chunking, awkward_to_flat_numpy, + dynamic_chunking_two_outputs, ) from ...plugin import FuseBasePlugin @@ -25,18 +28,33 @@ # Remove the path and file name option from the config and do this with the run_number?? @export class ChunkInput(FuseBasePlugin): - """Plugin to read XENONnT Geant4 root or csv files. + """Plugin to read XENONnT Geant4 root or csv files for the TPC and NV + detector. The plugin can distribute the events in time based on a source rate and will create multiple chunks of data if needed. """ - __version__ = "0.3.4" + __version__ = "0.4.0" depends_on: Tuple = tuple() - provides = "geant4_interactions" + provides = ("geant4_interactions", "nv_pmthits") + + dtype_tpc = deposit_positions_fields + g4_fields + primary_positions_fields + strax.time_fields + + dtype_nv = [ + ("pmthitEnergy", np.float32), + ("pmthitTime", np.float32), + ("pmthitID", np.int16), + ("evtid", np.int32), + ] + dtype_nv = dtype_nv + strax.time_fields - dtype = deposit_positions_fields + g4_fields + primary_positions_fields + strax.time_fields + dtype = dict() + dtype["geant4_interactions"] = dtype_tpc + dtype["nv_pmthits"] = dtype_nv + + data_kind = {"geant4_interactions": "geant4_interactions", "nv_pmthits": "nv_pmthits"} save_when = strax.SaveWhen.TARGET @@ -86,10 +104,10 @@ class ChunkInput(FuseBasePlugin): help="All interactions happening after this time (including the event time) will be cut", ) - n_interactions_per_chunk = straxen.URLConfig( - default=1e5, + file_size_limit = straxen.URLConfig( + default=500, type=(int, float), - help="Minimum number of interaction per chunk", + help="Target file size limit in MB", ) entry_start = straxen.URLConfig( @@ -128,6 +146,12 @@ class ChunkInput(FuseBasePlugin): help="Time length of the last chunk", ) + nv_output = straxen.URLConfig( + default=True, + type=bool, + help="Decide if you want to have NV output or not", + ) + def setup(self): super().setup() @@ -141,7 +165,7 @@ def setup(self): cut_delayed=self.cut_delayed, first_chunk_left=self.first_chunk_left, last_chunk_length=self.last_chunk_length, - n_interactions_per_chunk=self.n_interactions_per_chunk, + file_size_limit=self.file_size_limit, arg_debug=self.debug, outer_cylinder=None, # This is not running entry_start=self.entry_start, @@ -150,17 +174,39 @@ def setup(self): cut_nr_only=self.nr_only, fixed_event_spacing=self.fixed_event_spacing, log=self.log, + neutron_veto_output=self.nv_output, ) self.file_reader_iterator = self.file_reader.output_chunk() def compute(self): try: - chunk_data, chunk_left, chunk_right, source_done = next(self.file_reader_iterator) - chunk_data["endtime"] = chunk_data["time"] + + if self.nv_output: + chunk_data, chunk_left, chunk_right, source_done, nv_chunk_data = next( + self.file_reader_iterator + ) + chunk_data["endtime"] = chunk_data["time"] + nv_chunk_data["endtime"] = nv_chunk_data["time"] + + else: + chunk_data, chunk_left, chunk_right, source_done = next(self.file_reader_iterator) + chunk_data["endtime"] = chunk_data["time"] + nv_chunk_data = np.zeros(0, dtype=self.dtype["nv_pmthits"]) self.source_done = source_done - return self.chunk(start=chunk_left, end=chunk_right, data=chunk_data) + result = { + "geant4_interactions": self.chunk( + start=chunk_left, + end=chunk_right, + data=chunk_data, + data_type="geant4_interactions", + ), + "nv_pmthits": self.chunk( + start=chunk_left, end=chunk_right, data=nv_chunk_data, data_type="nv_pmthits" + ), + } + return result except StopIteration: raise RuntimeError("Bug in chunk building!") @@ -190,6 +236,7 @@ def __init__( random_number_generator, separation_scale=1e8, event_rate=1, + file_size_limit=500, subtract_first_interaction_time=True, n_interactions_per_chunk=500, cut_delayed=4e12, @@ -204,6 +251,7 @@ def __init__( cut_nr_only=False, fixed_event_spacing=False, log=None, + neutron_veto_output=False, ): self.directory = directory self.file_name = file_name @@ -211,6 +259,7 @@ def __init__( self.rng = random_number_generator self.separation_scale = separation_scale self.event_rate = event_rate / 1e9 # Conversion to ns + self.file_size_limit = file_size_limit self.subtract_first_interaction_time = subtract_first_interaction_time self.n_interactions_per_chunk = n_interactions_per_chunk self.cut_delayed = cut_delayed @@ -225,6 +274,7 @@ def __init__( self.cut_nr_only = cut_nr_only self.fixed_event_spacing = fixed_event_spacing self.log = log + self.neutron_veto_output = neutron_veto_output self.file = os.path.join(self.directory, self.file_name) @@ -234,6 +284,18 @@ def __init__( self.columns.remove("eventid") self.dtype += primary_positions_fields + strax.time_fields + # Set the nveto dtype -> move to dtype file later!!!! + if self.neutron_veto_output: + self.nv_dtype = [ + ("pmthitEnergy", np.float32), + ("pmthitTime", np.float32), + ("pmthitID", np.int16), + ("evtid", np.int32), + ] + self.nv_dtype = self.nv_dtype + strax.time_fields + + self.neutron_veto_column_names = ["pmthitEnergy", "pmthitTime", "pmthitID"] + # Prepare cut for root and csv case if self.outer_cylinder: self.cut_string = ( @@ -248,7 +310,12 @@ def output_chunk(self): """Function to return one chunk of data from the root or csv file.""" if self.file_type == "root": - interactions, n_simulated_events, start, stop = self._load_root_file() + if self.neutron_veto_output: + interactions, n_simulated_events, start, stop, nv_interactions = ( + self._load_root_file() + ) + else: + interactions, n_simulated_events, start, stop = self._load_root_file() elif self.file_type == "csv": interactions, n_simulated_events, start, stop = self._load_csv_file() else: @@ -256,35 +323,31 @@ def output_chunk(self): f'Cannot load events from file "{self.file}": .root or .cvs file needed.' ) - # Removing all events with zero energy deposit - # m = interactions["ed"] > 0 + # Find the global start time of the interactions to keep TPC and NV in sync + min_time_TPC = ak.min(interactions["t"], axis=1) + if self.neutron_veto_output: + min_time_NV = ak.min(nv_interactions["pmthitTime"], axis=1) + global_min_time = ak.min([min_time_TPC, min_time_NV], axis=0) + else: + global_min_time = min_time_TPC - if self.cut_nr_only: - self.log.info("'nr_only' set to True, keeping only the NR events") - m = ((interactions["type"] == "neutron") & (interactions["edproc"] == "hadElastic")) | ( - interactions["edproc"] == "ionIoni" - ) - e_dep_er = ak.sum(interactions[~m]["ed"], axis=1) - e_dep_nr = ak.sum(interactions[m]["ed"], axis=1) - interactions = interactions[(e_dep_er < 10) & (e_dep_nr > 0)] + # Sort interactions in events by time and subtract time of the first interaction + interactions = interactions[ak.argsort(interactions["t"], stable=True)] - # Removing all events with no interactions: - m = ak.num(interactions["ed"]) > 0 - # and all events with no deposited energy - m = m & (ak.sum(interactions["ed"], axis=1) > 0) + # Remove the global min time from the interactions + if self.event_rate > 0: + interactions["t"] = interactions["t"] - global_min_time - interactions = interactions[m] + if self.neutron_veto_output: + nv_interactions = nv_interactions[ak.argsort(nv_interactions["pmthitTime"])] + nv_interactions["pmthitTime"] = nv_interactions["pmthitTime"] - global_min_time - # Sort interactions in events by time and subtract time of the first interaction - interactions = interactions[ak.argsort(interactions["t"], stable=True)] + # Now lets distribute the interactions in time making sure that the + # detectors stay in sync with each other - if self.event_rate > 0 and self.subtract_first_interaction_time: - interactions["t"] = interactions["t"] - interactions["t"][:, 0] + # Right now this only works if there are TPC interactions. For NV only we have to update this part... - # Adjust event times if necessary - if self.event_rate > 0: num_interactions = len(interactions["t"]) - if self.fixed_event_spacing: self.log.info("Using fixed event spacing.") event_times = ( @@ -307,6 +370,10 @@ def output_chunk(self): interactions["time"] = interactions["t"] + event_times + if self.neutron_veto_output: + # Add the same event times to the NV interactions + nv_interactions["time"] = nv_interactions["pmthitTime"] + event_times + elif self.event_rate == 0: self.log.info("Using event times from provided input file.") if self.file_type == "root": @@ -317,14 +384,40 @@ def output_chunk(self): self.log.warning(msg) interactions["time"] = interactions["t"] + if self.neutron_veto_output: + nv_interactions["time"] = nv_interactions["pmthitTime"] + else: raise ValueError("Source rate cannot be negative!") + # Now that all events are distributed in time and the two detectors are still in sync + # we can apply some cuts to the data + + if self.cut_nr_only: + self.log.info("'nr_only' set to True, keeping only the NR events") + m = ((interactions["type"] == "neutron") & (interactions["edproc"] == "hadElastic")) | ( + interactions["edproc"] == "ionIoni" + ) + e_dep_er = ak.sum(interactions[~m]["ed"], axis=1) + e_dep_nr = ak.sum(interactions[m]["ed"], axis=1) + interactions = interactions[(e_dep_er < 10) & (e_dep_nr > 0)] + + # Removing all events with no interactions: + m = ak.num(interactions["ed"]) > 0 + # and all events with no deposited energy + m = m & (ak.sum(interactions["ed"], axis=1) > 0) + + interactions = interactions[m] + + # TPC delay cut interactions = interactions[interactions["t"] < self.cut_delayed] + if self.neutron_veto_output: + # Add all NV specific cuts here, I will just start with the delay cut for now + nv_interactions = nv_interactions[nv_interactions["pmthitTime"] < self.cut_delayed] + # Make into a flat numpy array interaction_time = awkward_to_flat_numpy(interactions["time"]) - # First caclulate sort index for the interaction times sort_idx = stable_argsort(interaction_time) # and now make it an integer for strax time field @@ -332,74 +425,203 @@ def output_chunk(self): # Sort the interaction times interaction_time = interaction_time[sort_idx] - chunk_idx = dynamic_chunking( - interaction_time, scale=self.separation_scale, n_min=self.n_interactions_per_chunk - ) + # Do the same for the NV interactions\ + if self.neutron_veto_output: + nv_interaction_time = awkward_to_flat_numpy(nv_interactions["time"]) + nv_sort_idx = stable_argsort(nv_interaction_time) + nv_interaction_time = nv_interaction_time.astype(np.int64) + nv_interaction_time = nv_interaction_time[nv_sort_idx] - unique_chunk_index_values = np.unique(chunk_idx) + # Chunking the data! - chunk_start = np.array( - [interaction_time[chunk_idx == i][0] for i in unique_chunk_index_values] - ) - chunk_end = np.array( - [interaction_time[chunk_idx == i][-1] for i in unique_chunk_index_values] - ) + # Move this somewhere else? Is this correct? + n_bytes_per_interaction_TPC = 6 * 8 + 5 * 4 + 2 * 2 + 4 * 40 + n_bytes_per_interaction_NV = 2 * 8 + 3 * 4 + 1 * 2 - if (len(chunk_start) > 1) & (len(chunk_end) > 1): - gap_length = chunk_start[1:] - chunk_end[:-1] - gap_length = np.append(gap_length, gap_length[-1] + self.last_chunk_length) - chunk_bounds = chunk_end + np.int64(self.chunk_delay_fraction * gap_length) - self.chunk_bounds = np.append(chunk_start[0] - self.first_chunk_left, chunk_bounds) - - else: - self.log.warning( - "Only one Chunk created! Only a few events simulated? " - "If no, your chunking parameters might not be optimal. " - "Try to decrease the source_rate or decrease the n_interactions_per_chunk." - ) - self.chunk_bounds = [ - chunk_start[0] - self.first_chunk_left, - chunk_end[0] + self.last_chunk_length, - ] + if not self.neutron_veto_output: + # Start with the TPC only case - # We need to get the min and max times for each event - # to preselect events with interactions in the chunk bounds - times_min = ak.to_numpy(ak.min(interactions["time"], axis=1)).astype(np.int64) - times_max = ak.to_numpy(ak.max(interactions["time"], axis=1)).astype(np.int64) + time_gaps = interaction_time[1:] - interaction_time[:-1] + time_gaps = np.append(time_gaps, 0) # Add last gap - # Process and yield each chunk - source_done = False - self.log.info(f"Simulating data in {len(unique_chunk_index_values)} chunks.") - for c_ix, chunk_left, chunk_right in zip( - unique_chunk_index_values, self.chunk_bounds[:-1], self.chunk_bounds[1:] - ): + chunk_idx = dynamic_chunking( + time_gaps, + self.file_size_limit, + self.separation_scale, + n_bytes_per_interaction_TPC, + ) + unique_chunk_index_values = np.unique(chunk_idx) - # We do a preselection of the events that have interactions within the chunk - # before converting the full array to numpy (which is expensive in terms of memory) - m = (times_min <= chunk_right) & (times_max >= chunk_left) - current_chunk = interactions[m] + chunk_start = np.array( + [interaction_time[chunk_idx == i][0] for i in unique_chunk_index_values] + ) + chunk_end = np.array( + [interaction_time[chunk_idx == i][-1] for i in unique_chunk_index_values] + ) - if len(current_chunk) == 0: - current_chunk = np.empty(0, dtype=self.dtype) + if (len(chunk_start) > 1) & (len(chunk_end) > 1): + gap_length = chunk_start[1:] - chunk_end[:-1] + gap_length = np.append(gap_length, gap_length[-1] + self.last_chunk_length) + chunk_bounds = chunk_end + np.int64(self.chunk_delay_fraction * gap_length) + self.chunk_bounds = np.append(chunk_start[0] - self.first_chunk_left, chunk_bounds) else: - # Convert the chunk from awkward array to a numpy array - current_chunk = full_array_to_numpy(current_chunk, self.dtype) + self.log.warning( + "Only one Chunk created! Only a few events simulated? " + "If no, your chunking parameters might not be optimal. " + "Try to decrease the source_rate or decrease the n_interactions_per_chunk." + ) + self.chunk_bounds = [ + chunk_start[0] - self.first_chunk_left, + chunk_end[0] + self.last_chunk_length, + ] + + # We need to get the min and max times for each event + # to preselect events with interactions in the chunk bounds + times_min = ak.to_numpy(ak.min(interactions["time"], axis=1)).astype(np.int64) + times_max = ak.to_numpy(ak.max(interactions["time"], axis=1)).astype(np.int64) + + # Process and yield each chunk + source_done = False + self.log.info(f"Simulating data in {len(unique_chunk_index_values)} chunks.") + for c_ix, chunk_left, chunk_right in zip( + unique_chunk_index_values, self.chunk_bounds[:-1], self.chunk_bounds[1:] + ): + + # We do a preselection of the events that have interactions within the chunk + # before converting the full array to numpy (which is expensive in terms of memory) + m = (times_min <= chunk_right) & (times_max >= chunk_left) + current_chunk = interactions[m] + + if len(current_chunk) == 0: + current_chunk = np.empty(0, dtype=self.dtype) + + else: + # Convert the chunk from awkward array to a numpy array + current_chunk = full_array_to_numpy(current_chunk, self.dtype) + + # Now we have the chunk of data in strax/numpy format + # We can now filter only the interactions within the chunk + select_times = current_chunk["time"] >= chunk_left + select_times &= current_chunk["time"] <= chunk_right + current_chunk = current_chunk[select_times] + + # Sorting each chunk by time within the chunk + sort_chunk = stable_argsort(current_chunk["time"]) + current_chunk = current_chunk[sort_chunk] + + if c_ix == unique_chunk_index_values[-1]: + source_done = True + + yield current_chunk, chunk_left, chunk_right, source_done + + elif self.neutron_veto_output: + # Now TPC and NV making this a bit more complicated + + # We need to combine the times of TPC and NV interactions so they always end up in the same chunk + combined_times = np.append(interaction_time, nv_interaction_time) + combined_types = np.append( + np.zeros(len(interaction_time)), np.ones(len(nv_interaction_time)) + ) - # Now we have the chunk of data in strax/numpy format - # We can now filter only the interactions within the chunk - select_times = current_chunk["time"] >= chunk_left - select_times &= current_chunk["time"] <= chunk_right - current_chunk = current_chunk[select_times] + sort_idx = stable_argsort(combined_times) + combined_times = combined_times[sort_idx] + combined_types = combined_types[sort_idx] + combined_time_gaps = combined_times[1:] - combined_times[:-1] + combined_time_gaps = np.append(combined_time_gaps, 0) # Add last gap + + combined_chunk_index = dynamic_chunking_two_outputs( + combined_time_gaps, + combined_types, + self.file_size_limit, + self.separation_scale, + n_bytes_per_interaction_TPC, + n_bytes_per_interaction_NV, + ) - # Sorting each chunk by time within the chunk - sort_chunk = stable_argsort(current_chunk["time"]) - current_chunk = current_chunk[sort_chunk] + unique_combined_chunk_index_values = np.unique(combined_chunk_index) - if c_ix == unique_chunk_index_values[-1]: - source_done = True + chunk_start = np.array( + [ + combined_times[combined_chunk_index == i][0] + for i in unique_combined_chunk_index_values + ] + ) + chunk_end = np.array( + [ + combined_times[combined_chunk_index == i][-1] + for i in unique_combined_chunk_index_values + ] + ) - yield current_chunk, chunk_left, chunk_right, source_done + if (len(chunk_start) > 1) & (len(chunk_end) > 1): + gap_length = chunk_start[1:] - chunk_end[:-1] + gap_length = np.append(gap_length, gap_length[-1] + self.last_chunk_length) + chunk_bounds = chunk_end + np.int64(self.chunk_delay_fraction * gap_length) + self.chunk_bounds = np.append(chunk_start[0] - self.first_chunk_left, chunk_bounds) + + else: + self.log.warning( + "Only one Chunk created! Only a few events simulated? " + "If no, your chunking parameters might not be optimal. " + "Try to decrease the source_rate or decrease the n_interactions_per_chunk." + ) + self.chunk_bounds = [ + chunk_start[0] - self.first_chunk_left, + chunk_end[0] + self.last_chunk_length, + ] + + # We need to get the min and max times for each event + # to preselect events with interactions in the chunk bounds + times_min = ak.to_numpy(ak.min(interactions["time"], axis=1)).astype(np.int64) + times_max = ak.to_numpy(ak.max(interactions["time"], axis=1)).astype(np.int64) + + times_min_nv = ak.to_numpy(ak.min(nv_interactions["time"], axis=1)).astype(np.int64) + times_max_nv = ak.to_numpy(ak.max(nv_interactions["time"], axis=1)).astype(np.int64) + + # Process and yield each chunk + source_done = False + self.log.info(f"Simulating data in {len(unique_combined_chunk_index_values)} chunks.") + for c_ix, chunk_left, chunk_right in zip( + unique_combined_chunk_index_values, self.chunk_bounds[:-1], self.chunk_bounds[1:] + ): + m = (times_min <= chunk_right) & (times_max >= chunk_left) + current_chunk = interactions[m] + + m_nv = (times_min_nv <= chunk_right) & (times_max_nv >= chunk_left) + current_chunk_nv = nv_interactions[m_nv] + + if len(current_chunk) == 0: + current_chunk = np.empty(0, dtype=self.dtype) + else: + current_chunk = full_array_to_numpy(current_chunk, self.dtype) + + if len(current_chunk_nv) == 0: + current_chunk_nv = np.empty(0, dtype=self.nv_dtype) + else: + current_chunk_nv = full_array_to_numpy(current_chunk_nv, self.nv_dtype) + + # Now we have the chunk of data in strax/numpy format + # We can now filter only the interactions within the chunk + select_times = current_chunk["time"] >= chunk_left + select_times &= current_chunk["time"] <= chunk_right + current_chunk = current_chunk[select_times] + + select_times_nv = current_chunk_nv["time"] >= chunk_left + select_times_nv &= current_chunk_nv["time"] <= chunk_right + current_chunk_nv = current_chunk_nv[select_times_nv] + + # Sorting each chunk by time within the chunk + sort_chunk = stable_argsort(current_chunk["time"]) + current_chunk = current_chunk[sort_chunk] + + sort_chunk_nv = stable_argsort(current_chunk_nv["time"]) + current_chunk_nv = current_chunk_nv[sort_chunk_nv] + + if c_ix == unique_combined_chunk_index_values[-1]: + source_done = True + + yield current_chunk, chunk_left, chunk_right, source_done, current_chunk_nv def last_chunk_bounds(self): return self.chunk_bounds[-1] @@ -512,7 +734,32 @@ def _load_root_file(self): interactions["y_pri"] = ak.broadcast_arrays(xyz_pri["y_pri"], interactions["x"])[0] interactions["z_pri"] = ak.broadcast_arrays(xyz_pri["z_pri"], interactions["x"])[0] - return interactions, n_simulated_events, start_index, stop_index + if self.neutron_veto_output: + nv_interactions = ttree.arrays( + self.neutron_veto_column_names, + entry_start=start_index, + entry_stop=stop_index, + ) + + # Time conversion to ns + nv_interactions["pmthitTime"] = nv_interactions["pmthitTime"] * 10**9 + + # Add event numbers + nv_eventids = ttree.arrays( + "eventid", + entry_start=start_index, + entry_stop=stop_index, + ) + nv_eventids = ak.broadcast_arrays( + nv_eventids["eventid"], nv_interactions["pmthitTime"] + )[0] + nv_interactions["evtid"] = nv_eventids + + return interactions, n_simulated_events, start_index, stop_index, nv_interactions + + else: + + return interactions, n_simulated_events, start_index, stop_index def _get_ttree(self): """Function which searches for the correct ttree in MC root file. diff --git a/fuse/plugins/neutron_veto/Testing_nv_hitlet.ipynb b/fuse/plugins/neutron_veto/Testing_nv_hitlet.ipynb new file mode 100644 index 00000000..365d47a7 --- /dev/null +++ b/fuse/plugins/neutron_veto/Testing_nv_hitlet.ipynb @@ -0,0 +1,902 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Getting Started\n", + "In this notebook you will learn how to run your first simulation with fuse.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Imports" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "You specified _auto_append_rucio_local=True and you are not on dali compute nodes, so we will add the following rucio local path: /project/lgrandi/rucio/\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/digangi/cutax/cutax/cut_lists/science.py:28: UserWarning: Removing cut \n", + " warnings.warn(f\"Removing cut {cut}\", UserWarning)\n", + "/home/digangi/.local/lib/python3.9/site-packages/straxen/config/preprocessors.py:16: UserWarning: From straxen version 2.1.0 onward, URLConfig parameterswill be sorted alphabetically before being passed to the plugins, this will change the lineage hash for non-sorted URLs. To load data processed with non-sorted URLs, you will need to use an older version.\n", + " warnings.warn(\n" + ] + } + ], + "source": [ + "import sys\n", + "\n", + "sys.path.insert(0, \"/home/digangi/cutax/\") # on midway3\n", + "\n", + "import fuse\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import awkward as ak\n", + "from tqdm import tqdm\n", + "import tqdm.notebook as tq\n", + "import pandas as pd\n", + "from sklearn.cluster import DBSCAN" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Warning! elife not in context config, skipping...\n", + "Warning! electron_drift_velocity not in context config, skipping...\n", + "Warning! electron_drift_time_gate not in context config, skipping...\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n" + ] + } + ], + "source": [ + "st = fuse.context.full_chain_context(output_folder=\"./fuse_data\")\n", + "st.set_config(\n", + " {\n", + " \"path\": \"/project2/lgrandi/layos/\",\n", + " \"file_name\": \"output_n_Veto_neutron_AmBe_1.root\",\n", + " \"entry_stop\": 2000,\n", + " \"nv_output\": True, # On inclus les pmts du Neutron Veto (par defaut False)\n", + " \"debug\": True,\n", + " \"file_size_limit\": 500 / 1e4, # Poser question a Hening sur cela....\n", + " }\n", + ")\n", + "\n", + "run_number = \"00003\" # Attention à ce parametre, aujourd'hui triviale, mais qui tiens en compte du mapping pour la TPC en fonction du run." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Warning! elife not in context config, skipping...\n", + "Warning! electron_drift_velocity not in context config, skipping...\n", + "Warning! electron_drift_time_gate not in context config, skipping...\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n", + "/home/digangi/.local/lib/python3.9/site-packages/strax/context.py:380: UserWarning: Provides of multi-output plugins overlap, deregister old plugins .\n", + " warnings.warn(\n" + ] + } + ], + "source": [ + "g4dir = \"/project/lgrandi/xenonnt/simulations/nveto/ambe/geant4/ambe_nveto_gd0p02_2024-01-23/gd_ambe_sr1_top_cw7d8m_with_gamma/results/\"\n", + "filename = \"nT_mc_gdt7p8m_wc_00000.root\"\n", + "\n", + "st = fuse.context.full_chain_context(output_folder=\"./fuse_data\")\n", + "st.set_config(\n", + " {\n", + " \"path\": g4dir,\n", + " \"file_name\": filename,\n", + " \"entry_stop\": 2000,\n", + " \"nv_output\": True, # On inclus les pmts du Neutron Veto (par defaut False)\n", + " \"debug\": True,\n", + " \"file_size_limit\": 500 / 1e4, # Poser question a Hening sur cela....\n", + " }\n", + ")\n", + "\n", + "run_number = \"00003\" # Attention à ce parametre, aujourd'hui triviale, mais qui tiens en compte du mapping pour la TPC en fonction du run." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "st.set_config({\"deteministic_seed\": False, \"user_defined_random_seed\": 42})" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "WARNING:strax:Option gain_model_nv not taken by any registered plugin\n", + "WARNING:strax:Option gain_model_mv not taken by any registered plugin\n", + "WARNING:strax:Option deteministic_seed not taken by any registered plugin\n", + "WARNING:strax:Option user_defined_random_seed not taken by any registered plugin\n", + "WARNING:strax:Option gain_model_nv not taken by any registered plugin\n", + "WARNING:strax:Option gain_model_mv not taken by any registered plugin\n", + "WARNING:strax:Option deteministic_seed not taken by any registered plugin\n", + "WARNING:strax:Option user_defined_random_seed not taken by any registered plugin\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "216ccda11e494bbaa5c22626763cd7ac", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading nv_pmthits: | | 0.00 % [00:00" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "plt.hist(events_nv[\"area\"], range=[0, 200], bins=20, alpha=0.3, label=\"events_nv\")\n", + "plt.hist(events_nv_stacket[\"area\"], range=[0, 200], bins=20, alpha=0.3, label=\"events_nv_stacked\")\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": 44, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Output saved to file: /scratch/midway3/digangi/simulations/NeutronVeto/hitlet_sim/output_data/fuseHitSim_sr0config_seed42_nparraynT_mc_gdt7p8m_wc_00000.root\n" + ] + } + ], + "source": [ + "output_file = (\n", + " \"/scratch/midway3/digangi/simulations/NeutronVeto/hitlet_sim/output_data/\"\n", + " + \"fuseHitSim_sr\"\n", + " + str(sr)\n", + " + \"config_seed42_nparray\"\n", + " + filename\n", + ")\n", + "np.save(output_file, events_nv)\n", + "print(\"Output saved to file: \", output_file)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([(4.1504655, 3.7096519e+01, 2006, 0, 776003102, 776003102),\n", + " (7.0233135, 1.2525828e+05, 123, 0, 776128323, 776128323),\n", + " (7.020823 , 1.2526002e+05, 300, 0, 776128325, 776128325),\n", + " ...,\n", + " (4.4297585, 1.4902604e+13, 2111, 370, 15282534046465, 15282534046465),\n", + " (2.2376018, 1.4902604e+13, 2035, 370, 15282534046465, 15282534046465),\n", + " (3.5634096, 1.4902604e+13, 2103, 370, 15282534046465, 15282534046465)],\n", + " dtype=[('pmthitEnergy', '[{pmthitEnergy: 4.15, pmthitTime: 37.1, pmthitID: 2006, evtid: 0, ...},\n", + " {pmthitEnergy: 7.02, pmthitTime: 1.25e+05, pmthitID: 123, evtid: 0, ...},\n", + " {pmthitEnergy: 7.02, pmthitTime: 1.25e+05, pmthitID: 300, evtid: 0, ...},\n", + " {pmthitEnergy: 6.95, pmthitTime: 1.25e+05, pmthitID: 465, evtid: 0, ...},\n", + " {pmthitEnergy: 6.95, pmthitTime: 1.25e+05, pmthitID: 87, evtid: 0, ...},\n", + " {pmthitEnergy: 6.96, pmthitTime: 1.25e+05, pmthitID: 392, evtid: 0, ...},\n", + " {pmthitEnergy: 6.98, pmthitTime: 1.25e+05, pmthitID: 54, evtid: 0, ...},\n", + " {pmthitEnergy: 6.99, pmthitTime: 1.25e+05, pmthitID: 465, evtid: 0, ...},\n", + " {pmthitEnergy: 6.93, pmthitTime: 1.25e+05, pmthitID: 365, evtid: 0, ...},\n", + " {pmthitEnergy: 7.02, pmthitTime: 1.25e+05, pmthitID: 54, evtid: 0, ...},\n", + " ...,\n", + " {pmthitEnergy: 3.22, pmthitTime: 1.49e+13, pmthitID: 2113, evtid: 370, ...},\n", + " {pmthitEnergy: 2.13, pmthitTime: 1.49e+13, pmthitID: 2005, evtid: 370, ...},\n", + " {pmthitEnergy: 4.97, pmthitTime: 1.49e+13, pmthitID: 2112, evtid: 370, ...},\n", + " {pmthitEnergy: 2.24, pmthitTime: 1.49e+13, pmthitID: 2113, evtid: 370, ...},\n", + " {pmthitEnergy: 3.32, pmthitTime: 1.49e+13, pmthitID: 2035, evtid: 370, ...},\n", + " {pmthitEnergy: 4.66, pmthitTime: 1.49e+13, pmthitID: 2005, evtid: 370, ...},\n", + " {pmthitEnergy: 4.43, pmthitTime: 1.49e+13, pmthitID: 2111, evtid: 370, ...},\n", + " {pmthitEnergy: 2.24, pmthitTime: 1.49e+13, pmthitID: 2035, evtid: 370, ...},\n", + " {pmthitEnergy: 3.56, pmthitTime: 1.49e+13, pmthitID: 2103, evtid: 370, ...}]\n", + "-----------------------------------------------------------------------------\n", + "type: 179349 * {\n", + " pmthitEnergy: float32,\n", + " pmthitTime: float32,\n", + " pmthitID: int16,\n", + " evtid: int32,\n", + " time: int64,\n", + " endtime: int64\n", + "}" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ak_pmthits = ak.Array(nv_pmthits)\n", + "ak_pmthits" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "179349\n", + "177403\n" + ] + } + ], + "source": [ + "# awkward array with 'evtid', 'pmthitTime', 'pmthitEnergy', 'pmthitID'\n", + "# pmthits=ak.Array(nv_pmthits)\n", + "pmthits = nv_pmthits\n", + "print(len(pmthits))\n", + "# select NV PMTs (need to exclude MV PMTs?)\n", + "mask = pmthits[\"pmthitID\"] >= 2000\n", + "pmthits = pmthits[mask]\n", + "print(len(pmthits))" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Applying QE and CE\n", + "Loading hit survive\n", + "Sampling hitlets charge pe\n", + "Getting time hitlets\n", + "Looking for stacket hitlets\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "3cefca70540547c59fc4454c87355865", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + " 0%| | 0/1824 [00:00" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Stampa del numero di eventi\n", + "print(\"[AM HitSim] NV events generated:\", len(events_AM))\n", + "print(\"[fuse HitSim] NV events generated:\", len(events_NV))\n", + "print(\"[fuse HitSim] NV events generated:\", len(events_NV_nparray))\n", + "\n", + "print(\"fuse - AM:\", (len(events_NV) - len(events_AM)) / len(events_NV) * 100, \"%\")\n", + "\n", + "plt.figure(figsize=(12, 8))\n", + "\n", + "# Parametri comuni per l'istogramma\n", + "bin_range = [0, 200]\n", + "bins = 20\n", + "\n", + "# Istogramma per AM HitSim\n", + "# counts_AM, edges_AM, _ = plt.hist(events_AM['area'], histtype='step', range=bin_range, bins=bins, label='AM HitSim: ' + str(len(events_AM)) + ' events', color='b')\n", + "# bin_centers_AM = (edges_AM[:-1] + edges_AM[1:]) / 2\n", + "# errors_AM = np.sqrt(counts_AM)\n", + "\n", + "# Istogramma per fuse HitSim\n", + "counts_nv, edges_nv, _ = plt.hist(\n", + " events_NV[\"area\"],\n", + " histtype=\"step\",\n", + " range=bin_range,\n", + " bins=bins,\n", + " label=\"fuse HitSim \" + str(len(events_NV)) + \" events\",\n", + " color=\"orange\",\n", + ")\n", + "bin_centers_nv = (edges_nv[:-1] + edges_nv[1:]) / 2\n", + "errors_nv = np.sqrt(counts_nv)\n", + "\n", + "# Istogramma per fuse HitSim\n", + "counts_nv_, edges_nv_, _ = plt.hist(\n", + " events_NV_nparray[\"area\"],\n", + " histtype=\"step\",\n", + " range=bin_range,\n", + " bins=bins,\n", + " label=\"fuse HitSim (np.array) \" + str(len(events_NV_nparray)) + \" events\",\n", + " color=\"r\",\n", + ")\n", + "bin_centers_nv_ = (edges_nv_[:-1] + edges_nv_[1:]) / 2\n", + "errors_nv_ = np.sqrt(counts_nv_)\n", + "\n", + "\n", + "# Aggiunta delle barre d'errore con tacchette orizzontali\n", + "# plt.errorbar(bin_centers_AM, counts_AM, yerr=errors_AM, fmt='none', color='b',\n", + "# label='AM HitSim Error', capsize=4, capthick=1, alpha=0.5)\n", + "plt.errorbar(\n", + " bin_centers_nv,\n", + " counts_nv,\n", + " yerr=errors_nv,\n", + " fmt=\"none\",\n", + " color=\"orange\",\n", + " label=\"fuse HitSim Error\",\n", + " capsize=4,\n", + " capthick=1,\n", + " alpha=0.5,\n", + ")\n", + "plt.errorbar(\n", + " bin_centers_nv_,\n", + " counts_nv_,\n", + " yerr=errors_nv_,\n", + " fmt=\"none\",\n", + " color=\"r\",\n", + " label=\"fuse HitSim (np.array) Error\",\n", + " capsize=4,\n", + " capthick=1,\n", + " alpha=0.5,\n", + ")\n", + "\n", + "\n", + "# Personalizzazione del grafico\n", + "plt.xlabel(\"Area\")\n", + "plt.ylabel(\"Counts\")\n", + "plt.title(\"gd_ambe_top_cw7d8m_with_gamma - Config: SR1, eCE=0.87\")\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.20" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/fuse/plugins/neutron_veto/__init__.py b/fuse/plugins/neutron_veto/__init__.py new file mode 100644 index 00000000..4d62ff79 --- /dev/null +++ b/fuse/plugins/neutron_veto/__init__.py @@ -0,0 +1,2 @@ +from . import nvhitlets +from .nvhitlets import * diff --git a/fuse/plugins/neutron_veto/nvhitlets.py b/fuse/plugins/neutron_veto/nvhitlets.py new file mode 100644 index 00000000..18cb9bdf --- /dev/null +++ b/fuse/plugins/neutron_veto/nvhitlets.py @@ -0,0 +1,244 @@ +import strax +import straxen +import numpy as np +import scipy.constants as const +from sklearn.cluster import DBSCAN + + +from ...plugin import FuseBasePlugin + +from ...dtypes import neutron_veto_hitlet_dtype + +export, __all__ = strax.exporter() + + +@export +class NeutronVetoHitlets(FuseBasePlugin): + """Plugin to simulate the Neutron Veto to hitlets. + + @Experts: Please add a better description of the plugin here. + """ + + __version__ = "0.0.1" + + depends_on = "nv_pmthits" + provides = "nv_hitlets" + data_kind = "nv_hitlets" + + dtype = neutron_veto_hitlet_dtype + strax.interval_dtype + + # Fix these URL configs! + nveto_pmt_qe = straxen.URLConfig( + default="nveto_pmt_qe://resource://simulation_config://" + "SIMULATION_CONFIG_FILE.json?" + "&key=nveto_pmt_qe" + "&fmt=json", + help="Quantum efficiency of NV PMTs", + ) + + nveto_spe_parameters = straxen.URLConfig( + default="nveto_spe://resource://simulation_config://" + "SIMULATION_CONFIG_FILE.json?&key=nveto_pmt_spe&fmt=json", + help="SPE model of NV PMTs", + ) + + stack_hitlets = straxen.URLConfig( + default=False, + help="Option to enable or disable hitlet stacking", + ) + + ce_scaling = straxen.URLConfig( + type=(int, float), + default="take://resource://SIMULATION_CONFIG_FILE.json?fmt=json&take=nveto_eCE", + cache=True, + help="Add good description here", + ) + + def compute(self, nv_pmthits): + + if len(nv_pmthits) == 0: + return np.zeros(0, self.dtype) + + hitlets = self._nv_hitlets(nv_pmthits) + + result = np.zeros(len(hitlets), dtype=self.dtype) + result["time"] = hitlets["time"] + result["length"] = 1 + result["dt"] = 10 + result["channel"] = hitlets["pmthitID"] + result["area"] = hitlets["pe_area"] + + return strax.sort_by_time(result) + + def _nv_hitlets(self, pmthits): + + extra_fields = [ + ("pe_area", "= 2000) & (pmthits["pmthitID"] < 2121) + pmthits = pmthits[mask] + + self.log.debug("Applying QE") + # Applying Quantum efficiency for each pmt + # --- super super slow..... + # qe = 1e-2 * np.vectorize(self.QE_E)(pmthits["pmthitEnergy"], pmthits["pmthitID"]) + + # A faster approach + qe = np.zeros(len(pmthits), dtype=np.float32) + unique_ids = np.unique(pmthits["pmthitID"]) + NVeto_PMT_QE = self.nveto_pmt_qe + for uid in unique_ids: + mask_id = pmthits["pmthitID"] == uid + qe[mask_id] = 1e-2 * QE_E(pmthits["pmthitEnergy"][mask_id], uid, NVeto_PMT_QE) + + self.log.debug("Applying CE") + # Applying effective collection efficiency + qe *= self.ce_scaling + + # Applying acceptance per pmt: for the approach in which SPE PDF has already applied a threshold for low charges + self.log.debug("Applying per pmt acceptance") + # also very slow, think this is a bottleneck of the URLConfigs, it is not very efficient to call them in a loop many times. + # qe = qe * np.vectorize(self.get_acceptance)(pmthits["pmthitID"]) + + NV_SPE = self.nveto_spe_parameters + acceptance_dict = {k: v["acceptance"] for k, v in NV_SPE.items()} + qe = qe * np.vectorize(acceptance_dict.get)(pmthits["pmthitID"]) + + self.log.debug("Binomial sampling") + # Generate a photoelectron based on (binomial) conversion probability qe*eCE*spe_acc + pe = np.array([np.random.binomial(1, j, 1)[0] for j in qe]) + + maks_qe = pe > 0 + pmthits = pmthits[maks_qe] + + # 2. Sampling charge from SPE for each pmthit with a generated pe + self.log.debug("Sampling hitlets charge pe") + # Same performance problems as above. Lets try something faster + # pmthits["pe_area"] = np.vectorize(self.pe_charge_N)(pmthits["pmthitID"]) + spe_charge = np.zeros(len(pmthits), dtype=np.float32) + unique_ids = np.unique(pmthits["pmthitID"]) + for uid in unique_ids: + mask_id = pmthits["pmthitID"] == uid + + # Just call the random choice once per PMT ID + SPE_channel = NV_SPE.get(uid) + spe_charge[mask_id] = self.rng.choice( + SPE_channel["pe"], p=SPE_channel["SPE_values"], size=np.sum(mask_id) + ) + pmthits["pe_area"] = spe_charge + + # 3. Creating hitlet times + self.log.debug("Getting time hitlets") + times = [] + for i in np.unique(pmthits["evtid"]): + mask = pmthits["evtid"] == i + pmthits_evt = pmthits[mask] + cluster_times_ns = pmthits_evt["pmthitTime"] - min(pmthits_evt["pmthitTime"]) + times.append(cluster_times_ns) + pmthits["cluster_times_ns"] = np.concatenate(times) + + if not self.stack_hitlets: + return pmthits + + elif self.stack_hitlets: + self.log.debug("Looking for stacked hitlets") + + arr_c_evt = [] + for i in np.unique(pmthits["evtid"]): + arr_evt = pmthits[pmthits["evtid"] == i] + arr_c_pmt = [] + for j in np.unique(arr_evt["pmthitID"]): + arr_pmt = arr_evt[arr_evt["pmthitID"] == j] + labels = channel_cluster_nv(arr_pmt["cluster_times_ns"]) + arr_pmt["labels"] = labels + arr_c = np.concatenate( + [ + get_clusters_arrays(arr_pmt[arr_pmt["labels"] == l], new_dtype) + for l in np.unique(labels) + ] + ) + + arr_c_pmt.append(arr_c) + arr_c_evt.append(np.concatenate(arr_c_pmt)) + + return np.concatenate(arr_c_evt) + + # Get Quantum efficiency + + # def QE_E(self, E, ID): + # WL = energy_to_wavelenght(E) + # ind = ID - 2000 + # qe = self.nveto_pmt_qe[ind](WL) + # # qe = self.nveto_pmt_qe["QE"][ind](WL) + + # return qe + + # def get_acceptance(self, ID): + # acc = self.nveto_spe_parameters.get(ID)["acceptance"] + # return acc + + # Get acceptance threshold + def get_threshold_acc(self, ID): + ind = ID - 2000 + threshold = self.nveto_spe_parameters.threshold_pe.values[ind] + return threshold + + # Sampling charge from SPE + # def pe_charge_N(self, pmt_id): + # SPE_channel = self.nveto_spe_parameters.get(pmt_id) + + # charge = self.rng.choice(SPE_channel["pe"], SPE_channel["SPE_values"], k=1)[0] + + # return charge + + +def energy_to_wavelenght(E): + Joules_to_eV = 1.602 * 1e-19 + return 1e9 * const.h * const.c / (E * Joules_to_eV) + + +def QE_E(E, ID, nveto_pmt_qe): + WL = energy_to_wavelenght(E) + ind = ID - 2000 + qe = nveto_pmt_qe[ind](WL) + # qe = self.nveto_pmt_qe["QE"][ind](WL) + return qe + + +# Cluster for stacket hitlets +def channel_cluster_nv(t): + db_cluster = DBSCAN( + eps=8, min_samples=1 + ) # As a preliminar value we fix distance between two photons arriving in the same pmt 8ns + t_val = np.array(t) + clusters = np.array(db_cluster.fit_predict(t_val.reshape(-1, 1))) + return clusters + + +def get_clusters_arrays(arr, typ): + arr_nv_c = np.zeros(1, dtype=typ) + arr_nv_c["n_clusters_hits"] = len(arr) + + for i in arr.dtype.names: # <-- CORRETTO: usare dtype.names invece di fields + if i in ["time", "pmthitTime", "cluster_times_ns"]: + arr_nv_c[i] = np.min(arr[i]) + elif i == "endtime": + arr_nv_c[i] = np.max(arr[i]) + elif i in ["pe_area", "pmthitEnergy"]: + arr_nv_c[i] = np.sum(arr[i]) + elif i in ["evtid", "pmthitID", "labels"]: + arr_nv_c[i] = np.unique(arr[i]) + + return arr_nv_c