diff --git a/HISTORY.md b/HISTORY.md index 5a8393d..cbc2f6d 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -3,11 +3,13 @@ * Fix calculation for cluster weighted average time (#65) * Add zenodo badge + 0.3.3 (2023-01-09) ================== * Specify not working numba in req. file (#62) * Install scikit-learn (#63) + 0.3.2 (2022-08-15) ================== * First implementation of BBF quanta generator (#57) diff --git a/epix/simulator/FastSimulator.py b/epix/simulator/FastSimulator.py new file mode 100644 index 0000000..2a25c6b --- /dev/null +++ b/epix/simulator/FastSimulator.py @@ -0,0 +1,125 @@ +import time +import numpy as np +import pandas as pd +import os +from .GenerateEvents import GenerateEvents +from .helpers import Helpers + +# 2023-02-19: configuration_files: +# 'nv_pmt_qe':'nveto_pmt_qe.json', +# 'photon_area_distribution':'XENONnT_SR0_spe_distributions_20210713_no_noise_scaled.csv', +# 's1_pattern_map': 'XENONnT_s1_xyz_patterns_LCE_MCvf051911_wires.pkl', +# 's2_pattern_map': 'XENONnT_s2_xy_patterns_GXe_LCE_corrected_qes_MCv4.3.0_wires.pkl', + +def monitor_time(prev_time, task): + t = time.time() + print(f'It took {(t - prev_time):.4f} sec to {task}') + return t + + +class FastSimulator: + """Simulator class for epix to go from epix instructions to fully processed data""" + + def __init__(self, instructions_epix, config, resource): + self.config = config + self.ge = GenerateEvents() + self.resource = resource + self.simulation_functions = sorted( + # get the functions of GenerateEvents in order to run through them all + [getattr(self.ge, field) for field in dir(self.ge) + if hasattr(getattr(self.ge, field), "order") + ], + # sort them by their order + key=(lambda field: field.order) + ) + self.instructions_epix = instructions_epix + + def cluster_events(self, ): + """Events have more than 1 s1/s2. Here we throw away all of them except the largest 2 + Take the position to be that of the main s2 + And strax wants some start and endtime, so we make something up + + First do the macro clustering. Clustered instructions will be flagged with amp=-1, + so we can safely throw those out""" + start_time = time.time() + + Helpers.macro_cluster_events(self.instructions_epix, self.config, self.resource) + self.instructions_epix = self.instructions_epix[self.instructions_epix['amp'] != -1] + + event_numbers = np.unique(self.instructions_epix['event_number']) + ins_size = len(event_numbers) + instructions = np.zeros(ins_size, dtype=Helpers.get_dtype()['events_tpc']) + + for ix in range(ins_size): + i = instructions[ix] + inst = self.instructions_epix[self.instructions_epix['event_number'] == event_numbers[ix]] + inst_s1 = inst[inst['type'] == 1] + inst_s2 = inst[inst['type'] == 2] + + s1 = np.argsort(inst_s1['amp']) + s2 = np.argsort(inst_s2['amp']) + + if len(s1) < 1 or len(s2) < 1: + continue + + i['s1_area'] = np.sum(inst_s1['amp']) + + i['s2_area'] = inst_s2[s2[-1]]['amp'] + i['e_dep'] = inst_s2[s2[-1]]['e_dep'] + + if len(s2) > 1: + i['alt_s2_area'] = inst_s2[s2[-2]]['amp'] + i['alt_e_dep'] = inst_s2[s2[-2]]['e_dep'] + i['alt_s2_x_true'] = inst_s2[s2[-2]]['x'] + i['alt_s2_y_true'] = inst_s2[s2[-2]]['y'] + i['alt_s2_z_true'] = inst_s2[s2[-2]]['z'] + + i['x_true'] = inst_s2[s2[-1]]['x'] + i['y_true'] = inst_s2[s2[-1]]['y'] + i['z_true'] = inst_s2[s2[-1]]['z'] + + i['x_pri'] = inst_s2[s2[-1]]['x_pri'] + i['y_pri'] = inst_s2[s2[-1]]['y_pri'] + i['z_pri'] = inst_s2[s2[-1]]['z_pri'] + + i['g4id'] = inst_s2[s2[-1]]['g4id'] + + # Strax wants begin and endtimes + i['time'] = ix * 1000 + i['endtime'] = ix * 1000 + 1 + + instructions = instructions[~((instructions['time'] == 0) & (instructions['endtime'] == 0))] + self.instructions = instructions + print(f'It took {(time.time() - start_time):.4f} sec to macro cluster events') + + def simulate(self, ): + for func in self.simulation_functions: + start = time.time() + func(i=self.instructions, + config=self.config, + resource=self.resource) + monitor_time(start, func.__name__) + + def run_simulator(self, ): + self.cluster_events() + # TODO this for debug purposes - not super nice. We should have a fastsim_truth data type + file_name = self.config['epix_config']['file_name'].split('/')[-1] + file_name = file_name.split('.')[0] + file_name += '_instruction_after_macro_clustering.csv' + save_epix = self.config['epix_config'].get('save_epix', False) + if save_epix: + if isinstance(save_epix, str): + # assume save epix as path to store + epix_path = os.path.join(self.config['epix_config']['save_epix'], file_name) + else: + # if save epix True store in normal path + epix_path = os.path.join(self.config['epix_config']['path'], file_name) + print(f'Saving epix instruction after macro clustering: {epix_path}') + df = pd.DataFrame(self.instructions) + df.to_csv(epix_path, index=False) + self.simulate() + return self.instructions + + + + diff --git a/epix/simulator/GenerateEvents.py b/epix/simulator/GenerateEvents.py index 32e348b..7f36215 100644 --- a/epix/simulator/GenerateEvents.py +++ b/epix/simulator/GenerateEvents.py @@ -1,145 +1,196 @@ from .helpers import Helpers import numpy as np +class GenerateEvents: + """Class to hold all the stuff to be applied to data. + The functions will be grouped together and executed by Simulator""" + + @staticmethod + @Helpers.assign_order(0) + def make_drift_time(i, config, resource): + """ + Calculate drift_time and alt_s2_interaction_drift_time + :params: i, numpy array with instructions of events_tpc dtype + :params: config, dict with configuration values for resolution + :params: resource, instance of wfsim Resource class + """ + for alt_s2, alt_s2_int in [('', ''), ('alt_s2_', 'alt_s2_interaction_')]: + i[f'{alt_s2_int}drift_time'] = Helpers.get_drift_time(i[f'{alt_s2}z_true'], + np.array([i[f'{alt_s2}x_true'], i[f'{alt_s2}y_true']]).T, + config, resource) + + @staticmethod + @Helpers.assign_order(1) + def get_true_polar_coordinates(i, config, resource): + for alt_s2, alt in [('', ''), ('alt_s2_', 'alt_')]: + i[f'{alt_s2}r_true'] = np.sqrt(i[f'{alt_s2}x_true'] ** 2 + i[f'{alt_s2}y_true'] ** 2) + i[f'{alt_s2}theta_true'] = np.arctan2(i[f'{alt_s2}y_true'], i[f'{alt_s2}x_true']) + + @staticmethod + @Helpers.assign_order(2) + def get_naive_positions(i, config, resource): + """ + Get the uncorrected observed positions s2_x, s2_y in gas gap based on true position + as well as z_naive from drift time + :params: i, numpy array with instructions of events_tpc dtype + :params: config, dict with configuration values for resolution + :params: resource, instance of wfsim Resource class + """ + v_drift = config['drift_velocity_liquid'] # cm/ns + for alt_s2_interaction, alt_s2, alt in [('', '', ''), ('alt_s2_interaction_', 'alt_s2_', 'alt_')]: + i[f'{alt_s2}r_naive'] = resource.fd_comsol(np.array([i[f'{alt_s2}r_true'], i[f'{alt_s2}z_true']]).T, + map_name='r_distortion_map') + i[f'{alt}s2_x'] = i[f'{alt_s2}r_naive'] * np.cos(i[f'{alt_s2}theta_true']) + i[f'{alt}s2_y'] = i[f'{alt_s2}r_naive'] * np.sin(i[f'{alt_s2}theta_true']) + + i[f'{alt_s2}z_naive'] = -(i[f'{alt_s2_interaction}drift_time'] - config['drift_time_gate']) * v_drift # in cm -class GenerateEvents(): - '''Class to hold all the stuff to be applied to data. - The functions will be grouped together and executed by Simulator''' @staticmethod - @Helpers.assignOrder(0) - def smear_positions(instructions, config, **kwargs): + @Helpers.assign_order(3) + def smear_positions(i, config, resource): """Take initial positions and apply gaussian smearing with some resolution to get the measured position :params: instructions, numpy array with instructions of events_tpc dtype :params: config, dict with configuration values for resolution """ - instructions['x'] = np.random.normal(instructions['x'], config['xy_resolution']) - instructions['y'] = np.random.normal(instructions['y'], config['xy_resolution']) - instructions['z'] = np.random.normal(instructions['z'], config['z_resolution']) - - instructions['alt_s2_x'] = np.random.normal(instructions['alt_s2_x'], config['xy_resolution']) - instructions['alt_s2_y'] = np.random.normal(instructions['alt_s2_y'], config['xy_resolution']) - instructions['alt_s2_z'] = np.random.normal(instructions['alt_s2_z'], config['z_resolution']) + # TODO check if this makes sense. Can we get the resolutions from somewhere? + for alt_s2, alt in [('', ''), ('alt_s2_', 'alt_')]: + i[f'{alt}s2_x'] = np.random.normal(i[f'{alt}s2_x'], config['xy_resolution']) + i[f'{alt}s2_y'] = np.random.normal(i[f'{alt}s2_y'], config['xy_resolution']) + i[f'{alt_s2}z_naive'] = np.random.normal(i[f'{alt_s2}z_naive'], config['z_resolution']) @staticmethod - @Helpers.assignOrder(1) - def make_s1(instructions, config, resource): - """Build the s1s. Takes number of quanta and calcultes the (alt) s1 area using wfsim + @Helpers.assign_order(4) + def get_corrected_positions(i, config, resource): + """ + Apply FDC to observed positions s2_x, s2_y to get x, y :params: instructions, numpy array with instructions of events_tpc dtype :params: config, dict with configuration values for resolution :params: resource, instance of wfsim Resource class """ - xyz = np.vstack([instructions['x'], instructions['y'], instructions['z']]).T + for alt_s2, alt, fdc in [('','', ''), ('alt_s2_', 'alt_', '_fdc')]: + + i[f'{alt_s2}r_field_distortion_correction'] = resource.fdc_map(np.array([i[f'{alt}s2_x'], i[f'{alt}s2_y'], + i[f'{alt_s2}z_naive']]).T) + with np.errstate(invalid='ignore', divide='ignore'): + r_cor = i[f'{alt_s2}r_naive'] + i[f'{alt_s2}r_field_distortion_correction'] + scale = np.divide(r_cor, i[f'{alt_s2}r_naive'], out=np.zeros_like(r_cor), where=i[f'{alt_s2}r_naive'] != 0) + + #i[f'{alt_s2}r{fdc}'] = i[f'{alt_s2}r_naive'] + i[f'{alt_s2}r_field_distortion_correction'] + #i[f'{alt_s2}x{fdc}'] = i[f'{alt_s2}r{fdc}']*np.cos(i[f'{alt_s2}theta_true']) + #i[f'{alt_s2}y{fdc}'] = i[f'{alt_s2}r{fdc}']*np.sin(i[f'{alt_s2}theta_true']) + i[f'{alt_s2}r{fdc}'] = r_cor + i[f'{alt_s2}x{fdc}'] = i[f'{alt}s2_x'] * scale + i[f'{alt_s2}y{fdc}'] = i[f'{alt}s2_y'] * scale + i[f'{alt_s2}theta'] = np.arctan2(i[f'{alt_s2}y{fdc}'], i[f'{alt_s2}x{fdc}']) + + + i[f'{alt_s2}z_dv_corr'] = resource.fd_comsol(np.array([i[f'{alt_s2}r_true'], i[f'{alt_s2}z_true']]).T, + map_name='z_distortion_map') + with np.errstate(invalid='ignore'): + z_cor = -(i[f'{alt_s2}z_naive'] ** 2 - i[f'{alt_s2}r_field_distortion_correction'] ** 2) ** 0.5 + invalid = np.abs(i[f'{alt_s2}z_naive']) < np.abs(i[f'{alt_s2}r_field_distortion_correction'] ) + # do not apply z correction above gate + invalid |= i[f'{alt_s2}z_naive'] >= 0 + z_cor[invalid] = i[f'{alt_s2}z_naive'][invalid] + i[f'{alt_s2}z_field_distortion_correction'] = z_cor - i[f'{alt_s2}z_naive'] + i[f'{alt_s2}z'] = i[f'{alt_s2}z_naive'] + - num_ph = instructions['s1_area'].astype(np.int64) + + + @staticmethod + @Helpers.assign_order(5) + def make_s1(i, config, resource): + """Build the s1s. Takes number of quanta and calculates the (alt) s1 area using wfsim + :params: i, numpy array with instructions of events_tpc dtype + :params: config, dict with configuration values for resolution + :params: resource, instance of wfsim Resource class + """ + xyz = np.vstack([i['x_true'], i['y_true'], i['z_true']]).T + + num_ph = i['s1_area'].astype(np.int64) # Here a WFsim function is called which remove the dpe # we have to introduce it again in fast simulator n_photons = Helpers.get_s1_light_yield(n_photons = num_ph, positions = xyz, - s1_lce_map = resource.s1_map, + s1_lce_map = resource.s1_lce_correction_map, config = config) * (1 + config['p_double_pe_emision']) - instructions['s1_area'] = Helpers.get_s1_area_with_spe(resource.photon_area_distribution, + i['s1_area'] = Helpers.get_s1_area_with_spe(resource.mean_photon_area_distribution, n_photons.astype(np.int64)) - - num_ph = instructions['alt_s1_area'].astype(np.int64) - alt_n_photons = Helpers.get_s1_light_yield(n_photons = num_ph, - positions = xyz, - s1_lce_map = resource.s1_map, - config = config) * (1 + config['p_double_pe_emision']) - - instructions['alt_s1_area'] = Helpers.get_s1_area_with_spe(resource.photon_area_distribution, - alt_n_photons.astype(np.int64)) @staticmethod - @Helpers.assignOrder(2) - def make_s2(instructions, config, resource): - """Call functions from wfsim to drift electrons. Since the s2 light yield implementation is a little bad how to do that? - Make sc_gain factor 11 to large to correct for this? Also whats the se gain variation? Lets take sqrt for now - :params: instructions, numpy array with instructions of events_tpc dtype - :params: config, dict with configuration values for resolution - :params: resource, instance of wfsim Resource class + @Helpers.assign_order(6) + def make_s2_area(i, config, resource): """ - xy = np.array([instructions['x'], instructions['y']]).T - alt_xy = np.array([instructions['alt_s2_x'], instructions['alt_s2_y']]).T - - n_el = instructions['s2_area'].astype(np.int64) - - n_electron = Helpers.get_s2_charge_yield(n_electron = n_el, - positions = xy, - z_obs = instructions['z'], - config = config, - resource = resource) - - n_el = instructions['alt_s2_area'].astype(np.int64) - alt_n_electron = Helpers.get_s2_charge_yield(n_electron=n_el, - positions=alt_xy, - z_obs=instructions['z'], - config=config, - resource=resource) - - sc_gain = Helpers.get_s2_light_yield(positions=xy, - config=config, - resource=resource) - sc_gain_sigma = np.sqrt(sc_gain) - - # Here a WFsim function is called which remove the dpe - # we have to introduce it again in fast simulator - instructions['s2_area'] = n_electron * np.random.normal(sc_gain, sc_gain_sigma) * (1 + config['p_double_pe_emision']) - instructions['drift_time'] = -instructions['z'] / config['drift_velocity_liquid'] + Call functions from wfsim to drift electrons. Since the s2 light yield implementation is a little bad how to do that? + Make sc_gain factor 11 too large to correct for this? Also, what's the se gain variation? Let's take sqrt for now + :params: i, numpy array with instructions of events_tpc dtype + :params: config, dict with configuration values for resolution + :params: resource, instance of wfsim Resource class + """ + for alt_s2, alt in [('', ''), ('alt_s2_', 'alt_')]: + xy_true = np.array([i[f'{alt_s2}x_true'], i[f'{alt_s2}y_true']]).T + xy = np.array([i[f'{alt}s2_x'], i[f'{alt}s2_y']]).T + n_el = i[f'{alt}s2_area'].astype(np.int64) + n_electron = Helpers.get_s2_charge_yield(n_electron = n_el, + xy = xy_true, + z = i[f'{alt_s2}z_true'], + config = config, + resource = resource) + # Here a WFsim function is called which removes the dpe + # we have to introduce it again in fast simulator + sc_gain = Helpers.get_s2_light_yield(positions=xy, + config=config, + resource=resource) + sc_gain *= (1 + config['p_double_pe_emision']) # intentional typo + i[f'{alt}s2_area'] = n_electron * sc_gain - instructions['alt_s2_area'] = alt_n_electron * np.random.normal(sc_gain, sc_gain_sigma) * (1 + config['p_double_pe_emision']) - instructions['alt_s2_drift_time'] = -instructions['alt_s2_z'] / config['drift_velocity_liquid'] @staticmethod - @Helpers.assignOrder(3) - def correction_s1(instructions, resource, **kwargs): - """" - Calculates (alt)cs1. Method taken from CorrectedAreas in straxen - :params: instructions, numpy array with instructions of events_tpc dtype + @Helpers.assign_order(7) + def correction_s1(i, config, resource): + """ + Calculates cs1. Method taken from CorrectedAreas in straxen + :params: i, numpy array with instructions of events_tpc dtype :params: config, dict with configuration values for resolution :params: resource, instance of wfsim Resource class """ - event_positions = np.vstack([instructions['x'], instructions['y'], instructions['z']]).T + event_positions = np.vstack([i['x_true'], i['y_true'], i['z_true']]).T - # Where does 0.15 come from ? Mean of s1_map in -130 < z < -20 and r < 50 + # Where does 0.15 come from ? Mean of s1_lce_correction_map in -130 < z < -20 and r < 50 # the_map = [] - # for xyz in resource.s1_map.coordinate_system: + # for xyz in resource.s1_lce_correction_map.coordinate_system: # r = np.sqrt(xyz[0]**2 + xyz[1]**2) # if (-130 < xyz[2] < -20) & ( r < 50): - # the_map.append(np.squeeze(resource.s1_map(xyz))[0]) + # the_map.append(np.squeeze(resource.s1_lce_correction_map(xyz))[0]) # print(np.mean(the_map)) - instructions['cs1'] = instructions['s1_area'] / (resource.s1_map(event_positions)[:, 0]/0.1581797073725071) - instructions['alt_cs1'] = instructions['alt_s1_area'] / (resource.s1_map(event_positions)[:, 0]/0.1581797073725071) + i['cs1'] = i['s1_area'] / (resource.s1_lce_correction_map(event_positions)[:, 0]/0.1581797073725071) @staticmethod - @Helpers.assignOrder(4) - def correction_s2(instructions, config, resource): + @Helpers.assign_order(8) + def correction_s2(i, config, resource): """" Calculates (alt)cs2. Method taken from CorrectedAreas in straxen - :params: instructions, numpy array with instructions of events_tpc dtype + :params: i, numpy array with instructions of events_tpc dtype :params: config, dict with configuration values for resolution :params: resource, instance of wfsim Resource class """ + for alt_s2_interaction, alt in [('', ''), ('alt_s2_interaction_', 'alt_')]: + lifetime_corr = np.exp(i[f'{alt_s2_interaction}drift_time'] / config['electron_lifetime_liquid']) - lifetime_corr = np.exp(instructions['drift_time'] / config['electron_lifetime_liquid']) - - alt_lifetime_corr = (np.exp((instructions['alt_s2_drift_time']) / config['electron_lifetime_liquid'])) - - # S2(x,y) corrections use the observed S2 positions - s2_positions = np.vstack([instructions['x'], instructions['y']]).T - alt_s2_positions = np.vstack([instructions['alt_s2_x'], instructions['alt_s2_y']]).T + # S2(x,y) corrections use the observed S2 positions + xy = np.vstack([i[f'{alt}s2_x'], i[f'{alt}s2_y']]).T - # Why S2 does not need the same threatment of S1 ? - instructions['cs2'] = (instructions['s2_area'] * lifetime_corr / resource.s2_map(s2_positions)) - instructions['alt_cs2'] = ( - instructions['alt_s2_area'] * alt_lifetime_corr / resource.s2_map(alt_s2_positions)) + # Why S2 does not need the same treatment of S1 ? + i[f'{alt}cs2'] = (i[f'{alt}s2_area'] * lifetime_corr / resource.s2_correction_map(xy)) - alt_s2_nan = instructions['alt_s2_area'] < 1e-6 - instructions['alt_s2_x'][alt_s2_nan] = 0.0 - instructions['alt_s2_y'][alt_s2_nan] = 0.0 - instructions['alt_s2_z'][alt_s2_nan] = 0.0 - instructions['alt_s2_drift_time'][alt_s2_nan] = 0.0 + #alt_s2_nan = i['alt_s2_area'] < 1e-6 + #i['alt_s2_x'][alt_s2_nan] = 0.0 + #i['alt_s2_y'][alt_s2_nan] = 0.0 + #i['alt_s2_z'][alt_s2_nan] = 0.0 + #i['alt_s2_drift_time'][alt_s2_nan] = 0.0 diff --git a/epix/simulator/StraxSimulator.py b/epix/simulator/StraxSimulator.py new file mode 100644 index 0000000..8bddae1 --- /dev/null +++ b/epix/simulator/StraxSimulator.py @@ -0,0 +1,181 @@ +import strax +import straxen +import wfsim.load_resource +import epix +import numpy as np +import pandas as pd +import os +from copy import deepcopy +from immutabledict import immutabledict +import inspect +import pickle +from .GenerateNveto import NVetoUtils +from .FastSimulator import FastSimulator +from .helpers import Helpers +import warnings + +# We should think a bit more to detector_config_override +# and tell fast_sim to look into epix_args +# also because one entry of self.config is epix_config +@strax.takes_config( + strax.Option('detector', default='XENONnT', help='Detector model'), + strax.Option('detector_config_override', default='', help='For the electric field, otherwise 200 V/cm'), + strax.Option('epix_config', default=dict(), help='Configuration file for epix', ), + strax.Option('configuration_files', default=dict(), help='Files required for simulating'), + strax.Option('fax_config', help='Fax configuration to load'), + strax.Option('fax_config_overrides', help='Fax configuration to override', default=None), + strax.Option('fax_config_override_from_cmt', default=None, infer_type=False, + help="Dictionary of fax parameter names (key) mapped to CMT config names (value) " + "where the fax parameter values will be replaced by CMT"), + strax.Option('xy_resolution', default=0, help='xy position resolution (cm)'), + strax.Option('z_resolution', default=0, help='xy position resolution (cm)'), + strax.Option('nv_spe_resolution', default=0.40, help='nVeto SPE resolution'), + strax.Option('nv_spe_res_threshold', default=0.50, help='nVeto SPE acceptance threshold'), + strax.Option('nv_max_time_ns', default=1e7, help='nVeto maximum time for the acceptance of PMT hits in event'), + strax.Option('s2_clustering_algorithm', default='mlp', help='clustering algorithm for S2, [ nsort | naive | mlp ]'), +) +class StraxSimulator(strax.Plugin): + provides = ('events_tpc', 'events_nveto') + depends_on = () + data_kind = immutabledict(zip(provides, provides)) + dtype = Helpers.get_dtype() + + def load_resources(self): + """First load the config through wfsim, then we add some things we'd like to have""" + self.resource = wfsim.load_resource.load_config(self.config) + # add mean photon area distribution with avg values. + self.resource.mean_photon_area_distribution = epix.Helpers.average_spe_distribution( + self.resource.photon_area_distribution) + self.resource.fdc_map = wfsim.load_resource.make_map(self.config['field_distortion_correction_map'], + fmt='json.gz') + working_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + self.resource.nn_weights = pickle.load(open(f'{working_dir}/nn_weights.p', 'rb+')) + if 'nv_pmt_qe' in self.config['configuration_files'].keys(): + self.resource.nv_pmt_qe = straxen.get_resource( + self.config['configuration_files']['nv_pmt_qe'], fmt='json') + else: + warnings.warn('The configuration_files should not exist!' + 'Everything should come by one config!' + 'Since nv_pmt_qe config is missing in configuration_files, ' + 'the default one nveto_pmt_qe.json will be used') + self.resource.nv_pmt_qe = straxen.get_resource('nveto_pmt_qe.json', fmt='json') + + def setup(self): + print('Setup') + overrides = self.config['fax_config_overrides'] + self.config.update(straxen.get_resource(self.config['fax_config'], fmt='json')) + if overrides is not None: + self.config.update(overrides) + if self.config['fax_config_override_from_cmt'] is not None: + for fax_field, cmt_option in self.config['fax_config_override_from_cmt'].items(): + if fax_field in ['fdc_3d', 's1_lce_correction_map'] and self.config.get('default_reconstruction_algorithm', False): + cmt_option = tuple(['suffix', self.config['default_reconstruction_algorithm'], *cmt_option]) + cmt_value = straxen.get_correction_from_cmt(self.run_id, cmt_option) + self.config[fax_field] = cmt_value + + # TODO: this should not be needed here. Where is this info normally taken from? + if 'gains' not in self.config.keys(): + self.config['gains'] = [1] * 494 + if 'n_top_pmts' not in self.config.keys(): + self.config['n_top_pmts'] = 253 + if 'n_tpc_pmts' not in self.config.keys(): + self.config['n_tpc_pmts'] = 494 + print('Loading resources') + self.load_resources() + + def get_nveto_data(self, ): + print('Getting nveto data') + file_loader = epix.io.file_loader(directory=self.config['epix_config']['path'], + file_name=self.config['epix_config']['file_name']) + file_tree, _ = file_loader._get_ttree() + + if 'pmthitID' in file_tree.keys(): + nv_hits = NVetoUtils.get_nv_hits(ttree=file_tree, + pmt_nv_json_dict=self.resource.nv_pmt_qe, + nveto_dtype=self.dtype['events_nveto'], + SPE_Resolution=self.config['nv_spe_resolution'], + SPE_ResThreshold=self.config['nv_spe_res_threshold'], + max_time_ns=self.config['nv_max_time_ns'], + batch_size=10000) + return nv_hits + + def save_epix_instruction(self, epix_ins, save_epix, config): + file_name = config['epix_config']['file_name'].split('/')[-1] + file_name = file_name.split('.')[0] + file_name += '_instructions.csv' + if isinstance(save_epix, str): + # assume save epix as path to store + epix_path = os.path.join(config['epix_config']['save_epix'], file_name) + else: + # if save epix True store in normal path + epix_path = os.path.join(config['epix_config']['path'], file_name) + print(f'Saving epix instruction: {epix_path}') + df = pd.DataFrame(epix_ins) + df.to_csv(epix_path, index=False) + + def get_epix_instructions(self, ): + epix_config = deepcopy(self.config['epix_config']) + fn = epix_config.get('file_name', '') + if fn.endswith('.csv'): + csv_file_path = os.path.join(epix_config['path'], epix_config['file_name']) + print(f'Loading epix instructions from csv-file from {csv_file_path}') + epix_ins = pd.read_csv(csv_file_path) + epix_ins = np.array(epix_ins.to_records(index=False)) + elif fn.endswith('.root'): + print('Generating epix instructions from root-file') + + detector = epix.init_detector(self.config['detector'].lower(), self.config['detector_config_override']) + epix_config['detector_config'] = detector + + outer_cylinder = getattr(epix.detectors, 'xenonnt') + _, outer_cylinder = outer_cylinder() + epix_config['outer_cylinder'] = outer_cylinder + epix_ins = epix.run_epix.main(epix_config, return_wfsim_instructions=True) + else: + print('No valid file_name! must be .root (Geant4 file) or .csv (epix instructions)') + return + return epix_ins + + def compute(self): + print('Compute') + #simulated_data_nveto = self.get_nveto_data() + simulated_data_nveto = None + self.epix_instructions = self.get_epix_instructions() + save_epix = self.config['epix_config'].get('save_epix', False) + if save_epix: + self.save_epix_instruction(self.epix_instructions, save_epix, self.config) + + self.Simulator = FastSimulator(instructions_epix=self.epix_instructions, + config=self.config, + resource=self.resource) + simulated_data = self.Simulator.run_simulator() + + simulated_data_chunk = self.chunk( + start=simulated_data['time'][0], + end=simulated_data['endtime'][-1], + data=simulated_data, + data_type='events_tpc') + + # write empty chunk if nveto data isn't there return {'events_tpc':simulated_data_chunk} + if simulated_data_nveto is None or len(simulated_data_nveto['time']) < 1: + simulated_data_nveto_chunk = self.chunk( + start=0, + end=1, + data=simulated_data_nveto, + data_type='events_nveto') + else: + simulated_data_nveto_chunk = self.chunk( + start=np.floor(simulated_data_nveto['time'][0]).astype(np.int64), + end=np.ceil(np.max(simulated_data_nveto['endtime'])).astype(np.int64), + data=simulated_data_nveto, + data_type='events_nveto') + + return {'events_tpc': simulated_data_chunk, + 'events_nveto': simulated_data_nveto_chunk} + + def is_ready(self, chunk): + # For this plugin we'll smash everything into 1 chunk, should be oke + return True if chunk == 0 else False + + def source_finished(self): + return True \ No newline at end of file diff --git a/epix/simulator/__init__.py b/epix/simulator/__init__.py index a0fdf66..90ac430 100644 --- a/epix/simulator/__init__.py +++ b/epix/simulator/__init__.py @@ -1,4 +1,4 @@ -from .fast_simulator import * +from .FastSimulator import * from .GenerateEvents import * from .GenerateNveto import * from .helpers import * \ No newline at end of file diff --git a/epix/simulator/fast_simulator.py b/epix/simulator/fast_simulator.py deleted file mode 100644 index 773a227..0000000 --- a/epix/simulator/fast_simulator.py +++ /dev/null @@ -1,309 +0,0 @@ -import strax -import straxen -from straxen import InterpolatingMap, get_resource -from wfsim.load_resource import load_config -import epix -import numpy as np -import pandas as pd -import os -from copy import deepcopy -from immutabledict import immutabledict -from .GenerateEvents import GenerateEvents -from .GenerateNveto import NVetoUtils -from .helpers import Helpers -import warnings -import pickle - -# 2023-02-19: configuration_files: -# 'nv_pmt_qe':'nveto_pmt_qe.json', -# 'photon_area_distribution':'XENONnT_SR0_spe_distributions_20210713_no_noise_scaled.csv', -# 's1_pattern_map': 'XENONnT_s1_xyz_patterns_LCE_MCvf051911_wires.pkl', -# 's2_pattern_map': 'XENONnT_s2_xy_patterns_GXe_LCE_corrected_qes_MCv4.3.0_wires.pkl', -# 's2_separation_bdt': 's2_separation_decision_tree_fast_sim.p' -# (FROM: /dali/lgrandi/jgrigat/s2_separation/s2_separation_decision_tree_fast_sim.p) - -class Simulator(): - '''Simulator class for epix to go from epix instructions to fully processed data''' - - def __init__(self, instructions_epix, config, resource): - self.config = config - self.ge = GenerateEvents() - self.resource = resource - self.simulation_functions = sorted( - # get the functions of GenerateEvents in order to run through them all - [getattr(self.ge, field) for field in dir(self.ge) - if hasattr(getattr(self.ge, field), "order") - ], - # sort them by their order - key=(lambda field: field.order) - ) - self.instructions_epix = instructions_epix - - #print(f"\n\nSimulator : Current directory [ {os.getcwd()} ]") - self.tree = pickle.load(open(self.config['configuration_files']['s2_separation_bdt'], 'rb+')) - - - def cluster_events(self, ): - """Events have more than 1 s1/s2. Here we throw away all of them except the largest 2 - Take the position to be that of the main s2 - And strax wants some start and endtime, so we make something up - - First do the macro clustering. Clustered instructions will be flagged with amp=-1 - so we can safely through those out""" - - Helpers.macro_cluster_events(self.tree, self.instructions_epix, self.config) - self.instructions_epix = self.instructions_epix[self.instructions_epix['amp'] != -1] - - # Why is it called for the second time?? - #Helpers.macro_cluster_events(self.tree, self.instructions_epix, self.config) - #self.instructions_epix = self.instructions_epix[self.instructions_epix['amp'] != -1] - - event_numbers = np.unique(self.instructions_epix['event_number']) - ins_size = len(event_numbers) - instructions = np.zeros(ins_size, dtype=StraxSimulator.dtype['events_tpc']) - - for ix in range(ins_size): - i = instructions[ix] - inst = self.instructions_epix[self.instructions_epix['event_number'] == event_numbers[ix]] - inst_s1 = inst[inst['type'] == 1] - inst_s2 = inst[inst['type'] == 2] - - s1 = np.argsort(inst_s1['amp']) - s2 = np.argsort(inst_s2['amp']) - - if len(s1) < 1 or len(s2) < 1: - continue - - - # why we consider only the larger s1 ? - #i['s1_area'] = inst_s1[s1[-1]]['amp'] - #if len(s1) > 1: - # I do not think it is correct - # i['alt_s1_area'] = inst_s1[s1[-2]]['amp'] - - i['s2_area'] = inst_s2[s2[-1]]['amp'] - i['e_dep'] = inst_s2[s2[-1]]['e_dep'] - - if len(s2) > 1: - i['alt_s2_area'] = inst_s2[s2[-2]]['amp'] - i['alt_e_dep'] = inst_s2[s2[-2]]['e_dep'] - i['alt_s2_x'] = inst_s2[s2[-2]]['x'] - i['alt_s2_y'] = inst_s2[s2[-2]]['y'] - i['alt_s2_z'] = inst_s2[s2[-2]]['z'] - - i['x'] = inst_s2[s2[-1]]['x'] - i['y'] = inst_s2[s2[-1]]['y'] - i['z'] = inst_s2[s2[-1]]['z'] - - i['x_pri'] = inst_s2[s2[-1]]['x_pri'] - i['y_pri'] = inst_s2[s2[-1]]['y_pri'] - i['z_pri'] = inst_s2[s2[-1]]['z_pri'] - - i['g4id'] = inst_s2[s2[-1]]['g4id'] - - # Strax wants begin and endtimes - i['time'] = ix * 1000 - i['endtime'] = ix * 1000 + 1 - - instructions = instructions[~((instructions['time'] == 0) & (instructions['endtime'] == 0))] - self.instructions = instructions - - def simulate(self, ): - for func in self.simulation_functions: - func(instructions=self.instructions, - config=self.config, - resource=self.resource) - - def run_simulator(self, ): - self.cluster_events() - - if isinstance(self.config['epix_config']['save_epix'], str): - epix_path = self.config['epix_config']['save_epix'] + self.config['epix_config']['file_name'][:-5] +'_epix_2' - print('Saving epix instruction: ', epix_path) - np.save(epix_path, self.instructions) - - self.simulate() - - return self.instructions - - -# We should think a bit more to detector_config_override -# and tell fast_sim to look into epix_args -# also because one entry of self.config is epix_config -@strax.takes_config( - strax.Option('detector', default='XENONnT', help='Detector model'), - strax.Option('detector_config_override', default='', help='For the electric field, otherwise 200 V/cm'), - strax.Option('g4_file', help='G4 file to simulate'), - strax.Option('epix_config', default=dict(), help='Configuration file for epix', ), - strax.Option('configuration_files', default=dict(), help='Files required for simulating'), - strax.Option('fax_config', help='Fax configuration to load'), - strax.Option('fax_config_overrides', help='Fax configuration to override', default=None), - strax.Option('xy_resolution', default=0, help='xy position resolution (cm)'), - strax.Option('z_resolution', default=0, help='xy position resolution (cm)'), - strax.Option('nv_spe_resolution', default=0.40, help='nVeto SPE resolution'), - strax.Option('nv_spe_res_threshold', default=0.50, help='nVeto SPE acceptance threshold'), - strax.Option('nv_max_time_ns', default=1e7, help='nVeto maximum time for the acceptance of PMT hits in event'), - strax.Option('s2_clustering_algorithm', default='bdt', help='Macroclustering algorithm for S2, [ nsort | naive | bdt ]'), -) -class StraxSimulator(strax.Plugin): - provides = ('events_tpc', 'events_nveto') - depends_on = () - data_kind = immutabledict(zip(provides, provides)) - dtype = dict(events_tpc=[('time', np.int64), - ('endtime', np.int64), - ('s1_area', np.float), - ('s2_area', np.float), - ('cs1', np.float), - ('cs2', np.float), - ('alt_s1_area', np.float), - ('alt_s2_area', np.float), - ('alt_cs1', np.float), - ('alt_cs2', np.float), - ('x', np.float), - ('y', np.float), - ('z', np.float), - ('alt_s2_x', np.float), - ('alt_s2_y', np.float), - ('alt_s2_z', np.float), - ('drift_time', np.float), - ('alt_s2_drift_time', np.float), - ('e_dep', np.float), - ('alt_e_dep', np.float), - ('x_pri', np.float), - ('y_pri', np.float), - ('z_pri', np.float), - ('g4id', np.int)], - events_nveto=[('time', np.float), - ('endtime', np.float), - ('event_id', np.int), - ('channel', np.int), ]) - - def load_config(self): - """First load the config through wfsim, then we add some things we'd like to have""" - self.resource = load_config(self.config) - - self.resource.s1_map = self.resource.s1_lce_correction_map - self.resource.s2_map = self.resource.s2_correction_map - - # Terrible way to handle the config - # we should have only one file - # instead of several config - if 'nv_pmt_qe' in self.config['configuration_files'].keys(): - self.resource.nv_pmt_qe = straxen.get_resource( - self.config['configuration_files']['nv_pmt_qe'], fmt='json') - else: - warnings.warn('The configuration_files should not exist!' - 'Everything should come by one config!' - 'Since nv_pmt_qe config is missing in configuration_files, ' - 'the default one nveto_pmt_qe.json will be used') - self.resource.nv_pmt_qe = straxen.get_resource('nveto_pmt_qe.json', fmt='json') - - if 'photon_area_distribution' in self.config['configuration_files'].keys(): - self.resource.photon_area_distribution = epix.Helpers.average_spe_distribution( - straxen.get_resource(self.config['configuration_files']['photon_area_distribution'], - fmt='csv')) - else: - warnings.warn('The configuration_files should not exist!' - 'Everything should come by one config!' - 'Since photon_area_distribution config is missing in configuration_files, ' - 'the one in fax_config will be used') - self.resource.photon_area_distribution = epix.Helpers.average_spe_distribution( - straxen.get_resource(self.config['photon_area_distribution'], fmt='csv')) - - def setup(self, ): - overrides = self.config['fax_config_overrides'] - self.config.update(straxen.get_resource(self.config['fax_config'], fmt='json')) - if overrides is not None: - self.config.update(overrides) - - if 'gains' not in self.config.keys(): - warnings.warn('Why we have to provide gains here? ' - 'No gains are passed in fax_config_overrides, default equal to 1') - self.config['gains'] = [1 for i in range(494)] - if 'n_top_pmts' not in self.config.keys(): - warnings.warn('This is a deault value, why we have to give it in fax_config_overrides? ' - 'No n_top_pmts are passed in fax_config_overrides, default equal to 253') - self.config['n_top_pmts'] = 253 - if 'n_tpc_pmts' not in self.config.keys(): - warnings.warn('This is a deault value, why we have to give it in fax_config_overrides? ' - 'No n_tpc_pmts are passed in fax_config_overrides, default equal to 494') - self.config['n_tpc_pmts'] = 494 - - self.load_config() - - def get_nveto_data(self, ): - file_loader = epix.io.file_loader(directory=self.config['epix_config']['path'], - file_name=self.config['epix_config']['file_name']) - file_tree, _ = file_loader._get_ttree() - - if 'pmthitID' in file_tree.keys(): - nv_hits = NVetoUtils.get_nv_hits(ttree=file_tree, - pmt_nv_json_dict=self.resource.nv_pmt_qe, - nveto_dtype=self.dtype['events_nveto'], - SPE_Resolution=self.config['nv_spe_resolution'], - SPE_ResThreshold=self.config['nv_spe_res_threshold'], - max_time_ns=self.config['nv_max_time_ns'], - batch_size=10000) - return nv_hits - - def get_epix_instructions(self, ): - detector = epix.init_detector(self.config['detector'].lower(), self.config['detector_config_override']) - - epix_config = deepcopy(self.config['epix_config']) - epix_config['detector_config'] = detector - - outer_cylinder = getattr(epix.detectors, 'xenonnt') - _, outer_cylinder = outer_cylinder() - epix_config['outer_cylinder'] = outer_cylinder - - epix_ins = epix.run_epix.main(epix_config, return_wfsim_instructions=True) - - if self.config['save_epix']: - epix_path = self.config['epix_config']['path'] + self.config['epix_config']['file_name'][:-5] +'_epix' - np.save(epix_path, epix_ins) - - return epix_ins - - def compute(self): - simulated_data_nveto = self.get_nveto_data() - self.epix_instructions = self.get_epix_instructions() - - if isinstance(self.config['epix_config']['save_epix'], str): - epix_path = self.config['epix_config']['save_epix'] + self.config['epix_config']['file_name'][:-5] +'_epix_1' - print('Saving epix instruction: ', epix_path) - np.save(epix_path, self.epix_instructions) - - self.Simulator = Simulator(instructions_epix=self.epix_instructions, - config=self.config, - resource=self.resource) - simulated_data = self.Simulator.run_simulator() - - simulated_data_chunk = self.chunk( - start=simulated_data['time'][0], - end=simulated_data['endtime'][-1], - data=simulated_data, - data_type='events_tpc') - - # write empty chunk if nveto data isn't there - if simulated_data_nveto is None or len(simulated_data_nveto['time']) < 1: - simulated_data_nveto_chunk = self.chunk( - start=0, - end=1, - data=simulated_data_nveto, - data_type='events_nveto') - else: - simulated_data_nveto_chunk=self.chunk( - start=np.floor(simulated_data_nveto['time'][0]).astype(np.int64), - end=np.ceil(np.max(simulated_data_nveto['endtime'])).astype(np.int64), - data=simulated_data_nveto, - data_type='events_nveto') - - return {'events_tpc':simulated_data_chunk, - 'events_nveto':simulated_data_nveto_chunk} - - def is_ready(self, chunk): - # For this plugin we'll smash everything into 1 chunk, should be oke - return True if chunk == 0 else False - - def source_finished(self): - return True diff --git a/epix/simulator/helpers.py b/epix/simulator/helpers.py index d764c7f..91296d3 100644 --- a/epix/simulator/helpers.py +++ b/epix/simulator/helpers.py @@ -2,13 +2,69 @@ from scipy.interpolate import interp1d import numpy as np import numba +import numpy.lib.recfunctions as rfn +DTYPE = dict(events_tpc=[('time', np.int64), + ('endtime', np.int64), + ('cs1', np.float), + ('cs2', np.float), + ('alt_cs2', np.float), + ('drift_time', np.float), + ('s1_area', np.float), + ('s2_area', np.float), + ('alt_s2_area', np.float), + ('alt_s2_interaction_drift_time', np.float), + ('s2_x', np.float), + ('s2_y', np.float), + ('alt_s2_x', np.float), + ('alt_s2_y', np.float), + ('x', np.float), + ('alt_s2_x_fdc', np.float), + ('y', np.float), + ('alt_s2_y_fdc', np.float), + ('r', np.float), + ('alt_s2_r_fdc', np.float), + ('z', np.float), + ('z_dv_corr', np.float), + ('alt_s2_z', np.float), + ('alt_s2_z_dv_corr', np.float), + ('r_naive', np.float), + ('alt_s2_r_naive', np.float), + ('z_naive', np.float), + ('alt_s2_z_naive', np.float), + ('r_field_distortion_correction', np.float), + ('alt_s2_r_field_distortion_correction', np.float), + ('z_field_distortion_correction', np.float32), + ('alt_s2_z_field_distortion_correction', np.float32), + ('alt_s2_theta', np.float32), + ('theta', np.float32), + # Truth values + ('x_true', np.float), + ('y_true', np.float), + ('z_true', np.float), + ('r_true', np.float), + ('theta_true', np.float), + ('alt_s2_x_true', np.float), + ('alt_s2_y_true', np.float), + ('alt_s2_z_true', np.float), + ('alt_s2_r_true', np.float), + ('alt_s2_theta_true', np.float), + ('e_dep', np.float), + ('alt_e_dep', np.float), + ('x_pri', np.float), + ('y_pri', np.float), + ('z_pri', np.float), + ('g4id', np.int)], + events_nveto=[('time', np.float), + ('endtime', np.float), + ('event_id', np.int), + ('channel', np.int), ]) # Numba and classes still are not a match made in heaven @numba.njit -def _merge_these_clusters_nsort(s2_area1, z1, s2_area2, z2, **kwargs): +def _merge_these_clusters_nsort(amp1, r1, z1, amp2, r2, z2, conf): sensitive_volume_ztop = 0 # it's the ground mesh, the top liquid level is at 2.7; // mm - max_s2_area = max(s2_area1, s2_area2) + max_s2_area = max(amp1, amp2) if max_s2_area > 5000: SeparationDistanceIntercept = 0.00024787 * 5000. + 3.4056346550312973 SeparationDistanceSlope = 5.5869678412887262e-07 * 5000. + 0.0044792968 @@ -19,32 +75,81 @@ def _merge_these_clusters_nsort(s2_area1, z1, s2_area2, z2, **kwargs): 5.5869678412887262e-07 * max_s2_area + 0.0044792968 SeparationDistance = \ SeparationDistanceIntercept - \ - SeparationDistanceSlope * (-sensitive_volume_ztop + \ - (z1 + z2) * 0.5) + SeparationDistanceSlope * (-sensitive_volume_ztop + (z1 + z2) * 0.5) return z1 - z2 < SeparationDistance @numba.njit -def _merge_these_clusters_nt_res_naive(s2_area1, z1, s2_area2, z2, **kwargs): +def _merge_these_clusters_nt_res_naive(amp1, r1, z1, amp2, r2, z2, conf): sensitive_volume_ztop = 0 # [cm] SeparationDistance = 1.6 # [cm], the worst case from [[weiss:analysis:he:zresoultion_zdependence]] return np.abs(z1 - z2) < SeparationDistance +@numba.njit +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_these_clusters_nt_res_jaron(s2_area1, z1, s2_area2, z2, tree): - return bool(tree.predict([[z1, z2-z1, s2_area1, s2_area2]])) -class Helpers(): - @staticmethod - def assignOrder(order): - """This is a trick to have the calculation function be executed in a particular order""" +def _merge_these_clusters_nt_res_mlp(inst1, inst2, config, resource): + v1 = Helpers.get_drift_velocity([inst1['z']], + np.array([[inst1['x']],[inst1['y']]]).T, config, resource)[0] * 1000 # cm/us + v2 = Helpers.get_drift_velocity([inst2['z']], + np.array([[inst2['x']],[inst2['y']]]).T, config, resource)[0] * 1000 # cm/us + dt1 = Helpers.get_drift_time(np.array([inst1['z']]), + np.array([[inst1['x']],[inst1['y']]]).T, config, resource)[0] / 1000 # us + dt2 = Helpers.get_drift_time(np.array([inst2['z']]), + np.array([[inst2['x']],[inst2['y']]]).T, config, resource)[0] / 1000 # us + diff1 = resource.diffusion_longitudinal_map([inst1['z']], + np.array([[inst1['x']], [inst1['y']]]).T)[0] * 1000 # cm2/us + diff2 = resource.diffusion_longitudinal_map([inst2['z']], + np.array([[inst2['x']], [inst2['y']]]).T)[0] * 1000 # cm2/us + survival1 = resource.field_dependencies_map([inst1['z']], np.array([[inst1['x']], [inst1['y']]]).T, + map_name='survival_probability_map')[0] + survival2 = resource.field_dependencies_map([inst2['z']], np.array([[inst2['x']], [inst2['y']]]).T, + map_name='survival_probability_map')[0] + width1 = 1.348*np.sqrt(2*diff1*dt1/v1**2) # us + width2 = 1.348*np.sqrt(2*diff2*dt2/v2**2) # us + delta_t = (dt2-dt1) # us + split_param = delta_t/(width1+width2) + + e_lifetime = config['electron_lifetime_liquid'] / 1000 # us + amp1_corr = int(config['electron_extraction_yield'] * survival1 * np.exp(-dt1/e_lifetime) * inst1['amp']) + amp2_corr = int(config['electron_extraction_yield'] * survival2 * np.exp(-dt2/e_lifetime) * inst2['amp']) + lower = min(amp1_corr, amp2_corr) + higher = max(amp1_corr, amp2_corr) + w0, w1, w2, w3, w4, w5, w6, w7, w8, w9 = resource.nn_weights + X = np.array((split_param, higher, lower), dtype=np.float32) + y = get_nn_prediction(X, w0, w1, w2, w3, w4, w5, w6, w7, w8, w9) + return y > 0.5 + + +class Helpers: + @staticmethod + def assign_order(order): + """ + This is a trick to have the calculation function be executed in a particular order + """ def do_assignment(to_func): to_func.order = order return to_func return do_assignment + @staticmethod + def get_dtype(): + return DTYPE + @staticmethod def average_spe_distribution(spe_shapes): """We take the spe distribution from all channels and take the average to be our spe distribution to draw photon areas from""" @@ -70,47 +175,40 @@ def average_spe_distribution(spe_shapes): return spe_distribution @staticmethod - def macro_cluster_events(tree, instructions, config): - """Loops over all instructions, checks if it's an s2 and if there is another s2 within the same event + def macro_cluster_events(instructions, config, resource): + """Loops over all instructions, checks if it's an S2 and if there is another s2 within the same event within the macro cluster distance, if it is they are merged.""" print(f"\n macro_cluster_events --> s2_clustering_algorithm == {config['s2_clustering_algorithm']} . . .") - - if config['s2_clustering_algorithm'] == 'bdt': - _merge_clusters = _merge_these_clusters_nt_res_jaron + if config['s2_clustering_algorithm'] == 'mlp': + _merge_clusters = _merge_these_clusters_nt_res_mlp elif config['s2_clustering_algorithm'] == 'naive': - tree = None _merge_clusters = _merge_these_clusters_nt_res_naive elif config['s2_clustering_algorithm'] == 'nsort': - tree = None _merge_clusters = _merge_these_clusters_nsort else: - _merge_clusters = None - + return for ix1, _ in enumerate(instructions): if instructions[ix1]['type'] != 2: continue - for ix2 in range(1, len(instructions[ix1:])): # why was + 1): ? - if instructions[ix1 + ix2]['type'] != 2: - continue + for ix2 in range(1, len(instructions[ix1:])): if instructions[ix1]['event_number'] != instructions[ix1 + ix2]['event_number']: break + if instructions[ix1 + ix2]['type'] != 2: + continue + if _merge_clusters(instructions[ix1], instructions[ix1 + ix2], config, resource): + amp1 = instructions[ix1]['amp'] + amp2 = instructions[ix1 + ix2]['amp'] + amp_total = int((instructions[ix1]['amp'] + instructions[ix1 + ix2]['amp'])) + instructions[ix1 + ix2]['x'] = (instructions[ix1]['x'] * amp1 + instructions[ix1 + ix2]['x'] * amp2) / amp_total + instructions[ix1 + ix2]['y'] = (instructions[ix1]['y'] * amp1 + instructions[ix1 + ix2]['y'] * amp2) / amp_total + instructions[ix1 + ix2]['z'] = (instructions[ix1]['z'] + instructions[ix1 + ix2]['z']) / 2 - # _nt_res - if _merge_clusters(instructions[ix1]['amp'], instructions[ix1]['z'], - instructions[ix1 + ix2]['amp'], instructions[ix1 + ix2]['z'], - tree): - - instructions[ix1 + ix2]['x'] = (instructions[ix1]['x'] + instructions[ix1 + ix2]['x']) * 0.5 - instructions[ix1 + ix2]['y'] = (instructions[ix1]['y'] + instructions[ix1 + ix2]['y']) * 0.5 - instructions[ix1 + ix2]['z'] = (instructions[ix1]['z'] + instructions[ix1 + ix2]['z']) * 0.5 - - # prymary position is one - instructions[ix1 + ix2]['x_pri'] = instructions[ix1]['x_pri'] - instructions[ix1 + ix2]['y_pri'] = instructions[ix1]['y_pri'] - instructions[ix1 + ix2]['z_pri'] = instructions[ix1]['z_pri'] + instructions[ix1 + ix2]['x_pri'] = instructions[ix1]['x_pri'] + instructions[ix1 + ix2]['y_pri'] = instructions[ix1]['y_pri'] + instructions[ix1 + ix2]['z_pri'] = instructions[ix1]['z_pri'] - instructions[ix1 + ix2]['amp'] = int((instructions[ix1]['amp'] + instructions[ix1 + ix2]['amp'])) + instructions[ix1 + ix2]['amp'] = amp_total instructions[ix1]['amp'] = -1 # flag to throw this instruction away later instructions[ix1 + ix2]['e_dep'] = (instructions[ix1]['e_dep'] + instructions[ix1 + ix2]['e_dep']) instructions[ix1]['e_dep'] = -1 # flag to throw this instruction away later @@ -145,10 +243,19 @@ def get_s2_light_yield(positions, config, resource): resource=resource) @staticmethod - def get_s2_charge_yield(n_electron, positions, z_obs, config, resource): + def get_s2_charge_yield(n_electron, xy, z, config, resource): """See wfsim.s2.get_electron_yield""" return wfsim.S2.get_electron_yield(n_electron=n_electron, - xy_int=positions, - z_int=z_obs, + xy_int=xy, + z_int=z, config=config, resource=resource) + @staticmethod + def get_drift_velocity(z, xy, config, resource): + return wfsim.S2.get_avg_drift_velocity(z, xy, config, resource) + + @staticmethod + def get_drift_time(z, xy, config, resource): + """See wfsim.S2.get_s2_drift_time_params""" + dt_mean, _ = wfsim.S2.get_s2_drift_time_params(z, xy, config, resource) + return dt_mean diff --git a/epix/simulator/nn_weights.p b/epix/simulator/nn_weights.p new file mode 100644 index 0000000..f906b49 Binary files /dev/null and b/epix/simulator/nn_weights.p differ