ML Data Loading¶

Canonical table of normalized streams available to ML loaders.

Attributes:

• frame (DataFrame) –

  Deterministically sorted Polars frame containing manifest identity, protocol, row counts, normalized segment paths, split assignments, and stable manifest/ML row identifiers.

Methods:

• filter_split –

  Return an index containing only one split.

filter_split(split: str) -> MlDatasetIndex

Group-aware split policy used while constructing an ML index.

Attributes:

• strategy (Literal['sample', 'provide', 'merge']) –

  "sample" hashes groups deterministically, "provide" selects explicit groups, and "merge" combines both.

• fraction (float) –

  Fraction of groups assigned to validation by sampling.

• seed (int) –

  Hash seed used by sampling.

• group_by (tuple[str | MappingSpec, ...]) –

  Manifest columns that define an indivisible group.

• provided (tuple[dict[str | MappingSpec, object], ...]) –

  Explicit partial or complete group selectors.

Examples:

Hold out one cell explicitly:

from batgrad.contracts.mapping import BaseColumns
from batgrad.ml.data.config import ValidationConfig

validation = ValidationConfig.provide(
    ({BaseColumns.set_id: "pozzato-2022", BaseColumns.cell_id: "Cell1"},)
)

Methods:

• sample –

  Create a deterministic sampled-group split.

• provide –

  Create a split containing only explicitly selected groups.

• merge –

  Create a split combining explicit and sampled groups.

sample(
    fraction: float = 0.2,
    seed: int = 69,
    group_by: tuple[str | MappingSpec, ...] = (
        set_id,
        cell_id,
        cidx,
    ),
) -> ValidationConfig

provide(
    provided: tuple[dict[str | MappingSpec, object], ...],
    group_by: tuple[str | MappingSpec, ...] = (
        set_id,
        cell_id,
        cidx,
    ),
) -> ValidationConfig

merge(
    provided: tuple[dict[str | MappingSpec, object], ...],
    fraction: float = 0.2,
    seed: int = 69,
    group_by: tuple[str | MappingSpec, ...] = (
        set_id,
        cell_id,
        cidx,
    ),
) -> ValidationConfig

Windowing settings for one protocol stream.

WindowConfig controls how a normalized manifest row is split into model windows before tensor materialization. It is protocol-specific because time series protocols such as cycling/HPPC and frequency-domain protocols such as EIS usually need different sequence lengths.

Attributes:

• batch_size (int) –

  Number of contiguous sequence lanes carved from each window.

• seq_len (int) –

  Number of input and shifted-target positions per lane.

• drop_incomplete (bool) –

  Drop a final source window that cannot fill every lane.

• step_rows (int | None) –

  Source-row distance between windows. The default advances by batch_size * seq_len.

Examples:

Configure cycling as long ordered time windows and EIS as one compact frequency window:

windows = {
    DatasetProtocolId.cycling: WindowConfig(batch_size=8, seq_len=1024),
    DatasetProtocolId.eis: WindowConfig(batch_size=16, seq_len=48),
}

step: int

window_rows: int

Runtime options for protocol-aware ML loading.

The canonical ML index stores facts about available manifest rows. Loader configuration describes how those rows are planned into windows and yielded. group_key identifies one protocol-specific stream/window source, while alignment_key identifies the shared physical context used by future multi-protocol schedules.

Usually group_key includes protocol and alignment_key does not:

group_key = (dataset id, cell id, cycle index, protocol)
alignment_key = (dataset id, cell id, cycle index)

With this setup cycling/HPPC/EIS remain separate streams but can later be bundled together for the same cell/cycle.

data_access controls the speed/RAM trade-off:

• "windowed" reads only the windows required for each batch from parquet. This has the lowest RAM footprint and is easiest to reason about, but it is IO-bound and can be orders of magnitude slower than memory-backed loading.
• "full_in_mem" preloads the full selected split, active protocol, and requested columns into CPU float32 tensors. RAM scales approximately as rows * selected_columns * 4 bytes, plus runtime and temporary conversion overhead. A synthetic normalized smoke with batch=64, seq=1024, three inputs and one target measured roughly 0.1-0.2M tokens/s for windowed and 30M+ tokens/s for full_in_mem after the cache is built.

Attributes:

• split (str) –

  Index split yielded by the loader.

• default_window (WindowConfig) –

  Window shape used unless a protocol override exists.

• window_by_protocol (dict[DatasetProtocolId, WindowConfig]) –

  Protocol-specific window overrides.

• seed (int) –

  Reproducible planning and epoch-phase seed.

• strategy (BatchStrategy) –

  Sequential windows or shuffled protocol-group streams.

• protocol_order (tuple[DatasetProtocolId, ...]) –

  Explicit protocol traversal order.

