PILArNetH5Dataset#

class PILArNetH5Dataset(data_root: str | None = None, split='train', transform=None, test_mode=False, test_cfg=None, loop=1, ignore_index=-1, energy_threshold=0.0, min_points=1024, max_len=-1, remove_low_energy_scatters=False, old_pid_mapping=False, revision: Literal['v1', 'v2', 'v3'] = 'v2', overlay_n_events=1, overlay_prob=1.0, overlay_allow_repeats=True)[source]#

Bases: PILArNetOverlayMixin, Dataset

PILArNet-M LArTPC dataset read directly from clustered HDF5 shards.

Loads events straight from point/cluster/cluster_extra HDF5 arrays (no per-event preprocessing), expands per-cluster truth to per-point arrays, and emits a flat dict. Standard keys per item: coord (N, 3), energy (N, 1, raw), segment_motif (N, 1, semantic class), segment_pid (N, 1, PID; v2/v3), momentum (N, 1; v2/v3), vertex (N, 3; v2/v3, interaction vertex in v3), is_primary (N, 1; v3 only), instance_particle and instance_interaction (N, 1, remapped contiguous ids), segment_interaction (N, 1, background flag), plus name/split/revision. After collation a batch adds offset. Registered as PILArNetH5Dataset – use as type under data.train/data.val/data.test.

Semantic classes (segment_motif): 0 shower, 1 track, 2 Michel, 3 delta, 4 low-energy deposit. PID classes (segment_pid): 0 photon, 1 electron, 2 muon, 3 pion, 4 proton, 5 none/LED (6 when old_pid_mapping).

Parameters:

  • data_root (str | None) – Root directory of the revision’s HDF5 shards. When None, falls back to $PILARNET_DATA_ROOT_V1/_V2/_V3 (by revision) then to ~/.cache/pimm/pilarnet/<revision>; raises if none exist. Defaults to None.

  • split (str | Sequence[str]) – Split name(s) used to glob *<split>/*.h5 under data_root. Defaults to "train".

  • transform (list[dict]) – List of transform configs (NOT a prebuilt Compose). Defaults to None.

  • 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.

  • loop (int) – Train-time epoch multiplier. Defaults to 1.

  • ignore_index (int) – Ignored-label value. Defaults to -1.

  • energy_threshold (float) – Drop points with energy at or below this value when positive. Defaults to 0.0.

  • min_points (int) – Minimum points per event; smaller events are excluded from the index. Defaults to 1024.

  • max_len (int) – Cap on event count before the loop multiplier (-1 = no cap). Defaults to -1.

  • remove_low_energy_scatters (bool) – Drop the first (LED scatter) cluster and its points. Defaults to False.

  • old_pid_mapping (bool) – Map LED PID to 6 instead of 5. Defaults to False.

  • revision ({"v1", "v2", "v3"}) – Dataset revision. v1 is the original PILArNet (no PID/momentum/vertex); v2 adds PID, momentum and particle vertices; v3 adds interaction-level vertices and primary-particle labels. Defaults to "v2".

  • overlay_n_events (int | tuple[int, int]) – Number (or inclusive range) of events to overlay into one point cloud; > 1 enables overlay. Defaults to 1.

  • overlay_prob (float) – Probability of applying overlay to a given sample. Defaults to 1.0.

  • overlay_allow_repeats (bool) – Allow the same event to be sampled more than once when overlaying. Defaults to True.

Note

Loader settings (batch_size, num_worker) live at the top level of the config, not on the dataset constructor. Split membership differs between v1 and v2/v3, so a v1-trained model evaluated on v2/v3 (or vice versa) is not seeing a comparable split. Event overlay deduplicates colliding voxels by semantic priority (track > shower > Michel > delta > LED) and rotates overlaid events by random 90-degree increments.

Example

>>> from pimm.datasets.builder import build_dataset
>>> ds = build_dataset(dict(type="PILArNetH5Dataset", revision="v2",
...                         split="train", transform=[], min_points=1024))
>>> sample = ds[0]
>>> sorted(sample)[:6]
['coord', 'energy', 'instance_interaction', 'instance_particle', 'momentum', 'name']
>>> sample["coord"].shape          # (N, 3) float32
(7366, 3)
>>> sample["segment_motif"].shape  # (N, 1) semantic class
(7366, 1)
>>> # in a config:
>>> # data = dict(train=dict(type="PILArNetH5Dataset", split="train",
>>> #             revision="v2", min_points=1024, transform=transform))
get_data(idx)[source]#

Load a point cloud from h5 file.

Output dictionary: - coord: (N, 3) array of coordinates - energy: (N, 1) array of energies - momentum: (N, 1) array of particle momentum (v2/v3 only) - vertex: (N, 3) array of vertices (v2/v3 only; interaction vertex for v3) - is_primary: (N, 1) array of primary-particle flags (v3 only) - segment_motif: (N, 1) array of motif labels - segment_pid: (N, 1) array of PID labels (v2/v3 only) - instance_particle: (N, 1) array of particle instance labels - instance_interaction: (N, 1) array of interaction instance labels - segment_interaction: (N, 1) array of interaction labels

get_data_name(idx)[source]#

Get name for the point cloud.

get_h5_files()[source]#

Get list of h5 files based on the split.

h5py_worker_init()[source]#

Initialize h5py files for each worker.

prepare_test_data(idx)[source]#

Prepare test data with test transforms.

prepare_train_data(idx)[source]#

Prepare training data with transforms.