nvalchemi.dynamics.SizeAwareSampler

class nvalchemi.dynamics.SizeAwareSampler(dataset, max_atoms=None, max_edges=None, max_batch_size=None, bin_width=1, shuffle=False, max_gpu_memory_fraction=0.8)

Size-aware sampler for inflight batching.

Manages dataset access, capacity budgets, and bin-packing logic for efficient GPU utilization during dynamics simulations. Ensures every replacement sample fits within the memory envelope of the graduated sample it replaces.

When CUDA is available, the sampler uses a heuristic to estimate the maximum number of atoms that fit in GPU memory. This estimate is combined with user-specified max_atoms — the more restrictive constraint wins. The GPU memory heuristic is best-effort and conservative; users who need tighter control should set max_atoms explicitly.

Parameters:

• dataset (Any) – Dataset with __len__, __getitem__, and get_metadata(idx) methods. get_metadata must return (num_atoms, num_edges).

• max_atoms (int | None) – Maximum total atoms across all samples in a batch. None disables the atom count constraint (GPU memory estimate may still apply). At least one of max_atoms or max_batch_size must be set.

• max_edges (int | None) – Maximum total edges across all samples in a batch. None disables the edge count constraint.

• max_batch_size (int | None) – Maximum number of samples (graphs) in a batch. None disables the graph-count constraint, letting max_atoms (and/or max_edges) alone control batch capacity. At least one of max_atoms or max_batch_size must be set.

• bin_width (int) – Atom-count bin width for grouping samples. Default 1.

• shuffle (bool) – Whether to shuffle within bins. Default False.

• max_gpu_memory_fraction (float) – Fraction of GPU memory to use when estimating atom capacity. Default 0.8 (80%), leaving 20% headroom for model parameters and CUDA context. Only used when CUDA is available.

Raises:

• RuntimeError – If any sample in the dataset has num_atoms > max_atoms or num_edges > max_edges — such samples can never be placed into any batch and indicate a configuration error.

• ValueError – If both max_atoms and max_batch_size are None, max_batch_size < 1, bin_width < 1, or max_gpu_memory_fraction is not in (0.0, 1.0].

Examples

>>> sampler = SizeAwareSampler(dataset, max_atoms=100, max_edges=500, max_batch_size=10)
>>> batch = sampler.build_initial_batch()
>>> replacement = sampler.request_replacement(num_atoms=5, num_edges=20)

build_initial_batch()

Build an initial batch using diverse round-robin bin packing.

Cycles across size bins in round-robin fashion, taking one sample per bin per round, until capacity constraints are reached. This produces a diverse mix of system sizes so that when systems converge and graduate, the freed atom budget can accommodate replacement systems of any size.

Returns:

A batch with status attribute initialized.

Return type:

Batch

Raises:

RuntimeError – If no samples can be added to the batch (e.g., all consumed or constraints too tight).

property exhausted: bool

Check if all samples have been consumed.

Returns:

True if all bins are empty or contain only consumed indices, False otherwise.

Return type:

bool

property max_atoms: int | None

Maximum total atoms per batch (user-specified constraint).

property max_batch_size: int | None

Maximum number of systems per batch (user-specified constraint).

property max_edges: int | None

Maximum total edges per batch (user-specified constraint).

request_replacement(num_atoms, num_edges)

Request a replacement sample that fits within the given constraints.

Searches for an unconsumed sample with at most num_atoms atoms and num_edges edges, starting from the target bin and progressively searching smaller bins.

Parameters:

• num_atoms (int) – Maximum number of atoms the replacement can have.

• num_edges (int) – Maximum number of edges the replacement can have.

Returns:

A replacement sample if found, or None if no suitable sample is available.

Return type:

AtomicData | None

request_replacements(node_counts, edge_counts)

Request replacement samples for multiple graduated systems using GPU-native constraint checking.

Eliminates the .tolist() D→H syncs from _refill_check. Constraint checking is fully vectorized on GPU. M scalar item() calls remain (unavoidable: Python dataset indexing requires CPU integers).

Parameters:

• node_counts (torch.Tensor) – Shape (M,) int64 tensor on GPU. Maximum atoms each replacement can have.

• edge_counts (torch.Tensor) – Shape (M,) int64 tensor on GPU. Maximum edges each replacement can have.

Returns:

Length-M list of replacement samples, or None where no suitable sample is available.

Return type:

list[AtomicData | None]

request_replacements_budget(atom_budget=None, edge_budget=None, max_count=None)

Request replacement samples that fit within a total atom/edge budget.

Searches bins from largest to smallest to maximize diversity — prefers filling the budget with one large system rather than many small ones. Multiple smaller systems can fill budget freed by one large graduated system, or vice versa.

Parameters:

• atom_budget (int | None) – Total atoms available for replacements. None = unconstrained.

• edge_budget (int | None) – Total edges available for replacements. None = unconstrained.

• max_count (int | None) – Maximum number of replacement samples. None = unconstrained.

Returns:

Replacement samples (may be empty if nothing fits or sampler is exhausted).

Return type:

list[AtomicData]