agedi.diffusion.samplers.pc
===========================

.. py:module:: agedi.diffusion.samplers.pc

.. autoapi-nested-parse::

   Predictor-corrector sampler — EM predictor with Langevin correctors.



Classes
-------

.. autoapisummary::

   agedi.diffusion.samplers.pc.PredictorCorrectorSampler


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

.. py:class:: 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: :py:obj:`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).

   :param score_fn: Score-model forward function.
   :type score_fn: callable
   :param noisers: Noisers in forward order.
   :type noisers: list of Noiser
   :param corrector_steps: Number of Langevin corrector passes per predictor step.  Default: 1.
   :type corrector_steps: int, optional
   :param corrector_step_size: Step size for each Langevin corrector step.  Default: 1e-3.
   :type corrector_step_size: float, optional

   .. rubric:: 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 :meth:`~agedi.diffusion.Diffusion.sample`
   uses the same convention via this class.


   .. py:attribute:: corrector_steps
      :value: 1



   .. py:attribute:: corrector_step_size
      :value: 0.001



   .. py:attribute:: _corrector_dt
      :type:  Optional[torch.Tensor]
      :value: None



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