diff --git a/src/proto_nd_flow/resources/geometry.py b/src/proto_nd_flow/resources/geometry.py index d2eea703..38bdbda4 100644 --- a/src/proto_nd_flow/resources/geometry.py +++ b/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) diff --git a/src/proto_nd_flow/util/lut.py b/src/proto_nd_flow/util/lut.py index db55f509..30e2e0ef 100644 --- a/src/proto_nd_flow/util/lut.py +++ b/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}") \ No newline at end of file + raise KeyError(f"No matching key found for the given value: {val}")