agedi.diffusion.samplers.pc

Predictor-corrector sampler — EM predictor with Langevin correctors.

Classes

PredictorCorrectorSampler

Euler-Maruyama predictor with Langevin corrector steps.

Module Contents

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

  1. EM predictor: evaluate score, apply noiser.denoise(), wrap & rebuild.

  2. Advance batch.time to t - dt.

  3. For each corrector iteration: evaluate score, apply Langevin step via noiser.langevin_step(), wrap & rebuild.