Skip to content

Commit

Permalink
for mpvmpr and no modular variations
Browse files Browse the repository at this point in the history
  • Loading branch information
Yifan Chen committed Feb 9, 2024
1 parent 719fb88 commit 458680c
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 78 deletions.
169 changes: 92 additions & 77 deletions src/proto_nd_flow/reco/charge/raw_event_generator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -148,11 +148,17 @@ def init(self):

self.packets_dtype = self.packets.dtype
if self.is_mc:
self.is_mc_neutrino = True
self.mc_assn = self.input_fh['mc_packets_assn']
self.mc_tracks = self.input_fh['segments']
self.mc_trajectories = self.input_fh['trajectories']
self.mc_events = self.input_fh['mc_hdr']
self.mc_stack = self.input_fh['mc_stack']
try:
self.mc_events = self.input_fh['mc_hdr']
self.mc_stack = self.input_fh['mc_stack']
except:
self.is_mc_neutrino = False
print("Hope you are not processing neutrino simulation! There is no information for neutrino interactions.")
pass

# initialize data objects
self.data_manager.create_dset(self.raw_event_dset_name, dtype=self.raw_event_dtype)
Expand All @@ -177,32 +183,35 @@ def init(self):
self.data_manager.set_attrs(self.raw_event_dset_name,
mc_tracks_dset_name=self.mc_tracks_dset_name,
mc_trajectories_dset_name=self.mc_trajectories_dset_name,
mc_events_dset_name=self.mc_events_dset_name,
mc_stack_dset_name=self.mc_stack_dset_name,
mc_packet_fraction_dset_name=self.mc_packet_fraction_dset_name)
if self.is_mc_neutrino:
self.data_manager.set_attrs(self.raw_event_dset_name,
mc_events_dset_name=self.mc_events_dset_name,
mc_stack_dset_name=self.mc_stack_dset_name)

self.data_manager.create_dset(self.mc_packet_fraction_dset_name, dtype=self.mc_assn['fraction'].dtype)
self.data_manager.create_ref(self.packets_dset_name, self.mc_packet_fraction_dset_name)

# copy datasets from source file
# MC interaction summary info
self.data_manager.create_dset(self.mc_events_dset_name, dtype=self.mc_events.dtype)
ninter = len(self.mc_events)
inter_sl = slice(
ceil(ninter / self.size * self.rank),
ceil(ninter / self.size * (self.rank + 1)))
self.data_manager.reserve_data(self.mc_events_dset_name, inter_sl)
self.data_manager.write_data(self.mc_events_dset_name, inter_sl,
self.mc_events[inter_sl])

# MC generator particle stack
self.data_manager.create_dset(self.mc_stack_dset_name, dtype=self.mc_stack.dtype)
nstack = len(self.mc_stack)
stack_sl = slice(
ceil(nstack / self.size * self.rank),
ceil(nstack / self.size * (self.rank + 1)))
self.data_manager.reserve_data(self.mc_stack_dset_name, stack_sl)
self.data_manager.write_data(self.mc_stack_dset_name, stack_sl, self.mc_stack[stack_sl])
if self.is_mc_neutrino:
# copy datasets from source file
# MC interaction summary info
self.data_manager.create_dset(self.mc_events_dset_name, dtype=self.mc_events.dtype)
ninter = len(self.mc_events)
inter_sl = slice(
ceil(ninter / self.size * self.rank),
ceil(ninter / self.size * (self.rank + 1)))
self.data_manager.reserve_data(self.mc_events_dset_name, inter_sl)
self.data_manager.write_data(self.mc_events_dset_name, inter_sl,
self.mc_events[inter_sl])

# MC generator particle stack
self.data_manager.create_dset(self.mc_stack_dset_name, dtype=self.mc_stack.dtype)
nstack = len(self.mc_stack)
stack_sl = slice(
ceil(nstack / self.size * self.rank),
ceil(nstack / self.size * (self.rank + 1)))
self.data_manager.reserve_data(self.mc_stack_dset_name, stack_sl)
self.data_manager.write_data(self.mc_stack_dset_name, stack_sl, self.mc_stack[stack_sl])

# edep-sim energy segments/deposits
self.data_manager.create_dset(self.mc_tracks_dset_name, dtype=self.mc_tracks.dtype)
Expand Down Expand Up @@ -233,18 +242,20 @@ def init(self):
self.mc_trajectories[traj_sl])

# set up references
self.data_manager.create_ref(self.raw_event_dset_name, self.mc_events_dset_name)
self.data_manager.create_ref(self.packets_dset_name, self.mc_tracks_dset_name)
self.data_manager.create_ref(self.mc_events_dset_name, self.mc_trajectories_dset_name)
self.data_manager.create_ref(self.mc_events_dset_name, self.mc_tracks_dset_name)
self.data_manager.create_ref(self.mc_events_dset_name, self.mc_stack_dset_name)
self.data_manager.create_ref(self.mc_stack_dset_name, self.mc_trajectories_dset_name)
self.data_manager.create_ref(self.mc_trajectories_dset_name, self.mc_tracks_dset_name)
if self.is_mc_neutrino:
self.data_manager.create_ref(self.raw_event_dset_name, self.mc_events_dset_name)
self.data_manager.create_ref(self.mc_events_dset_name, self.mc_trajectories_dset_name)
self.data_manager.create_ref(self.mc_events_dset_name, self.mc_tracks_dset_name)
self.data_manager.create_ref(self.mc_events_dset_name, self.mc_stack_dset_name)
self.data_manager.create_ref(self.mc_stack_dset_name, self.mc_trajectories_dset_name)