• stateful_n_windows (int) –

  Consecutive windows per stateful group, or -1 for whole-stream mode.

• cross_protocol_state_carry (CrossProtocolStateCarry | None) –

  "chain" groups aligned protocol streams for recurrent state carry.

• drop_incomplete_batches (bool) –

  Drop plans with fewer than the requested lanes.

• drop_incomplete_distributed (bool) –

  Equalize plan counts across distributed ranks.

• data_access (DataAccessMode) –

  Parquet-window or full-memory access.

• group_key (tuple[str | MappingSpec, ...]) –

  Columns identifying a protocol-specific stream.

• alignment_key (tuple[str | MappingSpec, ...]) –

  Columns identifying streams that may be aligned.

• num_workers (int) –

  PyTorch worker process count.

• prefetch_factor (int) –

  Batches prefetched by each worker.

• persistent_workers (bool) –

  Retain workers between iterations.

• pin_memory (bool) –

  Pin CPU tensors before device transfer.

• multiprocessing_context (MultiprocessingContext | None) –

  Worker start method.

• device (str) –

  Destination device used by optional prefetch.

• prefetch_to_device (bool) –

  Asynchronously move batches to CUDA.

• non_blocking (bool) –

  Use non-blocking tensor transfers where possible.

Methods:

• window_for –

  Return window settings for a protocol.

window_for(protocol: DatasetProtocolId) -> WindowConfig

Runtime scaling rule for one tensor or frame column.

Attributes:

• column (str | MappingSpec) –

  Canonical column name or mapping specification.

• input_min (float) –

  Lower physical-unit bound.

• input_max (float) –

  Upper physical-unit bound.

• output_min (float) –

  Lower model-space bound.

• output_max (float) –

  Upper model-space bound.

• clip (bool) –

  Clip forward-scaled values to output bounds.

• transform (ScalingTransform) –

  Linear scaling or log1p followed by linear scaling.

Protocol-agnostic time-series batch yielded by ML data loaders.

Targets are shifted one source row ahead of inputs. Incomplete source rows and null input values are represented by -2.0; null targets remain NaN so loss calculation can exclude them. mask identifies target positions backed by real source rows, while target finiteness provides feature-level validity.

Attributes:

• inputs (Tensor) –

  Model-space tensor shaped (B, T, C_in).

• targets (Tensor) –

  Model-space tensor shaped (B, T, C_out) for the following source rows.

• mask (Tensor) –

  Valid-target mask shaped (B, T) or (B, T, C_out).

• all_valid (bool) –

  Whether every target row is valid, allowing mask optimizations.

• state (BatchState) –

  Source and stateful-sequence metadata; remains on the CPU.

Examples:

Inspect a batch without assuming a specific protocol:

batch = next(iter(loader))
assert batch.inputs.shape[:2] == batch.targets.shape[:2]
print(batch.state.group_keys)

Methods:

• is_protocol –

  Return whether any batch lane uses the requested protocol.

• pin_memory –

  Return a batch with tensor fields copied to pinned CPU memory.

• to –

  Return a batch whose tensor fields are moved to device.

is_protocol(protocol: object) -> bool

pin_memory() -> Self

to(
    device: str | device | int | None,
    *,
    non_blocking: bool = False,
) -> Self

Traceability and recurrent-sequence metadata for a tensor batch.

Tuple-valued stream fields are lane-aligned. Stateful identifiers are set only when consecutive batches deliberately form one recurrent sequence.

Attributes:

• split (str) –

  Source index split.

• batch_idx (int) –

  Batch number within the loader iteration.

• protocols (tuple[DatasetProtocolId, ...]) –

  Protocol associated with each lane.

• manifest_paths (tuple[str, ...]) –

  Source manifest for each lane.

• manifest_row_ids (tuple[int, ...]) –

  Source manifest row identifier for each lane.

• group_keys (tuple[tuple[object, ...], ...]) –

  Protocol-specific stream keys.

• alignment_keys (tuple[tuple[object, ...], ...]) –

  Physical-context keys used for cross-protocol alignment.

• segments (tuple[BatchSegmentRef, ...]) –

  Normalized parquet ranges contributing to the batch.

• window_offsets (tuple[int, ...]) –

  Source-row start offset for each lane.

• stateful_group_idx (int | None) –

  Optional recurrent group identifier.

• stateful_step_idx (int | None) –

  Zero-based position in the recurrent group.

• stateful_steps (int | None) –

  Total number of steps in the recurrent group.

Traceable source range inside one normalized shard segment.

Attributes:

• path (str) –

  Normalized parquet segment path.

• row_start (int) –

  First source row read from the segment.

• row_count (int) –

  Number of source rows read.

• window_row_start (int) –

  First row within the logical materialized window.

• window_row_count (int) –

  Number of window rows supplied by this segment.