Skip to content

module adabmDCA.sampling

function sampling_profile_categorical

sampling_profile_categorical(
    params: dict[str, Tensor],
    nsamples: int,
    beta: float
) → Tensor

Sample integer states (N,L) from the uncoupled profile, without encoding.

function sampling_profile

sampling_profile(params: dict[str, Tensor], nsamples: int, beta: float) → Tensor

Samples from the profile model defined by the local biases only.

Args:

  • params (Dict[str, torch.Tensor]): Parameters of the model. - "bias": Tensor of shape (L, q) - local biases.
  • nsamples (int): Number of samples to generate.
  • beta (float): Inverse temperature.

Returns:

  • torch.Tensor: Sampled one-hot encoded sequences of shape (nsamples, L, q).

function gibbs_step_uniform_sites

gibbs_step_uniform_sites(
    chains: Tensor,
    params: dict[str, Tensor],
    beta: float = 1.0
) → Tensor

Performs a single mutation using the Gibbs sampler. In this version, the mutation is attempted at the same sites for all chains.

Args:

  • chains (torch.Tensor): One-hot encoded sequences of shape (batch_size, L, q).
  • params (Dict[str, torch.Tensor]): Parameters of the model. - "bias": Tensor of shape (L, q) - local biases. - "coupling_matrix": Tensor of shape (L, q, L, q) - coupling matrix.
  • beta (float, optional): Inverse temperature. Defaults to 1.0.

Returns:

  • torch.Tensor: Updated chains.

function gibbs_step_independent_sites

gibbs_step_independent_sites(
    chains: Tensor,
    params: dict[str, Tensor],
    beta: float = 1.0
) → Tensor

Performs a single mutation using the Gibbs sampler. This version selects different random sites for each chain. It is less efficient than the 'gibbs_step_uniform_sites' function, but it is more suitable for mutating starting from the same wild-type sequence since mutations are independent across chains.

Args:

  • chains (torch.Tensor): One-hot encoded sequences of shape (batch_size, L, q).
  • params (Dict[str, torch.Tensor]): Parameters of the model. - "bias": Tensor of shape (L, q) - local biases. - "coupling_matrix": Tensor of shape (L, q, L, q) - coupling matrix.
  • beta (float, optional): Inverse temperature. Defaults to 1.0.

Returns:

  • torch.Tensor: Updated chains.

function gibbs_sampling

gibbs_sampling(
    chains: Tensor,
    params: dict[str, Tensor],
    nsweeps: int,
    beta: float = 1.0
) → Tensor

Gibbs sampling. Attempts L * nsweeps mutations to each sequence in 'chains'.

Args:

  • chains (torch.Tensor): Initial one-hot encoded samples of size (batch_size, L, q).
  • params (Dict[str, torch.Tensor]): Parameters of the model. - "bias": Tensor of shape (L, q) - local biases. - "coupling_matrix": Tensor of shape (L, q, L, q) - coupling matrix.
  • nsweeps (int): Number of sweeps, where one sweep corresponds to attempting L mutations.
  • beta (float, optional): Inverse temperature. Defaults to 1.0.

Returns:

  • torch.Tensor: Updated chains.

function metropolis_step_uniform_sites

metropolis_step_uniform_sites(
    chains: Tensor,
    params: dict[str, Tensor],
    beta: float = 1.0
) → Tensor

Performs a single mutation using the Metropolis sampler. In this version, the mutation is attempted at the same sites for all chains.

Args:

  • chains (torch.Tensor): One-hot encoded sequences of shape (batch_size, L, q).
  • params (Dict[str, torch.Tensor]): Parameters of the model. - "bias": Tensor of shape (L, q) - local biases. - "coupling_matrix": Tensor of shape (L, q, L, q) - coupling matrix.
  • beta (float, optional): Inverse temperature. Defaults to 1.0.

Returns:

  • torch.Tensor: Updated chains.

function metropolis_step_independent_sites

metropolis_step_independent_sites(
    chains: Tensor,
    params: dict[str, Tensor],
    beta: float = 1.0
) → Tensor

Performs a single mutation using the Metropolis sampler. This version selects different random sites for each chain. It is less efficient than the 'metropolis_step_uniform_sites' function, but it is more suitable for mutating starting from the same wild-type sequence since mutations are independent across chains.

Args:

  • chains (torch.Tensor): One-hot encoded sequences of shape (batch_size, L, q).
  • params (Dict[str, torch.Tensor]): Parameters of the model. - "bias": Tensor of shape (L, q) - local biases. - "coupling_matrix": Tensor of shape (L, q, L, q) - coupling matrix.
  • beta (float, optional): Inverse temperature. Defaults to 1.0.

Returns:

  • torch.Tensor: Updated chains.

function metropolis_sampling

metropolis_sampling(
    chains: Tensor,
    params: dict[str, Tensor],
    nsweeps: int,
    beta: float = 1.0
) → Tensor

Metropolis sampling. Attempts L * nsweeps mutations to each sequence in 'chains'.

