Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

unfortunately all changes and bug fixes mangled into one #237

Merged
merged 1 commit into from
Jul 2, 2024
Merged

unfortunately all changes and bug fixes mangled into one #237

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 106 additions & 20 deletions cli/simulate_pixels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,8 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
results['unique_pix'],
results['current_fractions'],
results['track_pixel_map'],
results['traj_pixel_map'],
#results['segidx2trackid'],
output_filename, # defined earlier in script
uniq_event_times,
is_first_batch=is_first_batch,
Expand Down Expand Up @@ -454,17 +456,22 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
# First of all we load the edep-sim output
with h5py.File(input_filename, 'r') as f:
tracks = np.array(f['segments'])
if 'segment_id' in tracks.dtype.names:
segment_ids = tracks['segment_id']
else:
dtype = tracks.dtype.descr
dtype = [('segment_id','u4')] + dtype
new_tracks = np.empty(tracks.shape, dtype=np.dtype(dtype, align=True))
new_tracks['segment_id'] = np.arange(tracks.shape[0], dtype='u4')
for field in dtype[1:]:
new_tracks[field[0]] = tracks[field[0]]
tracks = new_tracks
segment_ids = tracks['segment_id']
#if 'segment_id' in tracks.dtype.names:
# segment_ids = tracks['segment_id']
# trajectory_ids = tracks['file_traj_id']
#else:
# dtype = tracks.dtype.descr
# dtype = [('segment_id','u4')] + dtype
# new_tracks = np.empty(tracks.shape, dtype=np.dtype(dtype, align=True))
# new_tracks['segment_id'] = np.arange(tracks.shape[0], dtype='u4')
# for field in dtype[1:]:
# new_tracks[field[0]] = tracks[field[0]]
# tracks = new_tracks
# segment_ids = tracks['segment_id']

segment_ids = tracks['segment_id']
trajectory_ids = tracks['file_traj_id']

try:
trajectories = np.array(f['trajectories'])
input_has_trajectories = True
Expand Down Expand Up @@ -506,6 +513,7 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
print(f'Selecting only the first {n_events} events for simulation.')
max_eventID = np.unique(tracks[sim.EVENT_SEPARATOR])[n_events-1]
segment_ids = segment_ids[tracks[sim.EVENT_SEPARATOR] <= max_eventID]
trajectory_ids = trajectory_ids[tracks[sim.EVENT_SEPARATOR] <= max_eventID]
tracks = tracks[tracks[sim.EVENT_SEPARATOR] <= max_eventID]

if input_has_trajectories:
Expand Down Expand Up @@ -604,6 +612,7 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
# First copy all tracks and segment_ids
all_mod_tracks = tracks
all_mod_segment_ids = segment_ids
all_mod_trajectory_ids = trajectory_ids
if mod2mod_variation == None or mod2mod_variation == False:
mod_ids = [-1]
# Sub-select only segments in active volumes
Expand All @@ -613,6 +622,7 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
active_tracks_mask = active_volume.select_active_volume(all_mod_tracks, detector.TPC_BORDERS)
tracks = all_mod_tracks[active_tracks_mask]
segment_ids = all_mod_segment_ids[active_tracks_mask]
trajectory_ids = all_mod_trajectory_ids[active_tracks_mask]
end_mask = time()
print(f" {end_mask-start_mask:.2f} s")
else:
Expand All @@ -622,6 +632,7 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o

# Convention module counting start from 1
for i_mod in mod_ids:
print(f'Simulating module {i_mod}')
if mod2mod_variation:
consts.detector.set_detector_properties(detector_properties, pixel_layout, i_mod)
# Currently shouln't be necessary to reload light props, but if
Expand All @@ -646,6 +657,7 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
module_tracks_mask = active_volume.select_active_volume(all_mod_tracks, module_borders)
tracks = all_mod_tracks[module_tracks_mask]
segment_ids = all_mod_segment_ids[module_tracks_mask]
trajectory_ids = all_mod_trajectory_ids[module_tracks_mask]
RangePop()

# find the module that triggers
Expand Down Expand Up @@ -763,6 +775,8 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o

track_ids = cp.asarray(np.arange(segment_ids.shape[0], dtype=int))
segment_ids_arr = cp.asarray(segment_ids)
trajectory_ids_arr = cp.asarray(trajectory_ids)

