PILArNetParquetIterableDataset#

class PILArNetParquetIterableDataset(repo_id: str | None = 'DeepLearnPhysics/PILArNet-M-mini', data_root: str | None = None, split='train', transform=None, revision: Literal['v1', 'v2', 'v3'] = 'v2', parquet_revision: str = 'refs/convert/parquet', config_name: str = 'default', energy_threshold=0.0, min_points=1024, max_len=-1, remove_low_energy_scatters=False, old_pid_mapping=False, shuffle=True, shuffle_buffer_size=1000, seed=0, loop=1, ignore_index=-1, test_mode=False, test_cfg=None)[source]#

Bases: IterableDataset

Streaming PILArNet reader for large-scale pretraining.

Streams parquet shards sequentially (no random single-row seeks – the access pattern parallel/network filesystems reward) and approximates a full shuffle with shard-order shuffling plus a reservoir buffer_size buffer. Shards are split across DDP ranks (split_dataset_by_node) and DataLoader workers automatically. Emits the same data_dict dictionaries as the map-style reader via the shared decode_event().

Note

This is a PyTorch IterableDataset – wiring it into pimm’s training engine (which currently assumes map-style datasets with samplers/length) is a separate integration step and has not yet been exercised end to end.

Parameters:

  • shuffle (bool) – Enable buffer + shard-order shuffle. Defaults to True.

  • shuffle_buffer_size (int) – Reservoir buffer size for streaming shuffle. Defaults to 1000.

  • seed (int) – Base shuffle seed. Defaults to 0.

See PILArNetParquetDataset for the shared loading/decoding args.

num_samples() → int[source]#

Count events with >= min_points, read cheaply from the parquet n_points column. The engine uses this to size iters_per_epoch when the config does not set it. Requires a local parquet root and the n_points column (written by the offline h5 -> parquet converter).

set_epoch(epoch: int) → None[source]#

Set the epoch so the next __iter__ reshuffles (engine contract).