# create references between trajectories and tracks
# eventID --> vertexID for latest production files
intr_evid = self.mc_events['vertex_id'][:]
stack_evid = self.mc_stack['vertex_id'][:]
if self.is_mc_neutrino:
intr_evid = self.mc_events['vertex_id'][:]
stack_evid = self.mc_stack['vertex_id'][:]
traj_evid = self.mc_trajectories['vertex_id'][:]
tracks_evid = self.mc_tracks['vertex_id'][:]
evs, ev_traj_start, ev_track_start = np.intersect1d(
Expand All @@ -257,7 +268,8 @@ def init(self):
ceil(len(evs) / self.size * self.rank),
ceil(len(evs) / self.size * (self.rank + 1)))

stack_trackid = self.mc_stack['traj_id'][:]
if self.is_mc_neutrino:
stack_trackid = self.mc_stack['traj_id'][:]
traj_trackid = self.mc_trajectories['traj_id'][:]
tracks_trackid = self.mc_tracks['traj_id'][:]
iter_ = tqdm(range(truth_slice.start, truth_slice.stop), smoothing=1, desc='generating truth references') if self.rank == 0 else range(truth_slice.start, truth_slice.stop)
Expand All @@ -278,42 +290,44 @@ def init(self):
ref[:, 1] += track_start
self.data_manager.write_ref(self.mc_trajectories_dset_name, self.mc_tracks_dset_name, ref)

# Create refs for interactions --> traj
intr_evid_block = np.expand_dims(intr_evid[:], 0) # Might need to modify for MPI running
ref = np.argwhere((ev == intr_evid_block) & (ev == traj_evid_block))
ref[:, 0] += traj_start
ref[:, 1] += 0 #i + inter_sl.start # Might need to modify for MPI running
self.data_manager.write_ref(self.mc_trajectories_dset_name, self.mc_events_dset_name, ref)

# Create refs for interactions --> tracks
intr_evid_block = np.expand_dims(intr_evid[:], -1) # Might need to modify for MPI running
ref = np.argwhere((ev == track_evid_block) & (ev == intr_evid_block))
ref[:, 0] += 0 #i + inter_sl.start # Might need to modify for MPI running
ref[:, 1] += track_start
self.data_manager.write_ref(self.mc_events_dset_name, self.mc_tracks_dset_name, ref)

# Create refs for interactions --> generator particle stack
stack_evid_block = np.expand_dims(stack_evid[:], 0) # Might need to modify for MPI running
ref = np.argwhere((ev == intr_evid_block) & (ev == stack_evid_block))
# ref[:, 0] += 0 # Placeholders for now.
# ref[:, 1] += 0 # This extra offset might be needed for future MPI running
self.data_manager.write_ref(self.mc_events_dset_name, self.mc_stack_dset_name, ref)

# Create refs for generator particle stack --> traj
stack_trackid_block = np.expand_dims(stack_trackid[:], -1) # Might need to modify for MPI running
traj_trackid_block = np.transpose(traj_trackid_block)
stack_evid_block = np.transpose(stack_evid_block) # Might need to modify for MPI running
traj_evid_block = np.transpose(traj_evid_block)
ref = np.argwhere((stack_trackid_block == traj_trackid_block) & (stack_evid_block == traj_evid_block))
ref[:, 0] += 0 # Might need to modify for MPI running
ref[:, 1] += traj_start
self.data_manager.write_ref(self.mc_stack_dset_name, self.mc_trajectories_dset_name, ref)
if self.is_mc_neutrino:
# Create refs for interactions --> traj
intr_evid_block = np.expand_dims(intr_evid[:], 0) # Might need to modify for MPI running
ref = np.argwhere((ev == intr_evid_block) & (ev == traj_evid_block))
ref[:, 0] += traj_start
ref[:, 1] += 0 #i + inter_sl.start # Might need to modify for MPI running
self.data_manager.write_ref(self.mc_trajectories_dset_name, self.mc_events_dset_name, ref)

# Create refs for interactions --> tracks
intr_evid_block = np.expand_dims(intr_evid[:], -1) # Might need to modify for MPI running
ref = np.argwhere((ev == track_evid_block) & (ev == intr_evid_block))
ref[:, 0] += 0 #i + inter_sl.start # Might need to modify for MPI running
ref[:, 1] += track_start
self.data_manager.write_ref(self.mc_events_dset_name, self.mc_tracks_dset_name, ref)