# We divide the sample in portions that can be processed by the GPU
is_first_batch = True
logger.start()
Expand All @@ -771,8 +785,16 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
i_trig = 0
sync_start = event_times[0] // (fee.CLOCK_RESET_PERIOD * fee.CLOCK_CYCLE) * (fee.CLOCK_RESET_PERIOD * fee.CLOCK_CYCLE) + (fee.CLOCK_RESET_PERIOD * fee.CLOCK_CYCLE)
det_borders = module_borders if mod2mod_variation else detector.TPC_BORDERS
#print('Detector boarders',det_borders)
i_batch=0
#for batch_mask in tqdm(batching.TPCBatcher(all_mod_tracks, tracks, sim.EVENT_SEPARATOR, tpc_batch_size=sim.EVENT_BATCH_SIZE, tpc_borders=det_borders),
# desc='Simulating batches...', ncols=80, smoothing=0):
# i_batch = i_batch+1
# print(i_batch,np.unique(all_mod_tracks[sim.EVENT_SEPARATOR]))
i_batch = 0
for batch_mask in tqdm(batching.TPCBatcher(all_mod_tracks, tracks, sim.EVENT_SEPARATOR, tpc_batch_size=sim.EVENT_BATCH_SIZE, tpc_borders=det_borders),
desc='Simulating batches...', ncols=80, smoothing=0):
#print('Batch index',i_batch)
i_batch = i_batch+1
# grab only tracks from current batch
track_subset = tracks[batch_mask]
Expand Down Expand Up @@ -804,13 +826,19 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
del null_light_results_acc['light_event_id']
# Nothing to simulate for charge readout?
continue

#print('Looping over',evt_tracks.shape[0],'tracks')
for itrk in tqdm(range(0, evt_tracks.shape[0], sim.BATCH_SIZE),
delay=1, desc=' Simulating event %i batches...' % ievd, leave=False, ncols=80):
if itrk > 0:
warnings.warn(f"Entered sub-batch loop, results may not be accurate! Consider increasing batch_size (currently {sim.BATCH_SIZE}) in the simulation_properties file.")

selected_tracks = evt_tracks[itrk:itrk+sim.BATCH_SIZE]
#segidx2trackid = []
#print('Tracks in batch',selected_tracks.dtype.fields.keys())
#for i,seg in enumerate(selected_tracks):
# #print(i,seg['file_traj_id'],seg['segment_id'])
# segidx2trackid.append(seg['file_traj_id'])
#segidx2trackid=np.array(segidx2trackid)

RangePush("event_id_map")
event_ids = selected_tracks[sim.EVENT_SEPARATOR]
Expand All @@ -835,20 +863,22 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o

active_pixels = cp.full((selected_tracks.shape[0], max_pixels[0]), -1, dtype=np.int32)
neighboring_pixels = cp.full((selected_tracks.shape[0], max_neighboring_pixels), -1, dtype=np.int32)
neighboring_radius = cp.full((selected_tracks.shape[0], max_neighboring_pixels), -1, dtype=np.int32)
n_pixels_list = cp.zeros(shape=(selected_tracks.shape[0]))

if not active_pixels.shape[1] or not neighboring_pixels.shape[1]:
if light.LIGHT_SIMULATED and (light.LIGHT_TRIG_MODE == 0 or light.LIGHT_TRIG_MODE == 1):
null_light_results_acc['light_event_id'].append(cp.full(1, ievd)) # one event
save_results(event_times, is_first_batch, null_light_results_acc, i_trig, i_mod, light_only=True)
i_trig += 1 # add to the trigger counter
i_trig += 1 # add to the trigger countern_max_pixels
del null_light_results_acc['light_event_id']
continue

RangePush("get_pixels")
pixels_from_track.get_pixels[BPG,TPB](selected_tracks,
active_pixels,
neighboring_pixels,
neighboring_radius,
n_pixels_list,
max_radius)
RangePop()
Expand All @@ -860,6 +890,34 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
unique_pix = unique_pix[(unique_pix != -1)]
RangePop()

def invert_array_map(in_map,pix_set):
'''
Invert the map of unique segment id => a set of unique pixel IDs to a map of unique
pixel index => a set of segment IDs.

Args:
in_map (:obj:`numpy.ndarray`): 2D array where segment index => list of pixel IDs
pix_set (:obj:`numpy.ndarray`): 1D array containing all unique pixel IDs
Returns:
ndarray: 2D array where pixel index => list of segment index
'''
pixids,counts=cp.unique(in_map[in_map>=0].flatten(),return_counts=True)

pix_id2idx = {val.item():i for i,val in enumerate(pix_set)}

mymap=cp.full(shape=(pix_set.shape[0],counts.max().item()),fill_value=-1,dtype=int)
curr_idx=cp.zeros(shape=(len(pix_id2idx),),dtype=int)
for seg_idx in range(in_map.shape[0]):
ass = in_map[seg_idx]
for pixid in ass:
if pixid<0: break
pix_idx = pix_id2idx[pixid.item()]
mymap[pix_idx][curr_idx[pix_idx]]=seg_idx
curr_idx[pix_idx] += 1
return mymap

