Use dynamic pixel batches based on num of segments

cli/simulate_pixels.py

@@ -1091,12 +1091,17 @@ def save_results(event_times, results, i_trig, i_mod=-1, light_only=False):
             assmap_pix2seg = invert_array_map(all_neighboring_pixels,all_unique_pix)
             RangePop() # invert_array_map
 
-            for ipix in tqdm(range(0, all_unique_pix.shape[0], sim.PIXEL_BATCH_SIZE),
-                             delay=1, desc='  Simulating event %i batches...' % ievd, leave=False, ncols=80):
+            pixel_ranges = batching.subbatch_pixel_ranges(assmap_pix2seg,
+                                                          sim.SEGMENT_BATCH_SIZE)
+
+            for start_pix, stop_pix in \
+                    tqdm(pixel_ranges, delay=1,
+                         desc='  Simulating event %i batches...' % ievd,
+                         leave=False, ncols=80):
                 RangePush("setup_pixel_batch")
-                selected_pix = all_unique_pix[ipix:ipix+sim.PIXEL_BATCH_SIZE]
+                selected_pix = all_unique_pix[start_pix:stop_pix]
 
-                selected_track_idcs = np.unique(assmap_pix2seg[ipix:ipix+sim.PIXEL_BATCH_SIZE])
+                selected_track_idcs = np.unique(assmap_pix2seg[start_pix:stop_pix])
                 selected_track_idcs = selected_track_idcs[selected_track_idcs != -1]
                 if selected_track_idcs.size == 0:
                     continue

larndsim/consts/sim.py

@@ -9,9 +9,9 @@
 
 from .units import mm, cm, V, kV
 
-PIXEL_BATCH_SIZE = 2000     # units = pixels
-EVENT_BATCH_SIZE = 1  # units = N tpcs
-WRITE_BATCH_SIZE = 1  # units = N batches
+SEGMENT_BATCH_SIZE = 10000    # units = track segments
+EVENT_BATCH_SIZE = 1          # units = N tpcs
+WRITE_BATCH_SIZE = 1          # units = N batches
 EVENT_SEPARATOR = 'event_id'  # 'spillID' or 'vertexID'
 
 # Filter out highly-delayed track segments
@@ -49,7 +49,7 @@ def set_simulation_properties(simprop_file):
             filename
         pixel_file (str): pixel layout YAML filename
     """
-    global PIXEL_BATCH_SIZE
+    global SEGMENT_BATCH_SIZE
     global EVENT_BATCH_SIZE
     global WRITE_BATCH_SIZE
     global EVENT_SEPARATOR
@@ -72,7 +72,7 @@ def set_simulation_properties(simprop_file):
         simprop = yaml.load(df, Loader=yaml.FullLoader)
 
     try:
-        PIXEL_BATCH_SIZE = simprop.get('pixel_batch_size', PIXEL_BATCH_SIZE)
+        SEGMENT_BATCH_SIZE = simprop.get('segment_batch_size', SEGMENT_BATCH_SIZE)
         EVENT_BATCH_SIZE = simprop.get('event_batch_size', EVENT_BATCH_SIZE)
         WRITE_BATCH_SIZE = simprop.get('write_batch_size', WRITE_BATCH_SIZE)
         EVENT_SEPARATOR = simprop.get('event_separator', EVENT_SEPARATOR)

larndsim/simulation_properties/2x2_NuMI_sim.yaml

@@ -1,4 +1,4 @@
-pixel_batch_size: 2000       # pixels
+segment_batch_size: 10000    # track segments
 event_batch_size: 2          # tpcs
 write_batch_size: 1          # batches
 is_spill_sim: 1              # boolean

larndsim/simulation_properties/2x2_NuMI_sim_no_modvar.yaml

@@ -1,4 +1,4 @@
-pixel_batch_size: 2000       # pixels
+segment_batch_size: 10000    # track segments
 event_batch_size: 8          # tpcs
 write_batch_size: 1          # batches
 is_spill_sim: 1              # boolean

larndsim/simulation_properties/NDLAr_LBNF_sim.yaml

@@ -1,6 +1,6 @@
 # Discussions of some of these parameters here:
 # https://dunescience.slack.com/archives/C05BMKCBW3B/p1692623569387309
-pixel_batch_size: 2000       # pixels
+segment_batch_size: 10000    # track segments
 event_batch_size: 2          # tpcs
 write_batch_size: 1          # batches
 is_spill_sim: 1              # boolean

larndsim/simulation_properties/cosmic_sim.yaml

@@ -1,4 +1,4 @@
-pixel_batch_size: 2000       # pixels
+segment_batch_size: 10000    # track segments
 event_batch_size: 2          # tpcs
 write_batch_size: 1          # batches
 is_spill_sim: 0              # boolean

larndsim/simulation_properties/singles_sim.yaml

@@ -1,4 +1,4 @@
-pixel_batch_size: 2000       # pixels
+segment_batch_size: 10000    # track segments
 event_batch_size: 2          # tpcs
 write_batch_size: 1          # batches
 is_spill_sim: 0              # boolean

larndsim/simulation_properties/singles_sim_ndlar.yaml

@@ -1,4 +1,4 @@
-pixel_batch_size: 2000       # pixels
+segment_batch_size: 10000    # track segments
 event_batch_size: 1          # tpcs
 write_batch_size: 1          # batches
 is_spill_sim: 0              # boolean

larndsim/simulation_properties/singles_sim_no_modvar.yaml

@@ -1,4 +1,4 @@
-pixel_batch_size: 2000       # pixels
+segment_batch_size: 10000    # track segments
 event_batch_size: 8          # tpcs
 write_batch_size: 1       # batches
 is_spill_sim: 0              # boolean

larndsim/util/batching.py

@@ -1,3 +1,4 @@
+import cupy as cp
 import numpy as np
 from math import ceil
 
@@ -65,3 +66,36 @@ def __next__(self):
         self._simulated = self._simulated | mask
 
         return self._events[self._curr_event], mask
+
+
+def subbatch_pixel_ranges(assmap_pix2seg: cp.ndarray, seg_batch_size: int) \
+        -> cp.ndarray:
+    """
+    Returns an array of disjoint ranges (i.e. pairs of indices) of the global
+    pixel array, such that, for each range, the number of associated segments is
+    approximately the targeted seg_batch_size.
+
+    Args:
+        assmap_pix2seg(:obj:`cp.ndarray`): Array of segments associated with
+            each pixel. Shape (n_unique_pix, max_associated_segments).
+        seg_batch_size(int): Targetted number of segments associated with each
+            range of pixels. The actual number of segments will be as large as
+            possible without exceeding seg_batch_size.
+
+    Returns:
+        array: Pairs of indices into the global pixel array. Inclusive in the
+            first index, exclusive in the second, so that they can be used
+            directly for slicing.
+    """
+
+    segs_per_pix = cp.sum(assmap_pix2seg != -1, axis=-1)
+    accum_segs = cp.cumsum(segs_per_pix)
+    total_segs = accum_segs[-1]
+    seg_divisions = cp.arange(0, total_segs, seg_batch_size)
+    indices = cp.searchsorted(accum_segs, seg_divisions)
+
+    result = cp.zeros(shape=(len(indices), 2), dtype=cp.uint32)
+    result[:, 0] = indices
+    result[:-1, 1] = indices[1:]
+    result[-1, 1] = total_segs
+    return result