"""LUCiDDataset: dataset for Water Cherenkov detector simulation output.
Loads PMT sensor data and/or 3D track segments from co-indexed HDF5 files.
Produces flat dicts compatible with pimm's transform/collation pipeline.
Example configs:
# PMT event classification (sensor response as fixed-geometry point cloud)
data = dict(train=dict(type="LUCiDDataset", data_root="dataset_wc",
modalities=("sensor",), dataset_name="wc", ...))
# Per-sensor instance separation (sparse per-particle entries)
data = dict(train=dict(type="LUCiDDataset", data_root="dataset_wc",
modalities=("sensor",), include_labels=True, ...))
# 3D track reconstruction
data = dict(train=dict(type="LUCiDDataset", data_root="dataset_wc",
modalities=("seg",), ...))
"""
import os
import numpy as np
from torch.utils.data import Dataset
from pimm.utils.logger import get_root_logger
from .builder import DATASETS
from .test_fragments import build_test_fragments
from .transform import Compose, TRANSFORMS
from .readers.lucid_seg_reader import LUCiDSegReader
from .readers.lucid_sensor_reader import LUCiDSensorReader
[docs]
@DATASETS.register_module()
class LUCiDDataset(Dataset):
"""Water Cherenkov detector dataset over co-indexed LUCiD HDF5 files.
Reads PMT sensor response and/or 3D track segments from event-aligned shard
families (``sensor/`` and/or ``seg/``) and emits flat dicts for the pimm
transform/collation pipeline. The public keys depend on ``output_mode`` for
the sensor modality: ``"response"`` emits one entry per PMT with ``coord``
(3D PMT position or a 1D sensor index), ``energy`` (total PE) and ``time``;
``"labels"`` emits sparse per-particle entries with ``coord``, ``energy``,
``segment`` (category) and ``instance`` (particle index); ``"separate"``
keeps the raw reader keys (``pmt_coord``, ``pmt_pe``, ``pp_*``). After
collation a batch adds ``offset``. Registered as ``LUCiDDataset`` -- use as
``type`` under ``data.train``/``data.val``/``data.test``.
Args:
data_root (str): Root directory holding ``seg/`` and/or ``sensor/``
subdirectories.
split (str): Split name used for shard discovery. Defaults to ``""``.
transform (list[dict]): List of transform configs (NOT a prebuilt
``Compose``). Defaults to ``None``.
modalities (tuple[str]): Which modalities to load, any of ``"sensor"``,
``"seg"``. Defaults to ``("sensor",)``.
dataset_name (str): Shard filename prefix (e.g. ``"wc"`` for
``wc_seg_0000.h5``). Defaults to ``"wc"``.
output_mode (str): Sensor output contract, one of ``"response"``,
``"labels"``, ``"separate"`` (see above). Defaults to ``"response"``.
include_labels (bool): Whether the sensor reader loads the per-particle
decomposition (needed for ``output_mode="labels"``). Defaults to
``True``.
pe_threshold (float): Minimum PE used to sparsify per-particle PE.
Defaults to ``0.0``.
min_segments (int): Minimum segments per event (seg reader filter).
Defaults to ``0``.
max_len (int): Cap on event count before the loop multiplier (-1 = no
cap). Defaults to ``-1``.
loop (int): Train-time epoch multiplier. Defaults to ``1``.
test_mode (bool): Emit voxelized/augmented test fragments and force
``loop = 1``. Defaults to ``False``.
test_cfg (object): Test config (``voxelize``, ``crop``, ``post_transform``,
``aug_transform``); required when ``test_mode``. Defaults to ``None``.
Note:
The dataset length is the minimum event count across the active readers
(they must be co-indexed). Loader settings (``batch_size``,
``num_worker``) live at the top level of the config.
Example:
.. code-block:: python
>>> from pimm.datasets.builder import build_dataset
>>> # data root not in this env -> shown with doctest +SKIP
>>> ds = build_dataset(dict(type="LUCiDDataset", data_root="dataset_wc",
... modalities=("sensor",), output_mode="response",
... transform=[])) # doctest: +SKIP
>>> sample = ds[0] # doctest: +SKIP
>>> # output_mode="response" sample keys: coord (N, 3 PMT pos, or N, 1
>>> # sensor index), energy (N, 1 total PE), time (N, 1), name, split
>>> # output_mode="labels": coord, energy, segment (category),
>>> # instance (particle idx), time, name, split
>>> # output_mode="separate": raw reader keys (pmt_coord, pmt_pe, pp_*)
>>> # + name, split
"""
def __init__(
self,
data_root,
split='',
transform=None,
modalities=('sensor',),
dataset_name='wc',
output_mode='response',
include_labels=True,
pe_threshold=0.0,
min_segments=0,
max_len=-1,
loop=1,
test_mode=False,
test_cfg=None,
):
"""Create modality readers and derive the co-indexed dataset length."""
super().__init__()
self.data_root = data_root
self.split = split
self.modalities = tuple(modalities)
self.dataset_name = dataset_name
self.output_mode = output_mode
self.max_len = max_len
self.loop = loop if not test_mode else 1
self.test_mode = test_mode
self.test_cfg = test_cfg if test_mode else None
self.transform = Compose(transform)
if test_mode and test_cfg is not None:
self.test_voxelize = TRANSFORMS.build(self.test_cfg.voxelize)
self.test_crop = (
TRANSFORMS.build(self.test_cfg.crop)
if self.test_cfg.crop else None)
self.post_transform = Compose(self.test_cfg.post_transform)
self.aug_transform = [
Compose(aug) for aug in self.test_cfg.aug_transform]
# Readers own raw HDF5 decoding; this wrapper owns output formatting.
self.seg_reader = None
self.sensor_reader = None
if 'seg' in self.modalities:
seg_root = self._modality_root('seg')
self.seg_reader = LUCiDSegReader(
data_root=seg_root, split=split,
dataset_name=dataset_name, min_segments=min_segments)
if 'sensor' in self.modalities:
sensor_root = self._modality_root('sensor')
self.sensor_reader = LUCiDSensorReader(
data_root=sensor_root, split=split,
dataset_name=dataset_name,
include_labels=include_labels,
pe_threshold=pe_threshold)
# Canonical reader and length
active_readers = [r for r in (self.seg_reader, self.sensor_reader)
if r is not None]
if not active_readers:
raise ValueError(f"Need 'seg' or 'sensor' in modalities, got {self.modalities}")
self._canonical_reader = active_readers[0]
self._n_events = min(len(r) for r in active_readers)
logger = get_root_logger()
logger.info(f"LUCiDDataset: {self._n_events} events, "
f"modalities={self.modalities}, output_mode={output_mode}")
def _modality_root(self, modality):
"""Resolve root directory for a modality shard family."""
mod_dir = os.path.join(self.data_root, modality)
if os.path.isdir(mod_dir):
return mod_dir
return self.data_root
[docs]
def get_data(self, idx):
"""Load one event and choose the public output contract by modality."""
data_dict = {}
# --- Seg (3D track segments) ---
if self.seg_reader is not None:
seg_data = self.seg_reader.read_event(idx)
if self.sensor_reader is not None and self.output_mode == 'separate':
# Prefix 3D keys to avoid collision with sensor coord
for k, v in seg_data.items():
data_dict[f'seg3d.{k}'] = v
else:
data_dict.update(seg_data)
# --- Sensor (PMT response + optional per-particle labels) ---
if self.sensor_reader is not None:
sensor_data = self.sensor_reader.read_event(idx)
if self.output_mode == 'response':
# PMT response: one entry per sensor.
n = len(sensor_data['pmt_pe'])
if 'pmt_coord' in sensor_data:
data_dict['coord'] = sensor_data['pmt_coord']
else:
# No 3D positions: use sensor index as 1D coord.
data_dict['coord'] = np.arange(n, dtype=np.float32)[:, None]
data_dict['energy'] = sensor_data['pmt_pe'][:, None]
data_dict['time'] = sensor_data['pmt_t'][:, None]
elif self.output_mode == 'labels':
# Sparse per-particle entries
if 'pp_sensor_idx' in sensor_data:
sidx = sensor_data['pp_sensor_idx']
if 'pmt_coord' in sensor_data:
data_dict['coord'] = sensor_data['pmt_coord'][sidx]
else:
data_dict['coord'] = sidx.astype(np.float32)[:, None]
data_dict['energy'] = sensor_data['pp_pe'][:, None]
data_dict['segment'] = sensor_data['pp_category']
data_dict['instance'] = sensor_data['pp_particle_idx']
if 'pp_t' in sensor_data:
data_dict['time'] = sensor_data['pp_t'][:, None]
elif self.output_mode == 'separate':
data_dict.update(sensor_data)
# Metadata
data_dict['name'] = self.get_data_name(idx)
data_dict['split'] = self.split if isinstance(self.split, str) else 'custom'
return data_dict
[docs]
def get_data_name(self, idx):
"""Return a stable shard/event name for logging and prediction files."""
reader = self._canonical_reader
file_idx = int(np.searchsorted(reader.cumulative_lengths, idx, side='right'))
local = idx - (int(reader.cumulative_lengths[file_idx - 1]) if file_idx > 0 else 0)
event_num = reader.indices[file_idx][local]
fname = os.path.basename(reader.h5_files[file_idx])
return f"{fname}_evt{event_num:03d}"
[docs]
def prepare_train_data(self, idx):
"""Load one event and apply the train transform pipeline."""
return self.transform(self.get_data(idx % len(self)))
[docs]
def prepare_test_data(self, idx):
"""Build augmented and voxelized fragments for test-time inference."""
data_dict = self.get_data(idx % len(self))
if self.transform is not None:
data_dict = self.transform(data_dict)
result_dict = dict(name=data_dict.pop("name"))
if "segment" in data_dict:
result_dict["segment"] = data_dict.pop("segment")
result_dict["fragment_list"] = build_test_fragments(
data_dict,
aug_transform=self.aug_transform,
test_voxelize=self.test_voxelize,
test_crop=self.test_crop,
post_transform=self.post_transform,
)
return result_dict
def __getitem__(self, idx):
"""Return a transformed train item or a fragmented test item."""
real_idx = idx % len(self)
if self.test_mode:
return self.prepare_test_data(real_idx)
return self.prepare_train_data(real_idx)
def __len__(self):
"""Return event count after max_len and loop are applied."""
n = self._n_events
if self.max_len > 0:
n = min(n, self.max_len)
return n * self.loop
def __del__(self):
"""Close any reader handles still owned by this dataset instance."""
for reader in (self.seg_reader, self.sensor_reader):
if reader is not None:
reader.close()
