From 09a22e52f4d7c2fb0a56f1e7382234b934928d56 Mon Sep 17 00:00:00 2001 From: jarongrigat Date: Thu, 12 Sep 2024 18:22:19 +0200 Subject: [PATCH] building fastsim --- fuse/context.py | 105 ++++++++++ fuse/plugins/__init__.py | 3 + fuse/plugins/fastsim/__init__.py | 11 ++ .../fastsim/fastsim_events_uncorrected.py | 180 ++++++++++++++++++ .../plugins/fastsim/fastsim_macro_clusters.py | 81 ++++++++ fuse/plugins/fastsim/fastsim_s1.py | 58 ++++++ fuse/plugins/fastsim/fastsim_s2.py | 121 ++++++++++++ 7 files changed, 559 insertions(+) create mode 100644 fuse/plugins/fastsim/__init__.py create mode 100644 fuse/plugins/fastsim/fastsim_events_uncorrected.py create mode 100644 fuse/plugins/fastsim/fastsim_macro_clusters.py create mode 100644 fuse/plugins/fastsim/fastsim_s1.py create mode 100644 fuse/plugins/fastsim/fastsim_s2.py diff --git a/fuse/context.py b/fuse/context.py index bda36abf..0654aa96 100644 --- a/fuse/context.py +++ b/fuse/context.py @@ -87,6 +87,13 @@ fuse.truth_information.ClusterTagging, ] +fastsim_plugins = [ + fuse.fastsim.MacroClusters, + fuse.fastsim.S1Areas, + fuse.fastsim.S2Areas, + fuse.fastsim.FastsimEventsUncorrected +] + def microphysics_context( output_folder="./fuse_data", simulation_config_file="fuse_config_nt_sr1_dev.json" @@ -224,6 +231,104 @@ def full_chain_context( return st +def fastsim_context( + output_folder="./fuse_data", + clustering_method="dbscan", + corrections_version=None, + simulation_config_file="fuse_config_nt_sr1_dev.json", + corrections_run_id="046477", + run_id_specific_config={ + "gain_model_mc": "gain_model", + "electron_lifetime_liquid": "elife", + "drift_velocity_liquid": "electron_drift_velocity", + "drift_time_gate": "electron_drift_time_gate", + }, + run_without_proper_corrections=False, + + +): + """Function to create a fuse fastsim context.""" + + if corrections_run_id is None: + raise ValueError("Specify a corrections_run_id to load the corrections") + if (corrections_version is None) & (not run_without_proper_corrections): + raise ValueError( + "Specify a corrections_version. If you want to run without proper " + "corrections for testing or just trying out fuse, " + "set run_without_proper_corrections to True" + ) + if simulation_config_file is None: + raise ValueError("Specify a simulation configuration file") + + if run_without_proper_corrections: + log.warning( + "Running without proper correction version. This is not recommended for production use." + "Take the context defined in cutax if you want to run XENONnT simulations." + ) + + st = strax.Context( + storage=strax.DataDirectory(output_folder), **straxen.contexts.xnt_common_opts + ) + + # Register microphysics plugins + if clustering_method == "dbscan": + for plugin in microphysics_plugins_dbscan_clustering: + st.register(plugin) + elif clustering_method == "lineage": + for plugin in microphysics_plugins_lineage_clustering: + st.register(plugin) + else: + raise ValueError(f"Clustering method {clustering_method} not implemented!") + + for plugin in remaining_microphysics_plugins: + st.register(plugin) + + # Register S1 plugins + for plugin in s1_simulation_plugins: + st.register(plugin) + + # Register S2 plugins + for plugin in s2_simulation_plugins: + st.register(plugin) + + # Register fastsim plugins + for plugin in fastsim_plugins: + st.register(plugin) + + # Register truth plugins + for plugin in truth_information_plugins: + st.register(plugin) + + if corrections_version is not None: + st.apply_xedocs_configs(version=corrections_version) + + set_simulation_config_file(st, simulation_config_file) + + local_versions = st.config + for config_name, url_config in local_versions.items(): + if isinstance(url_config, str): + if "run_id" in url_config: + local_versions[config_name] = straxen.URLConfig.format_url_kwargs( + url_config, run_id=corrections_run_id + ) + st.config = local_versions + + # Update some run specific config + for mc_config, processing_config in run_id_specific_config.items(): + if processing_config in st.config: + st.config[mc_config] = st.config[processing_config] + else: + print(f"Warning! {processing_config} not in context config, skipping...") + + # No blinding in simulations + st.config["event_info_function"] = "disabled" + + # Deregister plugins with missing dependencies + st.deregister_plugins_with_missing_dependencies() + + return st + + def set_simulation_config_file(context, config_file_name): """Function to loop over the plugin config and replace SIMULATION_CONFIG_FILE with the actual file name.""" diff --git a/fuse/plugins/__init__.py b/fuse/plugins/__init__.py index cc298898..e190986a 100644 --- a/fuse/plugins/__init__.py +++ b/fuse/plugins/__init__.py @@ -9,3 +9,6 @@ from . import truth_information from .truth_information import * + +from . import fastsim +from .fastsim import * diff --git a/fuse/plugins/fastsim/__init__.py b/fuse/plugins/fastsim/__init__.py new file mode 100644 index 00000000..6505aa27 --- /dev/null +++ b/fuse/plugins/fastsim/__init__.py @@ -0,0 +1,11 @@ +from . import fastsim_macro_clusters +from .fastsim_macro_clusters import * + +from . import fastsim_events_uncorrected +from .fastsim_events_uncorrected import * + +from . import fastsim_s1 +from .fastsim_s1 import * + +from . import fastsim_s2 +from .fastsim_s2 import * diff --git a/fuse/plugins/fastsim/fastsim_events_uncorrected.py b/fuse/plugins/fastsim/fastsim_events_uncorrected.py new file mode 100644 index 00000000..0226e0f4 --- /dev/null +++ b/fuse/plugins/fastsim/fastsim_events_uncorrected.py @@ -0,0 +1,180 @@ +import strax +import numpy as np +import straxen +import logging + +from ...plugin import FuseBasePlugin + +export, __all__ = strax.exporter() + +logging.basicConfig(handlers=[logging.StreamHandler()]) +log = logging.getLogger("fuse.fastsim.fastsim_s1") + + +@export +class FastsimEventsUncorrected(FuseBasePlugin): + """ + Plugin to simulate S1 and (alt) S2 areas from photon hits and electrons extracted + + """ + __version__ = "0.0.1" + + depends_on = ("fastsim_macro_clusters",) + provides = "fastsim_events_uncorrected" + data_kind = "fastsim_events" + dtype = [ + (("S1 area, uncorrected [PE]", "s1_area"), np.float32), + (("S2 area, uncorrected [PE]", "s2_area"), np.float32), + (("Alternate S2 area, uncorrected [PE]", "alt_s2_area"), np.float32), + (("Sum of S2 areas in event, uncorrected [PE]", "s2_sum"), np.float32), + (("Number of S2s in event", "multiplicity"), np.int32), + (("Drift time between main S1 and S2 [ns]", "drift_time"), np.float32), + (("Drift time using alternate S2 [ns]", "alt_s2_interaction_drift_time"), np.float32), + (("Main S2 reconstructed X position, uncorrected [cm]", "s2_x"), np.float32), + (("Main S2 reconstructed Y position, uncorrected [cm]", "s2_y"), np.float32), + (("Alternate S2 reconstructed X position, uncorrected [cm]", "alt_s2_x"), np.float32), + (("Alternate S2 reconstructed Y position, uncorrected [cm]", "alt_s2_y"), np.float32), + (("Main interaction r-position with observed position [cm]", "r_naive"), np.float32), + (("Alternate interaction r-position with observed position [cm]", "alt_s2_r_naive"), np.float32), + (("Main interaction z-position with observed position [cm]", "z_naive"), np.float32), + (("Alternate interaction z-position with observed position [cm]", "alt_s2_z_naive"), np.float32), + ] + strax.time_fields + + save_when = strax.SaveWhen.TARGET + + photon_area_distribution = straxen.URLConfig( + default="simple_load://resource://simulation_config://" + "SIMULATION_CONFIG_FILE.json?" + "&key=photon_area_distribution" + "&fmt=csv", + cache=True, + help="Photon area distribution", + ) + + s2_secondary_sc_gain_mc = straxen.URLConfig( + default="take://resource://" + "SIMULATION_CONFIG_FILE.json?&fmt=json" + "&take=s2_secondary_sc_gain", + type=(int, float), + cache=True, + help="Secondary scintillation gain [PE/e-]", + ) + + se_gain_from_map = straxen.URLConfig( + default="take://resource://" + "SIMULATION_CONFIG_FILE.json?&fmt=json" + "&take=se_gain_from_map", + cache=True, + help="Boolean indication if the secondary scintillation gain is taken from a map", + ) + + se_gain_map = straxen.URLConfig( + default="itp_map://resource://simulation_config://" + "SIMULATION_CONFIG_FILE.json?" + "&key=se_gain_map" + "&fmt=json", + cache=True, + help="Map of the single electron gain", + ) + + s2_correction_map = straxen.URLConfig( + default="itp_map://resource://simulation_config://" + "SIMULATION_CONFIG_FILE.json?" + "&key=s2_correction_map" + "&fmt=json", + cache=True, + help="S2 correction map", + ) + + p_double_pe_emision = straxen.URLConfig( + default="take://resource://" + "SIMULATION_CONFIG_FILE.json?&fmt=json" + "&take=p_double_pe_emision", + type=(int, float), + cache=True, + help="Probability of double photo-electron emission", + ) + + @staticmethod + def get_s1_area_with_spe(spe_distribution, num_photons): + """ + :params: spe_distribution, the spe distribution to draw photon areas from + :params: num_photons, number of photons to draw from spe distribution + """ + s1_area_spe = [] + for n_ph in num_photons: + s1_area_spe.append(np.sum(spe_distribution[ + (np.random.random(n_ph) * len(spe_distribution)).astype(np.int64)])) + + return np.array(s1_area_spe) + + def compute(self, fastsim_macro_clusters): + eventids = np.unique(fastsim_macro_clusters['eventid']) + result = np.zeros(len(eventids), dtype=self.dtype) + for i, eventid in enumerate(eventids): + these_clusters = fastsim_macro_clusters[fastsim_macro_clusters['eventid'] == eventid] + + result[i]["time"] = these_clusters[0]["time"] + result[i]["endtime"] = these_clusters[0]["endtime"] + + photons = np.sum(these_clusters['n_s1_photon_hits']) + result["s1_area"][i] = photons * 1.28 # TODO: replace 1.28 with correct spe value + + cluster_info = [] + for cluster in these_clusters: + pos = np.array([cluster["x"], cluster["y"]]).T # TODO: check if correct positions + ly = self.get_s2_light_yield(pos)[0] + s2_area = ly * cluster['n_electron_extracted'] + if s2_area > 0: + cluster_info.append((s2_area, cluster)) + + # Sort the clusters by s2_area in descending order + cluster_info_sorted = sorted(cluster_info, key=lambda x: x[0], reverse=True) + + # Assign the highest and second-highest s2_area and drift time values + if len(cluster_info_sorted) > 0: + s2_areas = [info[0] for info in cluster_info_sorted] + result[i]['s2_sum'] = np.sum(s2_areas) + result[i]['s2_area'] = cluster_info_sorted[0][0] + result[i]['drift_time'] = cluster_info_sorted[0][1]['drift_time_mean'] + result[i]['s2_x'] = cluster_info_sorted[0][1]['x_obs'] + result[i]['s2_y'] = cluster_info_sorted[0][1]['y_obs'] + result[i]['z_naive'] = cluster_info_sorted[0][1]['z_obs'] + + if len(cluster_info_sorted) > 1: + result[i]['alt_s2_area'] = cluster_info_sorted[1][0] + result[i]['alt_s2_interaction_drift_time'] = cluster_info_sorted[1][1]['drift_time_mean'] + result[i]['alt_s2_x'] = cluster_info_sorted[1][1]['x_obs'] + result[i]['alt_s2_y'] = cluster_info_sorted[1][1]['y_obs'] + result[i]['alt_s2_z_naive'] = cluster_info_sorted[1][1]['z_obs'] + + result[i]["multiplicity"] = len(cluster_info_sorted) + + result['r_naive'] = np.sqrt(result['s2_x'] ** 2 + result['s2_y'] ** 2) + result['alt_s2_r_naive'] = np.sqrt(result['alt_s2_x'] ** 2 + result['alt_s2_y'] ** 2) + + return result + + def get_s2_light_yield(self, positions): + """Calculate s2 light yield... + + Args: + positions: 2d array of positions (floats) returns array of floats (mean expectation) + """ + + if self.se_gain_from_map: + sc_gain = self.se_gain_map(positions) + else: + # calculate it from MC pattern map directly if no "se_gain_map" is given + sc_gain = self.s2_correction_map(positions) + sc_gain *= self.s2_secondary_sc_gain_mc + + # Depending on if you use the data driven or mc pattern map for light yield for S2 + # The shape of n_photon_hits will change. Mc needs a squeeze + if len(sc_gain.shape) != 1: + sc_gain = np.squeeze(sc_gain, axis=-1) + + # Data driven map contains nan, will be set to 0 here + sc_gain[np.isnan(sc_gain)] = 0 + + return sc_gain diff --git a/fuse/plugins/fastsim/fastsim_macro_clusters.py b/fuse/plugins/fastsim/fastsim_macro_clusters.py new file mode 100644 index 00000000..2ddc0549 --- /dev/null +++ b/fuse/plugins/fastsim/fastsim_macro_clusters.py @@ -0,0 +1,81 @@ +import strax +import numpy as np +from numba import njit +import straxen +import logging +import random + +from ...plugin import FuseBasePlugin + +export, __all__ = strax.exporter() + +logging.basicConfig(handlers=[logging.StreamHandler()]) +log = logging.getLogger("fuse.fastsim.fastsim_macro_clusters") + + +@njit(cache=True) +def get_nn_prediction(inp, w0, w1, w2, w3, w4, w5, w6, w7, w8, w9): + y = np.dot(inp, w0) + w1 + y[y < 0] = 0 + y = np.dot(y, w2) + w3 + y[y < 0] = 0 + y = np.dot(y, w4) + w5 + y[y < 0] = 0 + y = np.dot(y, w6) + w7 + y[y < 0] = 0 + y = np.dot(y, w8) + w9 + y = 1 / (1 + np.exp(-y)) + return y[0] + + +def merge_clusters(cluster1, cluster2): + return random.random() > 0.5 + + +@export +class MacroClusters(FuseBasePlugin): + """Plugin to simulate macro clusters for fastsim + + """ + __version__ = "0.0.1" + + depends_on = ("drifted_electrons", "extracted_electrons", "microphysics_summary", "s1_photon_hits") + provides = "fastsim_macro_clusters" + data_kind = "fastsim_macro_clusters" + + def infer_dtype(self): + return strax.merged_dtype([self.deps[d].dtype_for(d) for d in sorted(self.depends_on)]) + + def compute(self, interactions_in_roi): + for ix1, _ in enumerate(interactions_in_roi): + for ix2 in range(1, len(interactions_in_roi[ix1:])): + if interactions_in_roi[ix1]['eventid'] != interactions_in_roi[ix1 + ix2]['eventid']: + break + if merge_clusters(interactions_in_roi[ix1], interactions_in_roi[ix1 + ix2]): + ne1 = interactions_in_roi[ix1]['n_electron_extracted'] + ne2 = interactions_in_roi[ix1 + ix2]['n_electron_extracted'] + ne_total = ne1 + ne2 + interactions_in_roi[ix1 + ix2]['n_electron_extracted'] = ne_total + interactions_in_roi[ix1]['n_electron_extracted'] = -1 # flag to throw this instruction away later + + for quantity in ['photons', 'electrons', 'excitons', 'ed', 'n_electron_interface', + 'n_s1_photon_hits']: + interactions_in_roi[ix1 + ix2][quantity] += interactions_in_roi[ix1][quantity] + + for quantity in ['drift_time_mean', 'drift_time_spread']: + interactions_in_roi[ix1 + ix2][quantity] += interactions_in_roi[ix1][quantity] + interactions_in_roi[ix1 + ix2][quantity] /= 2 + + if ne_total > 0: + for coord in ['x', 'y', 'z']: + for obs in ['', '_obs']: + interactions_in_roi[ix1 + ix2][f'{coord}{obs}'] = \ + (interactions_in_roi[ix1][f'{coord}{obs}'] * ne1 + + interactions_in_roi[ix1 + ix2][f'{coord}{obs}'] * ne2) / ne_total + + interactions_in_roi[ix1 + ix2]['x_pri'] = interactions_in_roi[ix1]['x_pri'] + interactions_in_roi[ix1 + ix2]['y_pri'] = interactions_in_roi[ix1]['y_pri'] + interactions_in_roi[ix1 + ix2]['z_pri'] = interactions_in_roi[ix1]['z_pri'] + + break + return interactions_in_roi[interactions_in_roi['n_electron_extracted'] >= 0] diff --git a/fuse/plugins/fastsim/fastsim_s1.py b/fuse/plugins/fastsim/fastsim_s1.py new file mode 100644 index 00000000..012a8c0a --- /dev/null +++ b/fuse/plugins/fastsim/fastsim_s1.py @@ -0,0 +1,58 @@ +import strax +import numpy as np +import straxen +import logging + +from ...plugin import FuseBasePlugin + +export, __all__ = strax.exporter() + +logging.basicConfig(handlers=[logging.StreamHandler()]) +log = logging.getLogger("fuse.fastsim.fastsim_s1") + + +@export +class S1Areas(FuseBasePlugin): + """Plugin to simulate macro clusters for fastsim + + """ + __version__ = "0.0.1" + + depends_on = ("fastsim_macro_clusters",) + provides = "fastsim_s1" + data_kind = "fastsim_events" + dtype = [(("S1 area, uncorrected [PE]", "s1_area"), np.float32)] + strax.time_fields + + save_when = strax.SaveWhen.ALWAYS + + photon_area_distribution = straxen.URLConfig( + default="simple_load://resource://simulation_config://" + "SIMULATION_CONFIG_FILE.json?" + "&key=photon_area_distribution" + "&fmt=csv", + cache=True, + help="Photon area distribution", + ) + + @staticmethod + def get_s1_area_with_spe(spe_distribution, num_photons): + """ + :params: spe_distribution, the spe distribution to draw photon areas from + :params: num_photons, number of photons to draw from spe distribution + """ + s1_area_spe = [] + for n_ph in num_photons: + s1_area_spe.append(np.sum(spe_distribution[ + (np.random.random(n_ph) * len(spe_distribution)).astype(np.int64)])) + + return np.array(s1_area_spe) + + def compute(self, fastsim_macro_clusters): + eventids = np.unique(fastsim_macro_clusters['eventid']) + result = np.zeros(len(eventids), dtype=self.dtype) + result["time"] = fastsim_macro_clusters["time"] + result["endtime"] = fastsim_macro_clusters["endtime"] + for i, eventid in enumerate(eventids): + photons = np.sum(fastsim_macro_clusters[fastsim_macro_clusters['eventid'] == eventid]['n_s1_photon_hits']) + result["s1_area"][i] = photons * 1.28 + return result diff --git a/fuse/plugins/fastsim/fastsim_s2.py b/fuse/plugins/fastsim/fastsim_s2.py new file mode 100644 index 00000000..ff374edb --- /dev/null +++ b/fuse/plugins/fastsim/fastsim_s2.py @@ -0,0 +1,121 @@ +import strax +import numpy as np +import straxen +import logging + +from ...plugin import FuseBasePlugin + +export, __all__ = strax.exporter() + +logging.basicConfig(handlers=[logging.StreamHandler()]) +log = logging.getLogger("fuse.fastsim.fastsim_s2") + + +@export +class S2Areas(FuseBasePlugin): + """Plugin to simulate S2 areas from electrons extracted for fastsim + + """ + __version__ = "0.0.1" + + depends_on = ("fastsim_macro_clusters",) + provides = "fastsim_s2" + data_kind = "fastsim_events" + + dtype = [(("S2 area, uncorrected [PE]", "s2_area"), np.float32), + (("Alternate S2 area, uncorrected [PE]", "alt_s2_area"), np.float32), + (("Sum of S2 areas in event, uncorrected [PE]", "s2_sum"), np.float32), + (("Number of S2s in event", "multiplicity"), np.int32) + ] + strax.time_fields + + save_when = strax.SaveWhen.ALWAYS + + s2_secondary_sc_gain_mc = straxen.URLConfig( + default="take://resource://" + "SIMULATION_CONFIG_FILE.json?&fmt=json" + "&take=s2_secondary_sc_gain", + type=(int, float), + cache=True, + help="Secondary scintillation gain [PE/e-]", + ) + + se_gain_from_map = straxen.URLConfig( + default="take://resource://" + "SIMULATION_CONFIG_FILE.json?&fmt=json" + "&take=se_gain_from_map", + cache=True, + help="Boolean indication if the secondary scintillation gain is taken from a map", + ) + + se_gain_map = straxen.URLConfig( + default="itp_map://resource://simulation_config://" + "SIMULATION_CONFIG_FILE.json?" + "&key=se_gain_map" + "&fmt=json", + cache=True, + help="Map of the single electron gain", + ) + + s2_correction_map = straxen.URLConfig( + default="itp_map://resource://simulation_config://" + "SIMULATION_CONFIG_FILE.json?" + "&key=s2_correction_map" + "&fmt=json", + cache=True, + help="S2 correction map", + ) + + p_double_pe_emision = straxen.URLConfig( + default="take://resource://" + "SIMULATION_CONFIG_FILE.json?&fmt=json" + "&take=p_double_pe_emision", + type=(int, float), + cache=True, + help="Probability of double photo-electron emission", + ) + + def compute(self, fastsim_macro_clusters): + eventids = np.unique(fastsim_macro_clusters['eventid']) + result = np.zeros(len(eventids), dtype=self.dtype) + result["time"] = fastsim_macro_clusters["time"] + result["endtime"] = fastsim_macro_clusters["endtime"] + for i, eventid in enumerate(eventids): + clusters = fastsim_macro_clusters[fastsim_macro_clusters["eventid"] == eventid] + s2_areas = [] + for cluster in clusters: + pos = np.array([cluster["x"], cluster["y"]]).T # TODO: check if correct positions + ly = self.get_s2_light_yield(pos)[0] + area = ly * cluster['n_electron_extracted'] + if area > 0: + s2_areas.append(area) + if len(s2_areas): + result[i]["s2_area"] = max(s2_areas) + result[i]["s2_sum"] = sum(s2_areas) + if len(s2_areas) > 1: + result[i]["alt_s2_area"] = sorted(s2_areas, reverse=True)[1] + result[i]["multiplicity"] = len(s2_areas) + return result + + def get_s2_light_yield(self, positions): + """Calculate s2 light yield... + + Args: + positions: 2d array of positions (floats) returns array of floats (mean expectation) + """ + + if self.se_gain_from_map: + sc_gain = self.se_gain_map(positions) + else: + # calculate it from MC pattern map directly if no "se_gain_map" is given + sc_gain = self.s2_correction_map(positions) + sc_gain *= self.s2_secondary_sc_gain_mc + + # Depending on if you use the data driven or mc pattern map for light yield for S2 + # The shape of n_photon_hits will change. Mc needs a squeeze + if len(sc_gain.shape) != 1: + sc_gain = np.squeeze(sc_gain, axis=-1) + + # Data driven map contains nan, will be set to 0 here + sc_gain[np.isnan(sc_gain)] = 0 + + return sc_gain