diff --git a/wfsim/core/rawdata.py b/wfsim/core/rawdata.py index 5754e330..edfa97c8 100644 --- a/wfsim/core/rawdata.py +++ b/wfsim/core/rawdata.py @@ -2,7 +2,7 @@ from numba import njit import numpy as np from strax import exporter -from tqdm import tqdm +from strax.utils import tqdm from .afterpulse import PhotoIonization_Electron, PhotoElectric_Electron, PMT_Afterpulse from .pulse import Pulse from .s1 import S1 diff --git a/wfsim/strax_interface.py b/wfsim/strax_interface.py index 785d6ea4..b6a805cc 100644 --- a/wfsim/strax_interface.py +++ b/wfsim/strax_interface.py @@ -5,10 +5,11 @@ import pandas as pd from scipy.interpolate import interp1d import uproot - +import typing as ty import strax import straxen import wfsim +from strax.utils import tqdm export, __all__ = strax.exporter() logging.basicConfig(handlers=[logging.StreamHandler()]) @@ -59,28 +60,111 @@ def rand_instructions(c): """Random instruction generator function. This will be called by wfsim if you do not specify specific instructions. :params c: wfsim configuration dict""" - n = c['nevents'] = c['event_rate'] * c['chunk_size'] * c['nchunk'] - c['total_time'] = c['chunk_size'] * c['nchunk'] - - instructions = np.zeros(2 * n, dtype=instruction_dtype) - uniform_times = c['total_time'] * (np.arange(n) + 1.) / n - instructions['time'] = np.repeat(uniform_times, 2) * int(1e9) - instructions['event_number'] = np.digitize(instructions['time'], - 1e9 * np.arange(c['nchunk']) * c['chunk_size']) - 1 - instructions['type'] = np.tile([1, 2], n) - instructions['recoil'] = np.repeat(7, n * 2) # Use nest ids for ER - - r = np.sqrt(np.random.uniform(0, c['tpc_radius']**2, n)) - t = np.random.uniform(-np.pi, np.pi, n) - instructions['x'] = np.repeat(r * np.cos(t), 2) - instructions['y'] = np.repeat(r * np.sin(t), 2) - instructions['z'] = np.repeat(np.random.uniform(-c['tpc_length'], 0, n), 2) - - nphotons = np.random.uniform(2000, 2050, n) - nelectrons = 10 ** (np.random.uniform(3, 4, n)) - instructions['amp'] = np.vstack([nphotons, nelectrons]).T.flatten().astype(int) - - return instructions + log.warn('rand_instructions is deprecated, please use wfsim.random_instructions') + if 'drift_field' not in c: + log.warn('drift field not specified!') + # Do the conversion for now, don't break everything all at once + return _rand_instructions(event_rate=c.get('event_rate', 10), + chunk_size=c.get('chunk_size', 5), + n_chunk=c.get('n_chunk', 2), + energy_range=[1, 100], # keV + drift_field=c.get('drift_field', 100), # V/cm + tpc_radius=c.get('tpc_radius', straxen.tpc_r), + tpc_length=c.get('tpc_length', straxen.tpc_z), + nest_inst_types=[7] + ) + + +def random_instructions(**kwargs) -> np.ndarray: + """ + Generate instructions to run WFSim + :param event_rate: # events per second + :param chunk_size: the size of each chunk + :param n_chunk: the number of chunks + :param drift_field: + :param energy_range: the energy range (in keV) of the recoil type + :param tpc_length: the max depth of the detector + :param tpc_radius: the max radius of the detector + :param nest_inst_types: the nestpy type of the interaction see: + github.com/NESTCollaboration/nestpy/blob/8eb79414e5f834eb6cf6ddba5d6c433c6b0cbc70/src/nestpy/helpers.py#L27 + :return: Instruction array that can be used for simulation in wfsim + """ + return _rand_instructions(**kwargs) + + +def _rand_instructions( + event_rate: int, + chunk_size: int, + n_chunk: int, + drift_field: float, + energy_range: ty.Union[tuple, list, np.ndarray], + tpc_length: float = straxen.tpc_z, + tpc_radius: float = straxen.tpc_r, + nest_inst_types: ty.Union[ty.List[int], ty.Tuple[ty.List], np.ndarray, None] = None, +) -> np.ndarray: + import nestpy + if nest_inst_types is None: + nest_inst_types = [7] + + n_events = event_rate * chunk_size * n_chunk + total_time = chunk_size * n_chunk + + inst = np.zeros(2 * n_events, dtype=instruction_dtype) + inst[:] = -1 + + uniform_times = total_time * (np.arange(n_events) + 0.5) / n_events + + inst['time'] = np.repeat(uniform_times, 2) * int(1e9) + inst['event_number'] = np.digitize(inst['time'], + 1e9 * np.arange(n_chunk) * + chunk_size) - 1 + inst['type'] = np.tile([1, 2], n_events) + + r = np.sqrt(np.random.uniform(0, tpc_radius ** 2, n_events)) + t = np.random.uniform(-np.pi, np.pi, n_events) + inst['x'] = np.repeat(r * np.cos(t), 2) + inst['y'] = np.repeat(r * np.sin(t), 2) + inst['z'] = np.repeat(np.random.uniform(-tpc_length, 0, n_events), 2) + + # Here we'll define our XENON-like detector + nest_calc = nestpy.NESTcalc(nestpy.VDetector()) + nucleus_A = 131.293 + nucleus_Z = 54. + lxe_density = 2.862 # g/cm^3 #SR1 Value + + energy = np.random.uniform(*energy_range, n_events) + quanta = [] + exciton = [] + recoil = [] + e_dep = [] + for energy_deposit in tqdm(energy, desc='generating instructions from nest'): + interaction_type = np.random.choice(nest_inst_types) + interaction = nestpy.INTERACTION_TYPE(interaction_type) + y = nest_calc.GetYields(interaction, + energy_deposit, + lxe_density, + drift_field, + nucleus_A, + nucleus_Z, + ) + q = nest_calc.GetQuanta(y, lxe_density) + quanta.append(q.photons) + quanta.append(q.electrons) + exciton.append(q.excitons) + exciton.append(0) + # both S1 and S2 + recoil += [interaction_type, interaction_type] + e_dep += [energy_deposit, energy_deposit] + + inst['amp'] = quanta + inst['local_field'] = drift_field + inst['n_excitons'] = exciton + inst['recoil'] = recoil + inst['e_dep'] = e_dep + for field in inst.dtype.names: + if np.any(inst[field] == -1): + log.warn(f'{field} is not (fully) filled') + return inst def _read_optical_nveto(config, events, mask):