diff --git a/src/proto_nd_flow/reco/charge/raw_event_generator.py b/src/proto_nd_flow/reco/charge/raw_event_generator.py index e122d8fb..f96c2508 100644 --- a/src/proto_nd_flow/reco/charge/raw_event_generator.py +++ b/src/proto_nd_flow/reco/charge/raw_event_generator.py @@ -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) @@ -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) @@ -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( @@ -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) @@ -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 @@ -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 diff --git a/src/proto_nd_flow/reco/light/mc_event_generator.py b/src/proto_nd_flow/reco/light/mc_event_generator.py index a87f3e67..034ff60e 100644 --- a/src/proto_nd_flow/reco/light/mc_event_generator.py +++ b/src/proto_nd_flow/reco/light/mc_event_generator.py @@ -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