agedi.diffusion.samplers.ode

Probability-flow ODE samplers — deterministic reverse diffusion.

Classes

ProbabilityFlowODESampler

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

HeunODESampler

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

Module Contents

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

  1. Evaluate score model.

  2. For each SDE-based noiser: apply the ODE update (half diffusion, no noise).

  3. For other noisers: apply a deterministic EM step.

  4. 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.ode.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:

  1. First score at (x_t, t): s_1 = s_θ(x_t, t)

  2. ODE predictor: x̃ = x_t + dt·(0.5·g(t)²·s_1 - f(x_t, t))

  3. Second score at (x̃, t - dt): s_2 = s_θ(x̃, t - dt)

  4. Averaged score: s_avg = 0.5·(s_1 + s_2)

  5. 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.