agedi.diffusion.samplers.em

Euler-Maruyama sampler — one score call per reverse step.

Classes

EulerMaruyamaSampler

Standard Euler-Maruyama reverse-SDE sampler.

Module Contents

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

  1. Evaluate score model.

  2. Apply each noiser’s denoising update in reverse order.

  3. Wrap positions and rebuild the neighbour list.