mc_truth/interactions

data/proto_nd_flow/multi_tile_layout-2.5.16.yaml

@@ -19200,7 +19200,7 @@ chip_channel_to_position:
   - 76
   - 7
 multitile_layout_version: 2.5.16
-pixel_pitch: 3.8
+pixel_pitch: 3.87975
 tile_chip_to_io:
   1:
     11: 1001

src/proto_nd_flow/reco/charge/raw_event_generator.py

@@ -3,6 +3,7 @@
 from numpy.lib import recfunctions as rfn
 import h5py
 import logging
+import warnings
 from math import ceil
 from tqdm import tqdm
 
@@ -154,12 +155,28 @@ def init(self):
             self.mc_trajectories = self.input_fh['trajectories']
             try:
                 self.mc_events = self.input_fh['mc_hdr']
+            except:
+                self.mc_events = self.input_fh['vertices']
+            try:
                 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
 
+        # set up attribute name for vertex_id and traj_id
+        if 'file_vertex_id' in self.input_fh['vertices'].dtype.names:
+            self.vertex_id_name = 'file_vertex_id'
+        else:
+            self.vertex_id_name = 'vertex_id'
+            print("Using 'vertex_id'(unique for beam simulation, but not for mpvmpr) instead of 'file_vertex_id'.")
+
+        if 'file_traj_id' in self.input_fh['trajectories'].dtype.names:
+            self.traj_id_name = 'file_traj_id'
+        else:
+            self.traj_id_name = 'traj_id'
+            warnings.warn("Using 'traj_id' instead of 'file_traj_id'. 'traj_id' is not unique across the file and will cause reference issues.")
+
         # initialize data objects
         self.data_manager.create_dset(self.raw_event_dset_name, dtype=self.raw_event_dtype)
         self.data_manager.create_dset(self.packets_dset_name, dtype=self.packets_dtype)
@@ -183,27 +200,17 @@ 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_packet_fraction_dset_name=self.mc_packet_fraction_dset_name)
+                                        mc_packet_fraction_dset_name=self.mc_packet_fraction_dset_name,
+                                        mc_events_dset_name=self.mc_events_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
             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)
@@ -213,6 +220,16 @@ def init(self):
                 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])
 
+            # 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])
+
             # edep-sim energy segments/deposits
             self.data_manager.create_dset(self.mc_tracks_dset_name, dtype=self.mc_tracks.dtype)
             ntracks = len(self.mc_tracks)
@@ -244,34 +261,34 @@ def init(self):
             # set up references
             self.data_manager.create_ref(self.packets_dset_name, self.mc_tracks_dset_name)
             self.data_manager.create_ref(self.mc_trajectories_dset_name, self.mc_tracks_dset_name)
+            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)
             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
             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'][:]
+                stack_evid = self.mc_stack[self.vertex_id_name][:]
+            intr_evid = self.mc_events[self.vertex_id_name][:]
+            traj_evid = self.mc_trajectories[self.vertex_id_name][:]
+            tracks_evid = self.mc_tracks[self.vertex_id_name][:]
             evs, ev_traj_start, ev_track_start = np.intersect1d(
                 traj_evid, tracks_evid, return_indices=True)
             evs, ev_traj_end, ev_track_end = np.intersect1d(
                 traj_evid[::-1], tracks_evid[::-1], return_indices=True)
-            ev_traj_end = len(self.mc_trajectories['vertex_id']) - ev_traj_end
-            ev_track_end = len(self.mc_tracks['vertex_id']) - ev_track_end
+            ev_traj_end = len(self.mc_trajectories[self.vertex_id_name]) - ev_traj_end
+            ev_track_end = len(self.mc_tracks[self.vertex_id_name]) - ev_track_end
             truth_slice = slice(
                 ceil(len(evs) / self.size * self.rank),
                 ceil(len(evs) / self.size * (self.rank + 1)))
 
             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'][:]
+                stack_trackid = self.mc_stack[self.traj_id_name][:]
+            traj_trackid = self.mc_trajectories[self.traj_id_name][:]
+            tracks_trackid = self.mc_tracks[self.traj_id_name][:]
             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)
             for i in iter_:
                 if i < len(evs):
@@ -290,21 +307,21 @@ def init(self):
                     ref[:, 1] += track_start
                     self.data_manager.write_ref(self.mc_trajectories_dset_name, self.mc_tracks_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 --> 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)
+
+                    if self.is_mc_neutrino:
                         # 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))
@@ -323,9 +340,9 @@ def init(self):
                         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)))
                     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)))
 
@@ -461,19 +478,18 @@ 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])
 
-            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)
+            # 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
 

yamls/proto_nd_flow/reco/charge/RawEventGenerator.yaml

@@ -7,7 +7,7 @@ params:
 
   # configuration parameters
   mc_tracks_dset_name: 'mc_truth/segments'
-  buffer_size: 384000
+  buffer_size: 1152000 #38400 #1152000 #384000
   nhit_cut: 1
   sync_noise_cut: [1000000, 11000000] # 1e6 cut based on Brooke's study
   sync_noise_cut_enabled: True