agedi.diffusion.samplers.base
=============================

.. py:module:: agedi.diffusion.samplers.base

.. autoapi-nested-parse::

   Abstract base class for reverse-diffusion samplers.



Classes
-------

.. autoapisummary::

   agedi.diffusion.samplers.base.Sampler


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

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

   Bases: :py:obj:`abc.ABC`


   Abstract base class for reverse-diffusion samplers.

   A sampler encapsulates one complete reverse-diffusion outer step, from
   time ``t`` to ``t - dt``.  It may call the score model (and optionally
   the force-field model) any number of times within that outer step —
   enabling predictor-corrector schemes, second-order methods, and
   deterministic ODE integration.

   The outer loop in :meth:`~agedi.diffusion.Diffusion._sample_batch` sets
   ``batch.time`` before each call to :meth:`step` and handles trajectory
   saving, force-field guidance, and post-diffusion relaxation.  The sampler
   is responsible only for the diffusion update itself.

   :param score_fn: A function ``score_fn(batch) -> batch`` that runs the score model and
                    stores per-key scores in ``batch.{key}_score``.
   :type score_fn: callable
   :param noisers: The noisers attached to the diffusion model, in forward order.  The
                   sampler iterates them in reverse when denoising.
   :type noisers: list of Noiser
   :param Registry:
   :param --------:
   :param Use :meth:`register` to make a sampler accessible by string name so that:
   :param it can be selected via ``Diffusion.sample(sampler="name")``.:
   :param Built-in aliases (populated in ``samplers/__init__.py``):
   :param \* ``"em"``       — :class:`~agedi.diffusion.samplers.EulerMaruyamaSampler`:
   :param \* ``"pc"``       — :class:`~agedi.diffusion.samplers.PredictorCorrectorSampler`:
   :param \* ``"heun"``     — :class:`~agedi.diffusion.samplers.HeunSampler`:
   :param \* ``"ddim"``     — :class:`~agedi.diffusion.samplers.ProbabilityFlowODESampler`:
   :param \* ``"heun_ode"`` — :class:`~agedi.diffusion.samplers.HeunODESampler`:
   :param \* ``"ffpc"``     — :class:`~agedi.diffusion.samplers.ForcefieldCorrectorSampler`:


   .. py:attribute:: _registry
      :type:  ClassVar[Dict[str, Callable[Ellipsis, Sampler]]]


   .. py:attribute:: uses_force_field
      :type:  ClassVar[bool]
      :value: False



   .. py:attribute:: score_fn


   .. py:attribute:: noisers


   .. py:method:: register(name: str, factory: Callable[Ellipsis, Sampler]) -> None
      :classmethod:


      Register a sampler factory under a string alias.

      The factory is called with ``score_fn`` and ``noisers`` as keyword
      arguments and may accept additional sampler-specific keyword arguments
      (e.g. ``corrector_steps``).

      :param name: Alias used to look up this sampler, e.g. ``"heun"``.
      :type name: str
      :param factory: ``factory(score_fn, noisers, **kwargs) -> Sampler``.
      :type factory: callable

      .. rubric:: Examples

      Register a custom sampler so it can be selected by name::

          from agedi.diffusion.samplers import Sampler

          class MySampler(Sampler):
              ...

          Sampler.register("my_sampler",
              lambda score_fn, noisers, **kw: MySampler(score_fn, noisers))



   .. py:method:: _check_finite(batch: agedi.data.AtomsGraph, stage: str) -> None
      :staticmethod:


      Raise ``RuntimeError`` if any atom positions are non-finite.

      Call this immediately before every ``batch.update_graph()`` call.
      The neighbour-list C/CUDA kernel does not validate its inputs; passing
      NaN or infinite positions crashes the process with no Python traceback.

      :param batch: Batch whose ``pos`` tensor will be inspected.
      :type batch: AtomsGraph
      :param stage: Short description of where the check is called, included in the
                    error message to help locate the divergence.
      :type stage: str



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


      Perform one complete reverse-diffusion step from ``t`` to ``t - dt``.

      ``batch.time`` is set to the current time ``t`` by ``_sample_batch``
      before this method is called.  The implementation must return a batch
      with a valid neighbour list — i.e. it must call
      ``batch.wrap_positions()`` and ``batch.update_graph()`` (or equivalent)
      before returning.

      :param batch: Current state of the batch.  ``batch.time`` is the current time.
      :type batch: AtomsGraph
      :param dt: Positive step size: ``dt = t_i - t_{i+1}``.
      :type dt: torch.Tensor
      :param last: Whether this is the final reverse-diffusion step.  When ``True``,
                   stochastic samplers suppress the noise term so the final sample is
                   deterministic (matches the convention in the underlying noisers).
      :type last: bool

      :returns: Updated batch with a valid neighbour list.
      :rtype: AtomsGraph