From e7b2615764a176d008829123fb27b47ebb7a55a0 Mon Sep 17 00:00:00 2001 From: drinkingkazu Date: Fri, 25 Sep 2026 06:18:11 -0700 Subject: [PATCH] change the group size type from uint8 to uint16 so that it won't roll over beyond 255. --- docs/production/data-formats.md | 2 +- production/README.md | 2 +- production/save.py | 18 +++++++++++-- tests/test_production_io.py | 48 +++++++++++++++++++++++++++++++++ viewer/h5_worker.js | 4 +-- 5 files changed, 68 insertions(+), 6 deletions(-) diff --git a/docs/production/data-formats.md b/docs/production/data-formats.md index a8b3984..b1bfeae 100644 --- a/docs/production/data-formats.md +++ b/docs/production/data-formats.md @@ -221,7 +221,7 @@ enabled (it is on by default; `--no-track-hits` skips this file). See group_to_track (G,) int32 per-group Geant4 track_id (a label) {plane}/ one subgroup per active readout plane group_ids (G_p,) int32 active groups on this plane - group_sizes (G_p,) uint8 entries per group (CSR row lengths) + group_sizes (G_p,) uint16 entries per group (CSR row lengths) center_wires (G_p,) int16 peak-charge wire idx (wire) center_py (G_p,) int16 peak pixel-y idx (pixel) center_pz (G_p,) int16 peak pixel-z idx (pixel) diff --git a/production/README.md b/production/README.md index 3ea370d..2275b6b 100644 --- a/production/README.md +++ b/production/README.md @@ -280,7 +280,7 @@ entries (which groups contributed to which pixels). each group. {plane_label}/ one subgroup per readout plane group_ids (G_p,) int32 active groups on this plane - group_sizes (G_p,) uint8 entries per group + group_sizes (G_p,) uint16 entries per group center_wires (G_p,) int16 wire idx of peak-charge pixel (wire readout) center_py (G_p,) int16 pixel y idx of peak (pixel readout) center_pz (G_p,) int16 pixel z idx of peak (pixel readout) diff --git a/production/save.py b/production/save.py index 485b4cb..6f48ddf 100644 --- a/production/save.py +++ b/production/save.py @@ -433,6 +433,20 @@ def save_event_step(f, event_key, deposits, source_event_idx, pos_step_mm=0.3, vg.create_dataset('photons', data=np.asarray(vol.photons[:n]).astype(np.float32), **_COMPRESSION) +def _group_sizes(group_counts): + """CSR row lengths as uint16, refusing any that would wrap. + + A wrapped size misaligns every later group on the plane, so it is an + error, not a truncation. + """ + limit = np.iinfo(np.uint16).max + if len(group_counts) and int(group_counts.max()) > limit: + raise ValueError( + f"a correspondence group has {int(group_counts.max())} entries; " + f"group_sizes is uint16 and holds at most {limit}") + return group_counts.astype(np.uint16) + + def encode_correspondence_csr(gp_pk, gp_gid, gp_ch, gp_count, num_time_steps, threshold=0.0): """Convert flat correspondence arrays to CSR + delta + uint16/peak format. @@ -489,7 +503,7 @@ def encode_correspondence_csr(gp_pk, gp_gid, gp_ch, gp_count, num_time_steps, return { 'group_ids': unique_gids.astype(np.int32), - 'group_sizes': group_counts.astype(np.uint8), + 'group_sizes': _group_sizes(group_counts), 'center_wires': center_wires, 'center_times': center_times, 'peak_charges': peak_charges, 'delta_wires': delta_wires, 'delta_times': delta_times, @@ -572,7 +586,7 @@ def encode_correspondence_csr_pixel(gp_sk, gp_tk, gp_gid, gp_ch, gp_count, return { 'group_ids': unique_gids.astype(np.int32), - 'group_sizes': group_counts.astype(np.uint8), + 'group_sizes': _group_sizes(group_counts), 'center_py': center_py, 'center_pz': center_pz, 'center_times': center_times, 'peak_charges': peak_charges, 'delta_py': delta_py, 'delta_pz': delta_pz, diff --git a/tests/test_production_io.py b/tests/test_production_io.py index ef04a9a..4fb6192 100644 --- a/tests/test_production_io.py +++ b/tests/test_production_io.py @@ -321,6 +321,54 @@ def test_pixel_i16_roundtrip_signed(self): # signs follow the data, not the (negative) peak assert recovered[0] < 0 and recovered[1] > 0 and recovered[2] < 0 + def test_pixel_group_larger_than_255_roundtrip(self): + """A group with more than 255 entries keeps its size. group_sizes was + once uint8, so a 300-entry group was stored as 44 and every later + group on the plane was decoded from the wrong CSR offset.""" + import h5py + from production.load import _decode_plane_hits_pixel + + num_pz = 100 + # Group 1: 300 entries (5x5 pixels x 12 ticks); group 2: 3 entries. + py, pz, t = np.meshgrid(np.arange(5), np.arange(5), np.arange(12), + indexing='ij') + sk = np.concatenate([(10 + py.ravel()) * num_pz + (20 + pz.ravel()), + [60 * num_pz + 70] * 3]) + tk = np.concatenate([100 + t.ravel(), [500, 501, 502]]).astype(np.int32) + gid = np.array([1] * 300 + [2] * 3, dtype=np.int32) + ch = np.linspace(1.0, 50.0, 303).astype(np.float32) + + csr = encode_correspondence_csr_pixel(sk, tk, gid, ch, 303, num_pz=num_pz) + np.testing.assert_array_equal(csr['group_sizes'], [300, 3]) + + with tempfile.NamedTemporaryFile(suffix='.h5', delete=False) as tmp: + tmp_path = tmp.name + try: + with h5py.File(tmp_path, 'w') as f: + g = f.create_group('Pixel') + for key, val in csr.items(): + g.create_dataset(key, data=val) + with h5py.File(tmp_path, 'r') as f: + d_py, d_pz, d_t, d_gid, _, n = _decode_plane_hits_pixel(f['Pixel']) + + assert n == 303 + got = sorted(zip(d_gid.tolist(), (d_py * num_pz + d_pz).tolist(), + d_t.tolist())) + want = sorted(zip(gid.tolist(), sk.tolist(), tk.tolist())) + assert got == want + finally: + os.unlink(tmp_path) + + def test_group_size_overflow_raises(self): + """A group too large for uint16 group_sizes raises instead of wrapping.""" + n = 65536 + sk = np.zeros(n, dtype=np.int64) + tk = np.arange(n, dtype=np.int32) % 100 + gid = np.ones(n, dtype=np.int32) + ch = np.ones(n, dtype=np.float32) + with pytest.raises(ValueError, match='uint16'): + encode_correspondence_csr_pixel(sk, tk, gid, ch, n, num_pz=100) + def test_csr_encode_decode_roundtrip_via_hdf5(self): """Full CSR encode → HDF5 write → HDF5 read → decode roundtrip.""" import h5py diff --git a/viewer/h5_worker.js b/viewer/h5_worker.js index dc786c8..2619149 100644 --- a/viewer/h5_worker.js +++ b/viewer/h5_worker.js @@ -173,7 +173,7 @@ function decodeEvent(idx) { const g = iVol.get('Pixel'); if (g && g.get('group_ids')) { const grpIds = new Int32Array(g.get('group_ids').value); - const grpSz = new Uint8Array(g.get('group_sizes').value); + const grpSz = new Uint16Array(g.get('group_sizes').value); const cpy = new Int16Array(g.get('center_py').value); const cpz = new Int16Array(g.get('center_pz').value); const ct = new Int16Array(g.get('center_times').value); @@ -235,7 +235,7 @@ function decodeEvent(idx) { if (!g || !g.get('group_ids')) continue; const grpIds = new Int32Array(g.get('group_ids').value); - const grpSz = new Uint8Array(g.get('group_sizes').value); + const grpSz = new Uint16Array(g.get('group_sizes').value); const cw = new Int16Array(g.get('center_wires').value); const ct = new Int16Array(g.get('center_times').value); const pc = new Float32Array(g.get('peak_charges').value);