agedi.diffusion.samplers¶

Pluggable reverse-diffusion samplers for AGeDi.

Each sampler encapsulates one complete reverse-diffusion outer step, from time t to t - dt. Samplers may call the score model (and optionally the force-field model) any number of times within that outer step.

Available samplers¶

EulerMaruyamaSampler: Standard Euler-Maruyama reverse-SDE step (the default). One score call per step. String alias: "em".
PredictorCorrectorSampler: EM predictor followed by N Langevin corrector steps at the new noise level. String alias: "pc".
HeunSampler: Second-order stochastic sampler (Karras et al. 2022). Two score calls per step; the averaged score is used in the final stochastic update. String alias: "heun".
ProbabilityFlowODESampler: Deterministic probability-flow ODE sampler (DDIM / Anderson 1982). One score call per step; applies half the diffusion coefficient with no noise. String alias: "ddim".
HeunODESampler: Second-order deterministic ODE sampler (Heun’s method on the PF-ODE). Two score calls per step; fully deterministic. String alias: "heun_ode".

Usage¶

Pass a sampler by string alias or instance to sample():

structures = model.sample(10, steps=200, sampler="heun")
structures = model.sample(10, steps=500,
    sampler=PredictorCorrectorSampler(model.score_model, model.noisers,
                                      corrector_steps=3))

from agedi.diffusion.samplers import Sampler

class MySampler(Sampler):
    def step(self, batch, dt, last):
        ...

Sampler.register("my_sampler",
    lambda score_fn, noisers, **kw: MySampler(score_fn, noisers))

Submodules¶

Classes¶

