agedi.diffusion.samplers.heun
=============================

.. py:module:: agedi.diffusion.samplers.heun

.. autoapi-nested-parse::

   Stochastic Heun sampler — second-order reverse-SDE integration.



Classes
-------

.. autoapisummary::

   agedi.diffusion.samplers.heun.HeunSampler


Module Contents
---------------

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

   Bases: :py:obj:`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:

   1. First score at ``(x_t, t)``:
      ``s_1 = s_θ(x_t, t)``
   2. Deterministic EM predictor (no noise) for SDE noisers:
      ``x̃ = x_t + dt·(g(t)²·s_1 - f(x_t, t))``
   3. Second score at ``(x̃, t - dt)``:
      ``s_2 = s_θ(x̃, t - dt)``
   4. Average scores:
      ``s_avg = 0.5·(s_1 + s_2)``
   5. 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 :class:`~agedi.diffusion.noisers.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.

   :param score_fn: Score-model forward function.
   :type score_fn: callable
   :param noisers: Noisers in forward order.
   :type noisers: list of Noiser

   .. rubric:: 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.


   .. py:method:: step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) -> agedi.data.AtomsGraph

      Second-order stochastic Heun step.