agedi.diffusion.samplers.base

Abstract base class for reverse-diffusion samplers.

Classes

Sampler

Abstract base class for reverse-diffusion samplers.

Module Contents

class agedi.diffusion.samplers.base.Sampler(score_fn: Callable[[agedi.data.AtomsGraph], agedi.data.AtomsGraph], noisers: List[agedi.diffusion.noisers.Noiser])

Bases: abc.ABC

Abstract base class for reverse-diffusion samplers.

A sampler encapsulates one complete reverse-diffusion outer step, from time t to t - dt. It may call the score model (and optionally the force-field model) any number of times within that outer step — enabling predictor-corrector schemes, second-order methods, and deterministic ODE integration.

The outer loop in _sample_batch() sets batch.time before each call to step() and handles trajectory saving, force-field guidance, and post-diffusion relaxation. The sampler is responsible only for the diffusion update itself.

Parameters:

  • score_fn (callable) – A function score_fn(batch) -> batch that runs the score model and stores per-key scores in batch.{key}_score.

  • noisers (list of Noiser) – The noisers attached to the diffusion model, in forward order. The sampler iterates them in reverse when denoising.

  • Registry

  • --------

:param Use register() to make a sampler accessible by string name so that: :param it can be selected via Diffusion.sample(sampler="name").: :param Built-in aliases (populated in samplers/__init__.py): :param * "em" — EulerMaruyamaSampler: :param * "pc" — PredictorCorrectorSampler: :param * "heun" — HeunSampler: :param * "ddim" — ProbabilityFlowODESampler: :param * "heun_ode" — HeunODESampler: :param * "ffpc" — ForcefieldCorrectorSampler:

_registry: ClassVar[Dict[str, Callable[Ellipsis, Sampler]]]
uses_force_field: ClassVar[bool] = False
score_fn
noisers
classmethod register(name: str, factory: Callable[Ellipsis, Sampler]) → None

Register a sampler factory under a string alias.

The factory is called with score_fn and noisers as keyword arguments and may accept additional sampler-specific keyword arguments (e.g. corrector_steps).

Parameters:

  • name (str) – Alias used to look up this sampler, e.g. "heun".

  • factory (callable) – factory(score_fn, noisers, **kwargs) -> Sampler.

Examples

Register a custom sampler so it can be selected by name:

from agedi.diffusion.samplers import Sampler

class MySampler(Sampler):
    ...

Sampler.register("my_sampler",
    lambda score_fn, noisers, **kw: MySampler(score_fn, noisers))
static _check_finite(batch: agedi.data.AtomsGraph, stage: str) → None

Raise RuntimeError if any atom positions are non-finite.

Call this immediately before every batch.update_graph() call. The neighbour-list C/CUDA kernel does not validate its inputs; passing NaN or infinite positions crashes the process with no Python traceback.

Parameters:

  • batch (AtomsGraph) – Batch whose pos tensor will be inspected.

  • stage (str) – Short description of where the check is called, included in the error message to help locate the divergence.

abstractmethod step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) → agedi.data.AtomsGraph

Perform one complete reverse-diffusion step from t to t - dt.

batch.time is set to the current time t by _sample_batch before this method is called. The implementation must return a batch with a valid neighbour list — i.e. it must call batch.wrap_positions() and batch.update_graph() (or equivalent) before returning.

Parameters:

  • batch (AtomsGraph) – Current state of the batch. batch.time is the current time.

  • dt (torch.Tensor) – Positive step size: dt = t_i - t_{i+1}.

  • last (bool) – Whether this is the final reverse-diffusion step. When True, stochastic samplers suppress the noise term so the final sample is deterministic (matches the convention in the underlying noisers).

Returns:

Updated batch with a valid neighbour list.

Return type:

AtomsGraph