`Sampler`	Abstract base class for reverse-diffusion samplers.
`EulerMaruyamaSampler`	Standard Euler-Maruyama reverse-SDE sampler.
`ForcefieldCorrectorSampler`	EM predictor with force-field augmented Langevin corrector.
`HeunSampler`	Second-order stochastic sampler (Karras et al. 2022).
`HeunODESampler`	Second-order deterministic ODE sampler (Heun's method).
`ProbabilityFlowODESampler`	Deterministic probability-flow ODE sampler (DDIM / Anderson 1982).
`PredictorCorrectorSampler`	Euler-Maruyama predictor with Langevin corrector steps.

Package Contents¶

class agedi.diffusion.samplers.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¶

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

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

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

Bases: agedi.diffusion.samplers.base.Sampler

Standard Euler-Maruyama reverse-SDE sampler.

Performs one score-model evaluation per reverse step and delegates the position update to each noiser’s denoise() method. The update formula used (EM or DDPM posterior mean) is controlled by the sampler attribute of each PositionsNoiser.

This is the default sampler and exactly reproduces the behaviour of reverse_step() (minus guidance and timings).

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.

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

Euler-Maruyama reverse step.

Evaluate score model.
Apply each noiser’s denoising update in reverse order.
Wrap positions and rebuild the neighbour list.

class agedi.diffusion.samplers.ForcefieldCorrectorSampler(score_fn: Callable[[agedi.data.AtomsGraph], agedi.data.AtomsGraph], noisers: List[agedi.diffusion.noisers.Noiser], regressor_fn: Callable[[agedi.data.AtomsGraph], agedi.data.AtomsGraph] | None = None, corrector_steps: int = 1, corrector_step_size: float = 0.001, mixing_zeta: float = 1.0, temperature: float | None = None, terminal_steps: int = 0, terminal_step_size: float | None = None, terminal_dynamics: Literal['overdamped', 'langevin_md'] = 'overdamped', terminal_friction: float | None = None)¶

Bases: agedi.diffusion.samplers.base.Sampler

EM predictor with force-field augmented Langevin corrector.

Follows the standard predictor-corrector scheme from Song et al. (2021) but uses a force-field augmented score in the corrector steps:

Predictor — standard EM step using the neural score only:

\[x_{t-\Delta t} = x_t + \Delta t\,(g(t)^2\, s_\theta(x_t) - f(x_t,\,t)) + \sqrt{\Delta t}\, g(t)\, z\]

Corrector — corrector_steps Langevin steps at the new noise level t - \Delta t, using a score that blends the neural model with the force-field gradient:

\[\tilde{s}(x, t) = (1 - f(t))\, s_\theta(x) + f(t)\, F(x)\]

where \(f(t) = (1 - t)^\zeta\) increases as \(t \to 0\), so the force field has no influence at the start of sampling and full influence at the end. Temperature is not applied here; the relative weight of the force field is controlled entirely by mixing_zeta and corrector_step_size.

Terminal phase (optional) — applies terminal_steps additional refinement steps after the last diffusion step using only the force field. Two dynamics modes are supported via terminal_dynamics:

overdamped (default) — overdamped Langevin (no momenta):

\[x_{n+1} = x_n + \varepsilon\, F(x_n) + \sqrt{2\varepsilon T}\, z\]

where \(\varepsilon = \texttt{terminal\_step\_size}\) is a reduced step \(\varepsilon = \Delta t / T\). This parameterization gives T-independent stability (condition: \(\varepsilon < 2/k_{\max}\)) while preserving the correct Boltzmann stationary distribution \(\propto \exp(-U / T)\) through the \(\sqrt{T}\) factor in the noise amplitude.

langevin_md — standard Langevin MD (BAOAB integrator) with atomic masses from ASE. Momenta are initialised from the Maxwell-Boltzmann distribution at temperature T:

\[p_i \sim \mathcal{N}\!\left(0,\, m_i T\right)\]

followed by BAOAB integration:

\[\begin{split}p &\leftarrow p + \tfrac{\Delta t}{2}\, F \\ x &\leftarrow x + \tfrac{\Delta t}{2}\, p / m \\ p &\leftarrow e^{-\gamma \Delta t}\, p + \sqrt{T\, m\,(1 - e^{-2\gamma \Delta t})}\, z \\ x &\leftarrow x + \tfrac{\Delta t}{2}\, p / m \\ F &\leftarrow F(x) \\ p &\leftarrow p + \tfrac{\Delta t}{2}\, F\end{split}\]

When no regressor is attached (regressor_fn=None), the corrector falls back to a pure Langevin step with the neural score, and terminal steps are skipped.

Parameters:

score_fn (callable) – Neural score model: score_fn(batch) -> batch with pos_score set.
noisers (list[Noiser]) – Noisers in forward order; iterated in reverse for denoising.
regressor_fn (callable or None) – Force-field model: regressor_fn(batch) -> batch with forces_prediction set. Passed by _resolve_sampler().
corrector_steps (int) – Number of Langevin corrector steps per predictor step. Default: 1.
corrector_step_size (float) – Step size for each Langevin corrector step. Subject to the Langevin stability bound corrector_step_size < 2·var(t_{i-1}). Default: 1e-3.
mixing_zeta (float) – Exponent for the mixing schedule f(t) = (1 - t)**mixing_zeta. 1.0 (default) gives linear mixing; higher values concentrate force-field influence near the end of the trajectory.
temperature (float) –
Temperature T used only in the terminal phase (not in the corrector steps, where dividing forces by T causes blowup at low T).
- overdamped: appears as \(\sqrt{T}\) in the noise amplitude (see the update rule above) to preserve the correct Boltzmann distribution. The step-size stability bound is T-independent.
- langevin_md: the thermal energy \(k_B T\) in the same units as the model forces (typically eV). Sets the noise amplitude in the Ornstein-Uhlenbeck step and the Maxwell-Boltzmann velocity initialisation.
Default: None — falls back to 1.0 with a UserWarning reminding you to set an explicit value.
terminal_steps (int) – Number of refinement steps applied after the last diffusion step. 0 (default) disables them.
terminal_step_size (float or None) –
Step size for each terminal step.
- overdamped: the reduced step \(\varepsilon = \Delta t / T\). Stability requires \(\varepsilon < 2 / k_{\max}\) where \(k_{\max}\) is the largest force-constant eigenvalue — a bound that is independent of T. None (default) auto-selects 1e-3, which is conservative for most force fields.
- langevin_md: physical time step in units consistent with the model forces and ASE masses. When forces are in eV/Å and masses in amu the unit is femtoseconds. None (default) auto-selects 1.0 (1 fs), a safe starting point for typical ML potentials.
terminal_dynamics ("overdamped" or "langevin_md") – Dynamics used for the terminal phase. "overdamped" (default) uses overdamped Langevin (no momenta). "langevin_md" uses standard Langevin MD (BAOAB) with real atomic masses from ASE and momenta initialised from the Maxwell-Boltzmann distribution.
terminal_friction (float or None) – Friction coefficient \(\gamma\) for the Langevin thermostat in langevin_md dynamics (ignored for overdamped). Units are the inverse of terminal_step_size; when terminal_step_size is in fs a physically reasonable range is 0.001–0.1 fs⁻¹ (1–100 ps⁻¹). None (default) auto-selects \(\gamma = 0.1 / \Delta t\), giving a dimensionless reduced friction \(\gamma \Delta t = 0.1\) (moderately damped).
alias (String)
------------

:param "ffpc" — registered in agedi.diffusion.samplers.:

Notes

This sampler calls regressor_fn directly, not force_field_guidance_step. The LBFGS step sizer is therefore not required; uses_force_field is False.

uses_force_field: ClassVar[bool] = False¶

regressor_fn = None¶

corrector_steps = 1¶

corrector_step_size = 0.001¶

mixing_zeta = 1.0¶

temperature = None¶

terminal_steps = 0¶

terminal_step_size = None¶

terminal_dynamics = 'overdamped'¶

terminal_friction = None¶

_pos_noiser¶

_corrector_dt: torch.Tensor | None = None¶

_pending_frames: List = []¶

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

Predictor-corrector step with force-field augmented corrector.

EM predictor with neural score.
Advance batch.time to t - dt.
For each corrector iteration: evaluate neural score, blend with force-field gradient, apply Langevin step.
If last and terminal_steps > 0: terminal phase with the chosen dynamics. Intermediate frames are stored in _pending_frames for trajectory capture.

_run_terminal(batch: agedi.data.AtomsGraph) → None¶

Dispatch to the selected terminal dynamics.

Prepends a bridge frame (the denoised state before terminal dynamics start) unless corrector frames already captured it as their last entry — which is the case when corrector_steps > 0.

_terminal_overdamped(batch: agedi.data.AtomsGraph) → None¶

Overdamped Langevin terminal steps (no momenta).

terminal_step_size is the reduced step ε = dt/T. Update: x += ε·F + sqrt(2εT)·z, giving stationary ∝ exp(-U/T). Stability condition: ε < 2/k_max, independent of T.

_terminal_langevin_md(batch: agedi.data.AtomsGraph) → None¶

Standard Langevin MD terminal steps (BAOAB, real atomic masses).

Masses are looked up from ASE using batch.x (atomic numbers). Temperature T is the thermal energy \(k_B T\) in model force units. Time step units must be consistent with forces and masses (e.g. fs when forces are in eV/Å and masses in amu).

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

Bases: agedi.diffusion.samplers.base.Sampler

Second-order stochastic sampler (Karras et al. 2022).

Uses two score evaluations per step to achieve second-order accuracy for the stochastic reverse SDE:

First score at (x_t, t): s_1 = s_θ(x_t, t)
Deterministic EM predictor (no noise) for SDE noisers: x̃ = x_t + dt·(g(t)²·s_1 - f(x_t, t))
Second score at (x̃, t - dt): s_2 = s_θ(x̃, t - dt)
Average scores: s_avg = 0.5·(s_1 + s_2)
Stochastic EM step using averaged score: x_{t-dt} = x_t + dt·(g(t)²·s_avg - f(x_t, t)) + √dt·g(t)·z (no noise when last=True)

The score averaging provides second-order accuracy in the SDE sense, analogous to the stochastic Heun method from Karras et al. (2022) “Elucidating the Design Space of Diffusion-Based Generative Models”.

For noisers with a .sde attribute (continuous SDE-based noisers such as PositionsNoiser), the two-step procedure is applied. For noisers without .sde (e.g. discrete types noiser), a single EM step is applied using the first score evaluation, since score averaging is less principled for discrete diffusion.

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.

Notes

noiser.denoise() reads positions and scores directly from tensor attributes and does not require an up-to-date neighbour list. This allows the original positions x_t to be restored (step 5) without rebuilding the graph before the final denoising call in step 5 — the graph is only rebuilt once at the very end of the step.

step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) → agedi.data.AtomsGraph¶: Second-order stochastic Heun step.

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

