Samplers#

This module contains various sampling strategies and utilities for data processing.

class atomworks.ml.samplers.DistributedMixedSampler(datasets_info: list[dict[str, any]], num_replicas: int, rank: int, n_examples_per_epoch: int | None, shuffle: bool = True, drop_last: bool = True)[source]#

Bases: Sampler

Custom DistributedSampler implementation that samples from an arbitrary list of samplers with specified probabilities.

Child samplers can be any type of non-distributed sampler, including a MixedSampler. After gathering all indices, shards the samples across nodes, ensuring each node receives a unique slice of the dataset.

Example

Imagine we have the following sampling tree:

DistributedMixedSampler

|

0.8 0.2

Sampler1 MixedSampler

/ 0.9 0.1

Sampler2 Sampler3

If we initialized DistributedMixedSampler with n_examples_per_epoch=100 and num_replicas=2, it would collect 80 samples from Sampler1 and 20 samples from the MixedSampler. The MixedSampler would in turn collect 18 samples from Sampler2 and 2 samples from Sampler3. After collecting those 100 samples, the DistributedMixedSampler would shard the samples across the two nodes, ensuring each node receives a unique slice of 50 examples.

If any of the child samplers were distributed samples, then the DistributedMixedSampler would not receive n_examples_per_epoch indices, and we would raise an error.

NOTE: The order of the datasets in datasets_info MUST match the order of the datasets in the ConcatDataset associated with this MixedSampler.

Parameters:

  • datasets_info – List of dictionaries, where each dictionary must contain at a minimum: - “sampler”: Sampler object for the dataset - “dataset”: Dataset object associated with the sampler - “probability”: Probability of sampling from this dataset

  • num_replicas – Number of replicas (nodes) in the distributed setting

  • rank – Rank of the current node

  • n_examples_per_epoch – Number of examples in an epoch. Effectively, the “length” of the sampler (since we often sample with replacement). May be None, in which case the number of examples per epoch must be set dynamically by a parent sampler.

  • shuffle – Whether to shuffle the indices. If False, the iterator will return all sampled indices from the first dataset, then the second, etc.

  • drop_last – Whether to drop the last incomplete batch if the dataset size is not divisible by the batch size

Returns:

An iterator over indices of the dataset for the current process (of length n_samples, not n_examples_per_epoch)

Return type:

iter

Reference:

PyTorch DistributedSampler

set_epoch(epoch: int) → None[source]#
class atomworks.ml.samplers.FallbackSamplerWrapper(sampler: Sampler, fallback_sampler: Sampler, n_fallback_retries: int = 2)[source]#

Bases: Sampler

A wrapper around a sampler that allows for a fallback sampler to be used when an error occurs.

Meant to be used with a FallbackDatasetWrapper.

set_epoch(epoch: int) → None[source]#
class atomworks.ml.samplers.LazyWeightedRandomSampler(weights: Sequence[float], num_samples: int, replacement: bool = True, generator: Generator | None = None, prefetch_buffer_size: int = 1)[source]#

Bases: WeightedRandomSampler

class atomworks.ml.samplers.LoadBalancedDistributedSampler(dataset: Dataset, key_to_balance: str, num_replicas: int | None = None, rank: int | None = None, drop_last: bool = False)[source]#

Bases: DistributedSampler

DistributedSampler that balances large examples across replicas.

Helpful for validation, where we don’t want GPUs to be idle while waiting for the slowest replica to finish.

For example, we may want to avoid the scenario where one GPU receives many large examples that are slow to process, while another GPU receives many small examples that are quick to process.

NOTE: Only useful for validation, as the order of the examples is deterministic.

Parameters:

  • dataset – Dataset used for sampling.

  • key_to_balance – Key in the dataset data dataframe that contains the length (size) of each example. The dataset must have a data attribute that can be accessed like a dataframe. For example, if the dataset has a data attribute that is a pandas DataFrame, the key_to_balance should be a column in that DataFrame (i.e., “n_tokens”).

  • num_replicas (int, optional) – Number of processes participating in distributed training. By default, world_size is retrieved from the current distributed group.

  • rank (int, optional) – Rank of the current process within num_replicas. By default, rank is retrieved from the current distributed group.

  • drop_last (bool, optional) – if True, then the sampler will drop the tail of the data to make it evenly divisible across the number of replicas. If False, the sampler will add extra indices to make the data evenly divisible across the replicas. Default: False.

class atomworks.ml.samplers.MixedSampler(datasets_info: list[dict[str, any]], n_examples_per_epoch: int | None = None, shuffle: bool = True)[source]#

Bases: DistributedMixedSampler

A non-distributed sampler that samples from an arbitrary list of samplers with specified probabilities.

This class acts like a DistributedMixedSampler with rank=0 and num_replicas=1.

Parameters:

  • datasets_info – List of dictionaries, where each dictionary must contain at a minimum: - “sampler”: Sampler object for the dataset - “dataset”: Dataset object associated with the sampler - “probability”: Probability of sampling from this dataset

  • n_examples_per_epoch – Number of examples in an epoch. Effectively, the “length” of the sampler.

  • shuffle – Whether to shuffle the indices. If False, the iterator will return all sampled indices from the first dataset, then the second, etc.

atomworks.ml.samplers.calculate_weights_by_inverse_cluster_size(dataset_df: Table | DataFrame, cluster_column: str = 'cluster', altloc_seed_column: str | None = None, pdb_id_column: str = 'pdb_id') → Tensor[source]#

Calculate weights as the inverse of cluster size, optionally adjusted for altlocs.

Base weight per row: 1 / cluster_size. If altloc_seed_column is provided, weights are additionally multiplied by 1 / num_altlocs for uniform three-level sampling: cluster → example → altloc.

Parameters:

  • dataset_df – DataFrame or PyArrow Table containing the PN unit or interface data.

  • cluster_column – Column identifying sequence clusters. Defaults to "cluster".

  • altloc_seed_column – If provided, weights are divided by the number of unique altloc seeds per pdb_id. Defaults to None (no altloc adjustment).

  • pdb_id_column – Column identifying structures, used only when altloc_seed_column is set. Defaults to "pdb_id".

Returns:

Tensor of per-row weights with shape (len(dataset_df),).

atomworks.ml.samplers.calculate_weights_for_pdb_dataset_df(dataset_df: Table | DataFrame, alphas: dict[str, float], beta: float, cluster_column: str = 'cluster', altloc_seed_column: str | None = None, pdb_id_column: str = 'pdb_id') → Tensor[source]#

Calculate weights based on the AF-3 methodology, optionally adjusted for altlocs.

Base weight per row: (beta / cluster_size) * (a_prot * n_prot + ...). If altloc_seed_column is provided, weights are additionally multiplied by 1 / num_altlocs so that each altloc variant of a structure is sampled uniformly.

Parameters:

  • dataset_df – DataFrame or PyArrow Table containing the PN unit or interface data.

  • alphas – Alpha hyperparameters for the AF-3 weighting formula.

  • beta – Beta hyperparameter (distinct for interfaces vs. chains).

  • cluster_column – Column identifying sequence clusters. Defaults to "cluster".

  • altloc_seed_column – If provided, weights are divided by the number of unique altloc seeds per pdb_id. Defaults to None (no altloc adjustment).

  • pdb_id_column – Column identifying structures, used only when altloc_seed_column is set. Defaults to "pdb_id".

Returns:

Tensor of per-row weights with shape (len(dataset_df),).

atomworks.ml.samplers.set_sampler_epoch(sampler: Sampler, epoch: int, add_random_offset: bool = False) → None[source]#

Control the random seed for a sampler.