agedi.diffusion.samplers.ode
============================

.. py:module:: agedi.diffusion.samplers.ode

.. autoapi-nested-parse::

   Probability-flow ODE samplers — deterministic reverse diffusion.



Classes
-------

.. autoapisummary::

   agedi.diffusion.samplers.ode.ProbabilityFlowODESampler
   agedi.diffusion.samplers.ode.HeunODESampler


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

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

   Bases: :py:obj:`agedi.diffusion.samplers.base.Sampler`


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

   Integrates the reverse-time ODE:

   .. math::

       \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:

   .. math::

       \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 :class:`EulerMaruyamaSampler`, this sampler is fully deterministic:
   repeated calls with identical inputs produce identical outputs.

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


   .. py:method:: 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.



   .. py:method:: _ode_step(batch: agedi.data.AtomsGraph, noiser: agedi.diffusion.noisers.Noiser, dt: torch.Tensor) -> agedi.data.AtomsGraph
      :staticmethod:


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

      Computes:

      .. math::

          \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)

      :param batch: Current batch state.
      :type batch: AtomsGraph
      :param noiser: An SDE-based noiser with a ``.sde`` attribute.
      :type noiser: Noiser
      :param dt: Step size.
      :type dt: torch.Tensor

      :returns: Batch with the updated attribute.
      :rtype: AtomsGraph



.. py:class:: HeunODESampler(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 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 :class:`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.

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


   .. py:method:: 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.