Bases: agedi.diffusion.samplers.base.Sampler

Second-order deterministic ODE sampler (Heun’s method).

Applies Heun’s method to the probability-flow ODE, achieving second-order accuracy with two score evaluations per step:

First score at (x_t, t): s_1 = s_θ(x_t, t)
ODE predictor: x̃ = x_t + dt·(0.5·g(t)²·s_1 - f(x_t, t))
Second score at (x̃, t - dt): s_2 = s_θ(x̃, t - dt)
Averaged score: s_avg = 0.5·(s_1 + s_2)
ODE corrector: x_{t-dt} = x_t + dt·(0.5·g(t)²·s_avg - f(x_t, t))

Like ProbabilityFlowODESampler, this sampler is fully deterministic. It uses two score evaluations per step vs one for the first-order ODE sampler, but achieves better quality at lower step counts.

For noisers without .sde (discrete types), a single deterministic EM step is applied using the first score evaluation.

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.

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

Second-order Heun ODE step.

The last parameter is accepted for API compatibility but has no effect — this sampler is always deterministic.

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

Bases: agedi.diffusion.samplers.base.Sampler

Deterministic probability-flow ODE sampler (DDIM / Anderson 1982).

Integrates the reverse-time ODE:

\[\frac{d\mathbf{x}}{dt} = f(\mathbf{x}, t) - \tfrac{1}{2}\, g(t)^2\, \nabla_{\mathbf{x}} \log p_t(\mathbf{x})\]