assmap_pix2seg = invert_array_map(active_pixels,unique_pix)

if not unique_pix.shape[0]:
if light.LIGHT_SIMULATED and (light.LIGHT_TRIG_MODE == 0 or light.LIGHT_TRIG_MODE == 1):
null_light_results_acc['light_event_id'].append(cp.full(1, ievd)) # one event
Expand Down Expand Up @@ -900,10 +958,18 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o

RangePush("track_pixel_map")
# Mapping between unique pixel array and track array index
track_pixel_map = cp.full((unique_pix.shape[0], detsim.MAX_TRACKS_PER_PIXEL), -1)
max_segments_to_trace = max(assmap_pix2seg.shape[1],detsim.MAX_TRACKS_PER_PIXEL)
track_pixel_map = cp.full((unique_pix.shape[0], max_segments_to_trace), -1)
TPB = 32
#BPG = max(ceil(unique_pix.shape[0] / TPB),1)
BPG = max(ceil(unique_pix.shape[0] / TPB),1)
detsim.get_track_pixel_map[BPG, TPB](track_pixel_map, unique_pix, neighboring_pixels)
detsim.get_track_pixel_map2[BPG, TPB](track_pixel_map,
unique_pix,
#active_pixels,
neighboring_pixels,
neighboring_radius,
neighboring_radius.max().item()+1,
)
RangePop()

RangePush("sum_pixels_signals")
Expand All @@ -915,16 +981,32 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
BPG = (BPG_X, BPG_Y, BPG_Z)
pixels_signals = cp.zeros((len(unique_pix), len(detector.TIME_TICKS)))
pixels_tracks_signals = cp.zeros((len(unique_pix),
len(detector.TIME_TICKS),
track_pixel_map.shape[1]))
len(detector.TIME_TICKS),
track_pixel_map.shape[1],)
)
detsim.sum_pixel_signals[BPG,TPB](pixels_signals,
signals,
track_starts,
pixel_index_map,
track_pixel_map,
pixels_tracks_signals)
pixels_tracks_signals,
)
RangePop()

#np.savez('wf.npz',
# actives=cp.asnumpy(active_pixels),
# signals=cp.asnumpy(signals),
# pix2idx=cp.asnumpy(pixel_index_map),
# trk2pix=cp.asnumpy(track_pixel_map),
# neighbors=cp.asnumpy(neighboring_pixels),
# distances=cp.asnumpy(neighboring_radius),
# uniquepix=cp.asnumpy(unique_pix),
# track_starts=cp.asnumpy(track_starts/detector.TIME_SAMPLING),
# pixels_tracks_signals=cp.asnumpy(pixels_tracks_signals),
# assmap_pix2seg=cp.asnumpy(assmap_pix2seg),
# )


RangePush("get_adc_values")
# Here we simulate the electronics response (the self-triggering cycle) and the signal digitization
time_ticks = cp.linspace(0, len(unique_eventIDs) * detector.TIME_INTERVAL[1], pixels_signals.shape[1]+1)
Expand All @@ -948,7 +1030,6 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
rng_states,
current_fractions,
pixel_thresholds)

# get list of adc values
if pixel_gains_file is not None:
pixel_gains = cp.array(pixel_gains_lut[unique_pix.ravel()])
Expand All @@ -965,10 +1046,15 @@ def save_results(event_times, is_first_batch, results, i_trig, i_mod=-1, light_o
results_acc['adc_tot_ticks'].append(adc_ticks_list)
results_acc['unique_pix'].append(unique_pix)
results_acc['current_fractions'].append(current_fractions)
#track_pixel_map[track_pixel_map != -1] += first_trk_id + itrk
#results_acc['segidx2trackid'].append(segidx2trackid)
traj_pixel_map = cp.full(track_pixel_map.shape,-1)
traj_pixel_map[:] = track_pixel_map
traj_pixel_map[traj_pixel_map != -1] = trajectory_ids_arr[batch_mask][traj_pixel_map[traj_pixel_map != -1] + itrk]
track_pixel_map[track_pixel_map != -1] = segment_ids_arr[batch_mask][track_pixel_map[track_pixel_map != -1] + itrk]
results_acc['traj_pixel_map'].append(traj_pixel_map)
results_acc['track_pixel_map'].append(track_pixel_map)


# ~~~ Light detector response simulation ~~~
if light.LIGHT_SIMULATED:
RangePush("sum_light_signals")
Expand Down
Loading