# Create refs for interactions --> generator particle stack
stack_evid_block = np.expand_dims(stack_evid[:], 0) # Might need to modify for MPI running
ref = np.argwhere((ev == intr_evid_block) & (ev == stack_evid_block))
# ref[:, 0] += 0 # Placeholders for now.
# ref[:, 1] += 0 # This extra offset might be needed for future MPI running
self.data_manager.write_ref(self.mc_events_dset_name, self.mc_stack_dset_name, ref)

# Create refs for generator particle stack --> traj
stack_trackid_block = np.expand_dims(stack_trackid[:], -1) # Might need to modify for MPI running
traj_trackid_block = np.transpose(traj_trackid_block)
stack_evid_block = np.transpose(stack_evid_block) # Might need to modify for MPI running
traj_evid_block = np.transpose(traj_evid_block)
ref = np.argwhere((stack_trackid_block == traj_trackid_block) & (stack_evid_block == traj_evid_block))
ref[:, 0] += 0 # Might need to modify for MPI running
ref[:, 1] += traj_start
self.data_manager.write_ref(self.mc_stack_dset_name, self.mc_trajectories_dset_name, ref)
else:
self.data_manager.write_ref(self.mc_trajectories_dset_name, self.mc_tracks_dset_name, np.empty((0,2)))
self.data_manager.write_ref(self.mc_trajectories_dset_name, self.mc_events_dset_name, np.empty((0,2)))
self.data_manager.write_ref(self.mc_events_dset_name, self.mc_tracks_dset_name, np.empty((0,2)))
self.data_manager.write_ref(self.mc_events_dset_name, self.mc_stack_dset_name, np.empty((0,2)))
self.data_manager.write_ref(self.mc_stack_dset_name, self.mc_trajectories_dset_name, np.empty((0,2)))
if self.is_mc_neutrino:
self.data_manager.write_ref(self.mc_trajectories_dset_name, self.mc_events_dset_name, np.empty((0,2)))
self.data_manager.write_ref(self.mc_events_dset_name, self.mc_tracks_dset_name, np.empty((0,2)))
self.data_manager.write_ref(self.mc_events_dset_name, self.mc_stack_dset_name, np.empty((0,2)))
self.data_manager.write_ref(self.mc_stack_dset_name, self.mc_trajectories_dset_name, np.empty((0,2)))

# if self.is_mc:
# # copy meta-data from input file
Expand Down Expand Up @@ -447,18 +461,19 @@ def next(self):
self.data_manager.write_data(self.mc_packet_fraction_dset_name, sl, event_packet_fraction.compressed())
self.data_manager.write_ref(self.packets_dset_name, self.mc_packet_fraction_dset_name, np.c_[ref[:,0], sl])

# find events associated with tracks
if H5FLOW_MPI:
self.comm.barrier()
ref_dset, ref_dir = self.data_manager.get_ref(self.mc_tracks_dset_name, self.mc_events_dset_name)
ref_region = self.data_manager.get_ref_region(self.mc_tracks_dset_name, self.mc_events_dset_name)
mc_evs = dereference(ref[:, 1], ref_dset, region=ref_region,
ref_direction=ref_dir, indices_only=True)

ev_idcs = np.broadcast_to(np.expand_dims(ev_idcs, axis=-1), event_tracks.shape)
ref = np.c_[ev_idcs[~event_tracks.mask].ravel(), mc_evs.ravel()]
ref = np.unique(ref, axis=0) if len(ref) else ref
self.data_manager.write_ref(self.raw_event_dset_name, self.mc_events_dset_name, ref)
if self.is_mc_neutrino:
# find events associated with tracks
if H5FLOW_MPI:
self.comm.barrier()
ref_dset, ref_dir = self.data_manager.get_ref(self.mc_tracks_dset_name, self.mc_events_dset_name)
ref_region = self.data_manager.get_ref_region(self.mc_tracks_dset_name, self.mc_events_dset_name)
mc_evs = dereference(ref[:, 1], ref_dset, region=ref_region,
ref_direction=ref_dir, indices_only=True)

ev_idcs = np.broadcast_to(np.expand_dims(ev_idcs, axis=-1), event_tracks.shape)
ref = np.c_[ev_idcs[~event_tracks.mask].ravel(), mc_evs.ravel()]
ref = np.unique(ref, axis=0) if len(ref) else ref
self.data_manager.write_ref(self.raw_event_dset_name, self.mc_events_dset_name, ref)

return raw_event_slice if nevents else H5FlowGenerator.EMPTY

Expand Down
5 changes: 4 additions & 1 deletion src/proto_nd_flow/reco/light/mc_event_generator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,10 @@ def init(self):

# copy and remap the truth information
if type(self.light_dat) is h5py.Group: # We have four 96-column matrices
light_dat = self._bloat_light_dat(self.light_dat) # Now have 384 col
if 'light_dat_allmodules' in self.light_dat:
light_dat = self.light_dat['light_dat_allmodules']
else:
light_dat = self._bloat_light_dat(self.light_dat) # Now have 384 col
else: # We have the old 384-column matrix
light_dat = self.light_dat

Expand Down

0 comments on commit 458680c

Please sign in to comment.