Args:

  • chains (torch.Tensor): One-hot encoded sequences of shape (batch_size, L, q).
  • params (Dict[str, torch.Tensor]): Parameters of the model. - "bias": Tensor of shape (L, q) - local biases. - "coupling_matrix": Tensor of shape (L, q, L, q) - coupling matrix.
  • nsweeps (int): Number of sweeps to be performed, where one sweep corresponds to attempting L mutations.
  • beta (float, optional): Inverse temperature. Defaults to 1.0.

Returns:

  • torch.Tensor: Updated chains.

function metropolized_gibbs_step_uniform_sites

metropolized_gibbs_step_uniform_sites(
    chains: Tensor,
    params: dict[str, Tensor],
    beta: float = 1.0
) → Tensor

Performs a single Metropolized Gibbs update at the same site for all chains.

A new residue b different from the current one a is proposed from the site conditional restricted to the other residues and accepted with min(1, (1 - p_a) / (1 - p_b)) (Liu 1996). Both normalizers are summed directly so that dominant residues do not cancel.

Args:

  • chains (torch.Tensor): One-hot encoded sequences of shape (batch_size, L, q).
  • params (Dict[str, torch.Tensor]): Parameters of the model. - "bias": Tensor of shape (L, q) - local biases. - "coupling_matrix": Tensor of shape (L, q, L, q) - coupling matrix.
  • beta (float, optional): Inverse temperature. Defaults to 1.0.

Returns:

  • torch.Tensor: Updated chains.

function metropolized_gibbs_sampling

metropolized_gibbs_sampling(
    chains: Tensor,
    params: dict[str, Tensor],
    nsweeps: int,
    beta: float = 1.0
) → Tensor

Metropolized Gibbs sampling. Attempts L * nsweeps updates to each sequence in 'chains'.

Args:

  • chains (torch.Tensor): One-hot encoded sequences of shape (batch_size, L, q).
  • params (Dict[str, torch.Tensor]): Parameters of the model. - "bias": Tensor of shape (L, q) - local biases. - "coupling_matrix": Tensor of shape (L, q, L, q) - coupling matrix.
  • nsweeps (int): Number of sweeps, where one sweep corresponds to attempting L updates.
  • beta (float, optional): Inverse temperature. Defaults to 1.0.

Returns:

  • torch.Tensor: Updated chains.

function metropolized_gibbs_sampling_categorical

metropolized_gibbs_sampling_categorical(
    states: Tensor,
    params: dict[str, Tensor],
    nsweeps: int,
    beta: float = 1.0
) → Tensor

Metropolized Gibbs sampling of integer states of shape (N, L).

Each update draws one site, shared by all chains, proposes a different state b from the site conditional restricted to the other states and accepts it with min(1, (1 - p_a) / (1 - p_b)) (Liu 1996). This dominates Gibbs sampling in the Peskun order. Reference implementation for the Triton replica kernel; couplings must have zero within-site blocks.

function get_sampler

get_sampler(sampling_method: str) → Callable

Returns the sampling function corresponding to the chosen method.

Args:

  • sampling_method (str): String indicating the sampling method. Choose between 'metropolis', 'gibbs' and 'metropolized_gibbs'.

Raises:

  • KeyError: Unknown sampling method.

Returns:

  • Callable: Sampling function.

function prepare_sampler

prepare_sampler(sampling_method: str, device: device) → Callable

Select the fastest sampler for device.

CUDA uses the fused Triton kernels when Triton is installed. MPS uses the fused Metal kernels when torch.mps.compile_shader is available, unless disabled with ADABMDCA_MPS=0 (unsupported shapes/dtypes use PyTorch). CPU uses the multithreaded Numba kernels of adabmDCA.numba_kernels when Numba is installed (pip install adabmDCA[cpu]) and not disabled with ADABMDCA_NUMBA=0. Otherwise the TorchScript samplers of this module run. All of them perform the same random-site updates and sample the same distribution.

function prepare_fixed_model_sampler

prepare_fixed_model_sampler(
    sampling_method: str,
    device: device,
    dtype: str,
    params: dict[str, Tensor]
) → tuple[Callable, dict[str, Tensor]]

Prepare a sampler and parameters for a model that will not be updated.

BF16 mode keeps biases, chains, statistics, and energy calculations in float32. Only the fixed coupling matrix is rounded to BF16, once, before sampling. This is the inference counterpart of prepare_training_sampler, where the coupling copy must instead be refreshed after every parameter update.

function prepare_training_sampler

prepare_training_sampler(
    sampling_method: str,
    device: device,
    dtype: str = 'float32'
) → Callable

Prepare sampling for training with optional BF16 coupling storage.

In bfloat16 mode, parameters, chains and statistics outside the sampler stay float32. A fresh BF16 coupling copy is made after each parameter update; biases and all field/acceptance arithmetic remain float32. The rounded coupling matrix approximates the master model, so trajectories can change.

This file was automatically generated via lazydocs.