Speed up charge geometry loading

src/proto_nd_flow/resources/geometry.py

@@ -13,6 +13,8 @@
 from proto_nd_flow.util.compat import assert_compat_version
 import proto_nd_flow.util.units as units
 
+from tqdm import tqdm
+
 class Geometry(H5FlowResource):
     '''
         Provides helper functions for looking up geometric properties. 
@@ -388,9 +390,10 @@ def _get_module_RO_bounds(self):
 
         self._module_RO_bounds = []
 
+        io_group, io_channel, chip_id, channel_id = self.pixel_coordinates_2D.keys()
+
         # Loop through modules
-        for module_id in module_to_io_groups:  
-            io_group, io_channel, chip_id, channel_id = self.pixel_coordinates_2D.keys()
+        for module_id in module_to_io_groups:
             min_coord = np.finfo(self.pixel_coordinates_2D.dtype).min
             max_coord = np.finfo(self.pixel_coordinates_2D.dtype).max
             min_x, max_x = min_coord, max_coord
@@ -784,7 +787,8 @@ def _load_charge_geometry(self):
         self._pixel_pitch = [0.]*n_modules
 
         # Loop through modules
-        for module_id in module_to_io_groups:
+        for module_id in tqdm(module_to_io_groups, desc='Module geometries'):
+            iog_per_mod = len(det_geometry_yaml['module_to_io_groups'][module_id])
             geometry_yaml = geometry_yamls[self.crs_geometry_to_module[module_id-1]]
             pixel_pitch = geometry_yaml['pixel_pitch'] / units.cm # convert mm -> cm
             self._pixel_pitch[module_id-1] = pixel_pitch
@@ -802,7 +806,7 @@ def _load_charge_geometry(self):
 
                 for chip in tile_chip_to_io[tile]:
                     io_group_io_channel = tile_chip_to_io[tile][chip]
-                    io_group = io_group_io_channel//1000 + (module_id-1)*len(det_geometry_yaml['module_to_io_groups'][module_id])
+                    io_group = io_group_io_channel//1000 + (module_id-1)*iog_per_mod
                     io_channel = io_group_io_channel % 1000
                     self._tile_id[([io_group], [io_channel])] = tile+(module_id-1)*len(tile_chip_to_io)
 
@@ -826,9 +830,9 @@ def _load_charge_geometry(self):
                         else:
                             continue
 
-                    io_group = io_group_io_channel // 1000 + (module_id-1)*len(det_geometry_yaml['module_to_io_groups'][module_id])
+                    io_group = io_group_io_channel // 1000 + (module_id-1)*iog_per_mod
                     io_channel = io_group_io_channel % 1000
-                    
+
                     z = chip_channel_to_position[chip_channel][0] * \
                         pixel_pitch - z_size / 2 + pixel_pitch / 2
                     y = chip_channel_to_position[chip_channel][1] * \
@@ -849,7 +853,7 @@ def _load_charge_geometry(self):
 
                     for ioc in io_channels:
                         try:
-                            self._pixel_coordinates_2D[(io_group, ioc, chip, channel)] = z, y
+                            self._pixel_coordinates_2D.set_scalar((io_group, ioc, chip, channel), (z, y))
                         except:
                             print(io_group, ioc, chip, channel)
 

src/proto_nd_flow/util/lut.py

@@ -59,6 +59,7 @@ def __init__(self, dtype, *min_max_keys, default=None, shape=None):
         self.dtype = dtype
         self.min_max_keys = np.array(min_max_keys, dtype='i8')
         self.lengths = np.array([max_ - min_ + 1 for min_, max_ in self.min_max_keys])
+        self.strides = np.cumprod(self.lengths) # used by set_scalar
         self.max_hash = self._hash(*[max_ for min_, max_ in min_max_keys])[0]
         shape = (self.max_hash + 1,) + shape if shape else (self.max_hash + 1,)
         self._data = np.zeros(shape, dtype=self.dtype)
@@ -278,6 +279,13 @@ def keys(self):
     def __getitem__(self, keys):
         return self._data[self.hash(*keys)]
 
+    def set_scalar(self, keys, val):
+        idx = 1 + keys[0] - self.min_max_keys[0][0]
+        for i, key in enumerate(keys[1:]):
+            idx += (key - self.min_max_keys[i + 1][0]) * self.strides[i]
+        self._data[idx] = val
+        self._filled[idx] = True
+
     def __setitem__(self, keys, val):
         idx = self.hash(*keys)
 
@@ -309,4 +317,4 @@ def get_keys_from_val(self, val):
             if np.array_equal(stored_val, val):
                 return np.ravel(self.reverse_hash(idx))
 
-        raise KeyError(f"No matching key found for the given value: {val}")
+        raise KeyError(f"No matching key found for the given value: {val}")