From 0d75a6aba071420cbc5c6f8a29b913ee99cc1489 Mon Sep 17 00:00:00 2001 From: soleti Date: Thu, 25 Aug 2022 14:01:48 -0700 Subject: [PATCH] cleaning up some code --- README.md | 112 --------------------------------------- cli/simulate_pixels.py | 15 ++++-- larndsim/consts/light.py | 20 +++---- larndsim/detsim.py | 99 +++++++++------------------------- larndsim/fee.py | 20 +++---- larndsim/light_sim.py | 97 +++++++++++++++++---------------- 6 files changed, 104 insertions(+), 259 deletions(-) diff --git a/README.md b/README.md index be1a0cb7..b4000bf7 100644 --- a/README.md +++ b/README.md @@ -61,115 +61,3 @@ simulate_pixels.py \ The `response_38.npy` is a file containing an array of induced currents for several $(x,y)$ positions on a pixel with a 38 mm pitch. It is calculated externally to `larnd-sim`. Two versions, one with 44 mm pitch and one with 38 mm pitch, are available in the `larndsim/bin` directory. The output file will contain the datasets described in the [LArPix HDF5 documentation](https://larpix-control.readthedocs.io/en/stable/api/format/hdf5format.html), plus a dataset `tracks` containing the _true_ energy depositions in the detector, and a dataset `mc_packets_assn`, which has a list of indeces corresponding to the true energy deposition associated to each packet. - -## Step-by-step simulation - -Here we will describe, step-by-step, how we perform the simulation. A full example for the ND-LAr detector is available in the `examples/NDLAr example.ipynb` notebook. - -### Quenching and drifting stage - -The particles that interact in the TPC ionize the argon atoms. Some of the resulting electrons will immediately recombine with the atoms. This effect is simulated in the `quenching` module. - -The remaining electrons travel towards the anode and their spatial distribution is affected by longitudinal and transverse diffusion. The presence of impurities reduces the amount of electrons that reach the anode. These effects are simulated in the `drifting` module. - -These two modules modify in place the input array, here called `input_tracks`. - -```python -from larndsim import quenching, drifting - -TPB = 256 -BPG = ceil(tracks.shape[0] / TPB) -quenching.quench[BPG,TPB](input_tracks, consts.birks) -drifting.drift[BPG,TPB](input_tracks) -``` - -### Pixel simulation stage - -Once we have calculated the number and the position of the electrons reaching the anode, we can calculate the current induced on each pixel. -First, we find the pixels interesected by the projection of each track segment on the anode plane using the [Bresenham's line algorithm](https://en.wikipedia.org/wiki/Bresenham%27s_line_algorithm) implented in the `pixels_from_track` module. Due to diffusion, we consider also the neighboring pixels. - -```python -from larndsim import pixels_from_track -... -pixels_from_track.get_pixels[BPG,TPB](input_tracks, - active_pixels, - neighboring_pixels, - n_pixels_list, - max_radius) -``` - -Finally, we calculate the current induced on each pixel using the `tracks_current` function in the `detsim` module. The induced current is stored in the `signals` array, which is a three-dimensional array, where the dimensions correspond to the track segment, the pixel, and the time tick, respectively: `signals[0][1][2]` will contain the current induced by the track `0`, for the pixel `1`, at the time tick `2`. - -```python -from larndsim import detsim -... -detsim.tracks_current[BPG,TPB](signals, - neighboring_pixels, - input_tracks, - response, - rng_states) -``` - -Here, `response` is a Numpy array containing a look-up table with a pre-calculated field response. The file valid for Module0 and SingleCube LArPix tiles is availabe at `larndsim/response-44.npy`. For ND-LAr the file is `larndsim/response_38.npy`. - -### Accessing the signals - -The three-dimensional array can contain more than one signal for each pixel at different times. If we want to plot the full induced signal on the pixel, we need to join the signals corresponding to the same pixel. First, we find the start time of each signal with `time_intervals`: - -```python -from larndsim import detsim -... -detsim.time_intervals[blockspergrid,threadsperblock](track_starts, - max_length, - input_tracks) -``` - -Then, we join them using `sum_pixel_signals`: - -```python -from larndsim import detsim -... -detsim.sum_pixel_signals[blockspergrid,threadsperblock](pixels_signals, - signals, - track_starts, - pixel_index_map, - track_pixel_map, - pixels_tracks_signals) -``` - -### Electronics simulation - -Once we have the induced current for each active pixel in the detector we can apply our electronics simulation, which will calculate the ADC values for each pixel: - -```python -from larndsim import fee -from numba.cuda.random import create_xoroshiro128p_states -... -fee.get_adc_values[BPG,TPB](pixels_signals, - pixels_tracks_signals, - time_ticks, - integral_list, - adc_ticks_list, - 0, - rng_states, - current_fractions, - pixel_thresholds) -``` - -where the random states `rng_states` are needed for the noise simulation. - -### Export - -The final output is then exported to the [LArPix HDF5 format](https://larpix-control.readthedocs.io/en/stable/api/format/hdf5format.html) after each event in the input file: - -```python -fee.export_to_hdf5(event_id_list_batch, - adc_tot_list_batch, - adc_tot_ticks_list_batch, - cp.asnumpy(unique_pix_tot_batch), - cp.asnumpy(current_fractions_tot_batch), - cp.asnumpy(track_pixel_map_tot_batch), - output_filename, - t0=t0, - bad_channels=bad_channels) -``` diff --git a/cli/simulate_pixels.py b/cli/simulate_pixels.py index 9b9bb5c5..bb3ed53f 100755 --- a/cli/simulate_pixels.py +++ b/cli/simulate_pixels.py @@ -68,13 +68,16 @@ def swap_coordinates(tracks): def maybe_create_rng_states(n, seed=0, rng_states=None): """Create or extend random states for CUDA kernel""" + if rng_states is None: return create_xoroshiro128p_states(n, seed=seed) - elif n > len(rng_states): + + if n > len(rng_states): new_states = device_array(n, dtype=rng_states.dtype) new_states[:len(rng_states)] = rng_states new_states[len(rng_states):] = create_xoroshiro128p_states(n - len(rng_states), seed=seed) return new_states + return rng_states def run_simulation(input_filename, @@ -273,7 +276,8 @@ def run_simulation(input_filename, # pre-allocate some random number states rng_states = maybe_create_rng_states(1024*256, seed=0) last_time = 0 - for ievd in tqdm(range(0, tot_evids.shape[0], step), desc='Simulating events...', ncols=80, smoothing=0): + for ievd in tqdm(range(0, tot_evids.shape[0], step), + desc='Simulating events...', ncols=80, smoothing=0): first_event = tot_evids[ievd] last_event = tot_evids[min(ievd+step, tot_evids.shape[0]-1)] @@ -288,7 +292,8 @@ def run_simulation(input_filename, evt_tracks = track_subset[event_mask] first_trk_id = np.where(track_subset['eventID'] == evt_tracks['eventID'][0])[0][0] + min(start_idx[ievd:ievd + step]) - for itrk in tqdm(range(0, evt_tracks.shape[0], BATCH_SIZE), desc=' Simulating event %i batches...' % ievd, leave=False, ncols=80): + for itrk in tqdm(range(0, evt_tracks.shape[0], BATCH_SIZE), + desc=' Simulating event %i batches...' % ievd, leave=False, ncols=80): selected_tracks = evt_tracks[itrk:itrk+BATCH_SIZE] RangePush("event_id_map") event_ids = selected_tracks['eventID'] @@ -441,7 +446,7 @@ def run_simulation(input_filename, TPB = (1,64) BPG = (ceil(light_sample_inc.shape[0] / TPB[0]), - ceil(light_sample_inc.shape[1] / TPB[1])) + ceil(light_sample_inc.shape[1] / TPB[1])) light_sim.sum_light_signals[BPG, TPB]( selected_tracks, track_light_voxel[track_slice][event_mask][itrk:itrk+BATCH_SIZE], selected_track_id, light_inc, lut, light_t_start, light_sample_inc, light_sample_inc_true_track_id, @@ -512,7 +517,7 @@ def run_simulation(input_filename, light_waveforms_true_track_id_list_batch = np.concatenate(light_waveforms_true_track_id_list, axis=0) light_waveforms_true_photons_list_batch = np.concatenate(light_waveforms_true_photons_list, axis=0) light_trigger_modules = cp.array([detector.TPC_TO_MODULE[tpc] for tpc in light.OP_CHANNEL_TO_TPC[light_op_channel_idx_list_batch.get()][:,0]]) - + fee.export_to_hdf5(event_id_list_batch, adc_tot_list_batch, adc_tot_ticks_list_batch, diff --git a/larndsim/consts/light.py b/larndsim/consts/light.py index 5feea48b..cfee34f4 100644 --- a/larndsim/consts/light.py +++ b/larndsim/consts/light.py @@ -80,11 +80,11 @@ def set_light_properties(detprop_file): global ENABLE_LUT_SMEARING global LIGHT_TICK_SIZE global LIGHT_WINDOW - + global SINGLET_FRACTION global TAU_S global TAU_T - + global LIGHT_GAIN global SIPM_RESPONSE_MODEL global LIGHT_RESPONSE_TIME @@ -106,27 +106,27 @@ def set_light_properties(detprop_file): LUT_VOX_DIV = np.array(detprop['lut_vox_div']) N_OP_CHANNEL = detprop['n_op_channel'] OP_CHANNEL_EFFICIENCY = np.array(detprop['op_channel_efficiency']) - + tpc_to_op_channel = detprop['tpc_to_op_channel'] OP_CHANNEL_TO_TPC = np.zeros((N_OP_CHANNEL,), int) TPC_TO_OP_CHANNEL = np.zeros((len(tpc_to_op_channel), len(tpc_to_op_channel[0])), int) - for itpc in range(len(tpc_to_op_channel)): - TPC_TO_OP_CHANNEL[itpc] = np.array(tpc_to_op_channel[itpc]) - for idet in tpc_to_op_channel[itpc]: + for itpc, tpc_to_op in enumerate(tpc_to_op_channel): + TPC_TO_OP_CHANNEL[itpc] = np.array(tpc_to_op) + for idet in tpc_to_op: OP_CHANNEL_TO_TPC[idet] = itpc ENABLE_LUT_SMEARING = bool(detprop.get('enable_lut_smearing', ENABLE_LUT_SMEARING)) LIGHT_TICK_SIZE = float(detprop.get('light_tick_size', LIGHT_TICK_SIZE)) LIGHT_WINDOW = tuple(detprop.get('light_window', LIGHT_WINDOW)) assert len(LIGHT_WINDOW) == 2 - SINGLET_FRACTION = float(detprop.get('singlet_fraction', SINGLET_FRACTION)) TAU_S = float(detprop.get('tau_s', TAU_S)) TAU_T = float(detprop.get('tau_t', TAU_T)) - LIGHT_GAIN = np.array(detprop.get('light_gain', np.full(OP_CHANNEL_EFFICIENCY.shape, LIGHT_GAIN))) + if LIGHT_GAIN.size == 1: LIGHT_GAIN = np.full(OP_CHANNEL_EFFICIENCY.shape, LIGHT_GAIN) + assert LIGHT_GAIN.shape == OP_CHANNEL_EFFICIENCY.shape SIPM_RESPONSE_MODEL = int(detprop.get('sipm_response_model', SIPM_RESPONSE_MODEL)) assert SIPM_RESPONSE_MODEL in (0,1) @@ -134,6 +134,7 @@ def set_light_properties(detprop_file): LIGHT_RESPONSE_TIME = float(detprop.get('light_response_time', LIGHT_RESPONSE_TIME)) LIGHT_OSCILLATION_PERIOD = float(detprop.get('light_oscillation_period', LIGHT_OSCILLATION_PERIOD)) impulse_model_filename = str(detprop.get('impulse_model', '')) + if impulse_model_filename and SIPM_RESPONSE_MODEL == 1: print('Light impulse model:', impulse_model_filename) try: @@ -145,6 +146,7 @@ def set_light_properties(detprop_file): IMPULSE_MODEL = np.load(os.path.join(os.path.dirname(__file__), '../../') + impulse_model_filename) except FileNotFoundError: print("Impulse model file not found:", impulse_model_filename) + IMPULSE_TICK_SIZE = float(detprop.get('impulse_tick_size', IMPULSE_TICK_SIZE)) OP_CHANNEL_PER_TRIG = int(detprop.get('op_channel_per_det', OP_CHANNEL_PER_TRIG)) @@ -154,7 +156,7 @@ def set_light_properties(detprop_file): LIGHT_DIGIT_SAMPLE_SPACING = float(detprop.get('light_digit_sample_spacing', LIGHT_DIGIT_SAMPLE_SPACING)) LIGHT_NBIT = int(detprop.get('light_nbit', LIGHT_NBIT)) - + except KeyError: LIGHT_SIMULATED = False diff --git a/larndsim/detsim.py b/larndsim/detsim.py index eb762c4b..02e60c68 100644 --- a/larndsim/detsim.py +++ b/larndsim/detsim.py @@ -10,9 +10,8 @@ from numba import cuda from numba.cuda.random import xoroshiro128p_normal_float32 -from . import consts from .consts.detector import TPC_BORDERS, TIME_INTERVAL -from .consts import detector, physics +from .consts import detector from .pixels_from_track import id2pixel MAX_TRACKS_PER_PIXEL = 5 @@ -162,57 +161,6 @@ def rho(point, q, start, sigmas, segment): return expo -@nb.njit -def truncexpon(x, loc=0, scale=1): - """ - A truncated exponential distribution. - To shift and/or scale the distribution use the `loc` and `scale` parameters. - - Args: - x (float): where to evaluate the function - loc (float, optional): shift position of the distribution. Defaults to 0. - scale (float, optional): scale the distribution. Defaults to 1. - - Returns: - float: exp(-(x-loc)/scale)/scale - """ - y = (x-loc)/scale - - return exp(-y)/scale if y > 0 else 0 - -@nb.njit -def current_model(t, t0, x, y): - """ - Parametrization of the induced current on the pixel, which depends - on the of arrival at the anode (:math:`t_0`) and on the position - on the pixel pad. - - Args: - t (float): time where we evaluate the current - t0 (float): time of arrival at the anode - x (float): distance between the point on the pixel and the pixel center - on the :math:`x` axis - y (float): distance between the point on the pixel and the pixel center - on the :math:`y` axis - - Returns: - float: the induced current at time :math:`t` - """ - B_params = (1.060, -0.909, -0.909, 5.856, 0.207, 0.207) - C_params = (0.679, -1.083, -1.083, 8.772, -5.521, -5.521) - D_params = (2.644, -9.174, -9.174, 13.483, 45.887, 45.887) - t0_params = (2.948, -2.705, -2.705, 4.825, 20.814, 20.814) - - a = B_params[0] + B_params[1]*x+B_params[2]*y+B_params[3]*x*y+B_params[4]*x*x+B_params[5]*y*y - b = C_params[0] + C_params[1]*x+C_params[2]*y+C_params[3]*x*y+C_params[4]*x*x+C_params[5]*y*y - c = D_params[0] + D_params[1]*x+D_params[2]*y+D_params[3]*x*y+D_params[4]*x*x+D_params[5]*y*y - shifted_t0 = t0 + t0_params[0] + t0_params[1]*x + t0_params[2]*y + \ - t0_params[3]*x*y + t0_params[4]*x*x + t0_params[5]*y*y - - a = min(a, 1) - - return a * truncexpon(-t, -shifted_t0, b) + (1-a) * truncexpon(-t, -shifted_t0, c) - @nb.njit def track_point(start, direction, z): """ @@ -277,21 +225,21 @@ def overlapping_segment(x, y, start, end, radius): """ Calculates the segment of the track defined by start, end that overlaps with a circle centered at x,y - + """ dxy = x - start[0], y - start[1] v = end[0] - start[0], end[1] - start[1] l = sqrt(v[0]**2 + v[1]**2) v = v[0]/l, v[1]/l s = (dxy[0] * v[0] + dxy[1] * v[1])/l # position of point of closest approach - + r = sqrt((dxy[0] - v[0] * s * l)**2 + (dxy[1] - v[1] * s * l)**2) if r > radius: return start, start # no overlap - + s_plus = s + sqrt(radius**2 - r**2) / l s_minus = s - sqrt(radius**2 - r**2) / l - + if s_plus > 1: s_plus = 1 elif s_plus < 0: @@ -300,7 +248,7 @@ def overlapping_segment(x, y, start, end, radius): s_minus = 1 elif s_minus < 0: s_minus = 0 - + new_start = (start[0] * (1 - s_minus) + end[0] * s_minus, start[1] * (1 - s_minus) + end[1] * s_minus, start[2] * (1 - s_minus) + end[2] * s_minus) @@ -329,7 +277,7 @@ def tracks_current_mc(signals, pixels, tracks, response, rng_states): """ itrk, ipix, it = cuda.grid(3) ntrk, _, _ = cuda.gridsize(3) - + if itrk < signals.shape[0] and ipix < signals.shape[1] and it < signals.shape[2]: t = tracks[itrk] pID = pixels[itrk][ipix] @@ -348,7 +296,7 @@ def tracks_current_mc(signals, pixels, tracks, response, rng_states): else: end = (t["x_start"], t["y_start"], t["z_start"]) start = (t["x_end"], t["y_end"], t["z_end"]) - + t_start = max(TIME_INTERVAL[0], round((t["t_start"]-detector.TIME_PADDING) / detector.TIME_SAMPLING) * detector.TIME_SAMPLING) time_tick = t_start + it * detector.TIME_SAMPLING @@ -357,21 +305,21 @@ def tracks_current_mc(signals, pixels, tracks, response, rng_states): direction = (segment[0]/length, segment[1]/length, segment[2]/length) sigmas = (t["tran_diff"], t["tran_diff"], t["long_diff"]) - - impact_factor = sqrt(response.shape[0]**2 + + + impact_factor = sqrt(response.shape[0]**2 + response.shape[1]**2) * detector.RESPONSE_BIN_SIZE subsegment_start, subsegment_end = overlapping_segment(x_p, y_p, start, end, impact_factor) - subsegment = (subsegment_end[0]-subsegment_start[0], - subsegment_end[1]-subsegment_start[1], + subsegment = (subsegment_end[0]-subsegment_start[0], + subsegment_end[1]-subsegment_start[1], subsegment_end[2]-subsegment_start[2]) subsegment_length = sqrt(subsegment[0]**2 + subsegment[1]**2 + subsegment[2]**2) if subsegment_length == 0: return - + nstep = max(round(subsegment_length / MIN_STEP_SIZE), 1) step = subsegment_length / nstep # refine step size - + charge = t["n_electrons"] * (subsegment_length/length) / (nstep*MC_SAMPLE_MULTIPLIER) total_current = 0 rng_state = (rng_states[itrk + ntrk * ipix],) @@ -380,17 +328,17 @@ def tracks_current_mc(signals, pixels, tracks, response, rng_states): x = subsegment_start[0] + step * (istep + 0.5) * direction[0] y = subsegment_start[1] + step * (istep + 0.5) * direction[1] z = subsegment_start[2] + step * (istep + 0.5) * direction[2] - + z += xoroshiro128p_normal_float32(rng_state, 0) * sigmas[2] t0 = abs(z - TPC_BORDERS[t["pixel_plane"]][2][0]) / detector.V_DRIFT - detector.TIME_WINDOW if not t0 < time_tick < t0 + detector.TIME_WINDOW: continue - + x += xoroshiro128p_normal_float32(rng_state, 0) * sigmas[0] y += xoroshiro128p_normal_float32(rng_state, 0) * sigmas[1] x_dist = abs(x_p - x) y_dist = abs(y_p - y) - + if x_dist > detector.RESPONSE_BIN_SIZE * response.shape[0]: continue if y_dist > detector.RESPONSE_BIN_SIZE * response.shape[1]: @@ -449,7 +397,6 @@ def tracks_current(signals, pixels, tracks, response): sqrt(detector.PIXEL_PITCH**2 + detector.PIXEL_PITCH**2)/2)*2 z_poca, z_start, z_end = z_interval(start, end, x_p, y_p, impact_factor) -# print(z_poca,z_start,z_end,start[2]) if z_poca != 0: z_start_int = z_start - 4*sigmas[2] @@ -508,10 +455,10 @@ def tracks_current(signals, pixels, tracks, response): def sign(x): """ Sign function - + Args: x (float): input number - + Returns: int: 1 if x>=0 else -1 """ @@ -557,8 +504,12 @@ def sum_pixel_signals(pixels_signals, signals, track_starts, pixel_index_map, tr if itick < signals.shape[2]: itime = start_tick + itick if itime < pixels_signals.shape[1]: - cuda.atomic.add(pixels_signals, (pixel_index, itime), signals[itrk][ipix][itick]) - cuda.atomic.add(pixels_tracks_signals, (pixel_index, itime, counter), signals[itrk][ipix][itick]) + cuda.atomic.add(pixels_signals, + (pixel_index, itime), + signals[itrk][ipix][itick]) + cuda.atomic.add(pixels_tracks_signals, + (pixel_index, itime, counter), + signals[itrk][ipix][itick]) @cuda.jit def get_track_pixel_map(track_pixel_map, unique_pix, pixels): diff --git a/larndsim/fee.py b/larndsim/fee.py index c713dd44..e31cf103 100644 --- a/larndsim/fee.py +++ b/larndsim/fee.py @@ -91,18 +91,18 @@ def rotate_tile(pixel_id, tile_id): def gen_event_times(nevents, t0): """ Generate sequential event times assuming events are uncorrelated - + Args: nevents(int): number of event times to generate t0(int): offset to apply [microseconds] - + Returns: array: shape `(nevents,)`, sequential event times [microseconds] """ event_start_time = cp.random.exponential(scale=EVENT_RATE, size=int(nevents)) event_start_time = cp.cumsum(event_start_time) event_start_time += t0 - + return event_start_time @@ -141,20 +141,20 @@ def export_to_hdf5(event_id_list, Returns: tuple: a tuple containing the list of LArPix packets and the list of entries for the `mc_packets_assn` dataset """ - dtype = np.dtype([('track_ids', '(%i,)i8' % track_ids.shape[1]), ('fraction', '(%i,)f8' % current_fractions.shape[2])]) + dtype = np.dtype([('track_ids', f'({track_ids.shape[1]},)i8'), ('fraction', f'({current_fractions.shape[2]},)f8')]) io_groups = np.unique(np.array(list(detector.MODULE_TO_IO_GROUPS.values()))) packets = [] packets_mc = [] packets_frac = [] - + if is_first_event: for io_group in io_groups: packets.append(TimestampPacket(timestamp=0)) packets[-1].chip_key = Key(io_group,0,0) packets_mc.append([-1] * track_ids.shape[1]) packets_frac.append([0] * current_fractions.shape[2]) - + packets.append(SyncPacket(sync_type=b'S', timestamp=0, io_group=io_group)) packets_mc.append([-1] * track_ids.shape[1]) packets_frac.append([0] * current_fractions.shape[2]) @@ -197,7 +197,7 @@ def export_to_hdf5(event_id_list, event = event_id_list[itick,iadc] event_t0 = event_start_time_list[itick] time_tick = int(np.floor(t / CLOCK_CYCLE + event_t0)) - + if event_t0 > ROLLOVER_CYCLES-1 or time_tick > ROLLOVER_CYCLES-1: # 31-bit rollover rollover_count += 1 @@ -217,13 +217,13 @@ def export_to_hdf5(event_id_list, if event != last_event: for io_group in io_groups: packets.append(TimestampPacket(timestamp=event_start_times[unique_events_inv[itick]] * units.mus / units.s)) - packets[-1].chip_key = Key(io_group,0,0) + packets[-1].chip_key = Key(io_group,0,0) packets_mc.append([-1] * track_ids.shape[1]) packets_frac.append([0] * current_fractions.shape[2]) trig_mask = light_trigger_event_id == event if any(trig_mask): - for t_trig, event_id_trig, module_trig in zip(light_trigger_times[trig_mask], light_trigger_event_id[trig_mask], light_trigger_modules[trig_mask]): + for t_trig, module_trig in zip(light_trigger_times[trig_mask], light_trigger_modules[trig_mask]): t_trig = int(np.floor(t_trig / CLOCK_CYCLE + event_t0)) % ROLLOVER_CYCLES for io_group in detector.MODULE_TO_IO_GROUPS[int(module_trig)]: packets.append(TriggerPacket(io_group=io_group, trigger_type=b'\x02', timestamp=t_trig)) @@ -238,7 +238,7 @@ def export_to_hdf5(event_id_list, pix_y % detector.N_PIXELS_PER_TILE[1]), tile_id)] except KeyError: - logger.warning("Pixel ID not valid %i" % pixel_id) + logger.warning(f"Pixel ID not valid {pixel_id}") continue p.dataword = int(adc) diff --git a/larndsim/light_sim.py b/larndsim/light_sim.py index 97c3be27..8cf2b863 100644 --- a/larndsim/light_sim.py +++ b/larndsim/light_sim.py @@ -13,10 +13,9 @@ import h5py -from .consts import light -from .consts.light import LIGHT_TICK_SIZE, LIGHT_WINDOW, SINGLET_FRACTION, TAU_S, TAU_T, LIGHT_GAIN, LIGHT_OSCILLATION_PERIOD, LIGHT_RESPONSE_TIME, LIGHT_DET_NOISE_SAMPLE_SPACING, LIGHT_TRIG_THRESHOLD, LIGHT_TRIG_WINDOW, LIGHT_DIGIT_SAMPLE_SPACING, LIGHT_NBIT, OP_CHANNEL_TO_TPC, OP_CHANNEL_PER_TRIG, TPC_TO_OP_CHANNEL, OP_CHANNEL_PER_TRIG, SIPM_RESPONSE_MODEL, IMPULSE_TICK_SIZE, IMPULSE_MODEL, MC_TRUTH_THRESHOLD, ENABLE_LUT_SMEARING +from .consts.light import LIGHT_TICK_SIZE, LIGHT_WINDOW, SINGLET_FRACTION, TAU_S, TAU_T, LIGHT_GAIN, LIGHT_OSCILLATION_PERIOD, LIGHT_RESPONSE_TIME, LIGHT_DET_NOISE_SAMPLE_SPACING, LIGHT_TRIG_THRESHOLD, LIGHT_TRIG_WINDOW, LIGHT_DIGIT_SAMPLE_SPACING, LIGHT_NBIT, OP_CHANNEL_PER_TRIG, TPC_TO_OP_CHANNEL, OP_CHANNEL_PER_TRIG, SIPM_RESPONSE_MODEL, IMPULSE_TICK_SIZE, IMPULSE_MODEL, MC_TRUTH_THRESHOLD, ENABLE_LUT_SMEARING -from .consts.detector import TPC_BORDERS, MODULE_TO_TPCS +from .consts.detector import MODULE_TO_TPCS from .consts import units as units from .fee import CLOCK_CYCLE, ROLLOVER_CYCLES @@ -39,7 +38,7 @@ def get_nticks(light_incidence): end_time = np.max(light_incidence['t0_det'][mask]) + LIGHT_WINDOW[1] return int(np.ceil((end_time - start_time)/LIGHT_TICK_SIZE)), start_time return int((LIGHT_WINDOW[1] + LIGHT_WINDOW[0])/LIGHT_TICK_SIZE), 0 - + @cuda.jit def sum_light_signals(segments, segment_voxel, segment_track_id, light_inc, lut, start_time, light_sample_inc, light_sample_inc_true_track_id, light_sample_inc_true_photons): @@ -124,12 +123,12 @@ def sum_light_signals(segments, segment_voxel, segment_track_id, light_inc, lut, @nb.njit def scintillation_model(time_tick): """ - Calculates the fraction of scintillation photons emitted + Calculates the fraction of scintillation photons emitted during time interval `time_tick` to `time_tick + 1` - + Args: time_tick(int): time tick relative to t0 - + Returns: float: fraction of scintillation photons """ @@ -143,7 +142,7 @@ def calc_scintillation_effect(light_sample_inc, light_sample_inc_true_track_id, """ Applies a smearing effect due to the liquid argon scintillation time profile using a two decay component scintillation model. - + Args: light_sample_inc(array): shape `(ndet, ntick)`, light incident on each detector light_sample_inc_scint(array): output array, shape `(ndet, ntick)`, light incident on each detector after accounting for scintillation time @@ -153,7 +152,7 @@ def calc_scintillation_effect(light_sample_inc, light_sample_inc_true_track_id, if idet < light_sample_inc.shape[0]: if itick < light_sample_inc.shape[1]: conv_ticks = ceil((LIGHT_WINDOW[1] - LIGHT_WINDOW[0])/LIGHT_TICK_SIZE) - + for jtick in range(max(itick - conv_ticks, 0), itick+1): tick_weight = scintillation_model(itick-jtick) light_sample_inc_scint[idet,itick] += tick_weight * light_sample_inc[idet,jtick] @@ -162,7 +161,7 @@ def calc_scintillation_effect(light_sample_inc, light_sample_inc_true_track_id, for itrue in range(light_sample_inc_true_track_id.shape[-1]): if light_sample_inc_true_track_id[idet,jtick,itrue] == -1: break - + if tick_weight * light_sample_inc_true_photons[idet,jtick,itrue] < MC_TRUTH_THRESHOLD: continue @@ -177,13 +176,13 @@ def calc_scintillation_effect(light_sample_inc, light_sample_inc_true_track_id, @nb.njit def xoroshiro128p_poisson_int32(mean, states, index): """ - Return poisson distributed int32 and advance `states[index]`. For efficiency, + Return poisson distributed int32 and advance `states[index]`. For efficiency, if `mean > 100`, returns a gaussian distributed int32 with `mean == mean` and `std = sqrt(mean)` truncated at 0 (approximately equivalent to a poisson- distributed number) [DOI:10.1007/978-1-4613-8643-8_10] - + Args: mean(float): mean of poisson distribution states(array): array of RNG states @@ -206,13 +205,13 @@ def xoroshiro128p_poisson_int32(mean, states, index): break return x return max(int(cuda.random.xoroshiro128p_normal_float32(states, index) * sqrt(mean) + mean),0) - - + + @cuda.jit def calc_stat_fluctuations(light_sample_inc, light_sample_inc_disc, rng_states): """ Simulates Poisson fluctuations in the number of PE per time tick. - + Args: light_sample_inc(array): shape `(ndet, ntick)`, light incident on each detector light_sample_inc_disc(array): output array, shape `(ndet, ntick)`, number PE in each time interval @@ -260,17 +259,17 @@ def interp(idx, arr, low, high): v1 = arr[i1] return v0 + (v1 - v0) * (idx - i0) - + @nb.njit def sipm_response_model(idet, time_tick): """ Calculates the SiPM response from a PE at `time_tick` relative to the PE time - + Args: idet(int): SiPM index time_tick(int): time tick relative to t0 - + Returns: float: response """ @@ -289,13 +288,13 @@ def sipm_response_model(idet, time_tick): # normalize to 1 impulse /= IMPULSE_TICK_SIZE/LIGHT_TICK_SIZE return impulse - + @cuda.jit def calc_light_detector_response(light_sample_inc, light_sample_inc_true_track_id, light_sample_inc_true_photons, light_response, light_response_true_track_id, light_response_true_photons): """ Simulates the SiPM reponse and digit - + Args: light_sample_inc(array): shape `(ndet, ntick)`, PE produced on each SiPM at each time tick light_response(array): shape `(ndet, ntick)`, ADC value at each time tick @@ -305,16 +304,16 @@ def calc_light_detector_response(light_sample_inc, light_sample_inc_true_track_i if idet < light_sample_inc.shape[0]: if itick < light_sample_inc.shape[1]: conv_ticks = ceil((LIGHT_WINDOW[1] - LIGHT_WINDOW[0])/LIGHT_TICK_SIZE) - + for jtick in range(max(itick - conv_ticks, 0), itick+1): tick_weight = sipm_response_model(idet, itick-jtick) light_response[idet,itick] += LIGHT_GAIN[idet] * tick_weight * light_sample_inc[idet,jtick] - + # loop over convolution tick truth for itrue in range(light_sample_inc_true_track_id.shape[-1]): if light_sample_inc_true_track_id[idet,jtick,itrue] == -1: break - + if abs(tick_weight * light_sample_inc_true_photons[idet,jtick,itrue]) < MC_TRUTH_THRESHOLD: continue @@ -325,32 +324,32 @@ def calc_light_detector_response(light_sample_inc, light_sample_inc_true_track_i light_response_true_track_id[idet,itick,jtrue] = light_sample_inc_true_track_id[idet,itick,itrue] light_response_true_photons[idet,itick,jtrue] += tick_weight * light_sample_inc_true_photons[idet,jtick,itrue] break - + def gen_light_detector_noise(shape, light_det_noise): """ Generates uncorrelated noise with a defined frequency spectrum - + Args: shape(tuple): desired shape of output noise, `shape[0]` must equal `light_det_noise.shape[0]` light_det_noise(array): FFT of noise, `light_det_noise.ndim == 2` - + Returns: array: shape `(shape[0], shape[1])`, randomly generated sample noise """ if not shape[0]: return cp.empty(shape) - + noise_freq = cp.fft.rfftfreq((light_det_noise.shape[-1]-1)*2, d=LIGHT_DET_NOISE_SAMPLE_SPACING) desired_freq = cp.fft.rfftfreq(shape[-1], d=LIGHT_TICK_SIZE) - + bin_size = cp.diff(desired_freq).mean() noise_spectrum = cp.zeros((shape[0], desired_freq.shape[0])) for idet in range(shape[0]): noise_spectrum[idet] = cp.interp(desired_freq, noise_freq, light_det_noise[idet], left=0, right=0) # rescale noise spectrum to have constant noise power with digitizer sample spacing noise_spectrum *= cp.sqrt(cp.diff(noise_freq, axis=-1).mean()/bin_size) * LIGHT_DIGIT_SAMPLE_SPACING / LIGHT_TICK_SIZE - + # generate an FFT with the same frequency power, but with random phase noise = noise_spectrum * cp.exp(2j * cp.pi * cp.random.uniform(size=noise_spectrum.shape)) @@ -371,10 +370,10 @@ def gen_light_detector_noise(shape, light_det_noise): def get_triggers(signal): """ Identifies each simulated ticks that would initiate a trigger taking into account the ADC digitization window - + Args: signal(array): shape `(ndet, nticks)`, simulated signal on each channel - + Returns: tuple: array of tick indices at each trigger (shape `(ntrigs,)`) and array of op channel index (shape `(ntrigs, ndet_module)`) """ @@ -416,7 +415,7 @@ def get_triggers(signal): def digitize_signal(signal, trigger_idx, op_channel_idx, signal_true_track_id, signal_true_photons, digit_signal, digit_signal_true_track_id, digit_signal_true_photons): """ Interpolate signal to the appropriate sampling frequency - + Args: signal(array): shape `(ndet, nticks)`, simulated signal on each channel trigger_idx(array): shape `(ntrigs,)`, tick index for each trigger @@ -424,7 +423,7 @@ def digitize_signal(signal, trigger_idx, op_channel_idx, signal_true_track_id, s digit_signal(array): output array, shape `(ntrigs, ndet_module, nsamples)`, digitized signal """ itrig,idet_module,isample = cuda.grid(3) - + if itrig < digit_signal.shape[0]: if idet_module < digit_signal.shape[1]: if isample < digit_signal.shape[2]: @@ -434,7 +433,7 @@ def digitize_signal(signal, trigger_idx, op_channel_idx, signal_true_track_id, s itick0 = int(floor(sample_tick)) itick1 = int(ceil(sample_tick)) - + itrue = 0 # loop over previous tick truth for jtrue in range(signal_true_track_id.shape[-1]): @@ -442,7 +441,7 @@ def digitize_signal(signal, trigger_idx, op_channel_idx, signal_true_track_id, s break if signal_true_track_id[idet,itick0,jtrue] == -1: break - + photons0, photons1 = 0, 0 # if matches the current sample track or we have empty truth slot, add truth info @@ -451,7 +450,7 @@ def digitize_signal(signal, trigger_idx, op_channel_idx, signal_true_track_id, s itrue += 1 # interpolate true photons photons0 = signal_true_photons[idet,itick0,jtrue] - + if abs(photons0) < MC_TRUTH_THRESHOLD: continue @@ -473,7 +472,7 @@ def digitize_signal(signal, trigger_idx, op_channel_idx, signal_true_track_id, s def sim_triggers(bpg, tpb, signal, signal_true_track_id, signal_true_photons, trigger_idx, op_channel_idx, digit_samples, light_det_noise): """ Generates digitized waveforms at specified simulation tick indices - + Args: bpg(tuple): blocks per grid used to generate digitized waveforms, `len(bpg) == 3`, `prod(bpg) * prod(tpb) >= digit_samples.size` tpb(tuple): threads per grid used to generate digitized waveforms, `len(bpg) == 3`, `bpg[i] * tpb[i] >= digit_samples.shape[i]` @@ -484,7 +483,7 @@ def sim_triggers(bpg, tpb, signal, signal_true_track_id, signal_true_photons, tr op_channel_idx(array): shape `(ntrigs, ndet_module)`, optical channel indices for each trigger digit_samples(int): number of digitizations per waveform light_det_noise(array): shape `(ndet, nnoise_bins)`, noise spectrum for each channel (only used if waveforms extend past simulated signal) - + Returns: array: shape `(ntrigs, ndet_module, digit_samples)`, digitized waveform on each channel for each trigger """ @@ -494,7 +493,7 @@ def sim_triggers(bpg, tpb, signal, signal_true_track_id, signal_true_photons, tr # exit if no triggers if digit_signal.shape[0] == 0: return digit_signal, digit_signal_true_track_id, digit_signal_true_photons - + padded_trigger_idx = trigger_idx.copy() # pad front of simulation with noise, if trigger close to start of simulation window @@ -506,30 +505,30 @@ def sim_triggers(bpg, tpb, signal, signal_true_track_id, signal_true_photons, tr signal_true_track_id = cp.concatenate([cp.full(pad_shape + signal_true_track_id.shape[-1:], -1, dtype=signal_true_track_id.dtype), signal_true_track_id], axis=-2) signal_true_photons = cp.concatenate([cp.zeros(pad_shape + signal_true_photons.shape[-1:], signal_true_photons.dtype), signal_true_photons], axis=-2) padded_trigger_idx += pad_shape[1] - + # pad end of simulation with noise, if trigger close to end of simulation window post_digit_ticks = int(ceil(LIGHT_TRIG_WINDOW[1]/LIGHT_TICK_SIZE)) if post_digit_ticks + padded_trigger_idx.max() > signal.shape[1]: pad_shape = (signal.shape[0], int(post_digit_ticks + padded_trigger_idx.max() - signal.shape[1])) - print('posttrigger', pad_shape, padded_trigger_idx.min(), post_digit_ticks, signal.shape) + print('posttrigger', pad_shape, padded_trigger_idx.min(), post_digit_ticks, signal.shape) signal = cp.concatenate([signal, gen_light_detector_noise(pad_shape, light_det_noise)], axis=-1) signal_true_track_id = cp.concatenate([signal_true_track_id, cp.full(pad_shape + signal_true_track_id.shape[-1:], -1, dtype=signal_true_track_id.dtype)], axis=-2) signal_true_photons = cp.concatenate([signal_true_photons, cp.zeros(pad_shape + signal_true_photons.shape[-1:], dtype=signal_true_photons.dtype)], axis=-2) - + digitize_signal[bpg,tpb](signal, padded_trigger_idx, op_channel_idx, signal_true_track_id, signal_true_photons, digit_signal, digit_signal_true_track_id, digit_signal_true_photons) # truncate to correct number of bits digit_signal = cp.round(digit_signal / 2**(16-LIGHT_NBIT)) * 2**(16-LIGHT_NBIT) - + return digit_signal, digit_signal_true_track_id, digit_signal_true_photons def export_to_hdf5(event_id, start_times, trigger_idx, op_channel_idx, waveforms, output_filename, event_times, waveforms_true_track_id, waveforms_true_photons): """ Saves waveforms to output file - + Args: event_id(array): shape `(ntrigs,)`, event id for each trigger start_times(array): shape `(ntrigs,)`, simulation time offset for each trigger [microseconds] @@ -540,15 +539,15 @@ def export_to_hdf5(event_id, start_times, trigger_idx, op_channel_idx, waveforms event_times(array): shape `(nevents,)`, global event t0 for each unique event [microseconds] waveforms_true_track_id(array): shape `(ntrigs, ndet, nsamples, ntruth)`, segment ids contributing to each sample waveforms_true_photons(array): shape `(ntrigs, ndet, nsamples, ntruth)`, true photocurrent at each sample - + """ if event_id.shape[0] == 0: return - + unique_events, unique_events_inv = np.unique(event_id, return_inverse=True) event_start_times = event_times[unique_events_inv] event_sync_times = (event_times[unique_events_inv] / CLOCK_CYCLE).astype(int) % ROLLOVER_CYCLES - + with h5py.File(output_filename, 'a') as f: trig_data = np.empty(trigger_idx.shape[0], dtype=np.dtype([('op_channel','i4',(op_channel_idx.shape[-1])), ('ts_s','f8'), ('ts_sync','u8')])) trig_data['op_channel'] = op_channel_idx.get() @@ -561,7 +560,7 @@ def export_to_hdf5(event_id, start_times, trigger_idx, op_channel_idx, waveforms truth_data = np.empty(waveforms_true_track_id.shape[:-1], dtype=truth_dtype) truth_data['track_ids'] = waveforms_true_track_id.get() truth_data['pe_current'] = waveforms_true_photons.get() - + if 'light_wvfm' not in f: f.create_dataset('light_wvfm', data=waveforms.get(), maxshape=(None,None,None)) f.create_dataset('light_trig', data=trig_data, maxshape=(None,)) @@ -578,4 +577,4 @@ def export_to_hdf5(event_id, start_times, trigger_idx, op_channel_idx, waveforms f['light_wvfm_mc_assn'].resize(f['light_wvfm_mc_assn'].shape[0] + truth_data.shape[0], axis=0) f['light_wvfm_mc_assn'][-truth_data.shape[0]:] = truth_data - +