which is the deterministic counterpart of the reverse-time SDE. The factor of 0.5 on the diffusion term (vs 1.0 in Euler-Maruyama) removes the stochastic component while preserving the marginal distributions.

The update for one step t → t - dt is:

\[\mathbf{x}_{t - \Delta t} = \mathbf{x}_t + \Delta t \bigl( \tfrac{1}{2}\, g(t)^2 \, s_\theta(\mathbf{x}_t, t) - f(\mathbf{x}_t, t) \bigr)\]

For noisers with a .sde attribute (continuous SDE-based noisers, such as positions), the ODE update is applied directly using the SDE’s drift and diffusion methods. For other noisers (e.g. discrete types), this sampler falls back to a deterministic EM step (noiser.denoise(batch, dt, last=True)).

Unlike EulerMaruyamaSampler, this sampler is fully deterministic: repeated calls with identical inputs produce identical outputs.

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.

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

Probability-flow ODE step.

Evaluate score model.
For each SDE-based noiser: apply the ODE update (half diffusion, no noise).
For other noisers: apply a deterministic EM step.
Wrap positions and rebuild the neighbour list.

The last parameter is accepted for API compatibility but has no effect — this sampler is always deterministic.

static _ode_step(batch: agedi.data.AtomsGraph, noiser: agedi.diffusion.noisers.Noiser, dt: torch.Tensor) → agedi.data.AtomsGraph¶

Apply the probability-flow ODE update for one SDE-based noiser.

Computes:

\[\mathbf{r}_{t-\Delta t} = \mathbf{r}_t + \Delta t \bigl( \tfrac{1}{2}\, g(t)^2\, s_\theta - f(\mathbf{r}_t, t) \bigr)\]

Parameters:

batch (AtomsGraph) – Current batch state.
noiser (Noiser) – An SDE-based noiser with a .sde attribute.
dt (torch.Tensor) – Step size.

Returns:

Batch with the updated attribute.

Return type:

AtomsGraph

class agedi.diffusion.samplers.PredictorCorrectorSampler(score_fn: Callable[[agedi.data.AtomsGraph], agedi.data.AtomsGraph], noisers: List[agedi.diffusion.noisers.Noiser], corrector_steps: int = 1, corrector_step_size: float = 0.001)¶

Bases: agedi.diffusion.samplers.base.Sampler

Euler-Maruyama predictor with Langevin corrector steps.

After the standard EM predictor step advances the state from t to t - dt, this sampler applies corrector_steps Langevin corrector steps at the new noise level t - dt. This is the conventional predictor-corrector scheme from Song et al. (2021).

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.
corrector_steps (int, optional) – Number of Langevin corrector passes per predictor step. Default: 1.
corrector_step_size (float, optional) – Step size for each Langevin corrector step. Default: 1e-3.

Notes

The corrector runs at the new time t_{i-1} = t_i - dt (after the predictor), which is the standard convention. The legacy corrector_steps integer parameter on sample() uses the same convention via this class.

corrector_steps = 1¶

corrector_step_size = 0.001¶

_corrector_dt: torch.Tensor | None = None¶

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

Predictor-corrector step.

EM predictor: evaluate score, apply noiser.denoise(), wrap & rebuild.
Advance batch.time to t - dt.
For each corrector iteration: evaluate score, apply Langevin step via noiser.langevin_step(), wrap & rebuild.