Sampler

A sampler decides which elements of a dataset appear in each batch during training or evaluation. Given the dataset table-of-contents (toc) and an active split partition, index() returns Awkward Array indices at the scene, frame, molecule, residue, or graph level.

sample() then groups those indices into batches according to batch_size, batch_on (e.g. "molecules" or "scenes"), and distributed-training settings (world_size, rank). Different samplers implement different sampling strategies — uniform over molecules, frame-based windows, mutation-aware grouping, etc.

See Samplers.

class bioverse.sampler.Sampler[source]

Bases: ABC

Select which dataset elements form each training or evaluation batch.

Samplers map the dataset table-of-contents (toc) and an active split partition to Awkward index arrays at the scene, frame, molecule, or residue level. sample() groups those indices into batches according to batch_size, batch_on, and distributed-training settings.

Subclasses implement index(). Common strategies include sampling every molecule (MoleculeSampler) or every frame (FrameSampler).

Examples

from bioverse.samplers import MoleculeSampler

sampler = MoleculeSampler()
batch_indices = sampler.sample(
    dataset, partition="train", split="default", batch_size=32
)
abstract index(toc: Array, mask: Array) → Array[source]

Return row indices for elements in the active split partition.

Parameters:

  • toc – Table-of-contents array describing dataset size at each level.

  • mask – Boolean mask selecting scenes in the current partition.

Returns:

Structured index with fields such as scene, frame, molecule.

Return type:

ak.Array