agedi.diffusion.samplers.em
===========================

.. py:module:: agedi.diffusion.samplers.em

.. autoapi-nested-parse::

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



Classes
-------

.. autoapisummary::

   agedi.diffusion.samplers.em.EulerMaruyamaSampler


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

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

   Bases: :py:obj:`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 :meth:`~agedi.diffusion.noisers.Noiser.denoise`
   method.  The update formula used (EM or DDPM posterior mean) is controlled
   by the ``sampler`` attribute of each :class:`~agedi.diffusion.noisers.PositionsNoiser`.

   This is the default sampler and exactly reproduces the behaviour of
   :meth:`~agedi.diffusion.Diffusion.reverse_step` (minus guidance and timings).

   :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

      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.